thds.core 0.0.1__py3-none-any.whl → 1.31.20250116223856__py3-none-any.whl

thds/core/sqlite/functions.py

@@ -0,0 +1,63 @@
+import hashlib
+import inspect
+import sqlite3
+import typing as ty
+
+
+def _slowish_hash_str_function(__string: str) -> int:
+    # Raymond Hettinger says that little-endian is slightly faster, though that was 2021.
+    # I have also tested this myself and found it to be true.
+    # https://bugs.python.org/msg401661
+    return int.from_bytes(hashlib.md5(__string.encode()).digest()[:7], byteorder="little")
+
+
+_THE_HASH_FUNCTION = _slowish_hash_str_function
+
+
+# If you need it to be faster, you can 'replace' this with your own implementation.
+def set_hash_function(f: ty.Callable[[str], int]) -> None:
+    global _THE_HASH_FUNCTION
+    _THE_HASH_FUNCTION = f
+
+
+def _pyhash_values(*args) -> int:
+    _args = (x if isinstance(x, str) else str(x) for x in args)
+    concatenated = "".join(_args)
+    hash_value = _THE_HASH_FUNCTION(concatenated)
+    return hash_value
+
+
+def _has_param_kind(signature, kind) -> bool:
+    return any(p.kind == kind for p in signature.parameters.values())
+
+
+def _num_parameters(f: ty.Callable) -> int:
+    signature = inspect.signature(f)
+    if _has_param_kind(signature, inspect.Parameter.VAR_KEYWORD):
+        raise NotImplementedError("**kwargs in sqlite functions is not supported")
+    elif _has_param_kind(signature, inspect.Parameter.VAR_POSITIONAL):
+        return -1
+    else:
+        return len(signature.parameters.keys())
+
+
+_FUNCTIONS = [_pyhash_values]
+
+
+def register_functions_on_connection(
+    conn: sqlite3.Connection,
+    *,
+    functions: ty.Collection[ty.Callable] = _FUNCTIONS,
+) -> sqlite3.Connection:
+    """By default registers our default functions.
+
+    Returns the connection itself, for chaining.
+    """
+    for f in _FUNCTIONS:
+        narg = _num_parameters(f)
+        # SPOOKY: we're registering a function here with SQLite, and the SQLite name will
+        # be the same as its name in Python.  Be very careful that you do not register two
+        # different functions with the same name - you can read their docs on what will
+        # happen, but it would be far better to just not ever do this.
+        conn.create_function(name=f.__name__, narg=narg, func=f, deterministic=True)
+    return conn

thds/core/sqlite/index.py

@@ -0,0 +1,22 @@
+import typing as ty
+
+from .connect import Connectable, autoconn_scope, autoconnect
+
+
+@autoconn_scope.bound
+def create(connectable: Connectable, table_name: str, columns: ty.Collection[str], unique: bool = False):
+    """Is idempotent, but does not verify that your index DDL matches what you're asking for."""
+    conn = autoconnect(connectable)
+    """Create an index on a table in a SQLite database using only sqlite3 and SQL."""
+    colnames = "_".join(colname for colname in columns).replace("-", "_")
+
+    sql_create_index = (
+        f"CREATE {'UNIQUE' if unique else ''} INDEX IF NOT EXISTS "
+        f"[{table_name}_{colnames}_idx] ON [{table_name}] ({', '.join(columns)})"
+    )
+    try:
+        conn.execute(sql_create_index)
+        # do not commit - let the caller decide when to commit, or allow autoconnect to do its job
+    except Exception:
+        print(sql_create_index)
+        raise

thds/core/sqlite/insert_utils.py

