diff --git a/tests/conftest.py b/tests/conftest.py index 42013fc..6f89fe1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,6 @@ """tests/conftest.py""" +import shutil from collections.abc import Iterator import boto3 @@ -265,6 +266,26 @@ def timdex_metadata_with_deltas( return TIMDEXDatasetMetadata(timdex_dataset_with_runs.location) +@pytest.fixture +def timdex_metadata_merged_deltas( + tmp_path, timdex_metadata_with_deltas, timdex_dataset_with_runs +): + """TIMDEXDatasetMetadata after merging append deltas to static database file.""" + # copy directory of a dataset with runs + dataset_location = str(tmp_path / "cloned_dataset_with_runs/") + shutil.copytree(timdex_metadata_with_deltas.location, dataset_location) + + # 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() + + return metadata + + # ================================================================================ # Utility Fixtures # ================================================================================ diff --git a/tests/test_metadata.py b/tests/test_metadata.py index f3bae4a..d63144c 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -8,6 +8,18 @@ from timdex_dataset_api import TIMDEXDatasetMetadata +ORDERED_METADATA_COLUMN_NAMES = [ + "timdex_record_id", + "source", + "run_date", + "run_type", + "action", + "run_id", + "run_record_offset", + "run_timestamp", + "filename", +] + def test_tdm_init_no_metadata_file_warning_success(caplog, timdex_dataset_with_runs): TIMDEXDatasetMetadata(timdex_dataset_with_runs.location) @@ -266,6 +278,49 @@ def test_tdm_current_records_most_recent_version(timdex_metadata_with_deltas): assert current_version.iloc[0]["run_id"] == most_recent.iloc[0]["run_id"] +def test_tdm_merge_append_deltas_static_counts_match_records_count_before_merge( + timdex_metadata_with_deltas, timdex_metadata_merged_deltas +): + static_count_merged_deltas = timdex_metadata_merged_deltas.conn.query( + """select count(*) as count from static_db.records;""" + ).fetchone()[0] + assert static_count_merged_deltas == timdex_metadata_with_deltas.records_count + + +def test_tdm_merge_append_deltas_adds_records_to_static_db( + timdex_metadata_with_deltas, timdex_metadata_merged_deltas +): + append_deltas = timdex_metadata_with_deltas.conn.query( + f""" + select + {','.join(ORDERED_METADATA_COLUMN_NAMES)} + from metadata.append_deltas + """ + ).to_df() + + merged_static_db = timdex_metadata_merged_deltas.conn.query( + f""" + select + {','.join(ORDERED_METADATA_COLUMN_NAMES)} + from static_db.records + """ + ).to_df() + + assert set(map(tuple, append_deltas.to_numpy())).issubset( + set(map(tuple, merged_static_db.to_numpy())) + ) + + +def test_tdm_merge_append_deltas_deletes_append_deltas( + timdex_metadata_with_deltas, timdex_metadata_merged_deltas +): + assert timdex_metadata_with_deltas.append_deltas_count != 0 + assert os.listdir(timdex_metadata_with_deltas.append_deltas_path) + + assert timdex_metadata_merged_deltas.append_deltas_count == 0 + 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( monkeypatch, tmp_path_factory, timdex_dataset_with_runs ): diff --git a/tests/test_s3client.py b/tests/test_s3client.py index bf440b4..4a83634 100644 --- a/tests/test_s3client.py +++ b/tests/test_s3client.py @@ -42,6 +42,22 @@ def test_split_s3_uri_invalid(): client._split_s3_uri("timdex/path/to/file.txt") +def test_list_objects(s3_bucket_mocked, tmp_path): + client = S3Client() + + # Create a test file + test_file = tmp_path / "test.txt" + test_file.write_text("test content") + + # Upload the file + s3_uri = "s3://timdex/metadata/append_deltas/test.txt" + client.upload_file(test_file, s3_uri) + + # Verify list of objects + s3_prefix = "s3://timdex/metadata/append_deltas" + assert client.list_objects(s3_prefix) == ["metadata/append_deltas/test.txt"] + + def test_upload_download_file(s3_bucket_mocked, tmp_path): """Test upload_file and download_file methods.""" client = S3Client() diff --git a/timdex_dataset_api/metadata.py b/timdex_dataset_api/metadata.py index e92a34b..227f3e0 100644 --- a/timdex_dataset_api/metadata.py +++ b/timdex_dataset_api/metadata.py @@ -396,7 +396,8 @@ def _create_append_deltas_view(self, conn: DuckDBPyConnection) -> None: create or replace view metadata.append_deltas as ( select * from read_parquet( - '{self.append_deltas_path}/*.parquet' + '{self.append_deltas_path}/*.parquet', + filename = 'append_delta_filename' ) ); """ @@ -414,14 +415,17 @@ def _create_append_deltas_view(self, conn: DuckDBPyConnection) -> None: def _create_records_union_view(self, conn: DuckDBPyConnection) -> None: logger.debug("creating view of unioned records") + conn.execute( - """ + f""" create or replace view metadata.records as ( - select * + select + {','.join(ORDERED_METADATA_COLUMN_NAMES)} from static_db.records union all - select * + select + {','.join(ORDERED_METADATA_COLUMN_NAMES)} from metadata.append_deltas ); """ @@ -463,6 +467,79 @@ def _create_current_records_view(self, conn: DuckDBPyConnection) -> None: """ conn.execute(query) + def merge_append_deltas(self) -> None: + """Merge append deltas into the static metadata database file.""" + logger.info("merging append deltas into static metadata database file") + + start_time = time.perf_counter() + + s3_client = S3Client() + + # get filenames of append deltas + append_delta_filenames = ( + self.conn.query( + """ + select distinct(append_delta_filename) + from metadata.append_deltas + """ + ) + .to_df()["append_delta_filename"] + .to_list() + ) + + if len(append_delta_filenames) == 0: + logger.info("no append deltas found") + return + + logger.debug(f"{len(append_delta_filenames)} append deltas found") + + with tempfile.TemporaryDirectory() as temp_dir: + # create local copy of the static metadata database (static db) file + local_db_path = str(Path(temp_dir) / self.metadata_database_filename) + if self.location_scheme == "s3": + s3_client.download_file( + s3_uri=self.metadata_database_path, local_path=local_db_path + ) + else: + shutil.copy(src=self.metadata_database_path, dst=local_db_path) + + # attach to local static db + self.conn.execute(f"""attach '{local_db_path}' AS local_static_db;""") + + # insert records from append deltas to local static db + self.conn.execute( + f""" + insert into local_static_db.records + select + {','.join(ORDERED_METADATA_COLUMN_NAMES)} + from metadata.append_deltas + """ + ) + + # detach from local static db + self.conn.execute("""detach local_static_db;""") + + # overwrite static db file with local version + if self.location_scheme == "s3": + s3_client.upload_file( + local_db_path, + self.metadata_database_path, + ) + else: + shutil.copy(src=local_db_path, dst=self.metadata_database_path) + + # delete append deltas + for append_delta_filename in append_delta_filenames: + if self.location_scheme == "s3": + s3_client.delete_file(s3_uri=append_delta_filename) + else: + os.remove(append_delta_filename) + + logger.debug( + "append deltas merged into the static metadata database file: " + f"{self.metadata_database_path}, {time.perf_counter()-start_time}s" + ) + def write_append_delta_duckdb(self, filepath: str) -> None: """Write an append delta for an ETL parquet file. diff --git a/timdex_dataset_api/utils.py b/timdex_dataset_api/utils.py index 538ea05..06f6ad8 100644 --- a/timdex_dataset_api/utils.py +++ b/timdex_dataset_api/utils.py @@ -59,6 +59,12 @@ def object_exists(self, s3_uri: str) -> bool: return False raise + def list_objects(self, s3_prefix: str) -> list[str]: + bucket, _ = self._split_s3_uri(s3_prefix) + objects = [obj.key for obj in self.resource.Bucket(bucket).objects.all()] + logger.debug(f"Found {len(objects)} objects in {s3_prefix}: {objects}") + return objects + 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)