From 2fd22187aede39986bf7a21239aef03c8e63e985 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Sun, 3 Aug 2025 08:57:13 -0400 Subject: [PATCH 1/5] Begin rebuild of TIMDEXDatasetMetadata Why these changes are being introduced: The current overarching work is to support the creation and reading of a static metadata database file and append deltas. To get there, very little of the original TIMDEXDatasetMetadata class is needed or wanted. This commit begins the process of rebuilding TIMDEXDatasetMetadata, oriented around managing a static metadata database file, and providing a readonly projection over that and append delta paqruet files. How this addresses that need: TIMDEXDatasetMetadata is almost completely rebuilt, with the first functionality being the creation of the static metadata file by scanning the ETL records. Then, the ability to remotely attach in readonly mode to this metadata database file for reading. Note: these changes are breaking. TIMDEXDataset cannot provide "current" records and many unit tests are broken. This will be addressed in future commits as we build this class back up with new functionality. Side effects of this change: * TIMDEXDataset cannot provide current records * Unit tests are either temporarily skipped or failing Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-530 --- tests/test_metadata.py | 57 +---- timdex_dataset_api/metadata.py | 379 ++++++++++----------------------- timdex_dataset_api/utils.py | 55 ++++- 3 files changed, 178 insertions(+), 313 deletions(-) diff --git a/tests/test_metadata.py b/tests/test_metadata.py index 9f9d74c..8af6c2d 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -39,51 +39,14 @@ def test_tdm_get_duckdb_connection(timdex_dataset_metadata): assert isinstance(conn, duckdb.DuckDBPyConnection) -def test_tdm_set_threads(timdex_dataset_metadata): - # set to 64 - timdex_dataset_metadata.set_database_thread_usage(64) - sixty_four_thread_count = timdex_dataset_metadata.conn.query( - """SELECT current_setting('threads');""" - ).fetchone()[0] - assert sixty_four_thread_count == 64 - - # set to 12 - timdex_dataset_metadata.set_database_thread_usage(12) - sixty_four_thread_count = timdex_dataset_metadata.conn.query( - """SELECT current_setting('threads');""" - ).fetchone()[0] - assert sixty_four_thread_count == 12 - - -def test_tdm_init_sets_up_database(timdex_dataset_metadata): - df = timdex_dataset_metadata.conn.query("show tables;").to_df() - assert set(df.name) == {"current_records", "records"} - - -def test_tdm_get_current_parquet_files(timdex_dataset_metadata): - parquet_files = timdex_dataset_metadata.get_current_parquet_files() - # assert 5 total parquet files in dataset - # but only 3 contain current records - assert len(timdex_dataset_metadata.timdex_dataset.dataset.files) == 5 - assert len(parquet_files) == 3 - - -def test_tdm_get_record_to_run_mapping(timdex_dataset_metadata): - record_map = timdex_dataset_metadata.get_current_record_to_run_map() - - assert len(record_map) == 75 - assert record_map["alma:0"] == "run-5" - assert record_map["alma:5"] == "run-4" - assert record_map["alma:19"] == "run-4" - assert "run-3" not in record_map.values() - assert record_map["alma:20"] == "run-2" - - -def test_tdm_current_records_subset_of_all_records(timdex_dataset_metadata): - records_df = timdex_dataset_metadata.conn.query("select * from records;").to_df() - current_records_df = timdex_dataset_metadata.conn.query( - "select * from current_records;" +def test_tdm_connection_has_static_database_attached(timdex_dataset_metadata): + assert set( + timdex_dataset_metadata.conn.query("""show databases;""").to_df().database_name + ) == {"memory", "static_db"} + + +def test_tdm_connection_static_database_records_table_exists(timdex_dataset_metadata): + records_df = timdex_dataset_metadata.conn.query( + """select * from static_db.records;""" ).to_df() - assert set(current_records_df.timdex_record_id).issubset( - set(records_df.timdex_record_id) - ) + assert len(records_df) > 0 diff --git a/timdex_dataset_api/metadata.py b/timdex_dataset_api/metadata.py index def7533..dca957d 100644 --- a/timdex_dataset_api/metadata.py +++ b/timdex_dataset_api/metadata.py @@ -1,316 +1,165 @@ """timdex_dataset_api/metadata.py""" import os +import tempfile import time -from typing import TYPE_CHECKING, Unpack +from pathlib import Path from urllib.parse import urlparse import duckdb +from duckdb import DuckDBPyConnection from timdex_dataset_api.config import configure_logger - -if TYPE_CHECKING: - from timdex_dataset_api.dataset import DatasetFilters, TIMDEXDataset +from timdex_dataset_api.utils import S3Client, configure_duckdb_s3_secret logger = configure_logger(__name__) +ORDERED_DATASET_COLUMN_NAMES = [ + "timdex_record_id", + "source", + "run_date", + "run_type", + "action", + "run_id", + "run_record_offset", + "run_timestamp", + "filename", +] -class TIMDEXDatasetMetadata: - """Collect and provide access to metadata about the parquet dataset. - - The ETL parquet dataset is essentially parquet files in S3. This class utilizes - DuckDB to generate metadata about the parquet dataset, down to the individual record - layer. This is somewhat similar to how other data lakes like Apache Iceberg or - DuckDB DuckLake provide a metadata layer over the stored, large, raw files. - Because this metadata is somewhat infrequently needed, e.g. only for bulk operations - or analysis, the architectural decision has been made to pay an initial time penalty - of crawling the dataset to generate metadata which is then used to dramatically - speed up and simplify other operations. In the event this dataset-wide metadata - is needed more often, it may be worth exploring storing it in S3 alongside the data - and updating it for each write; very much mirroring other data lake frameworks. - """ +class TIMDEXDatasetMetadata: def __init__( self, - timdex_dataset: "TIMDEXDataset", - db_path: str = ":memory:", - ): - """Initialize TIMDEXDatasetMetadata. + location: str, + ) -> None: + """Init TIMDEXDatasetMetadata. Args: - timdex_dataset: The TIMDEX dataset instance to extract metadata from - db_path: Path to the DuckDB database file. Defaults to ":memory:" for - in-memory database + location: root location of TIMDEX dataset, e.g. 's3://timdex/dataset' """ - self.timdex_dataset = timdex_dataset - self.db_path = db_path - - self.conn = self.get_connection() - self._setup_database() - - @classmethod - def from_dataset_location( - cls, - timdex_dataset_location: str, - **kwargs: str, - ) -> "TIMDEXDatasetMetadata": - """Factory method to init TIMDEXDatasetMetadata from a dataset location. - - This first instantiates and loads a TIMDEXDataset instance, then instantiates this - class using that. While this class will likely most commonly be used by - TIMDEXDataset to limit to current records, it is hoped and expected this dataset - metadata client will be increasingly useful in its own right, thus this method. - - Args: - timdex_dataset_location: S3 path or local path to the TIMDEX dataset - **kwargs: Additional keyword arguments passed to the class constructor, - such as db_path + self.location = location + self.conn: None | DuckDBPyConnection = self.setup_duckdb_context() + + @property + def metadata_root(self) -> str: + return f"{self.location.removesuffix('/')}/metadata" + + @property + def metadata_database_filename(self) -> str: + return "metadata.duckdb" + + @property + def metadata_database_path(self) -> str: + return f"{self.metadata_root}/{self.metadata_database_filename}" + + @property + def append_deltas_path(self) -> str: + return f"{self.metadata_root}/append_deltas" + + def database_exists(self) -> bool: + """Check if static metadata database file exists.""" + if urlparse(self.metadata_database_path).scheme == "s3": + s3_client = S3Client() + return s3_client.object_exists(self.metadata_database_path) + return os.path.exists(self.metadata_database_path) + + def recreate_static_database_file(self) -> None: + """Create/recreate the static metadata database file. + + The following work is performed: + 1. Create a local working directory + 2. Open a DuckDB connection with a database file in this local working dir + 3. Create tables and views by scanning ETL data in dataset/data/records + 4. Close DuckDB connection ensuring a fully formed, local database file + 5. Upload DuckDB database file to target destination, making that the new + static metadata database file """ - # avoids circular import dependency - from .dataset import TIMDEXDataset # noqa: PLC0415 - - timdex_dataset = TIMDEXDataset(timdex_dataset_location) - timdex_dataset.load() - return cls(timdex_dataset, **kwargs) + s3_client = S3Client() - def get_connection(self) -> duckdb.DuckDBPyConnection: - """Get a DuckDB connection to the metadata database.""" - return duckdb.connect(self.db_path) - - def set_database_thread_usage(self, thread_count: int) -> None: - """Set the number of threads for DuckDB operations.""" - self.conn.execute(f"""SET threads = {thread_count};""") - - def _setup_database(self) -> None: - """Initialize DuckDB database with AWS credentials and base tables and views.""" - start_time = time.perf_counter() + # remove any append deltas that may exist at this time of database recreation + s3_client.delete_folder(self.append_deltas_path) - # bump threads for high parallelization of lightweight data calls for metadata - self.set_database_thread_usage(64) + # build database locally + with tempfile.TemporaryDirectory() as temp_dir: + local_db_path = str(Path(temp_dir) / self.metadata_database_filename) - # configure s3 connection - self._configure_s3_connection() + with duckdb.connect(local_db_path) as conn: + conn.execute("""SET threads = 64;""") + configure_duckdb_s3_secret(conn) - # create a table of metadata about all rows in dataset - self._create_full_dataset_table() + self._create_full_dataset_table(conn) - # create a view for current records - self._create_current_records_view() - - logger.info( - f"metadata database setup elapsed: {time.perf_counter()-start_time}, " - f"path: '{self.db_path}'" - ) - - def _configure_s3_connection(self) -> None: - """Configure S3 connection for DuckDB access. - - If the env var 'MINIO_S3_ENDPOINT_URL' is present, assume a local MinIO S3 - instance and configure accordingly, otherwise assume normal AWS S3 and setup a - credentials chain in DuckDB. - """ - logger.info("configuring S3 connection") - - if os.getenv("MINIO_S3_ENDPOINT_URL"): - self.conn.execute( - f""" - create or replace secret minio_s3_secret ( - type s3, - endpoint '{urlparse(os.environ["MINIO_S3_ENDPOINT_URL"]).netloc}', - key_id '{os.environ["MINIO_USERNAME"]}', - secret '{os.environ["MINIO_PASSWORD"]}', - region 'us-east-1', - url_style 'path', - use_ssl false - ); - """ + # copy local database file to remote location + s3_client.upload_file( + local_db_path, + self.metadata_database_path, ) - else: - self.conn.execute( - """ - create or replace secret aws_s3_secret ( - type s3, - provider credential_chain, - chain 'sso;env;config', - refresh true - ); - """ - ) + # refresh DuckDB connection + self.conn = self.setup_duckdb_context() - def _create_full_dataset_table(self) -> None: - """Create a table of metadata about all records in the parquet dataset. + def _create_full_dataset_table(self, conn: DuckDBPyConnection) -> None: + """Create a table of metadata for all records in the ETL parquet dataset. - While this table will obviously have a high number of rows, the data is small. - Testing has shown around 20 million records results in 1gb in memory or ~150mb on - disk. + This is one of the few times we fully materialize data in a DuckDB connection. + This is most commonly used when recreating the baseline static metadata database + file. """ start_time = time.perf_counter() logger.info("creating table of full dataset metadata") - parquet_glob_pattern = self._prepare_parquet_file_glob_pattern() query = f""" - create or replace table records as ( - select - timdex_record_id, - source, - run_date, - run_type, - run_id, - action, - run_record_offset, - run_timestamp, - filename, - from read_parquet( - {parquet_glob_pattern}, - hive_partitioning=true, - filename=true - ) - ); - """ - self.conn.execute(query) - - row_count = self.conn.query("""select count(*) from records;""").fetchone()[0] # type: ignore[index] + create or replace table records as ( + select + {','.join(ORDERED_DATASET_COLUMN_NAMES)} + from read_parquet( + '{self.location}/data/records/**/*.parquet', + hive_partitioning=true, + filename=true + ) + ); + """ + conn.execute(query) + + row_count = conn.query("""select count(*) from records;""").fetchone()[0] # type: ignore[index] logger.info( f"'records' table created - rows: {row_count}, " f"elapsed: {time.perf_counter() - start_time}" ) - def _prepare_parquet_file_glob_pattern(self) -> str: - """Prepare a parquet file glob pattern suitable for DuckDB read_parquet().""" - if isinstance(self.timdex_dataset.location, list): - return ",".join([f"'{file}'" for file in self.timdex_dataset.location]) + def setup_duckdb_context(self) -> DuckDBPyConnection | None: + """Create a DuckDB connection that provides full dataset metadata information. - prefix = self.timdex_dataset.location.removesuffix("/") - return f"'{prefix}/**/*.parquet'" + The following work is performed: + 1. Attach to static metadata database file. + 2. Create views that union static metadata with any append deltas. + 3. Create additional metadata views as needed. - def _create_current_records_view(self) -> None: - """Create a view of current records. - - This view builds on the table `records`. - - This view includes only the most current version of each record in the dataset. - Because it includes the `timdex_record_id` and `run_id`, it makes yielding the - current version of a record via a TIMDEXDataset instance trivial: for any given - `timdex_record_id` if the `run_id` doesn't match, it's not the current version. + The resulting, in-memory DuckDB connection is used for all metadata queries. """ - start_time = time.perf_counter() - logger.info("creating view of current records metadata") - - query = """ - create or replace view 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 - where r.run_timestamp >= ( - select max(r2.run_timestamp) - from records r2 - where r2.source = r.source - and r2.run_type = 'full' + if not self.database_exists(): + logger.warning( + f"Static metadata database not found @ '{self.metadata_database_path}'. " + "Please recreate via TIMDEXDatasetMetadata.recreate_database_file()." ) - ) - select - timdex_record_id, - source, - run_date, - run_type, - run_id, - action, - run_record_offset, - run_timestamp, - filename - from ranked_records - where rn = 1; - """ - self.conn.execute(query) + return None - row_count = self.conn.query( # type: ignore[index] - """select count(*) from current_records;""" - ).fetchone()[0] - logger.info( - f"'current_records' view created - rows: {row_count}, " - f"elapsed: {time.perf_counter() - start_time}" - ) - - def get_current_parquet_files( - self, - *, - strip_protocol_prefix: bool = True, - **filters: Unpack["DatasetFilters"], - ) -> list[str]: - """Provide a list of parquet files that contain one or more current records. - - Args: - - strip_protocol_prefix: boolean if the file protocol should be removed, - e.g. "s3://" - - **filters: keyword dataset filters like `source="alma"` or - `run_date="2025-05-01"` - """ - where_clause = self._prepare_where_clause_from_dataset_filters(**filters) - - query = f""" - select distinct - filename as parquet_filename - from current_records - {where_clause} - order by run_timestamp desc; - """ - parquet_files_df = self.conn.query(query).to_df() + conn = duckdb.connect() + configure_duckdb_s3_secret(conn) - if strip_protocol_prefix: - parquet_files_df["parquet_filename"] = parquet_files_df[ - "parquet_filename" - ].apply(lambda x: x.removeprefix("s3://")) + self._attach_database_file(conn) - return list(parquet_files_df["parquet_filename"]) + return conn - def get_current_record_to_run_map(self, **filters: Unpack["DatasetFilters"]) -> dict: - """Provide a dictionary of timdex_record_id --> run_id for current records. + def _attach_database_file(self, conn: DuckDBPyConnection) -> None: + """Readonly attach to static metadata database. - This dictionary is all that read methods in TIMDEXDataset would require to ensure - they only yield the current version of a record. - - Args: - - **filters: keyword dataset filters like `source="alma"` or - `run_date="2025-05-01"` - """ - start_time = time.perf_counter() - - where_clause = self._prepare_where_clause_from_dataset_filters(**filters) - - query = f""" - select - timdex_record_id, - run_id - from current_records - {where_clause} - ; + Attaching to a remote DuckDB database file is supported, but only in readonly + mode: https://duckdb.org/docs/stable/sql/statements/attach.html, though it does + support multiple, concurrent attachments. """ - mapper_df = self.conn.query(query).to_df() - mapper_dict = mapper_df.set_index("timdex_record_id")["run_id"].to_dict() - logger.info( - f"Record-to-run mapper dict created elapsed: {time.perf_counter()-start_time}" + logger.debug(f"Attaching to static database file: {self.metadata_database_path}") + conn.execute( + f"""attach '{self.metadata_database_path}' AS static_db (READ_ONLY);""" ) - return mapper_dict - - def _prepare_where_clause_from_dataset_filters( - self, **filters: Unpack["DatasetFilters"] - ) -> str: - """Given keyword filters from DatasetFilters, provide a SQL WHERE clause. - - Note: this implementation of translating TIMDEXDataset DatasetFilters to a single - SQL WHERE clause is quite naive. This does the trick for now, supporting filters - like `source` or `run_date`, but this should be revisited if more robust filtering - is needed. - """ - conditions = [f"{column} = '{value}'" for column, value in filters.items()] - - if conditions: - return f"where {' and '.join(conditions)}" - return "" diff --git a/timdex_dataset_api/utils.py b/timdex_dataset_api/utils.py index 4e71419..31473f6 100644 --- a/timdex_dataset_api/utils.py +++ b/timdex_dataset_api/utils.py @@ -6,6 +6,7 @@ from urllib.parse import urlparse import boto3 +from duckdb import DuckDBPyConnection from mypy_boto3_s3.service_resource import S3ServiceResource logger = logging.getLogger(__name__) @@ -34,6 +35,16 @@ def _create_resource(self) -> S3ServiceResource: ) return boto3.resource("s3") + def object_exists(self, s3_uri: str) -> bool: + bucket, key = self._split_s3_uri(s3_uri) + try: + self.resource.Object(bucket, key).load() + return True # noqa: TRY300 + except self.resource.meta.client.exceptions.ClientError as e: + if e.response["Error"]["Code"] == "404": + return False + raise + def download_file(self, s3_uri: str, local_path: str | pathlib.Path) -> None: bucket, key = self._split_s3_uri(s3_uri) local_path = pathlib.Path(local_path) @@ -61,7 +72,7 @@ def delete_folder(self, s3_uri: str) -> list[str]: deleted_keys = [] for request in receipt: deleted_keys.extend([item["Key"] for item in request["Deleted"]]) - logger.info(f"Deleted {deleted_keys}") + logger.info(f"Deleted objects with prefix '{s3_uri}': {deleted_keys}") return deleted_keys @staticmethod @@ -74,3 +85,45 @@ def _split_s3_uri(s3_uri: str) -> tuple[str, str]: bucket = parsed.netloc key = parsed.path.lstrip("/") # strip leading slash from /key return bucket, key + + +def configure_duckdb_s3_secret( + conn: DuckDBPyConnection, + scope: str | None = None, +) -> None: + """Configure a secret in a DuckDB connection for S3 access. + + If a scope is provided, e.g. an S3 URI prefix like 's3://timdex', set a scope + parameter in the config. Else, leave it blank. + """ + # establish scope string + scope_str = f", scope '{scope}'" if scope else "" + + if os.getenv("MINIO_S3_ENDPOINT_URL"): + conn.execute( + f""" + create or replace secret minio_s3_secret ( + type s3, + endpoint '{urlparse(os.environ["MINIO_S3_ENDPOINT_URL"]).netloc}', + key_id '{os.environ["MINIO_USERNAME"]}', + secret '{os.environ["MINIO_PASSWORD"]}', + region 'us-east-1', + url_style 'path', + use_ssl false + {scope_str} + ); + """ + ) + + else: + conn.execute( + f""" + create or replace secret aws_s3_secret ( + type s3, + provider credential_chain, + chain 'sso;env;config', + refresh true + {scope_str} + ); + """ + ) From ff2aff07e194faad012891a21b2c8dd382eb807a Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Sun, 3 Aug 2025 09:08:35 -0400 Subject: [PATCH 2/5] Remove current records functionality in TIMDEXDataset Why these changes are being introduced: While the TIMDEXDatasetMetadata class is rebuilt, TIMDEXDataset itself can no longer provide "current" records from the dataaset as it has no metadata to work with. This is temporary until TIMDEXDatasetMetadata is rebuilt, and TIMDEXDataset gets new functionality based on *that* new metadata. How this addresses that need: * Any reference to "current records" is removed Side effects of this change: * TIMDEXDataset cannot provide current records Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-530 --- ...25_consistent_run_timestamp_per_etl_run.py | 9 +- timdex_dataset_api/dataset.py | 82 +------------------ 2 files changed, 11 insertions(+), 80 deletions(-) diff --git a/migrations/002_2025_06_25_consistent_run_timestamp_per_etl_run.py b/migrations/002_2025_06_25_consistent_run_timestamp_per_etl_run.py index 2d009ce..c741a23 100644 --- a/migrations/002_2025_06_25_consistent_run_timestamp_per_etl_run.py +++ b/migrations/002_2025_06_25_consistent_run_timestamp_per_etl_run.py @@ -1,4 +1,7 @@ -# ruff: noqa: BLE001, D212, TRY300, TRY400 +# ruff: noqa: PGH004 +# ruff: noqa +# type: ignore + """ Date: 2025-06-25 @@ -29,6 +32,10 @@ pipenv run python migrations/002_2025_06_25_consistent_run_timestamp_per_etl_run.py \ \ --dry-run + +Update: 2025-08-04 + +This migration is no longer functional given changes to TIMDEXDataset. """ import argparse diff --git a/timdex_dataset_api/dataset.py b/timdex_dataset_api/dataset.py index f835074..fbc7366 100644 --- a/timdex_dataset_api/dataset.py +++ b/timdex_dataset_api/dataset.py @@ -20,11 +20,11 @@ from timdex_dataset_api.config import configure_logger from timdex_dataset_api.exceptions import DatasetNotLoadedError -from timdex_dataset_api.metadata import TIMDEXDatasetMetadata if TYPE_CHECKING: from timdex_dataset_api.record import DatasetRecord # pragma: nocover + logger = configure_logger(__name__) TIMDEX_DATASET_SCHEMA = pa.schema( @@ -126,10 +126,6 @@ def __init__( # writing self._written_files: list[ds.WrittenFile] = None # type: ignore[assignment] - # reading - self._current_records: bool = False - self.metadata: TIMDEXDatasetMetadata = None # type: ignore[assignment] - @property def row_count(self) -> int: """Get row count from loaded dataset.""" @@ -139,8 +135,6 @@ def row_count(self) -> int: def load( self, - *, - current_records: bool = False, **filters: Unpack[DatasetFilters], ) -> None: """Lazy load a pyarrow.dataset.Dataset and set to self.dataset. @@ -161,21 +155,12 @@ def load( - filters: kwargs typed via DatasetFilters TypedDict - Filters passed directly in method call, e.g. source="alma", run_date="2024-12-20", etc., but are typed according to DatasetFilters. - - current_records: bool - - if True, all records yielded from this instance will be the current - version of the record in the dataset. """ start_time = time.perf_counter() # reset paths from original location before load _, self.paths = self.parse_location(self.location) - # read dataset metadata if only current records are requested - self._current_records = current_records - if current_records: - self.metadata = TIMDEXDatasetMetadata(timdex_dataset=self) - self.paths = self.metadata.get_current_parquet_files(**filters) - # perform initial load of full dataset self.dataset = self._load_pyarrow_dataset() @@ -465,10 +450,6 @@ def read_batches_iter( While batch_size will limit the max rows per batch, filtering may result in some batches having less than this limit. - If the flag self._current_records is set, this method leans on - self._yield_current_record_deduped_batches() to apply deduplication of records to - ensure only current versions of the record are ever yielded. - Args: - columns: list[str], list of columns to return from the dataset - filters: pairs of column:value to filter the dataset @@ -479,13 +460,6 @@ def read_batches_iter( ) dataset = self._get_filtered_dataset(**filters) - # if current records, add required columns for deduplication - if self._current_records: - if not columns: - columns = TIMDEX_DATASET_SCHEMA.names - columns.extend(["timdex_record_id", "run_id"]) - columns = list(set(columns)) - batches = dataset.to_batches( columns=columns, batch_size=self.config.read_batch_size, @@ -493,59 +467,9 @@ def read_batches_iter( fragment_readahead=self.config.fragment_read_ahead, ) - if self._current_records: - yield from self._yield_current_record_batches(batches, **filters) - else: - for batch in batches: - if len(batch) > 0: - yield batch - - def _yield_current_record_batches( - self, - batches: Iterator[pa.RecordBatch], - **filters: Unpack[DatasetFilters], - ) -> Iterator[pa.RecordBatch]: - """Method to yield only the most recent version of each record. - - When multiple versions of a record (same timdex_record_id) exist in the dataset, - this method ensures only the most recent version is returned. If filtering is - applied that removes this most recent version of a record, that timdex_record_id - will not be yielded at all. - - This method uses TIMDEXDatasetMetadata to provide a mapping of timdex_record_id to - run_id for the current ETL run for that record. While yielding records, only when - the timdex_record_id + run_id match the mapping is a record yielded. - - Args: - - batches: batches of records to actually yield from - - filters: pairs of column:value to filter the dataset metadata required - """ - # get map of timdex_record_id to run_id for current version of that record - record_to_run_map = self.metadata.get_current_record_to_run_map(**filters) - - # loop through batches, yielding only current records for batch in batches: - - if batch.num_rows == 0: - continue - - to_yield_indices = [] - - record_ids = batch.column("timdex_record_id").to_pylist() - run_ids = batch.column("run_id").to_pylist() - - for i, (record_id, run_id) in enumerate( - zip( - record_ids, - run_ids, - strict=True, - ) - ): - if record_to_run_map.get(record_id) == run_id: - to_yield_indices.append(i) - - if to_yield_indices: - yield batch.take(pa.array(to_yield_indices)) # type: ignore[arg-type] + if len(batch) > 0: + yield batch def read_dataframes_iter( self, From 37c9275276a7ae4c31dba4271d97f88dd7188146 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Mon, 4 Aug 2025 09:47:48 -0400 Subject: [PATCH 3/5] Property for ETL records data Why these changes are being introduced: This is a small change now, that will lead to a larger change later. The TIMDEX dataset is getting more structure, and this means we will want to initialize a TIMDEXDataset instance with the root of the dataset, but then internally there will be more opinionation about where files should be read and written to. How this addresses that need: A new property 'data_records_root' is added to TIMDEXDataset that mirrors similar properties in TIMDEXDatasetMetadata. This informs any operations that need to read or write ETL records where precisely they are in the dataset. At this time only .write() utilizes it, but in a future ticket the load method will be heavily reworked (if not outright removed) and this property will be fully integrated. This is needed now to continue updates to TIMDEXMetadataDataset for TIMX-530. Side effects of this change: * Initialization of TIMDEXDataset should provide the true dataset root, not point to /data/records. The pipeline lambda currently does this, but will be updated in TIMX-531. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-530 * https://mitlibraries.atlassian.net/browse/TIMX-531 --- timdex_dataset_api/dataset.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/timdex_dataset_api/dataset.py b/timdex_dataset_api/dataset.py index fbc7366..0b20923 100644 --- a/timdex_dataset_api/dataset.py +++ b/timdex_dataset_api/dataset.py @@ -126,6 +126,10 @@ def __init__( # writing self._written_files: list[ds.WrittenFile] = None # type: ignore[assignment] + @property + def data_records_root(self) -> str: + return f"{self.location.removesuffix('/')}/data/records" # type: ignore[union-attr] + @property def row_count(self) -> int: """Get row count from loaded dataset.""" @@ -370,7 +374,7 @@ def write( start_time = time.perf_counter() self._written_files = [] - dataset_filesystem, dataset_path = self.parse_location(self.location) + dataset_filesystem, dataset_path = self.parse_location(self.data_records_root) if isinstance(dataset_path, list): raise TypeError( "Dataset location must be the root of a single dataset for writing" From 7f7800dd7868513a560e08ee4caffb6bedbbb682 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Mon, 4 Aug 2025 10:59:56 -0400 Subject: [PATCH 4/5] Begin rebuilding of data and metadata tests Why these changes are being introduced: With the big changes to TIMDEXMetadataDataset comes the need to virtually rewrite the test suite for that class. The changes too TIMDEXMetadataDataset are also influencing tests for TIMDEXDataset, both how its loaded and tested for 'current' record reading. How this addresses that need: This begins with some basic tests around the loading, creating, and attaching of a static database file for TIMDEXMetadataDataset. Future tests will more fully exercise the final views and tables created. This commit also *temporarily* skips a bunch of tests for TIMDEXDataset that will not pass until the ability to limit to 'current' records is reinsated with the updated TIMDEXMetadataDataset. Side effects of this change: * Test suite passes, but multiple tests are temporarily skipped. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-530 --- tests/conftest.py | 21 ++++++++---- tests/test_dataset.py | 60 ++++++++++++++++++++-------------- tests/test_metadata.py | 48 ++++++++++++--------------- tests/test_s3client.py | 10 +++--- tests/test_write.py | 3 ++ timdex_dataset_api/metadata.py | 38 ++++++++++++++++----- 6 files changed, 107 insertions(+), 73 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 9fc7c6e..2c35bd1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -137,7 +137,7 @@ def dataset_with_runs_location(tmp_path) -> str: @pytest.fixture -def local_dataset_with_runs(dataset_with_runs_location) -> TIMDEXDataset: +def dataset_with_runs(dataset_with_runs_location) -> TIMDEXDataset: return TIMDEXDataset(dataset_with_runs_location) @@ -195,19 +195,26 @@ def dataset_with_same_day_runs(tmp_path) -> TIMDEXDataset: return timdex_dataset -@pytest.fixture -def timdex_dataset_metadata(dataset_with_same_day_runs): - return TIMDEXDatasetMetadata(timdex_dataset=dataset_with_same_day_runs) - - @pytest.fixture def timdex_bucket(): return "timdex" @pytest.fixture -def mock_s3_resource(timdex_bucket): +def mocked_timdex_bucket(timdex_bucket): with moto.mock_aws(): conn = boto3.resource("s3", region_name="us-east-1") conn.create_bucket(Bucket=timdex_bucket) yield conn + + +@pytest.fixture +def timdex_dataset_metadata_empty(dataset_with_runs_location): + return TIMDEXDatasetMetadata(dataset_with_runs_location) + + +@pytest.fixture +def timdex_dataset_metadata(dataset_with_runs_location): + tdm = TIMDEXDatasetMetadata(dataset_with_runs_location) + tdm.recreate_static_database_file() + return tdm diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 098f228..18061dd 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -137,6 +137,7 @@ def test_dataset_load_with_multi_nonpartition_filters_success(fixed_local_datase assert fixed_local_dataset.row_count == 1 +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") def test_dataset_load_current_records_all_sources_success(dataset_with_runs_location): timdex_dataset = TIMDEXDataset(dataset_with_runs_location) @@ -149,6 +150,7 @@ def test_dataset_load_current_records_all_sources_success(dataset_with_runs_loca assert len(timdex_dataset.dataset.files) == 12 +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") def test_dataset_load_current_records_one_source_success(dataset_with_runs_location): timdex_dataset = TIMDEXDataset(dataset_with_runs_location) timdex_dataset.load(current_records=True, source="alma") @@ -346,9 +348,9 @@ def test_dataset_local_dataset_row_count_missing_dataset_raise_error(local_datas _ = td.row_count -def test_dataset_all_records_not_current_and_not_deduped(local_dataset_with_runs): - local_dataset_with_runs.load() - all_records_df = local_dataset_with_runs.read_dataframe() +def test_dataset_all_records_not_current_and_not_deduped(dataset_with_runs): + dataset_with_runs.load() + all_records_df = dataset_with_runs.read_dataframe() # assert counts reflect all records from dataset, no deduping assert all_records_df.source.value_counts().to_dict() == {"alma": 254, "dspace": 194} @@ -358,9 +360,10 @@ def test_dataset_all_records_not_current_and_not_deduped(local_dataset_with_runs assert all_records_df.run_date.max() == date(2025, 2, 5) -def test_dataset_all_current_records_deduped(local_dataset_with_runs): - local_dataset_with_runs.load(current_records=True) - all_records_df = local_dataset_with_runs.read_dataframe() +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") +def test_dataset_all_current_records_deduped(dataset_with_runs): + dataset_with_runs.load(current_records=True) + all_records_df = dataset_with_runs.read_dataframe() # assert both sources have accurate record counts for current records only assert all_records_df.source.value_counts().to_dict() == {"dspace": 90, "alma": 100} @@ -373,9 +376,10 @@ def test_dataset_all_current_records_deduped(local_dataset_with_runs): assert all_records_df.run_date.max() == date(2025, 2, 5) # dspace -def test_dataset_source_current_records_deduped(local_dataset_with_runs): - local_dataset_with_runs.load(current_records=True, source="alma") - alma_records_df = local_dataset_with_runs.read_dataframe() +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") +def test_dataset_source_current_records_deduped(dataset_with_runs): + dataset_with_runs.load(current_records=True, source="alma") + alma_records_df = dataset_with_runs.read_dataframe() # assert only alma records present and correct count assert alma_records_df.source.value_counts().to_dict() == {"alma": 100} @@ -388,36 +392,40 @@ def test_dataset_source_current_records_deduped(local_dataset_with_runs): assert alma_records_df.run_date.max() == date(2025, 1, 5) +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") def test_dataset_all_read_methods_get_deduplication( - local_dataset_with_runs, + dataset_with_runs, ): - local_dataset_with_runs.load(current_records=True, source="alma") + dataset_with_runs.load(current_records=True, source="alma") - full_df = local_dataset_with_runs.read_dataframe() - all_records = list(local_dataset_with_runs.read_dicts_iter()) - transformed_records = list(local_dataset_with_runs.read_transformed_records_iter()) + full_df = dataset_with_runs.read_dataframe() + all_records = list(dataset_with_runs.read_dicts_iter()) + transformed_records = list(dataset_with_runs.read_transformed_records_iter()) assert len(full_df) == len(all_records) == len(transformed_records) +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") def test_dataset_current_records_no_additional_filtering_accurate_records_yielded( - local_dataset_with_runs, + dataset_with_runs, ): - local_dataset_with_runs.load(current_records=True, source="alma") - df = local_dataset_with_runs.read_dataframe() + dataset_with_runs.load(current_records=True, source="alma") + df = dataset_with_runs.read_dataframe() assert df.action.value_counts().to_dict() == {"index": 99, "delete": 1} +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") def test_dataset_current_records_action_filtering_accurate_records_yielded( - local_dataset_with_runs, + dataset_with_runs, ): - local_dataset_with_runs.load(current_records=True, source="alma") - df = local_dataset_with_runs.read_dataframe(action="index") + dataset_with_runs.load(current_records=True, source="alma") + df = dataset_with_runs.read_dataframe(action="index") assert df.action.value_counts().to_dict() == {"index": 99} +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") def test_dataset_current_records_index_filtering_accurate_records_yielded( - local_dataset_with_runs, + dataset_with_runs, ): """This is a somewhat complex test, but demonstrates that only 'current' records are yielded when .load(current_records=True) is applied. @@ -437,14 +445,14 @@ def test_dataset_current_records_index_filtering_accurate_records_yielded( "influenced" what records we would see as we continue backwards in time. """ # with current_records=False, we get all 25 records from run-5 - local_dataset_with_runs.load(current_records=False, source="alma") - df = local_dataset_with_runs.read_dataframe(run_id="run-5") + dataset_with_runs.load(current_records=False, source="alma") + df = dataset_with_runs.read_dataframe(run_id="run-5") assert len(df) == 25 # with current_records=True, we only get 15 records from run-5 # because newer run-6 influenced what records are current for older run-5 - local_dataset_with_runs.load(current_records=True, source="alma") - df = local_dataset_with_runs.read_dataframe(run_id="run-5") + dataset_with_runs.load(current_records=True, source="alma") + df = dataset_with_runs.read_dataframe(run_id="run-5") assert len(df) == 15 assert list(df.timdex_record_id) == [ "alma:10", @@ -465,6 +473,7 @@ def test_dataset_current_records_index_filtering_accurate_records_yielded( ] +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") def test_dataset_load_current_records_gets_correct_same_day_full_run( dataset_with_same_day_runs, ): @@ -477,6 +486,7 @@ def test_dataset_load_current_records_gets_correct_same_day_full_run( assert list(df.run_id.unique()) == ["run-2"] +@pytest.mark.skip(reason="All tests for 'current' records will be reworked.") def test_dataset_load_current_records_gets_correct_same_day_daily_runs_ordering( dataset_with_same_day_runs, ): diff --git a/tests/test_metadata.py b/tests/test_metadata.py index 8af6c2d..e5b5e75 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -1,42 +1,36 @@ -# ruff: noqa: PLR2004 +from duckdb import DuckDBPyConnection -import duckdb +from timdex_dataset_api import TIMDEXDatasetMetadata -from timdex_dataset_api import TIMDEXDataset, TIMDEXDatasetMetadata +def test_tdm_init_no_metadata_file_warning_success(caplog, dataset_with_runs_location): + tdm = TIMDEXDatasetMetadata(dataset_with_runs_location) -def test_tdm_init_from_timdex_dataset_instance_success(dataset_with_same_day_runs): - tdm = TIMDEXDatasetMetadata(timdex_dataset=dataset_with_same_day_runs) - assert isinstance(tdm.timdex_dataset, TIMDEXDataset) + assert tdm.conn is None + assert "Static metadata database not found" in caplog.text -def test_tdm_init_from_timdex_dataset_path_success(dataset_with_runs_location): - tdm = TIMDEXDatasetMetadata.from_dataset_location(dataset_with_runs_location) - assert isinstance(tdm.timdex_dataset, TIMDEXDataset) +def test_tdm_local_dataset_structure_properties(): + local_root = "/path/to/nothing" + tdm_local = TIMDEXDatasetMetadata(local_root) + assert tdm_local.location == local_root + assert tdm_local.location_scheme == "file" -def test_tdm_default_database_location_in_memory(timdex_dataset_metadata): - assert timdex_dataset_metadata.db_path == ":memory:" - result = timdex_dataset_metadata.conn.query("PRAGMA database_list;").fetchone() - assert result[1] == "memory" # name of database - assert result[2] is None # file associated with database, where None is memory +def test_tdm_s3_dataset_structure_properties(mocked_timdex_bucket): + 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_explicit_database_in_file(tmp_path, dataset_with_runs_location): - db_path = str(tmp_path / "tda.duckdb") - tdm = TIMDEXDatasetMetadata.from_dataset_location( - dataset_with_runs_location, - db_path=db_path, - ) - assert tdm.db_path == db_path - result = tdm.conn.query("PRAGMA database_list;").fetchone() - assert result[1] == "tda" # name of database - assert result[2] == db_path # filepath passed during init +def test_tdm_create_metadata_database_file_success(caplog, timdex_dataset_metadata_empty): + caplog.set_level("DEBUG") + timdex_dataset_metadata_empty.recreate_static_database_file() -def test_tdm_get_duckdb_connection(timdex_dataset_metadata): - conn = timdex_dataset_metadata.get_connection() - assert isinstance(conn, duckdb.DuckDBPyConnection) +def test_tdm_init_metadata_file_found_success(timdex_dataset_metadata): + assert isinstance(timdex_dataset_metadata.conn, DuckDBPyConnection) def test_tdm_connection_has_static_database_attached(timdex_dataset_metadata): diff --git a/tests/test_s3client.py b/tests/test_s3client.py index 31de7c1..0f8f045 100644 --- a/tests/test_s3client.py +++ b/tests/test_s3client.py @@ -42,7 +42,7 @@ def test_split_s3_uri_invalid(): client._split_s3_uri("timdex/path/to/file.txt") -def test_upload_download_file(mock_s3_resource, tmp_path): +def test_upload_download_file(mocked_timdex_bucket, tmp_path): """Test upload_file and download_file methods.""" client = S3Client() @@ -62,7 +62,7 @@ def test_upload_download_file(mock_s3_resource, tmp_path): assert download_path.read_text() == "test content" -def test_delete_file(mock_s3_resource, tmp_path): +def test_delete_file(mocked_timdex_bucket, tmp_path): """Test delete_file method.""" client = S3Client() @@ -76,12 +76,12 @@ def test_delete_file(mock_s3_resource, tmp_path): client.delete_file(s3_uri) # Verify the file is deleted - bucket = mock_s3_resource.Bucket("timdex") + bucket = mocked_timdex_bucket.Bucket("timdex") objects = list(bucket.objects.all()) assert len(objects) == 0 -def test_delete_folder(mock_s3_resource, tmp_path): +def test_delete_folder(mocked_timdex_bucket, tmp_path): """Test delete_folder method.""" client = S3Client() @@ -104,7 +104,7 @@ def test_delete_folder(mock_s3_resource, tmp_path): assert len(deleted_keys) == 3 assert all(key.startswith("folder/") for key in deleted_keys) - bucket = mock_s3_resource.Bucket("timdex") + bucket = mocked_timdex_bucket.Bucket("timdex") objects = list(bucket.objects.all()) assert len(objects) == 1 assert objects[0].key == "other.txt" diff --git a/tests/test_write.py b/tests/test_write.py index 5529be7..d5fc6b9 100644 --- a/tests/test_write.py +++ b/tests/test_write.py @@ -52,6 +52,9 @@ def test_dataset_write_record_batches_uses_batch_size( ) +@pytest.mark.skip( + reason="Test unneeded soon when list[str] not supported for dataset location." +) def test_dataset_write_to_multiple_locations_raise_error(sample_records_iter): timdex_dataset = TIMDEXDataset( location=["/path/to/records-1.parquet", "/path/to/records-2.parquet"] diff --git a/timdex_dataset_api/metadata.py b/timdex_dataset_api/metadata.py index dca957d..156a890 100644 --- a/timdex_dataset_api/metadata.py +++ b/timdex_dataset_api/metadata.py @@ -1,9 +1,11 @@ """timdex_dataset_api/metadata.py""" import os +import shutil import tempfile import time from pathlib import Path +from typing import Literal from urllib.parse import urlparse import duckdb @@ -41,6 +43,15 @@ def __init__( self.location = location self.conn: None | DuckDBPyConnection = self.setup_duckdb_context() + @property + def location_scheme(self) -> Literal["file", "s3"]: + scheme = urlparse(self.location).scheme + if scheme == "": + return "file" + if scheme == "s3": + return "s3" + raise ValueError(f"Location with scheme type '{scheme}' not supported.") + @property def metadata_root(self) -> str: return f"{self.location.removesuffix('/')}/metadata" @@ -59,7 +70,7 @@ def append_deltas_path(self) -> str: def database_exists(self) -> bool: """Check if static metadata database file exists.""" - if urlparse(self.metadata_database_path).scheme == "s3": + if self.location_scheme == "s3": s3_client = S3Client() return s3_client.object_exists(self.metadata_database_path) return os.path.exists(self.metadata_database_path) @@ -75,10 +86,11 @@ def recreate_static_database_file(self) -> None: 5. Upload DuckDB database file to target destination, making that the new static metadata database file """ - s3_client = S3Client() - - # remove any append deltas that may exist at this time of database recreation - s3_client.delete_folder(self.append_deltas_path) + if self.location_scheme == "s3": + s3_client = S3Client() + s3_client.delete_folder(self.append_deltas_path) + else: + shutil.rmtree(self.append_deltas_path, ignore_errors=True) # build database locally with tempfile.TemporaryDirectory() as temp_dir: @@ -91,10 +103,18 @@ def recreate_static_database_file(self) -> None: self._create_full_dataset_table(conn) # copy local database file to remote location - s3_client.upload_file( - local_db_path, - self.metadata_database_path, - ) + if self.location_scheme == "s3": + s3_client = S3Client() + s3_client.upload_file( + local_db_path, + self.metadata_database_path, + ) + else: + Path(self.metadata_database_path).parent.mkdir( + parents=True, + exist_ok=True, + ) + shutil.copy(local_db_path, self.metadata_database_path) # refresh DuckDB connection self.conn = self.setup_duckdb_context() From aaaadd03bc4841ea7c3f9fc804fa5b704457ed9c Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Wed, 6 Aug 2025 11:08:12 -0400 Subject: [PATCH 5/5] Set DuckDB secret refresh to auto --- timdex_dataset_api/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timdex_dataset_api/utils.py b/timdex_dataset_api/utils.py index 31473f6..36e0256 100644 --- a/timdex_dataset_api/utils.py +++ b/timdex_dataset_api/utils.py @@ -122,7 +122,7 @@ def configure_duckdb_s3_secret( type s3, provider credential_chain, chain 'sso;env;config', - refresh true + refresh auto {scope_str} ); """