promnesia 1.1.20230129__py3-none-any.whl → 1.2.20240810__py3-none-any.whl

promnesia/database/common.py

@@ -0,0 +1,66 @@
+from datetime import datetime
+from typing import Sequence, Tuple
+
+from sqlalchemy import (
+    Column,
+    Integer,
+    Row,
+    String,
+)
+
+# TODO maybe later move DbVisit here completely?
+# kinda an issue that it's technically an "api" because hook in config can patch up DbVisit
+from ..common import DbVisit, Loc
+
+
+def get_columns() -> Sequence[Column]:
+    # fmt: off
+    res: Sequence[Column] = [
+        Column('norm_url'     , String()),
+        Column('orig_url'     , String()),
+        Column('dt'           , String()),
+        Column('locator_title', String()),
+        Column('locator_href' , String()),
+        Column('src'          , String()),
+        Column('context'      , String()),
+        Column('duration'     , Integer())
+    ]
+    # fmt: on
+    assert len(res) == len(DbVisit._fields) + 1  # +1 because Locator is 'flattened'
+    return res
+
+
+def db_visit_to_row(v: DbVisit) -> Tuple:
+    # ugh, very hacky...
+    # we want to make sure the resulting tuple only consists of simple types
+    # so we can use dbengine directly
+    dt_s = v.dt.isoformat()
+    row = (
+        v.norm_url,
+        v.orig_url,
+        dt_s,
+        v.locator.title,
+        v.locator.href,
+        v.src,
+        v.context,
+        v.duration,
+    )
+    return row
+
+
+def row_to_db_visit(row: Sequence) -> DbVisit:
+    (norm_url, orig_url, dt_s, locator_title, locator_href, src, context, duration) = row
+    dt_s = dt_s.split()[0]  # backwards compatibility: previously it could be a string separated with tz name
+    dt = datetime.fromisoformat(dt_s)
+    return DbVisit(
+        norm_url=norm_url,
+        orig_url=orig_url,
+        dt=dt,
+        locator=Loc(
+            title=locator_title,
+            href=locator_href,
+        ),
+        src=src,
+        context=context,
+        duration=duration,
+    )

promnesia/database/dump.py