@@ -0,0 +1,23 @@
+# this works only if you have something to map the tuples to an attrs class.
+# it also does not currently offer any parallelism.
+import typing as ty
+
+from thds.core import progress
+from thds.core.sqlite import DbAndTable, TableSource, table_source
+
+T = ty.TypeVar("T")
+
+
+def tuples_to_table_source(
+    data: ty.Iterable[tuple],
+    upsert_many: ty.Callable[[ty.Iterable[T]], DbAndTable],
+    func: ty.Callable[[tuple], ty.Optional[T]],
+) -> TableSource:
+    """
+    Converts an iterable of tuples to a TableSource.
+    :param data: An iterable of tuples.
+    :param upsert_many: A function that accepts an iterable of T and returns a DbAndTable.
+    :param func: A function that processes each tuple and returns an optional T.
+    :return: The resulting table source after upserting data.
+    """
+    return table_source(*upsert_many(progress.report(filter(None, map(func, data)))))

thds/core/sqlite/merge.py

@@ -0,0 +1,84 @@
+import os
+import typing as ty
+from contextlib import closing
+from pathlib import Path
+from sqlite3 import connect
+from timeit import default_timer
+
+from thds.core import log, types
+
+from .meta import get_indexes, get_tables, pydd
+
+logger = log.getLogger(__name__)
+
+
+def merge_databases(
+    filenames: ty.Iterable[types.StrOrPath],
+    table_names: ty.Collection[str] = tuple(),
+    *,
+    replace: bool = False,
+    copy_indexes: bool = True,
+) -> Path:
+    """Merges the listed tables, if found in the other filenames,
+    into the first database in the list. If no table names are listed,
+    we will instead merge _all_ tables found in _every_ database into the first.
+
+    If the table already exists in the destination, it will not be altered.
+    If the table does not, it will be created using the SQL CREATE TABLE statement
+    found in the first database where it is encountered.
+
+    This mutates the first database in the list! If you don't want it
+    mutated, make a copy of the file first, or start with an empty database file!
+
+    Allows SQL injection via the table names - don't use this on untrusted inputs.
+
+    By default, also copies indexes associated with the tables where an index with the
+    same name does not already exist in the destination table.  You can disable this
+    wholesale and then create/copy the specific indexes that you want after the fact.
+    """
+    _or_replace_ = "OR REPLACE" if replace else ""
+    filenames = iter(filenames)
+    first_filename = next(filenames)
+    logger.info(f"Connecting to {first_filename}")
+    conn = connect(first_filename)
+    destination_tables = get_tables(conn)
+
+    merge_start = default_timer()
+    with closing(conn.cursor()) as cursor:
+        cursor.execute("pragma synchronous = off;")
+        cursor.execute("pragma journal_mode = off;")
+        for filename in filenames:
+            pydd(Path(filename))
+            start = default_timer()
+            to_merge = "to_merge"  # just a local/temporary constant alias
+            logger.info(f"Merging {filename} into {first_filename}")
+            cursor.execute(f"ATTACH '{os.fspath(filename)}' AS " + to_merge)
+
+            attached_tables = get_tables(conn, schema_name=to_merge)
+            cursor.execute("BEGIN")
+            for table_name in table_names or attached_tables.keys():
+                if table_name not in attached_tables:
+                    continue  # table doesn't exist in the database to merge, so we skip it!
+                if table_name not in destination_tables:
+                    cursor.execute(attached_tables[table_name])  # create the table in the destination
+                    destination_tables = get_tables(conn)  # refresh tables dict
+                cursor.execute(
+                    f"INSERT {_or_replace_} INTO {table_name} SELECT * FROM {to_merge}.[{table_name}]"
+                )
+                if copy_indexes:
+                    dest_indexes = get_indexes(table_name, conn)
+                    for idx_name, index_sql in get_indexes(
+                        table_name, conn, schema_name=to_merge
+                    ).items():
+                        if idx_name not in dest_indexes:
+                            cursor.execute(index_sql)
+            logger.info(
+                f"Committing merge of {filename} into {first_filename} after {default_timer() - start:.2f}s"
+            )
+            conn.commit()  # without a commit, DETACH DATABASE will error with Database is locked.
+            # https://stackoverflow.com/questions/56243770/sqlite3-detach-database-produces-database-is-locked-error
+            cursor.execute(f"DETACH DATABASE {to_merge}")
+
+    logger.info(f"Merge complete after {default_timer() - merge_start:.2f}s")
+    conn.close()
+    return Path(first_filename)

