thds.tabularasa 0.13.1__py3-none-any.whl → 0.14.1__py3-none-any.whl

thds/tabularasa/__main__.py

@@ -943,16 +943,18 @@ class ReferenceDataManager:
                         return None
                     raise IOError(table.name)
 
-        failed: List[str] = []
+        failed: list[tuple[str, Exception]] = []
         synced: List[str] = []
         for table_name, res in parallel.yield_all([(t.name, partial(inner, t)) for t in tables_to_sync]):
             if isinstance(res, parallel.Error):
-                failed.append(table_name)
+                failed.append((table_name, res.error))
             elif res is not None:
                 synced.append(table_name)
 
         if failed:
-            raise RuntimeError(f"Sync failed for tables {', '.join(failed)}")
+            first_exc = failed[0][1]
+            table_names = [name for name, _ in failed]
+            raise RuntimeError(f"Sync failed for tables {', '.join(table_names)}") from first_exc
 
         down_ = (
             f"to local build directory {pkg_resources.resource_filename(self.package, self.package_data_dir)}"

thds/tabularasa/data_dependencies/adls.py

@@ -6,6 +6,7 @@ from typing import List, Optional, Protocol, Union
 import attr
 
 from thds.adls import ADLSFileSystem, fqn
+from thds.core import fretry
 from thds.tabularasa.schema.files import ADLSDataSpec, RemoteBlobStoreSpec
 
 CACHE_DIR = ".cache/"
@@ -40,6 +41,7 @@ def adls_filesystem(account: str, filesystem: str, cache_dir: Optional[Union[Pat
     return ADLSFileSystem(account, filesystem, cache_dir=cache_dir)
 
 
+@fretry.retry_regular(fretry.is_exc(Exception), fretry.n_times(3))
 def sync_adls_data(
     adls_spec: ADLSDataSpec, cache_dir: Optional[Union[Path, str]] = CACHE_DIR
 ) -> List[ADLSDownloadResult]:

thds/tabularasa/data_dependencies/sqlite.py

@@ -222,7 +222,9 @@ def populate_sqlite_db(
     table_predicate: Callable[[Table], bool] = is_build_time_package_table,
     data_path_overrides: Optional[Mapping[str, Path]] = None,
 ):
-    """Populate a sqlite database with data for a set of tables from a `reference_data.schema.Schema`
+    """Populate a sqlite database with data for a set of tables from a `reference_data.schema.Schema`.
+    Note that this can safely be called concurrently in multiple processes on the same database file; a file lock
+    is acquired on the database file and only released when the data insertion is complete.
 
     :param schema: the `reference_data.schema.Schema` object defining the data to be inserted
     :param db_package: name of the package where the database file is stored, if any. In case `None` is
@@ -259,6 +261,9 @@ def populate_sqlite_db(
     # gather all tables before executing any I/O
     insert_tables = [table for table in schema.filter_tables(table_predicate) if table.has_indexes]
 
+    if not insert_tables:
+        return
+
     with bulk_write_connection(db_path, db_package, close=True) as con:
         for table in insert_tables:
             table_filename: Optional[str]

thds/tabularasa/loaders/sqlite_util.py

@@ -22,7 +22,6 @@ from thds.tabularasa.sqlite3_compat import sqlite3
 
 DEFAULT_ATTR_SQLITE_CACHE_SIZE = 100_000
 DEFAULT_MMAP_BYTES = int(os.environ.get("TABULA_RASA_DEFAULT_MMAP_BYTES", 8_589_934_592))  # 8 GB
-DISABLE_WAL_MODE = bool(os.environ.get("REF_D_DISABLE_SQLITE_WAL_MODE", False))
 
 PARAMETERIZABLE_BUILTINS = sys.version_info >= (3, 9)
 
@@ -159,11 +158,6 @@ def set_bulk_write_mode(con: sqlite3.Connection) -> sqlite3.Connection:
     logger.debug("Setting pragmas for bulk write optimization")
     # https://www.sqlite.org/pragma.html#pragma_synchronous
     _log_exec_sql(logger, con, "PRAGMA synchronous = 0")  # OFF
-    # https://www.sqlite.org/pragma.html#pragma_journal_mode
-    if not DISABLE_WAL_MODE:
-        _log_exec_sql(logger, con, "PRAGMA journal_mode = WAL")
-    # https://www.sqlite.org/pragma.html#pragma_locking_mode
-    _log_exec_sql(logger, con, "PRAGMA locking_mode = EXCLUSIVE")
 
     return con
 
@@ -171,16 +165,7 @@ def set_bulk_write_mode(con: sqlite3.Connection) -> sqlite3.Connection:
 def unset_bulk_write_mode(con: sqlite3.Connection) -> sqlite3.Connection:
     logger = logging.getLogger(__name__)
     logger.debug("Setting pragmas for bulk write optimization")
-    # https://www.sqlite.org/pragma.html#pragma_journal_mode
-    # resetting this to the default. This is a property of the database, rather than the connection.
-    # the other settings are connection-specific.
-    # according to the docs, the WAL journal mode should be disabled before the locking mode is restored,
-    # else any attempt to do so is a no-op.
-    _log_exec_sql(logger, con, "PRAGMA journal_mode = DELETE")
-    # https://www.sqlite.org/pragma.html#pragma_synchronous
     _log_exec_sql(logger, con, "PRAGMA synchronous = 2")  # FULL (default)
-    # https://www.sqlite.org/pragma.html#pragma_locking_mode
-    _log_exec_sql(logger, con, "PRAGMA locking_mode = NORMAL")
 
     return con
 
@@ -191,9 +176,10 @@ def bulk_write_connection(
 ) -> ty.Generator[sqlite3.Connection, None, None]:
     """Context manager to set/unset bulk write mode on a sqlite connection. Sets pragmas for efficient bulk writes,
     such as loosening synchronous and locking modes. If `close` is True, the connection will be closed on exit.
-    Since setting the pragmas may mutate the database file, and since by design this context manager exists to enable
-    bulk writes which intentionally mutate the database, if a `db_path` (and optionally `db_package`, if specified as
-    package data) is given, we also acquire a file lock on the database file on entry and release it on exit.
+    To avoid bulk insert routines being run by other processes concurrently, we also acquire a file lock on the
+    database file on entry and release it on exit. Other processes attempting to perform bulk writes to the same file
+    will block until the lock is released. In the case of tabularasa init-sqlite, the semantics then imply that those
+    workers will perform no writes at all, since metadata will indicate that the data in the file is up-to-date.
     """
     db_path_ = to_local_path(db_path, db_package).absolute()
     lock_path = db_path_.with_suffix(".lock")

thds/tabularasa/testing/__init__.py

@@ -0,0 +1,5 @@
+__all__ = [
+    "mock_sqlite_loader",
+]
+
+from .mock_sqlite import mock_sqlite_loader

thds/tabularasa/testing/mock_sqlite.py

@@ -0,0 +1,114 @@
+import contextlib
+import inspect
+import sqlite3
+import tempfile
+import typing as ty
+from pathlib import Path
+
+import attrs
+import pyarrow as pa
+import pyarrow.parquet
+
+from thds.core import scope
+from thds.core.types import StrOrPath
+from thds.tabularasa.data_dependencies import sqlite, util
+from thds.tabularasa.schema import load_schema
+
+
+class _GeneratedSqliteLoader(ty.Protocol):
+    def __init__(
+        self,
+        package: ty.Optional[str],
+        db_path: str,
+    ) -> None: ...
+
+
+_UNTIL_EXIT_SCOPE = scope.Scope("tabularasa.testing.mock_sqlite_loader")
+# this scope is for creating temporary sqlite database files that persist until program exit, in case the caller of
+# mock_sqlite_loader doesn't want to manage the database file themselves
+
+L = ty.TypeVar("L", bound=_GeneratedSqliteLoader)
+
+
+def mock_sqlite_loader(
+    loader_cls: ty.Type[L],
+    data: ty.Mapping[str, ty.Collection[attrs.AttrsInstance]],
+    tmp_db_path: ty.Optional[StrOrPath] = None,
+    *,
+    package: ty.Optional[str] = None,
+    schema_path: str = "schema.yaml",
+    validate: bool = False,
+) -> L:
+    """Construct an instance of your custom generated sqlite loader from mocked data. Note that this is guaranteed
+    typesafe because regardless of how you define your mock records, the resulting sqlite loader will be a true instance
+    of your generated loader class, and will have all the same lookup methods and will use all the same deserialization
+    logic for reading rows from the database and returning actual instances from your library's data model.
+
+    :param loader_cls: The generated sqlite loader class to instantiate.
+    :param data: A mapping from table names to collections of attrs records representing rows.
+    :param package: The root package name containing the schema and generated loader(s). If omitted, it will be inferred
+      from the loader class's `__module__` attribute by climbing up until a schema file is found.
+    :param schema_path: The path to the schema file within the package.
+    :param tmp_db_path: Optional path to a file to use for the sqlite database. If None, a temporary file is created.
+      Note that in this case the temporary file will not be cleaned up until program exit.
+    :param validate: Whether to validate data against the schema when inserting data into the database.
+    :return: An instance of the specified sqlite loader class populated with the provided mocked data, with empty
+      tables for any table names that were not included in the `data` mapping.
+    """
+    if package is None:
+        if package_ := inspect.signature(loader_cls).parameters["package"].default:
+            package_candidates = [package_]
+        else:
+            loader_module_path = loader_cls.__module__.split(".")
+            package_candidates = [
+                ".".join(loader_module_path[:i]) for i in range(len(loader_module_path), 0, -1)
+            ]
+    else:
+        package_candidates = [package]
+
+    for package_ in package_candidates:
+        try:
+            schema = load_schema(package_, schema_path)
+        except (ModuleNotFoundError, FileNotFoundError):
+            continue
+        else:
+            break
+    else:
+        raise ValueError(
+            f"Could not infer package containing schema from loader class {loader_cls.__qualname__}; "
+            "please specify the 'package' argument explicitly."
+        )
+
+    if tmp_db_path is None:
+        tmp_db_path = _UNTIL_EXIT_SCOPE.enter(tempfile.NamedTemporaryFile(suffix=".sqlite")).name
+
+    unknown_tables = set(data.keys()).difference(schema.tables.keys())
+    if unknown_tables:
+        raise ValueError(f"Data provided for unknown tables: {sorted(unknown_tables)}")
+
+    with (
+        tempfile.TemporaryDirectory() as tmpdir,
+        contextlib.closing(sqlite3.connect(str(tmp_db_path))) as con,
+    ):
+        # this tmpdir is only for staging parquet files before loading into sqlite; it's fine that they get deleted
+        # immediately after the database is populated
+        for name, table in schema.tables.items():
+            rows = data.get(name, [])
+            pa_table = pa.Table.from_pylist(
+                [attrs.asdict(row, recurse=True) for row in rows], schema=table.parquet_schema
+            )
+            filename = name + ".parquet"
+            pyarrow.parquet.write_table(
+                pa_table, Path(tmpdir) / filename, version=util.PARQUET_FORMAT_VERSION
+            )
+            sqlite.insert_table(
+                con,
+                table,
+                package=None,
+                data_dir=tmpdir,
+                filename=filename,
+                validate=validate,
+                cast=False if validate else True,
+            )
+
+    return loader_cls(package=None, db_path=str(tmp_db_path))

{thds_tabularasa-0.13.1.dist-info → thds_tabularasa-0.14.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: thds.tabularasa
-Version: 0.13.1
+Version: 0.14.1
 Summary: Trilliant Health reference data build system.
 Author-email: Trilliant Health <info@trillianthealth.com>
 Project-URL: Repository, https://github.com/TrilliantHealth/ds-monorepo

{thds_tabularasa-0.13.1.dist-info → thds_tabularasa-0.14.1.dist-info}/RECORD

@@ -1,5 +1,5 @@
 thds/tabularasa/__init__.py,sha256=jc6w1WD868MQ2t4wkRNYvRssojwXvPDcNyC8V5gwbl0,169
-thds/tabularasa/__main__.py,sha256=w10WQRwQmer4Hn3JmgHqjtVJ2WEjs9MtTiDvsAS9gog,47648
+thds/tabularasa/__main__.py,sha256=Ryfd7YogTE_qFjp8IJA-KTeTXXD9INS5GJGmdPVvWBw,47791
 thds/tabularasa/compat.py,sha256=j0313TPIXtkbfvRI0AH4if8GLrjQSrDJ9heayCIl9w8,1037
 thds/tabularasa/git_util.py,sha256=fBFhaCPi_5W2BpG2B3WiPcAWJvuVI_pG47rt73wLO6E,1388
 thds/tabularasa/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -7,9 +7,9 @@ thds/tabularasa/sqlite3_compat.py,sha256=67hldmiLFTjG0qL2TPX0USV7XNfTjEN3j8MneqN
 thds/tabularasa/sqlite_from_parquet.py,sha256=yJatUIAbgErHUOL5dhchWJwzKZCrDrx93SP0nGm7It8,1115
 thds/tabularasa/to_sqlite.py,sha256=5lcEUh38MNebxAJdLp2XGWOP_WQDIADtL1fyhOvi9UU,1715
 thds/tabularasa/data_dependencies/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-thds/tabularasa/data_dependencies/adls.py,sha256=smL-iRYr9aXFFpW4inQBQmB-ieBKXcBVe9AdwV2zisw,3161
+thds/tabularasa/data_dependencies/adls.py,sha256=vJAuc5Key-vO1N6DGo5dj9fIx_4hMALAVC17qhvkT7Y,3257
 thds/tabularasa/data_dependencies/build.py,sha256=6iYgw93sOF2Nlnb6WSmA9NvPNwOf_Yyi2wXUQpRVkJM,23382
-thds/tabularasa/data_dependencies/sqlite.py,sha256=eweuLdoxyGlG-PvQUANarlMe_mmZgA5cxuMbOYxcpsQ,12576
+thds/tabularasa/data_dependencies/sqlite.py,sha256=sMP_NInBEDoH5SScIRYxtOvcPUi9WXfE3_jCoOBduGo,12825
 thds/tabularasa/data_dependencies/tabular.py,sha256=oq9wFse235ikLEv8Zvol59ptRRojZbkbzXJyQeFfC9o,6529
 thds/tabularasa/data_dependencies/util.py,sha256=FQ9G1nIpqKh00z2lXOt0Y2R1mLQsEb-BC6Tka1z2egc,8489
 thds/tabularasa/diff/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -19,7 +19,7 @@ thds/tabularasa/diff/summary.py,sha256=gENtDwhSrDYeN-8fWr6Ug2zgdp584b0pZF9UBYzKF
 thds/tabularasa/loaders/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 thds/tabularasa/loaders/lazy_adls.py,sha256=jrWy5tTKDQfWEv6aHQ3UJhFzLrDPOlSGsArv9zcl1g8,1375
 thds/tabularasa/loaders/parquet_util.py,sha256=u75j3PkMSakO2zfq4zksWzXLYnaO--WizAgXTcSpXRY,13354
-thds/tabularasa/loaders/sqlite_util.py,sha256=GLP271SF-GVQwxaaB0p5kCG0eGay6j4mWabL00a9pz0,11954
+thds/tabularasa/loaders/sqlite_util.py,sha256=3Gi1Y4iTVCD9FXqylQw1eyFwVuplQUrjY1J0SC5FFWg,11099
 thds/tabularasa/loaders/util.py,sha256=XmsGkDdL6O8R6B4667Iqi5HoRgq0YMs6LP3VvPIqPVU,21369
 thds/tabularasa/schema/__init__.py,sha256=bowvNXrrDrWB3TAmwDxCeEAvVEe9z7iRfqRaNg1Qmo4,440
 thds/tabularasa/schema/constraints.py,sha256=V2vh01BhYR8OVQvgdujqSi0l_fMJvFKYSlBvWExZFG0,9744
@@ -39,8 +39,10 @@ thds/tabularasa/schema/compilation/pyarrow.py,sha256=pcNQ3a6UPJT1PBj6xHOl99UvZft
 thds/tabularasa/schema/compilation/sphinx.py,sha256=we5X-ZpCk6WH-8KCXAv6Nklg1JZmnkGPT3V2EHa2_rg,17491
 thds/tabularasa/schema/compilation/sqlite.py,sha256=wSrSlVCYeuTpOf9AOHAnp6gJHkjHZhx8UkgkYgfoQVw,2368
 thds/tabularasa/schema/compilation/util.py,sha256=YXFe1_yoBobED010hstKIoq-dwLHo6SBv1v1IAw6AYU,3886
-thds_tabularasa-0.13.1.dist-info/METADATA,sha256=flLUSZeccW-NUJgBBPwUrBeB94h1-TsjwdIUaNkdB8c,26786
-thds_tabularasa-0.13.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-thds_tabularasa-0.13.1.dist-info/entry_points.txt,sha256=PX4ShRonjv6lMsVjrGu8RkFzpyyvgM9EnZlNfMomd9k,61
-thds_tabularasa-0.13.1.dist-info/top_level.txt,sha256=LTZaE5SkWJwv9bwOlMbIhiS-JWQEEIcjVYnJrt-CriY,5
-thds_tabularasa-0.13.1.dist-info/RECORD,,
+thds/tabularasa/testing/__init__.py,sha256=XoLzB-DotxFw9KHt2vfH72k7pyAAFI2bW-qqq6nww1g,85
+thds/tabularasa/testing/mock_sqlite.py,sha256=xoV4w_GaDgtZf17iUux2-LA6Va1XRJdC2FU34dysh0o,4769
+thds_tabularasa-0.14.1.dist-info/METADATA,sha256=rqg7l_iBlrh7E8-iXCeZyLcnOxJIuFDBw_5QNe1A9V0,26786
+thds_tabularasa-0.14.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+thds_tabularasa-0.14.1.dist-info/entry_points.txt,sha256=PX4ShRonjv6lMsVjrGu8RkFzpyyvgM9EnZlNfMomd9k,61
+thds_tabularasa-0.14.1.dist-info/top_level.txt,sha256=LTZaE5SkWJwv9bwOlMbIhiS-JWQEEIcjVYnJrt-CriY,5
+thds_tabularasa-0.14.1.dist-info/RECORD,,