Merge pull request #163 from MITLibraries/TIMX-528-merge-append-deltas

Author: jonavellecuerdo

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
 # ================================================================================

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
 ):

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()

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.
 

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)