thds/core/sqlite/meta.py

@@ -0,0 +1,190 @@
+"""Read-only utilities for inspecting a SQLite database.
+
+Note that none of these are safe from SQL injection - you should probably
+not be allowing users to specify tables in an ad-hoc fashion.
+"""
+
+import contextlib
+import os
+import sqlite3
+import typing as ty
+from pathlib import Path
+
+from thds.core import log
+from thds.core.source import from_file
+
+from .connect import autoconn_scope, autoconnect
+from .types import Connectable, TableSource
+
+logger = log.getLogger(__name__)
+
+
+def fullname(table_name: str, schema_name: str = "") -> str:
+    if schema_name:
+        return f"{schema_name}.[{table_name}]"
+    return f"[{table_name}]"
+
+
+@autoconn_scope.bound
+def list_tables(connectable: Connectable, schema_name: str = "") -> ty.List[str]:
+    conn = autoconnect(connectable)
+    return [
+        row[0]
+        for row in conn.execute(
+            f"SELECT name FROM {fullname('sqlite_master', schema_name)} WHERE type='table'"
+        )
+    ]
+
+
+@autoconn_scope.bound
+def get_tables(connectable: Connectable, *, schema_name: str = "main") -> ty.Dict[str, str]:
+    """Keys of the returned dict are the names of tables in the database.
+
+    Values of the returned dict are the raw SQL that can be used to recreate the table.
+    """
+    conn = autoconnect(connectable)
+    return {
+        row[0]: row[1]
+        for row in conn.execute(
+            f"""
+            SELECT name, sql
+            FROM {fullname('sqlite_master', schema_name)}
+            WHERE type = 'table'
+            AND sql is not null
+            """
+        )
+    }
+
+
+def pydd(path: os.PathLike):
+    """Sometimes running this on a big sqlite file before starting a
+    query will make a big difference to overall query performance.
+    """
+    with open(path, "rb") as f:
+        while f.read(1024 * 1024):
+            pass
+
+
+def table_name_from_path(db_path: Path) -> str:
+    tables = list_tables(db_path)
+    assert len(tables) == 1, f"Expected exactly one table, got {tables}"
+    return tables[0]
+
+
+def table_source(db_path: Path, table_name: str = "") -> TableSource:
+    if not table_name:
+        table_name = table_name_from_path(db_path)
+    return TableSource(from_file(db_path), table_name)
+
+
+@autoconn_scope.bound
+def primary_key_cols(table_name: str, connectable: Connectable) -> ty.Tuple[str, ...]:
+    conn = autoconnect(connectable)
+    return tuple(
+        row[0]
+        for row in conn.execute(
+            f"""
+            SELECT name
+            FROM pragma_table_info('{table_name}')
+            WHERE pk <> 0
+            """
+        )
+    )
+
+
+@autoconn_scope.bound
+def column_names(table_name: str, connectable: Connectable) -> ty.Tuple[str, ...]:
+    conn = autoconnect(connectable)
+    return tuple([row[1] for row in conn.execute(f"PRAGMA table_info({table_name})")])
+
+
+def preload_sources(*table_srcs: ty.Optional[TableSource]) -> None:
+    for table_src in table_srcs:
+        if not table_src:
+            continue
+        assert isinstance(table_src, TableSource)
+        logger.debug("Preloading %s from %s", table_src.table_name, table_src.db_src.uri)
+        pydd(table_src.db_src)
+    logger.debug("Preloading complete")
+
+
+@autoconn_scope.bound
+def get_indexes(
+    table_name: str, connectable: Connectable, *, schema_name: str = "main"
+) -> ty.Dict[str, str]:
+    """Keys of the returned dict are the names of indexes belonging to the given table.
+
+    Values of the returned dict are the raw SQL that can be used to recreate the index(es).
+    """
+    conn = autoconnect(connectable)
+    return {
+        row[0]: row[1]
+        for row in conn.execute(
+            f"""
+            SELECT name, sql
+            FROM {fullname('sqlite_master', schema_name)}
+            WHERE type = 'index' AND tbl_name = '{table_name}'
+            AND sql is not null
+            """
+        )
+        if row[1]
+        # sqlite has internal indexes, usually prefixed with "sqlite_autoindex_",
+        # that do not have sql statements associated with them.
+        # we exclude these because they are uninteresting to the average caller of this function.
+    }
+
+
+@contextlib.contextmanager
+@autoconn_scope.bound
+def debug_errors(connectable: Connectable) -> ty.Iterator:
+    try:
+        yield
+    except Exception:
+        try:
+            conn = autoconnect(connectable)
+            cur = conn.cursor()
+            cur.execute("PRAGMA database_list")
+            rows = cur.fetchall()
+            databases = {row[1]: dict(num=row[0], file=row[2]) for row in rows}
+            logger.error(f"Database info: {databases}")
+            for db_name in databases:
+                logger.error(
+                    f"    In {db_name}, tables are: {get_tables(connectable, schema_name=db_name)}"
+                )
+        except Exception:
+            logger.error(f"SQLite database: {connectable} is not introspectable")
+
+        raise
+
+
+@autoconn_scope.bound
+def get_table_schema(
+    conn: ty.Union[sqlite3.Connection, Connectable], table_name: str
+) -> ty.Dict[str, str]:
+    """
+    Retrieve the schema of a given table.
+
+    Args:
+        conn: The database connection object or a Connectable.
+        table_name: The name of the table.
+
+    Returns: A dictionary with column names as keys and their types as values.
+    """
+    # Ensure we have a connection object
+    connection = autoconnect(conn)
+
+    # Fetch the table schema
+    cursor = connection.cursor()
+    cursor.execute(f"PRAGMA table_info('{table_name}')")
+    schema = {row[1]: row[2].lower() for row in cursor.fetchall()}
+    return schema
+
+
+@autoconn_scope.bound
+def attach(connectable: Connectable, db_path: os.PathLike, schema_name: str) -> None:
+    """ATTACH a database to the current connection, using your provided schema name.
+
+    It must be an actual file.
+    """
+    conn = autoconnect(connectable)
+    conn.execute(f"ATTACH DATABASE '{os.fspath(db_path)}' AS {schema_name}")