@@ -0,0 +1,187 @@
+from pathlib import Path
+import sqlite3
+from typing import Dict, Iterable, List, Optional, Set
+
+from more_itertools import chunked
+
+from sqlalchemy import (
+    Engine,
+    MetaData,
+    Table,
+    create_engine,
+    event,
+    exc,
+    func,
+    select,
+)
+from sqlalchemy.dialects import sqlite as dialect_sqlite
+
+from ..common import (
+    DbVisit,
+    Loc,
+    Res,
+    SourceName,
+    get_logger,
+    now_tz,
+)
+from .common import get_columns, db_visit_to_row
+from .. import config
+
+
+# NOTE: I guess the main performance benefit from this is not creating too many tmp lists and avoiding overhead
+# since as far as sql is concerned it should all be in the same transaction. only a guess
+# not sure it's the proper way to handle it
+# see test_index_many
+_CHUNK_BY = 10
+
+# I guess 1 hour is definitely enough
+_CONNECTION_TIMEOUT_SECONDS = 3600
+
+SRC_ERROR = 'error'
+
+
+# using WAL keeps database readable while we're writing in it
+# this is tested by test_query_while_indexing
+def enable_wal(dbapi_con, con_record) -> None:
+    dbapi_con.execute('PRAGMA journal_mode = WAL')
+
+
+def begin_immediate_transaction(conn):
+    conn.exec_driver_sql('BEGIN IMMEDIATE')
+
+
+Stats = Dict[Optional[SourceName], int]
+
+
+# returns critical warnings
+def visits_to_sqlite(
+    vit: Iterable[Res[DbVisit]],
+    *,
+    overwrite_db: bool,
+    _db_path: Optional[Path] = None,  # only used in tests
+) -> List[Exception]:
+    if _db_path is None:
+        db_path = config.get().db
+    else:
+        db_path = _db_path
+
+    logger = get_logger()
+
+    now = now_tz()
+
+    index_stats: Stats = {}
+
+    def vit_ok() -> Iterable[DbVisit]:
+        for v in vit:
+            ev: DbVisit
+            if isinstance(v, DbVisit):
+                ev = v
+            else:
+                # conform to the schema and dump. can't hurt anyway
+                ev = DbVisit(
+                    norm_url='<error>',
+                    orig_url='<error>',
+                    dt=now,
+                    locator=Loc.make('<errror>'),
+                    src=SRC_ERROR,
+                    # todo attach backtrace?
+                    context=repr(v),
+                )
+            index_stats[ev.src] = index_stats.get(ev.src, 0) + 1
+            yield ev
+
+    meta = MetaData()
+    table = Table('visits', meta, *get_columns())
+
+    def query_total_stats(conn) -> Stats:
+        query = select(table.c.src, func.count(table.c.src)).select_from(table).group_by(table.c.src)
+        return {src: cnt for (src, cnt) in conn.execute(query).all()}
+
+    def get_engine(*args, **kwargs) -> Engine:
+        # kwargs['echo'] = True  # useful for debugging
+        e = create_engine(*args, **kwargs)
+        event.listen(e, 'connect', enable_wal)
+        return e
+
+    ### use readonly database just to get stats
+    pengine = get_engine('sqlite://', creator=lambda: sqlite3.connect(f"file:{db_path}?mode=ro", uri=True))
+    stats_before: Stats
+    try:
+        with pengine.begin() as conn:
+            stats_before = query_total_stats(conn)
+    except exc.OperationalError as oe:
+        if oe.code == 'e3q8':
+            # db doesn't exist yet
+            stats_before = {}
+        else:
+            raise oe
+    pengine.dispose()
+    ###
+
+    # needtimeout, othewise concurrent indexing might not work
+    # (note that this also requires WAL mode)
+    engine = get_engine(f'sqlite:///{db_path}', connect_args={'timeout': _CONNECTION_TIMEOUT_SECONDS})
+
+    cleared: Set[str] = set()
+
+    # by default, sqlalchemy does some sort of BEGIN (implicit) transaction, which doesn't provide proper isolation??
+    # see https://docs.sqlalchemy.org/en/20/dialects/sqlite.html#serializable-isolation-savepoints-transactional-ddl
+    event.listen(engine, 'begin', begin_immediate_transaction)
+    # TODO to allow more concurrent indexing, maybe could instead write to a temporary table?
+    # or collect visits first and only then start writing to the db to minimize db access window.. not sure
+
+    # engine.begin() starts a transaction
+    # so everything inside this block will be atomic to the outside observers
+    with engine.begin() as conn:
+        table.create(conn, checkfirst=True)
+
+        if overwrite_db:
+            conn.execute(table.delete())
+
+        insert_stmt = table.insert()
+        # using raw statement gives a massive speedup for inserting visits
+        # see test_benchmark_visits_dumping
+        insert_stmt_raw = str(insert_stmt.compile(dialect=dialect_sqlite.dialect(paramstyle='qmark')))
+
+        for chunk in chunked(vit_ok(), n=_CHUNK_BY):
+            srcs = set(v.src or '' for v in chunk)
+            new = srcs.difference(cleared)
+
+            for src in new:
+                conn.execute(table.delete().where(table.c.src == src))
+                cleared.add(src)
+
+            bound = [db_visit_to_row(v) for v in chunk]
+            conn.exec_driver_sql(insert_stmt_raw, bound)
+
+        stats_after = query_total_stats(conn)
+    engine.dispose()
+
+    stats_changes = {}
+    # map str just in case some srcs are None
+    for k in sorted(map(str, {*stats_before.keys(), *stats_after.keys()})):
+        diff = stats_after.get(k, 0) - stats_before.get(k, 0)
+        if diff == 0:
+            continue
+        sdiff = ('+' if diff > 0 else '') + str(diff)
+        stats_changes[k] = sdiff
+
+    action = 'overwritten' if overwrite_db else 'updated'
+    total_indexed = sum(index_stats.values())
+    total_err = index_stats.get(SRC_ERROR, 0)
+    total_ok = total_indexed - total_err
+    logger.info(f'indexed (current run) : total: {total_indexed}, ok: {total_ok}, errors: {total_err} {index_stats}')
+    logger.info(f'database "{db_path}" : {action}')
+    logger.info(f'database stats before : {stats_before}')
+    logger.info(f'database stats after  : {stats_after}')
+
+    if len(stats_changes) == 0:
+        logger.info('database stats changes: no changes')
+    else:
+        for k, v in stats_changes.items():
+            logger.info(f'database stats changes: {k} {v}')
+
+    res: List[Exception] = []
+    if total_ok == 0:
+        res.append(RuntimeError('No visits were indexed, something is probably wrong!'))
+    return res

