Setup DuckDB context on TIMDEXDataset

Why these changes are being introduced:

TIMDEXDatasetMetadata (TDM) has a DuckDB context for metadata
attachments and views.  TIMDEXDataset (TD) will need one for
DuckDB queries that return actual ETL data, not just metadata.

How this addresses that need:
* TD reuses the TDM.conn DuckDB connection and builds upon it
* TDM DuckDB connection builds metadata related views in a
"metadata" schema
* TD will build views in a "data" schema

Side effects of this change:
* All TDM metadata views are under a "metadata" schema

Relevant ticket(s):
* https://mitlibraries.atlassian.net/browse/TIMX-529

Author: ghukill

tests/test_dataset.py

@@ -8,6 +8,7 @@
 
 import pyarrow as pa
 import pytest
+from duckdb.duckdb import DuckDBPyConnection
 from pyarrow import fs
 
 from timdex_dataset_api.dataset import (
@@ -291,3 +292,21 @@ def test_dataset_records_data_structure_is_idempotent(timdex_dataset_with_runs):
     assert os.path.exists(timdex_dataset_with_runs.data_records_root)
     end_file_count = glob.glob(f"{timdex_dataset_with_runs.data_records_root}/**/*")
     assert start_file_count == end_file_count
+
+
+def test_dataset_duckdb_context_crated_on_init(timdex_dataset):
+    assert isinstance(timdex_dataset.conn, DuckDBPyConnection)
+
+
+def test_dataset_duckdb_context_creates_data_schema(timdex_dataset):
+    assert (
+        timdex_dataset.conn.query(
+            """
+            select count(*)
+            from information_schema.schemata
+            where catalog_name = 'memory'
+            and schema_name = 'data';
+            """
+        ).fetchone()[0]
+        == 1
+    )

tests/test_metadata.py

@@ -36,6 +36,20 @@ def test_tdm_init_metadata_file_found_success(timdex_metadata):
     assert isinstance(timdex_metadata.conn, DuckDBPyConnection)
 
 
+def test_tdm_duckdb_context_creates_metadata_schema(timdex_metadata):
+    assert (
+        timdex_metadata.conn.query(
+            """
+            select count(*)
+            from information_schema.schemata
+            where catalog_name = 'memory'
+            and schema_name = 'metadata';
+            """
+        ).fetchone()[0]
+        == 1
+    )
+
+
 def test_tdm_connection_has_static_database_attached(timdex_metadata):
     assert set(
         timdex_metadata.conn.query("""show databases;""").to_df().database_name
@@ -71,7 +85,9 @@ def test_tdm_views_created_on_init(timdex_metadata):
 
 
 def test_tdm_records_view_structure(timdex_metadata):
-    records_df = timdex_metadata.conn.query("""select * from records limit 1;""").to_df()
+    records_df = timdex_metadata.conn.query(
+        """select * from metadata.records limit 1;"""
+    ).to_df()
     expected_columns = {
         "timdex_record_id",
         "source",
@@ -88,7 +104,7 @@ def test_tdm_records_view_structure(timdex_metadata):
 
 def test_tdm_current_records_view_structure(timdex_metadata):
     current_records_df = timdex_metadata.conn.query(
-        """select * from current_records limit 1;"""
+        """select * from metadata.current_records limit 1;"""
     ).to_df()
     expected_columns = {
         "timdex_record_id",
@@ -106,7 +122,7 @@ def test_tdm_current_records_view_structure(timdex_metadata):
 
 def test_tdm_append_deltas_view_empty_structure(timdex_metadata):
     append_deltas_df = timdex_metadata.conn.query(
-        """select * from append_deltas;"""
+        """select * from metadata.append_deltas;"""
     ).to_df()
     expected_columns = {
         "timdex_record_id",
@@ -127,7 +143,7 @@ def test_tdm_records_count_property(timdex_metadata):
     assert timdex_metadata.records_count > 0
 
     manual_count = timdex_metadata.conn.query(
-        """select count(*) from records;"""
+        """select count(*) from metadata.records;"""
     ).fetchone()[0]
     assert timdex_metadata.records_count == manual_count
 
@@ -136,7 +152,7 @@ def test_tdm_current_records_count_property(timdex_metadata):
     assert timdex_metadata.current_records_count > 0
 
     manual_count = timdex_metadata.conn.query(
-        """select count(*) from current_records;"""
+        """select count(*) from metadata.current_records;"""
     ).fetchone()[0]
     assert timdex_metadata.current_records_count == manual_count
 
@@ -150,7 +166,7 @@ def test_tdm_records_equals_static_without_deltas(timdex_metadata):
         """select count(*) from static_db.records;"""
     ).fetchone()[0]
     records_count = timdex_metadata.conn.query(
-        """select count(*) from records;"""
+        """select count(*) from metadata.records;"""
     ).fetchone()[0]
     assert static_count == records_count
 
@@ -198,7 +214,7 @@ def test_tdm_current_records_with_deltas_logic(timdex_metadata_with_deltas):
 
     # verify current records view returns unique timdex_record_id values
     current_records_df = timdex_metadata_with_deltas.conn.query(
-        """select timdex_record_id from current_records;"""
+        """select timdex_record_id from metadata.current_records;"""
     ).to_df()
 
     unique_count = len(current_records_df.timdex_record_id.unique())
@@ -210,7 +226,7 @@ def test_tdm_current_records_most_recent_version(timdex_metadata_with_deltas):
     multi_version_records = timdex_metadata_with_deltas.conn.query(
         """
         select timdex_record_id, count(*) as version_count
-        from records
+        from metadata.records
         group by timdex_record_id
         having count(*) > 1
         limit 1;
@@ -224,7 +240,7 @@ def test_tdm_current_records_most_recent_version(timdex_metadata_with_deltas):
         most_recent = timdex_metadata_with_deltas.conn.query(
             f"""
             select run_timestamp, run_id
-            from records
+            from metadata.records
             where timdex_record_id = '{record_id}'
             order by run_timestamp desc
             limit 1;
@@ -235,7 +251,7 @@ def test_tdm_current_records_most_recent_version(timdex_metadata_with_deltas):
         current_version = timdex_metadata_with_deltas.conn.query(
             f"""
             select run_timestamp, run_id
-            from current_records
+            from metadata.current_records
             where timdex_record_id = '{record_id}';
             """
         ).to_df()

timdex_dataset_api/dataset.py

@@ -18,6 +18,7 @@
 import pandas as pd
 import pyarrow as pa
 import pyarrow.dataset as ds
+from duckdb import DuckDBPyConnection
 from pyarrow import fs
 
 from timdex_dataset_api.config import configure_logger
@@ -128,6 +129,9 @@ def __init__(
         # dataset metadata
         self.metadata = TIMDEXDatasetMetadata(location)
 
+        # DuckDB context
+        self.conn = self.setup_duckdb_context()
+
     @property
     def location_scheme(self) -> Literal["file", "s3"]:
         scheme = urlparse(self.location).scheme
@@ -221,6 +225,24 @@ 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.
+
+        This relies on TIMDEXDatasetMetadata.setup_duckdb_context() to produce a DuckDB
+        connection that has all metadata already created.
+        """
+        start_time = time.perf_counter()
+
+        conn = self.metadata.conn
+
+        # create data schema
+        conn.execute("""create schema data;""")
+
+        logger.debug(
+            f"DuckDB data context created, {round(time.perf_counter()-start_time,2)}s"
+        )
+        return conn
+
     def write(
         self,
         records_iter: Iterator["DatasetRecord"],

timdex_dataset_api/metadata.py

@@ -91,17 +91,35 @@ def append_deltas_path(self) -> str:
     @property
     def records_count(self) -> int:
         """Count of all records in dataset."""
-        return self.conn.query("""select count(*) from records;""").fetchone()[0]  # type: ignore[index]
+        return self.conn.query(
+            """
+            select count(*) from metadata.records;
+            """
+        ).fetchone()[
+            0
+        ]  # type: ignore[index]
 
     @property
     def current_records_count(self) -> int:
         """Count of all current records in dataset."""
-        return self.conn.query("""select count(*) from current_records;""").fetchone()[0]  # type: ignore[index]
+        return self.conn.query(
+            """
+            select count(*) from metadata.current_records;
+            """
+        ).fetchone()[
+            0
+        ]  # type: ignore[index]
 
     @property
     def append_deltas_count(self) -> int:
         """Count of all append deltas."""
-        return self.conn.query("""select count(*) from append_deltas;""").fetchone()[0]  # type: ignore[index]
+        return self.conn.query(
+            """
+            select count(*) from metadata.append_deltas;
+            """
+        ).fetchone()[
+            0
+        ]  # type: ignore[index]
 
     def create_metadata_structure(self) -> None:
         """Ensure metadata structure exists in TIDMEX dataset.."""
@@ -280,13 +298,16 @@ def setup_duckdb_context(self) -> DuckDBPyConnection:
             )
             return conn
 
+        # create metadata schema
+        conn.execute("create schema metadata;")
+
         self._attach_database_file(conn)
         self._create_append_deltas_view(conn)
         self._create_records_union_view(conn)
         self._create_current_records_view(conn)
 
         logger.debug(
-            f"DuckDB context created, {round(time.perf_counter()-start_time,2)}s"
+            f"DuckDB metadata context created, {round(time.perf_counter()-start_time,2)}s"
         )
         return conn
 
@@ -327,7 +348,7 @@ def _create_append_deltas_view(self, conn: DuckDBPyConnection) -> None:
         # if deltas, create view projecting over those parquet files
         if append_delta_count > 0:
             query = f"""
-            create view append_deltas as (
+            create or replace view metadata.append_deltas as (
                 select *
                 from read_parquet(
                     '{self.append_deltas_path}/*.parquet'
@@ -338,7 +359,7 @@ def _create_append_deltas_view(self, conn: DuckDBPyConnection) -> None:
         # if not, create a view that mirrors the structure of static_db.records
         else:
             query = """
-            create view append_deltas as (
+            create or replace view metadata.append_deltas as (
                 select *
                 from static_db.records
                 where 1 = 0
@@ -350,13 +371,13 @@ def _create_records_union_view(self, conn: DuckDBPyConnection) -> None:
         logger.debug("creating view of unioned records")
         conn.execute(
             """
-            create view records as
+            create or replace view metadata.records as
             (
                 select *
                 from static_db.records
                 union all
                 select *
-                from append_deltas
+                from metadata.append_deltas
             );
             """
         )
@@ -374,18 +395,18 @@ def _create_current_records_view(self, conn: DuckDBPyConnection) -> None:
         logger.info("creating view of current records metadata")
 
         query = f"""
-        create or replace view current_records as
+        create or replace view metadata.current_records as
         with ranked_records as (
             select
                 r.*,
                 row_number() over (
                     partition by r.timdex_record_id
                     order by r.run_timestamp desc
                 ) as rn
-            from records r
+            from metadata.records r
             where r.run_timestamp >= (
                 select max(r2.run_timestamp)
-                from records r2
+                from metadata.records r2
                 where r2.source = r.source
                 and r2.run_type = 'full'
             )