thds/core/sqlite/read.py

@@ -0,0 +1,66 @@
+import typing as ty
+from sqlite3 import Connection, Row
+
+
+def matching_where(to_match_colnames: ty.Iterable[str]) -> str:
+    """Creates a where clause for these column names, with named @{col} placeholders for each column."""
+    qs = " AND ".join(f"{k} = @{k}" for k in to_match_colnames)
+    return f"WHERE {qs}" if qs else ""
+
+
+def matching_select(
+    table_name: str,
+    conn: Connection,
+    to_match: ty.Mapping[str, ty.Any],
+    columns: ty.Sequence[str] = tuple(),
+) -> ty.Iterator[ty.Mapping[str, ty.Any]]:
+    """Get a single row from a table by key.
+
+    This is susceptible to SQL injection because the keys are
+    formatted directly. Do _not_ give external users the ability to
+    call this function directly and specify any of its keys.
+    """
+    cols = ", ".join(columns) if columns else "*"
+
+    qs = " AND ".join(f"{k} = ?" for k in to_match.keys())
+    where = f"WHERE {qs}" if qs else ""
+    # because we control the whole query, we're matching on the 'dumb' ? placeholder.
+
+    old_row_factory = conn.row_factory
+    conn.row_factory = Row  # this is an optimized approach to getting 'mappings' (with key names)
+    for row in conn.execute(f"SELECT {cols} FROM {table_name} {where}", tuple(to_match.values())):
+        yield row
+    conn.row_factory = old_row_factory
+
+
+matching = matching_select  # alias
+
+
+def partition(
+    n: int, i: int, columns: ty.Optional[ty.Union[str, ty.Collection[str]]] = None
+) -> ty.Dict[str, int]:
+    """Can (surprisingly) be used directly with matching().
+
+    i should be zero-indexed, whereas N is a count (natural number).
+
+    columns is an optional parameter to specify column(s) to partition on
+    if no columns are specified partitioning will be based on rowid
+    """
+    assert 0 <= i < n
+    if not columns:
+        return {f"rowid % {n}": i}
+    hash_cols = columns if isinstance(columns, str) else ", ".join(columns)
+    # when SQLite uses this in a WHERE clause, as "hash(foo, bar) % 5 = 3",
+    # it hashes the _values_ in the row for those columns.
+    # Note that if we do have to fall back to this, the partitioning will be
+    # quite a bit slower because of the necessity of calling back into Python.
+    return {f"_pyhash_values({hash_cols}) % {n}": i}
+
+
+def maybe(
+    table_name: str, conn: Connection, to_match: ty.Mapping[str, ty.Any]
+) -> ty.Optional[ty.Mapping[str, ty.Any]]:
+    """Get a single row, if it exists, from a table by key"""
+    results = list(matching(table_name, conn, to_match))
+    assert len(results) == 0 or len(results) == 1
+    return results[0] if results else None