promnesia/{read_db.py → database/load.py}

@@ -1,32 +1,29 @@
 from pathlib import Path
 from typing import Tuple, List
 
-from cachew import NTBinder
 from sqlalchemy import (
     create_engine,
     exc,
+    Engine,
     MetaData,
     Index,
     Table,
 )
-from sqlalchemy.engine import Engine
 
-from .common import DbVisit
+from .common import DbVisit, get_columns, row_to_db_visit
 
 
-DbStuff = Tuple[Engine, NTBinder, Table]
+DbStuff = Tuple[Engine, Table]
 
 
 def get_db_stuff(db_path: Path) -> DbStuff:
     assert db_path.exists(), db_path
     # todo how to open read only?
     # actually not sure if we can since we are creating an index here
-    engine = create_engine(f'sqlite:///{db_path}') # , echo=True)
-
-    binder = NTBinder.make(DbVisit)
+    engine = create_engine(f'sqlite:///{db_path}')  # , echo=True)
 
     meta = MetaData()
-    table = Table('visits', meta, *binder.columns)
+    table = Table('visits', meta, *get_columns())
 
     idx = Index('index_norm_url', table.c.norm_url)
     try:
@@ -39,13 +36,15 @@ def get_db_stuff(db_path: Path) -> DbStuff:
             raise e
 
     # NOTE: apparently it's ok to open connection on every request? at least my comparisons didn't show anything
-    return engine, binder, table
+    return engine, table
 
 
 def get_all_db_visits(db_path: Path) -> List[DbVisit]:
     # NOTE: this is pretty inefficient if the DB is huge
     # mostly intended for tests
-    engine, binder, table = get_db_stuff(db_path)
+    engine, table = get_db_stuff(db_path)
     query = table.select()
     with engine.connect() as conn:
-        return [binder.from_row(row) for row in conn.execute(query)]
+        res = [row_to_db_visit(row) for row in conn.execute(query)]
+    engine.dispose()
+    return res

promnesia/extract.py

@@ -28,6 +28,7 @@ DEFAULT_FILTERS = (
 )
 
 
+# TODO maybe move these to configs?
 @lru_cache(1) #meh, not sure what would happen under tests?
 def filters() -> Sequence[Filter]:
     from . import config

promnesia/kjson.py

@@ -74,7 +74,7 @@ def test_json_processor():
     handled = []
     class Proc(JsonProcessor):
         def handle_dict(self, value: JDict, path):
-            if 'skipme' in self.kpath(path):
+            if 'skipme' in self.kpath(path):  # type: ignore[comparison-overlap]
                 return JsonProcessor.SKIP
 
         def handle_str(self, value: str, path):

promnesia/logging.py

@@ -1,13 +1,14 @@
 #!/usr/bin/env python3
 '''
-Default logger is a bit, see 'test'/run this file for a demo
+Default logger is a bit meh, see 'test'/run this file for a demo
 '''
 
 def test() -> None:
     import logging
     import sys
     from typing import Callable
