promnesia 1.2.20230515__py3-none-any.whl → 1.3.20241021__py3-none-any.whl

promnesia/database/dump.py

@@ -0,0 +1,188 @@
+from __future__ import annotations
+
+import sqlite3
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Optional
+
+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 .. import config
+from ..common import (
+    DbVisit,
+    Loc,
+    Res,
+    SourceName,
+    get_logger,
+    now_tz,
+)
+from .common import db_visit_to_row, get_columns
+
+# 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: Path | None = 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 dict(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 = {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 __future__ import annotations
+
 from pathlib import Path
-from typing import Tuple, List
 
-from cachew import NTBinder
 from sqlalchemy import (
-    create_engine,
-    exc,
-    MetaData,
+    Engine,
     Index,
+    MetaData,
     Table,
+    create_engine,
+    exc,
 )
-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]:
+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

@@ -1,20 +1,22 @@
-from functools import lru_cache
+from __future__ import annotations
+
 import re
-import traceback
-from typing import Set, Iterable, Sequence, Union
+from collections.abc import Iterable, Sequence
+from functools import lru_cache
 
 from .cannon import CanonifyException
 from .common import (
-    logger,
-    DbVisit, Visit,
-    Res,
-    SourceName, Source,
+    DbVisit,
     Filter,
+    Res,
+    Results,
+    Source,
+    SourceName,
     Url,
-    Results, Extractor,
+    Visit,
+    logger,
 )
 
-
 DEFAULT_FILTERS = (
     r'^chrome-\w+://',
     r'chrome://newtab',
@@ -28,6 +30,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
@@ -52,7 +55,7 @@ def extract_visits(source: Source, *, src: SourceName) -> Iterable[Res[DbVisit]]
         yield e
         return
 
-    handled: Set[Visit] = set()
+    handled: set[Visit] = set()
     try:
         for p in vit:
             if isinstance(p, Exception):
@@ -93,7 +96,7 @@ def filtered(url: Url) -> bool:
     return any(f(url) for f in filters())
 
 
-def make_filter(thing: Union[str, Filter]) -> Filter:
+def make_filter(thing: str | Filter) -> Filter:
     if isinstance(thing, str):
         rc = re.compile(thing)
         def filter_(u: str) -> bool:

promnesia/kjson.py

@@ -3,19 +3,19 @@ Some experimental ideas on JSON processing.
 This is a bit overengineered and I admit it!
 I'll make it more readable, but in the meantime feel free to open an issue if you're confused about something.
 """
+from __future__ import annotations
 
-from typing import Any, Dict, List, Union, Tuple, cast
+from typing import Any, Union, cast
 
-
-JDict = Dict[str, Any] # TODO not sure if we can do recursive..
-JList = List[Any]
+JDict = dict[str, Any] # TODO not sure if we can do recursive..
+JList = list[Any]
 JPrim = Union[str, int, float] # , type(None)]
 
 Json = Union[JDict, JList, JPrim]
 
-JPathPart = Tuple[Json, Union[str, int]]
+JPathPart = tuple[Json, Union[str, int]]
 
-JPath = Tuple[JPathPart, ...]
+JPath = tuple[JPathPart, ...]
 
 
 class JsonProcessor:
@@ -36,7 +36,7 @@ class JsonProcessor:
         if res is self.SKIP:
             return
         for k, v in js.items():
-            path = cast(JPath, jp + ((js, k), ))
+            path = cast(JPath, jp + ((js, k), ))  # noqa: RUF005
             self._do(v, path)
 
     def do_list(self, js: JList, jp: JPath) -> None:
@@ -45,7 +45,7 @@ class JsonProcessor:
         if res is self.SKIP:
             return
         for i, x in enumerate(js):
-            path = cast(JPath, jp + ((js, i), ))
+            path = cast(JPath, jp + ((js, i), ))  # noqa: RUF005
             self._do(x, path)
 
     def _do(self, js: Json, path: JPath) -> None:
@@ -65,7 +65,7 @@ class JsonProcessor:
         self._do(js, path)
 
     @classmethod
-    def kpath(cls, path: JPath) -> Tuple[JPathPart, ...]:
+    def kpath(cls, path: JPath) -> tuple[JPathPart, ...]:
         return tuple(x[1] for x in path) # type: ignore
 
 # TODO path is a sequence of jsons and keys?
@@ -73,9 +73,10 @@ class JsonProcessor:
 def test_json_processor():
     handled = []
     class Proc(JsonProcessor):
-        def handle_dict(self, value: JDict, path):
-            if 'skipme' in self.kpath(path):
+        def handle_dict(self, value: JDict, path):  # noqa: ARG002
+            if 'skipme' in self.kpath(path):  # type: ignore[comparison-overlap]
                 return JsonProcessor.SKIP
+            return None
 
         def handle_str(self, value: str, path):
             if 'http' in value:

promnesia/logging.py

@@ -29,9 +29,9 @@ def test() -> None:
 
 
 import logging
-from typing import Union, Optional, cast
 import os
 import warnings
+from typing import Optional, Union, cast
 
 Level = int
 LevelIsh = Optional[Union[Level, str]]
@@ -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()
@@ -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)
 
 

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,14 +1,12 @@
-#!/usr/bin/env python3
 from __future__ import annotations
 
 import argparse
 import os
+import platform
 import sys
 import time
 from pathlib import Path
-import platform
 from subprocess import check_call, run
-from typing import List
 
 SYSTEM = platform.system()
 UNSUPPORTED_SYSTEM = RuntimeError(f'Platform {SYSTEM} is not supported yet!')
@@ -58,7 +56,7 @@ def systemd(*args: str | Path, method=check_call) -> None:
     ])
 
 
-def install_systemd(name: str, out: Path, launcher: str, largs: List[str]) -> None:
+def install_systemd(name: str, out: Path, launcher: str, largs: list[str]) -> None:
     unit_name = name
 
     import shlex
@@ -80,7 +78,7 @@ def install_systemd(name: str, out: Path, launcher: str, largs: List[str]) -> No
         raise e
 
 
-def install_launchd(name: str, out: Path, launcher: str, largs: List[str]) -> None:
+def install_launchd(name: str, out: Path, launcher: str, largs: list[str]) -> None:
     service_name = name
     arguments = '\n'.join(f'<string>{a}</string>' for a in [launcher, *largs])
     out.write_text(LAUNCHD_TEMPLATE.format(
@@ -115,16 +113,16 @@ def install(args: argparse.Namespace) -> None:
     print(f"Writing launch script to {out}", file=sys.stderr)
 
     # ugh. we want to know whether we're invoked 'properly' as an executable or ad-hoc via scripts/promnesia
+    extra_exe: list[str] = []
     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
-        launcher = exe # older systemd wants absolute paths..
+        launcher = sys.executable
+        extra_exe = ['-m', 'promnesia']
 
     db = args.db
     largs = [
+        *extra_exe,
         'serve',
         *([] if db is None else ['--db', str(db)]),
         '--timezone', args.timezone,