thds/core/sqlite/sqlmap.py

@@ -0,0 +1,179 @@
+import inspect
+import shutil
+import typing as ty
+from collections import defaultdict
+from concurrent.futures import ProcessPoolExecutor
+from functools import partial
+from pathlib import Path
+from uuid import uuid4
+
+from thds.core import log, parallel, scope, thunks, tmp, types
+
+from .merge import merge_databases
+
+logger = log.getLogger(__name__)
+_tmpdir_scope = scope.Scope()
+
+
+class Partition(ty.NamedTuple):
+    partition: int  # 0 indexed
+    of: int  # count (1 indexed)
+
+
+ALL = Partition(0, 1)
+
+
+def _name(callable: ty.Callable) -> str:
+    if hasattr(callable, "func"):
+        return callable.func.__name__
+    return callable.__name__
+
+
+def _write_partition(
+    basename: str, writer: ty.Callable[[Partition, Path], ty.Any], base_dir: Path, partition: Partition
+) -> Path:
+    part_dir = base_dir / f"-{basename}-{partition.partition:02d}of{partition.of:02d}"
+    part_dir.mkdir(exist_ok=True, parents=True)
+    with log.logger_context(p=f"{partition.partition + 1:2d}/{partition.of:2d}"):
+        logger.info(f"Writing partition '{partition}' outputs into {part_dir}")
+        try:
+            writer(partition, part_dir)
+        except Exception:
+            logger.exception(f"Failed to write partition '{partition}' outputs into {part_dir}")
+            raise
+        logger.info(f"Finished writing partition {partition} outputs into {part_dir}")
+        return part_dir
+
+
+Merger = ty.Callable[[ty.Iterable[types.StrOrPath]], Path]
+
+
+def merge_sqlite_dirs(
+    merger: ty.Optional[Merger],
+    part_dirs: ty.Iterable[Path],
+    output_dir: Path,
+    max_cores: int = 2,
+) -> ty.Dict[str, Path]:
+    """Return a dictionary where the keys are the filenames, and the values are the Paths
+    of merged SQLite databases.
+
+    Any file found in any of the input part directories is assumed to be a SQLite
+    database, and will be merged, using `merger`, into all the other SQLite databases that
+    _bear the same name_ in any of the other directories.
+
+    Each final, merged SQLite database will then be _moved_ into the output_dir provided.
+
+    max_cores is the maximum number of _databases_ to merge in parallel;
+    since SQLite is doing almost all of the work, we don't imagine that we'd be able to get
+    much speedup by merging multiple databases using the same core. This has not been benchmarked.
+    """
+    _ensure_output_dir(output_dir)
+    sqlite_dbs_by_filename: ty.Dict[str, ty.List[Path]] = defaultdict(list)
+    for partition_dir in part_dirs:
+        if not partition_dir.exists():
+            continue  # a partition writer is allowed to write out nothing to a partition
+        if not partition_dir.is_dir():
+            # this may happen if people don't read the parallel_to_sqlite docstring and
+            # assume the Path is meant to be a file.
+            raise ValueError(
+                f"Partition directory {partition_dir} is not a directory!"
+                " Your code may have written directly to the provided Path as a file,"
+                " rather than writing SQLite database files into the directory as required."
+            )
+        for sqlite_db_path in partition_dir.iterdir():
+            if sqlite_db_path.is_file():
+                sqlite_dbs_by_filename[sqlite_db_path.name].append(sqlite_db_path)
+
+    thunking_merger = thunks.thunking(merger or _default_merge_databases)
+    for merged_db in parallel.yield_results(
+        [
+            thunking_merger(sqlite_db_paths)
+            for filename, sqlite_db_paths in sqlite_dbs_by_filename.items()
+        ],
+        # SQLite merge is CPU-intensive, so we use a Process Pool.
+        executor_cm=ProcessPoolExecutor(max_workers=max(min(max_cores, len(sqlite_dbs_by_filename)), 1)),
+    ):
+        logger.info(f"Moving merged database {merged_db} into {output_dir}")
+        shutil.move(str(merged_db), output_dir)
+
+    return {filename: output_dir / filename for filename in sqlite_dbs_by_filename}
+
+
+def _ensure_output_dir(output_directory: Path):
+    if output_directory.exists():
+        if not output_directory.is_dir():
+            raise ValueError("Output path must be a directory if it exists!")
+    else:
+        output_directory.mkdir(parents=True, exist_ok=True)
+    assert output_directory.is_dir()
+
+
+_default_merge_databases: Merger = partial(
+    merge_databases,
+    **{
+        k: v.default
+        for k, v in inspect.signature(merge_databases).parameters.items()
+        if v.default != inspect._empty
+    },
+)
+# TODO - feels like this suggests creating a general utility that creates partials where all the defaults are applied
+#   the typing seems a bit tricky though
+
+
+@_tmpdir_scope.bound
+def partitions_to_sqlite(
+    partition_writer: ty.Callable[[Partition, Path], ty.Any],
+    output_directory: Path,
+    partitions: ty.Sequence[Partition],
+    *,
+    custom_merger: ty.Optional[Merger] = None,
+    max_workers: int = 0,
+) -> ty.Dict[str, Path]:
+    """By default, will use one Process worker per partition provided."""
+    temp_dir = _tmpdir_scope.enter(tmp.tempdir_same_fs(output_directory))
+
+    part_directories = list(
+        parallel.yield_results(
+            [
+                thunks.thunking(_write_partition)(
+                    _name(partition_writer) + uuid4().hex[:20],
+                    partition_writer,
+                    temp_dir,
+                    partition,
+                )
+                for partition in partitions
+            ],
+            executor_cm=ProcessPoolExecutor(max_workers=max_workers or len(partitions)),
+            # executor_cm=contextlib.nullcontext(loky.get_reusable_executor(max_workers=N)),
+        )
+    )
+    return merge_sqlite_dirs(
+        custom_merger if custom_merger is not None else _default_merge_databases,
+        part_directories,
+        output_directory,
+        max_cores=max_workers,
+    )
+
+
+def parallel_to_sqlite(
+    partition_writer: ty.Callable[[Partition, Path], ty.Any],
+    output_directory: Path,
+    N: int = 8,
+    custom_merger: ty.Optional[Merger] = None,
+) -> ty.Dict[str, Path]:
+    """The partition_writer will be provided a partition number and a directory (as a Path).
+
+    It must write one or more sqlite databases to the directory provided, using (of
+    course) the given partition to filter/query its input.
+
+    Any files found in that output directory will be assumed to be a SQLite database, and
+    any _matching_ filenames across the set of partitions that are written will be merged
+    into each other. Therefore, there will be a single database in the output directory
+    for every unique filename found in any of the partition directories after writing.
+    """
+    return partitions_to_sqlite(
+        partition_writer,
+        output_directory,
+        [Partition(i, N) for i in range(N)],
+        custom_merger=custom_merger,
+    )