-    M: Callable[[str], None]  = lambda s: print(s, file=sys.stderr)
+
+    M: Callable[[str], None] = lambda s: print(s, file=sys.stderr)
 
     M("   Logging module's defaults are not great...'")
     l = logging.getLogger('test_logger')
@@ -20,7 +21,7 @@ def test() -> None:
     M("")
     M("    With LazyLogger you get a reasonable logging format, colours and other neat things")
 
-    ll = LazyLogger('test') # No need for basicConfig!
+    ll = LazyLogger('test')  # No need for basicConfig!
     ll.info("default level is INFO")
     ll.debug(".. so this shouldn't be displayed")
     ll.warning("warnings are easy to spot!")
@@ -37,10 +38,10 @@ LevelIsh = Optional[Union[Level, str]]
 
 
 def mklevel(level: LevelIsh) -> Level:
-    # todo do the same for Promnesia?
-    # glevel = os.environ.get('HPI_LOGS', None)
-    # if glevel is not None:
-    #     level = glevel
+    # todo put in some global file, like envvars.py
+    glevel = os.environ.get('PROMNESIA_LOGS', None)
+    if glevel is not None:
+        level = glevel
     if level is None:
         return logging.NOTSET
     if isinstance(level, int):
@@ -53,7 +54,6 @@ FORMAT_COLOR   = FORMAT.format(start='%(color)s', end='%(end_color)s')
 FORMAT_NOCOLOR = FORMAT.format(start='', end='')
 DATEFMT = '%Y-%m-%d %H:%M:%S'
 
-# NOTE: this is a bit experimental and temporary..
 COLLAPSE_DEBUG_LOGS = os.environ.get('COLLAPSE_DEBUG_LOGS', False)
 
 _init_done = 'lazylogger_init_done'
@@ -61,7 +61,7 @@ _init_done = 'lazylogger_init_done'
 def setup_logger(logger: logging.Logger, level: LevelIsh) -> None:
     lvl = mklevel(level)
     try:
