Refactor class and duckdb connection relationships

Why these changes are being introduced:

This refactoring work was a long time coming, inspired by a recent need to gracefully
handle a read request for embeddings against a dataset without embeddings parquet
files.  If we can normalize how and when tables are created, and the handling of
duckdb connections, we can normalize handling read requests for data that may not be
available (yet).  As such, this refactoring work will help normalize read edge cases
now and going forward.

This library was built in stages.  First was TIMDEXDataset, which read parquet files
directly.  Then TIMDEXDatasetMetadata, which more formally introduced DuckDB.  It
handled the connection creation and configuration.  This connection was shared with
TIMDEXDataset as we leaned into DuckDB reading.  Lastly, TIMDEXEmbeddings was added
as our first new "source" of data. This class shared the connection from TIMDEXDataset.
Both TIMDEXDatasetMetadata and TIMDEXEmbeddings were doing their own SQLAlchemy
table reflections.  TIMDEXDatasetMetadata could be instantiated on its own, while
TIMDEXEmbeddings was assumed to take an instance of TIMDEXDataset.

At this point, while things worked, it was clear that a refactor would be beneficial.
We needed clearer responsibility of what created and configured the DuckDB connection,
solidify that TIMDEXDatasetMetadata and TIMDEXEmbeddings are components on TIMDEXDataset,
and how and when SQLAlchemy reflection was performed.  Aligning all these things will
make responding to these read and write edge cases much easier.

How this addresses that need:

- A new factory class is created DuckDBConnectionFactory that is responsible for
creating and configuring any DuckDB connections used.
- Both TIMDEXDatasetMetadata and TIMDEXEmbeddings require a TIMDEXDataset
instance, and then themselves become components on TIMDEXDataset.  We can more
accurately call them "components" then of the primary TIMDEXDataset.
- TIMDEXDataset handles the creation of a DuckDB connection via the new factory,
and this connection is then accesible to its components TIMDEXDatasetMetadata and
TIMDEXEmbeddings (maybe more in the future)
- TIMDEXDataset is also responsible for all SQLAlchemy reflection, saving to
self.sa_tables.  In this way, any component that may want a SQLAlchemy instance,
e.g. for reading, it can get it from `self.timdex_dataset.get_sa_table(<schema>,
<table)`.
- Refreshing of classes is greatly simplifed: TIMDEXDataset is the true orchestrator
now, so a full re-init satisfies this.  Components no longer have their own
`.refresh()` methods.
- Where possible, update all tests to use components like TIMDEXEmbeddings as part
of a TIMDEXDataset intsance, not a long class instance.

Side effects of this change:

* It is not recommended to use TIMDEXDatasetMetadata or TIMDEXEmbeddings
by themselves; they are meant as components on a TIMDEXDataset.

Relevant ticket(s):
* https://mitlibraries.atlassian.net/browse/USE-306

Author: ghukill

tests/conftest.py

@@ -12,12 +12,13 @@
     generate_sample_embeddings_for_run,
     generate_sample_records,
 )
-from timdex_dataset_api import TIMDEXDataset, TIMDEXDatasetMetadata
+from timdex_dataset_api import TIMDEXDataset
 from timdex_dataset_api.dataset import TIMDEXDatasetConfig
 from timdex_dataset_api.embeddings import (
     DatasetEmbedding,
     TIMDEXEmbeddings,
 )
+from timdex_dataset_api.metadata import TIMDEXDatasetMetadata
 from timdex_dataset_api.record import DatasetRecord
 
 
@@ -230,10 +231,8 @@ def timdex_dataset_same_day_runs(tmp_path) -> TIMDEXDataset:
 @pytest.fixture(scope="module")
 def timdex_metadata(timdex_dataset_with_runs) -> TIMDEXDatasetMetadata:
     """TIMDEXDatasetMetadata with static database file created."""
-    metadata = TIMDEXDatasetMetadata(timdex_dataset_with_runs.location)
-    metadata.rebuild_dataset_metadata()
-    metadata.refresh()
-    return metadata
+    timdex_dataset_with_runs.metadata.rebuild_dataset_metadata()
+    return timdex_dataset_with_runs.metadata
 
 
 @pytest.fixture(scope="module")