-        import logzero # type: ignore[import]
+        import logzero  # type: ignore[import-not-found]
         formatter = logzero.LogFormatter(
             fmt=FORMAT_COLOR,
             datefmt=DATEFMT,
@@ -75,7 +75,7 @@ def setup_logger(logger: logging.Logger, level: LevelIsh) -> None:
     logger.addFilter(AddExceptionTraceback())
     if use_logzero and not COLLAPSE_DEBUG_LOGS: # all set, nothing to do
         # 'simple' setup
-        logzero.setup_logger(logger.name, level=lvl, formatter=formatter)
+        logzero.setup_logger(logger.name, level=lvl, formatter=formatter)  # type: ignore[possibly-undefined]
         return
 
     h = CollapseDebugHandler() if COLLAPSE_DEBUG_LOGS else logging.StreamHandler()
@@ -83,7 +83,7 @@ def setup_logger(logger: logging.Logger, level: LevelIsh) -> None:
     h.setLevel(lvl)
     h.setFormatter(formatter)
     logger.addHandler(h)
-    logger.propagate = False # ugh. otherwise it duplicates log messages
+    logger.propagate = False  # ugh. otherwise it duplicates log messages? not sure about it..
 
 
 class LazyLogger(logging.Logger):
@@ -92,7 +92,7 @@ class LazyLogger(logging.Logger):
 
         # this is called prior to all _log calls so makes sense to do it here?
         def isEnabledFor_lazyinit(*args, logger=logger, orig=logger.isEnabledFor, **kwargs) -> bool:
-            if not getattr(logger, _init_done, False):
+            if not getattr(logger, _init_done, False):  # init once, if necessary
                 setup_logger(logger, level=level)
                 setattr(logger, _init_done, True)
                 logger.isEnabledFor = orig # restore the callback
@@ -101,7 +101,7 @@ class LazyLogger(logging.Logger):
         # oh god.. otherwise might go into an inf loop
         if not hasattr(logger, _init_done):
             setattr(logger, _init_done, False) # will setup on the first call
-            logger.isEnabledFor = isEnabledFor_lazyinit  # type: ignore[assignment]
+            logger.isEnabledFor = isEnabledFor_lazyinit  # type: ignore[method-assign]
         return cast(LazyLogger, logger)
 
 
@@ -145,7 +145,7 @@ class CollapseDebugHandler(logging.StreamHandler):
             import os
             columns, _ = os.get_terminal_size(0)
             # ugh. the columns thing is meh. dunno I guess ultimately need curses for that
-            # TODO also would be cool to have a terminal post-processor? kinda like tail but aware of logging keyworkds (INFO/DEBUG/etc)
+            # TODO also would be cool to have a terminal post-processor? kinda like tail but aware of logging keywords (INFO/DEBUG/etc)
             self.stream.write(msg + ' ' * max(0, columns - len(msg)) + ('' if cur else '\n'))
             self.flush()
         except:

promnesia/misc/config_example.py

@@ -11,7 +11,6 @@ SOURCES = [
     Source(
         auto.index,
         # just some arbitrary directory with plaintext files
-        '/usr/include/c++/',
-        '/usr/local/include/c++/', # on apple they are here apparently..
+        '/usr/share/vim/',
     )
 ]

promnesia/misc/install_server.py

@@ -1,10 +1,13 @@
 #!/usr/bin/env python3
+from __future__ import annotations
+
 import argparse
 import os
 import sys
 import time
 from pathlib import Path
 import platform
+import shutil
 from subprocess import check_call, run
 from typing import List
 
@@ -50,7 +53,7 @@ LAUNCHD_TEMPLATE = '''
 '''
 
 
-def systemd(*args, method=check_call):
+def systemd(*args: str | Path, method=check_call) -> None:
     method([
         'systemctl', '--no-pager', '--user', *args,
     ])
@@ -116,9 +119,7 @@ def install(args: argparse.Namespace) -> None:
     if os.environ.get('DIRTY_RUN') is not None:
         launcher = str(root() / 'scripts/promnesia')
     else:
-        # must be installed, so available in PATH
-        import distutils.spawn
-        exe = distutils.spawn.find_executable('promnesia'); assert exe is not None
+        exe = shutil.which('promnesia'); assert exe is not None
         launcher = exe # older systemd wants absolute paths..
 
     db = args.db

promnesia/server.py

@@ -1,15 +1,16 @@
 #!/usr/bin/python3
-__package__ = 'promnesia'  # ugh. hacky way to make wsgi runner work properly...
+from __future__ import annotations
 
 import argparse
 from dataclasses import dataclass
-import os
-import json
 from datetime import timedelta
-from pathlib import Path
-import logging
 from functools import lru_cache
-from typing import List, NamedTuple, Dict, Optional, Any, Tuple
+import importlib.metadata
+import json
+import logging
+import os
+from pathlib import Path
+from typing import List, NamedTuple, Dict, Optional, Any, Tuple, Protocol
 
 
 import pytz
@@ -17,15 +18,15 @@ from pytz import BaseTzInfo
 
 import fastapi
 
-from sqlalchemy import MetaData, exists, literal, between, or_, and_, exc, select
+from sqlalchemy import literal, between, or_, and_, exc, select
 from sqlalchemy import Column, Table, func, types
 from sqlalchemy.sql.elements import ColumnElement
 from sqlalchemy.sql import text
 
 
 from .common import PathWithMtime, DbVisit, Url, setup_logger, default_output_dir, get_system_tz
-from .compat import Protocol
 from .cannon import canonify
+from .database.load import DbStuff, get_db_stuff, row_to_db_visit
 
 
 Json = Dict[str, Any]
@@ -50,8 +51,7 @@ def get_logger() -> logging.Logger:
 
 
 def get_version() -> str:
-    from pkg_resources import get_distribution
-    return get_distribution(__package__).version
+    return importlib.metadata.version(__package__)
 
 
 class ServerConfig(NamedTuple):
@@ -118,8 +118,6 @@ def get_db_path(check: bool=True) -> Path:
     return db
 
 
-from .read_db import DbStuff, get_db_stuff
-
 @lru_cache(1)
 # PathWithMtime aids lru_cache in reloading the sqlalchemy binder
 def _get_stuff(db_path: PathWithMtime) -> DbStuff:
@@ -135,7 +133,7 @@ def get_stuff(db_path: Optional[Path]=None) -> DbStuff: # TODO better name
 
 
 def db_stats(db_path: Path) -> Json:
-    engine, binder, table = get_stuff(db_path)
+    engine, table = get_stuff(db_path)
     query = select(func.count()).select_from(table)
     with engine.connect() as conn:
         total = list(conn.execute(query))[0][0]
@@ -150,8 +148,8 @@ class Where(Protocol):
 
 @dataclass
 class VisitsResponse:
-    original_url: Url
-    normalised_url: Url
+    original_url: str
+    normalised_url: str
     visits: Any
 
 
@@ -166,7 +164,7 @@ def search_common(url: str, where: Where) -> VisitsResponse:
         url = original_url
     logger.info('normalised url: %s', url)
 
-    engine, binder, table = get_stuff()
+    engine, table = get_stuff()
 
     query = table.select().where(where(table=table, url=url))
     logger.debug('query: %s', query)
@@ -174,7 +172,7 @@ def search_common(url: str, where: Where) -> VisitsResponse:
     with engine.connect() as conn:
         try:
             # TODO make more defensive here
-            visits: List[DbVisit] = [binder.from_row(row) for row in conn.execute(query)]
+            visits: List[DbVisit] = [row_to_db_visit(row) for row in conn.execute(query)]
         except exc.OperationalError as e:
             if getattr(e, 'msg', None) == 'no such table: visits':
                 logger.warn('you may have to run indexer first!')
@@ -231,6 +229,7 @@ def status() -> Json:
     try:
         version = get_version()
     except Exception as e:
+        logger.exception(e)
         version = None
 
     return {
@@ -240,10 +239,9 @@ def status() -> Json:
     }
 
 
-from dataclasses import dataclass
 @dataclass
 class VisitsRequest:
-    url: Url
+    url: str
 
 @app.get ('/visits', response_model=VisitsResponse)
 @app.post('/visits', response_model=VisitsResponse)
@@ -254,15 +252,17 @@ def visits(request: VisitsRequest) -> VisitsResponse:
         url=url,
         # odd, doesn't work just with: x or (y and z)
         where=lambda table, url: or_(
-            table.c.norm_url == url,  # exact match
-            and_(table.c.context != None, table.c.norm_url.startswith(url, autoescape=True)) # + child visits, but only 'interesting' ones
+            # exact match
+            table.c.norm_url == url,
+            # + child visits, but only 'interesting' ones
+            and_(table.c.context != None, table.c.norm_url.startswith(url, autoescape=True))  # noqa: E711
         ),
     )
 
 
 @dataclass
 class SearchRequest:
-    url: Url
+    url: str
 
 @app.get ('/search', response_model=VisitsResponse)
 @app.post('/search', response_model=VisitsResponse)
@@ -360,7 +360,7 @@ def visited(request: VisitedRequest) -> VisitedResponse:
     if len(snurls) == 0:
         return []
 
-    engine, binder, table = get_stuff()
+    engine, table = get_stuff()
 
     # sqlalchemy doesn't seem to support SELECT FROM (VALUES (...)) in its api
     # also doesn't support array binding...
@@ -388,7 +388,7 @@ SELECT queried, visits.*
     # brings down large queries to 50ms...
     with engine.connect() as conn:
         res = list(conn.execute(query))
-        present: Dict[str, Any] = {row[0]: binder.from_row(row[1:]) for row in res}
+        present: Dict[str, Any] = {row[0]: row_to_db_visit(row[1:]) for row in res}
     results = []
     for nu in nurls:
         r = present.get(nu, None)