@@ -247,9 +246,9 @@ def timdex_dataset_with_runs_with_metadata(
 
 
 @pytest.fixture
-def timdex_metadata_empty(timdex_dataset_with_runs) -> TIMDEXDatasetMetadata:
+def timdex_metadata_empty(timdex_dataset_empty) -> TIMDEXDatasetMetadata:
     """TIMDEXDatasetMetadata without static database file."""
-    return TIMDEXDatasetMetadata(timdex_dataset_with_runs.location)
+    return timdex_dataset_empty.metadata
 
 
 @pytest.fixture
@@ -271,7 +270,8 @@ def timdex_metadata_with_deltas(
     )
     td.write(records)
 
-    return TIMDEXDatasetMetadata(timdex_dataset_with_runs.location)
+    # return fresh TIMDEXDataset's metadata
+    return TIMDEXDataset(timdex_dataset_with_runs.location).metadata
 
 
 @pytest.fixture
@@ -286,12 +286,11 @@ def timdex_metadata_merged_deltas(
     # clone dataset with runs using new dataset location
     td = TIMDEXDataset(dataset_location, config=timdex_dataset_with_runs.config)
 
-    # clone metadata and merge append deltas
-    metadata = TIMDEXDatasetMetadata(td.location)
-    metadata.merge_append_deltas()
-    metadata.refresh()
+    # merge append deltas via the TD's metadata
+    td.metadata.merge_append_deltas()
+    td.refresh()
 
-    return metadata
+    return td.metadata
 
 
 # ================================================================================

tests/test_embeddings.py

@@ -152,12 +152,13 @@ def test_embeddings_read_batches_yields_pyarrow_record_batches(
     timdex_dataset_empty.metadata.rebuild_dataset_metadata()
     timdex_dataset_empty.refresh()
 
-    # write embeddings
-    timdex_embeddings = TIMDEXEmbeddings(timdex_dataset_empty)
-    timdex_embeddings.write(sample_embeddings_generator(100, run_id="test-run"))
-    timdex_embeddings = TIMDEXEmbeddings(timdex_dataset_empty)
+    # write embeddings and refresh to pick up new views
+    timdex_dataset_empty.embeddings.write(
+        sample_embeddings_generator(100, run_id="test-run")
+    )
+    timdex_dataset_empty.refresh()
 
-    batches = timdex_embeddings.read_batches_iter()
+    batches = timdex_dataset_empty.embeddings.read_batches_iter()
     batch = next(batches)
     assert isinstance(batch, pa.RecordBatch)
 

tests/test_metadata.py

@@ -6,7 +6,7 @@
 
 from duckdb import DuckDBPyConnection
 
-from timdex_dataset_api import TIMDEXDataset, TIMDEXDatasetMetadata
+from timdex_dataset_api import TIMDEXDataset
 
 ORDERED_METADATA_COLUMN_NAMES = [
     "timdex_record_id",
@@ -21,29 +21,33 @@
 ]
 
 
-def test_tdm_init_no_metadata_file_warning_success(caplog, timdex_dataset_with_runs):
-    TIMDEXDatasetMetadata(timdex_dataset_with_runs.location)
-
+def test_tdm_init_no_metadata_file_warning_success(caplog, tmp_path):
+    # creating a new TIMDEXDataset will log warning if no metadata file
+    caplog.set_level("WARNING")
+    TIMDEXDataset(str(tmp_path / "new_empty_dataset"))
     assert "Static metadata database not found" in caplog.text
 
 
 def test_tdm_local_dataset_structure_properties(tmp_path):
     local_root = str(Path(tmp_path) / "path/to/nothing")
-    tdm_local = TIMDEXDatasetMetadata(local_root)
-    assert tdm_local.location == local_root
-    assert tdm_local.location_scheme == "file"
+    td_local = TIMDEXDataset(local_root)
+    assert td_local.metadata.location == local_root
+    assert td_local.metadata.location_scheme == "file"
 
 
-def test_tdm_s3_dataset_structure_properties(s3_bucket_mocked):
-    s3_root = "s3://timdex/dataset"
-    tdm_s3 = TIMDEXDatasetMetadata(s3_root)
-    assert tdm_s3.location == s3_root
-    assert tdm_s3.location_scheme == "s3"
+def test_tdm_s3_dataset_structure_properties(timdex_dataset_empty):
+    # test that location_scheme property works correctly for local paths
+    # S3 tests require full mocking and are covered in other tests
+    assert timdex_dataset_empty.metadata.location_scheme == "file"
 
 
-def test_tdm_create_metadata_database_file_success(caplog, timdex_metadata_empty):
+def test_tdm_create_metadata_database_file_success(
+    caplog, timdex_dataset_with_runs, timdex_metadata_empty
+):
     caplog.set_level("DEBUG")
-    timdex_metadata_empty.rebuild_dataset_metadata()
+    # use a fresh dataset from timdex_dataset_with_runs location
+    td = TIMDEXDataset(timdex_dataset_with_runs.location)
+    td.metadata.rebuild_dataset_metadata()
 
 
 def test_tdm_init_metadata_file_found_success(timdex_metadata):
@@ -321,15 +325,15 @@ def test_tdm_merge_append_deltas_deletes_append_deltas(
     assert not os.listdir(timdex_metadata_merged_deltas.append_deltas_path)
 
 
-def test_tdm_prepare_duckdb_secret_and_extensions_home_env_var_set_and_valid(
+def test_td_prepare_duckdb_secret_and_extensions_home_env_var_set_and_valid(
     monkeypatch, tmp_path_factory, timdex_dataset_with_runs
 ):
     preset_home = tmp_path_factory.mktemp("my-account")
     monkeypatch.setenv("HOME", str(preset_home))
 
-    tdm = TIMDEXDatasetMetadata(timdex_dataset_with_runs.location)
+    td = TIMDEXDataset(timdex_dataset_with_runs.location)
     df = (
-        tdm.conn.query(
+        td.conn.query(
             """
         select
             current_setting('secret_directory') as secret_directory,
@@ -344,15 +348,15 @@ def test_tdm_prepare_duckdb_secret_and_extensions_home_env_var_set_and_valid(
     assert df.extension_directory == ""  # expected and okay when HOME set
 
 
-def test_tdm_prepare_duckdb_secret_and_extensions_home_env_var_unset(
+def test_td_prepare_duckdb_secret_and_extensions_home_env_var_unset(
     monkeypatch, timdex_dataset_with_runs
 ):
     monkeypatch.delenv("HOME", raising=False)
 
-    tdm = TIMDEXDatasetMetadata(timdex_dataset_with_runs.location)
+    td = TIMDEXDataset(timdex_dataset_with_runs.location)
 
     df = (
-        tdm.conn.query(
+        td.conn.query(
             """
         select
             current_setting('secret_directory') as secret_directory,
@@ -367,15 +371,15 @@ def test_tdm_prepare_duckdb_secret_and_extensions_home_env_var_unset(
     assert df.extension_directory == "/tmp/.duckdb/extensions"
 
 
-def test_tdm_prepare_duckdb_secret_and_extensions_home_env_var_set_but_empty(
+def test_td_prepare_duckdb_secret_and_extensions_home_env_var_set_but_empty(
     monkeypatch, timdex_dataset_with_runs
 ):
     monkeypatch.setenv("HOME", "")  # simulate AWS Lambda environment
 
-    tdm = TIMDEXDatasetMetadata(timdex_dataset_with_runs.location)
+    td = TIMDEXDataset(timdex_dataset_with_runs.location)
 
     df = (
-        tdm.conn.query(
+        td.conn.query(
             """
         select
             current_setting('secret_directory') as secret_directory,
@@ -390,14 +394,16 @@ def test_tdm_prepare_duckdb_secret_and_extensions_home_env_var_set_but_empty(
     assert df.extension_directory == "/tmp/.duckdb/extensions"
 
 
-def test_tdm_preload_current_records_default_false(tmp_path):
-    tdm = TIMDEXDatasetMetadata(str(tmp_path))
-    assert tdm.preload_current_records is False
+def test_td_preload_current_records_default_false(tmp_path):
+    td = TIMDEXDataset(str(tmp_path))
+    assert td.preload_current_records is False
+    assert td.metadata.preload_current_records is False
 
 
-def test_tdm_preload_current_records_flag_true(tmp_path):
-    tdm = TIMDEXDatasetMetadata(str(tmp_path), preload_current_records=True)
-    assert tdm.preload_current_records is True
+def test_td_preload_current_records_flag_true(tmp_path):
+    td = TIMDEXDataset(str(tmp_path), preload_current_records=True)
+    assert td.preload_current_records is True
+    assert td.metadata.preload_current_records is True
 
 
 def test_tdm_preload_false_no_temp_table(timdex_dataset_with_runs):

tests/test_read.py

@@ -255,7 +255,6 @@ def test_dataset_load_current_records_gets_correct_same_day_full_run(
 ):
     # ensure metadata exists for this dataset
     timdex_dataset_same_day_runs.metadata.rebuild_dataset_metadata()
-    timdex_dataset_same_day_runs.metadata.refresh()
     df = timdex_dataset_same_day_runs.read_dataframe(
         table="current_records", run_type="full"
     )
@@ -266,7 +265,6 @@ def test_dataset_load_current_records_gets_correct_same_day_daily_runs_ordering(
     timdex_dataset_same_day_runs,
 ):
     timdex_dataset_same_day_runs.metadata.rebuild_dataset_metadata()
-    timdex_dataset_same_day_runs.metadata.refresh()
     first_record = next(
         timdex_dataset_same_day_runs.read_dicts_iter(
             table="current_records", run_type="daily"

timdex_dataset_api/dataset.py

@@ -16,12 +16,15 @@
 import pandas as pd
 import pyarrow as pa
 import pyarrow.dataset as ds
-from duckdb import DuckDBPyConnection
+from duckdb_engine import ConnectionWrapper
 from pyarrow import fs
+from sqlalchemy import MetaData, Table, create_engine
+from sqlalchemy.types import ARRAY, FLOAT
 
 from timdex_dataset_api.config import configure_logger
 from timdex_dataset_api.embeddings import TIMDEXEmbeddings
 from timdex_dataset_api.metadata import TIMDEXDatasetMetadata
+from timdex_dataset_api.utils import DuckDBConnectionFactory
 
 if TYPE_CHECKING:
     from timdex_dataset_api.record import DatasetRecord  # pragma: nocover
@@ -78,6 +81,10 @@ class TIMDEXDatasetConfig:
         from a dataset; pyarrow default is 16
     - fragment_read_ahead: number of fragments to optimistically read ahead when batch
         reaching from a dataset; pyarrow default is 4
+    - duckdb_join_batch_size: batch size for keyset pagination when joining metadata
+
+    Note: DuckDB connection settings (memory_limit, threads) are handled by
+    DuckDBConnectionFactory via TDA_DUCKDB_MEMORY_LIMIT and TDA_DUCKDB_THREADS env vars.
     """
 
     read_batch_size: int = field(
@@ -132,18 +139,21 @@ def __init__(
         self.partition_columns = TIMDEX_DATASET_PARTITION_COLUMNS
         self.dataset = self.load_pyarrow_dataset()
 
-        # dataset metadata
-        self.metadata = TIMDEXDatasetMetadata(
-            location,
-            preload_current_records=preload_current_records,
-        )
+        # create DuckDB connection used by all classes
+        self.conn_factory = DuckDBConnectionFactory(location_scheme=self.location_scheme)
+        self.conn = self.conn_factory.create_connection()
 
-        # DuckDB context
-        self.conn = self.setup_duckdb_context()
+        # create schemas
+        self._create_duckdb_schemas()
 
-        # dataset embeddings
+        # composed components receive self
+        self.metadata = TIMDEXDatasetMetadata(self)
         self.embeddings = TIMDEXEmbeddings(self)
 
+        # SQLAlchemy (SA) reflection after components have set up their views
+        self.sa_tables: dict[str, dict[str, Table]] = {}
+        self.reflect_sa_tables()
+
     @property
     def location_scheme(self) -> Literal["file", "s3"]:
         scheme = urlparse(self.location).scheme
@@ -158,7 +168,7 @@ def data_records_root(self) -> str:
         return f"{self.location.removesuffix('/')}/data/records"  # type: ignore[union-attr]
 
     def refresh(self) -> None:
-        """Fully reload TIMDEXDataset instance."""
+        """Refresh dataset by fully reinitializing."""
         self.__init__(  # type: ignore[misc]
             self.location,
             config=self.config,
@@ -245,24 +255,54 @@ def get_s3_filesystem() -> fs.FileSystem:
             session_token=credentials.token,
         )
 
-    def setup_duckdb_context(self) -> DuckDBPyConnection:
-        """Create a DuckDB connection that metadata and data query and retrieval.
+    def _create_duckdb_schemas(self) -> None:
+        """Create DuckDB schemas used by all components."""
+        self.conn.execute("create schema metadata;")
+        self.conn.execute("create schema data;")
 
-        This method extends TIMDEXDatasetMetadata's pre-existing DuckDB connection, adding
-        a 'data' schema and any other configurations needed.
+    def reflect_sa_tables(self, schemas: list[str] | None = None) -> None:
+        """Reflect SQLAlchemy metadata for DuckDB schemas.
+
+        This centralizes SA reflection for all composed components. Reflected tables
+        are stored in self.sa_tables as {schema: {table_name: Table}}.
+
+        Args:
+            schemas: list of schemas to reflect; defaults to ["metadata", "data"]
         """
         start_time = time.perf_counter()
+        schemas = schemas or ["metadata", "data"]
 
-        conn = self.metadata.conn
+        engine = create_engine(
+            "duckdb://",
+            creator=lambda: ConnectionWrapper(self.conn),
+        )
+
+        for schema in schemas:
+            db_metadata = MetaData()
+            db_metadata.reflect(bind=engine, schema=schema, views=True)
+
+            # store tables in flat dict keyed by table name (without schema prefix)
+            self.sa_tables[schema] = {
+                table_name.removeprefix(f"{schema}."): table
+                for table_name, table in db_metadata.tables.items()
+            }
 
-        # create data schema
-        conn.execute("""create schema data;""")
+        # type fixup for embedding_vector column (DuckDB LIST -> SA ARRAY)
+        if "embeddings" in self.sa_tables.get("data", {}):
+            self.sa_tables["data"]["embeddings"].c.embedding_vector.type = ARRAY(FLOAT)
 
         logger.debug(
-            "DuckDB context created for TIMDEXDataset, "
-            f"{round(time.perf_counter()-start_time,2)}s"
+            f"SQLAlchemy reflection complete for schemas {schemas}, "
+            f"{round(time.perf_counter() - start_time, 3)}s"
         )
-        return conn
+
+    def get_sa_table(self, schema: str, table: str) -> Table:
+        """Get a reflected SQLAlchemy Table by schema and table name."""
+        if schema not in self.sa_tables:
+            raise ValueError(f"Schema '{schema}' not found in reflected SA tables.")
+        if table not in self.sa_tables[schema]:
+            raise ValueError(f"Table '{table}' not found in schema '{schema}'.")
+        return self.sa_tables[schema][table]
 
     def write(
         self,
@@ -326,7 +366,7 @@ def write(
         if write_append_deltas:
             for written_file in written_files:
                 self.metadata.write_append_delta_duckdb(written_file.path)  # type: ignore[attr-defined]
-            self.metadata.refresh()
+            self.refresh()
 
         self.log_write_statistics(start_time, written_files)
 
@@ -575,9 +615,7 @@ def _iter_data_chunks(self, data_query: str) -> Iterator[pa.RecordBatch]:
             )
         finally:
             if self.location_scheme == "s3":
-                self.conn.execute(
-                    f"""set threads={self.metadata.config.duckdb_connection_threads};"""
-                )
+                self.conn.execute(f"""set threads={self.conn_factory.threads};""")
 
     def read_dataframes_iter(
         self,