POC: updates from ruff formatter

Author: ghukill

migrations/001_2025_05_30_backfill_run_timestamp_column.py

@@ -83,7 +83,7 @@ def backfill_dataset(location: str, *, dry_run: bool = False) -> None:
         logger.info(json.dumps(result))
 
     logger.info(
-        f"Backfill complete. Elapsed: {time.perf_counter()-start_time}, "
+        f"Backfill complete. Elapsed: {time.perf_counter() - start_time}, "
         f"Success: {success_count}, Skipped: {skip_count}, Errors: {error_count}"
     )
 

migrations/002_2025_06_25_consistent_run_timestamp_per_etl_run.py

@@ -70,7 +70,6 @@ def fix_backfilled_run_timestamps(location: str, *, dry_run: bool = False) -> No
     error_count = 0
 
     for idx, row in parquet_to_run_timestamp_df.iterrows():
-
         if row.status == "OK":
             continue
 
@@ -94,7 +93,7 @@ def fix_backfilled_run_timestamps(location: str, *, dry_run: bool = False) -> No
         logger.info(json.dumps(result))
 
     logger.info(
-        f"Backfill complete. Elapsed: {time.perf_counter()-start_time}, "
+        f"Backfill complete. Elapsed: {time.perf_counter() - start_time}, "
         f"Success: {success_count}, Skipped: {skip_count}, Errors: {error_count}"
     )
 

tests/test_dataset.py

@@ -248,12 +248,15 @@ def test_dataset_duckdb_context_created_on_init(timdex_dataset):
 
 
 def test_dataset_duckdb_context_creates_data_schema(timdex_dataset):
-    assert timdex_dataset.conn.query("""
+    assert (
+        timdex_dataset.conn.query("""
             select count(*)
             from information_schema.schemata
             where catalog_name = 'memory'
             and schema_name = 'data';
-            """).fetchone()[0] == 1
+            """).fetchone()[0]
+        == 1
+    )
 
 
 def test_dataset_preload_current_records_default_false(timdex_dataset):

tests/test_metadata.py

@@ -55,12 +55,15 @@ def test_tdm_init_metadata_file_found_success(timdex_metadata):
 
 
 def test_tdm_duckdb_context_creates_metadata_schema(timdex_metadata):
-    assert timdex_metadata.conn.query("""
+    assert (
+        timdex_metadata.conn.query("""
             select count(*)
             from information_schema.schemata
             where catalog_name = 'memory'
             and schema_name = 'metadata';
-            """).fetchone()[0] == 1
+            """).fetchone()[0]
+        == 1
+    )
 
 
 def test_tdm_connection_has_static_database_attached(timdex_metadata):
@@ -285,13 +288,13 @@ def test_tdm_merge_append_deltas_adds_records_to_static_db(
 ):
     append_deltas = timdex_metadata_with_deltas.conn.query(f"""
             select
-            {','.join(ORDERED_METADATA_COLUMN_NAMES)}
+            {",".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)}
+            {",".join(ORDERED_METADATA_COLUMN_NAMES)}
             from static_db.records
         """).to_df()
 
@@ -317,12 +320,16 @@ def test_td_prepare_duckdb_secret_and_extensions_home_env_var_set_and_valid(
     monkeypatch.setenv("HOME", str(preset_home))
 
     td = TIMDEXDataset(timdex_dataset_with_runs.location)
-    df = td.conn.query("""
+    df = (
+        td.conn.query("""
         select
             current_setting('secret_directory') as secret_directory,
             current_setting('extension_directory') as extension_directory
         ;
-        """).to_df().iloc[0]
+        """)
+        .to_df()
+        .iloc[0]
+    )
     assert "my-account" in df.secret_directory
     assert df.extension_directory == ""  # expected and okay when HOME set
 
@@ -334,12 +341,16 @@ def test_td_prepare_duckdb_secret_and_extensions_home_env_var_unset(
 
     td = TIMDEXDataset(timdex_dataset_with_runs.location)
 
-    df = td.conn.query("""
+    df = (
+        td.conn.query("""
         select
             current_setting('secret_directory') as secret_directory,
             current_setting('extension_directory') as extension_directory
         ;
-        """).to_df().iloc[0]
+        """)
+        .to_df()
+        .iloc[0]
+    )
     assert df.secret_directory == "/tmp/.duckdb/secrets"
     assert df.extension_directory == "/tmp/.duckdb/extensions"
 
@@ -351,12 +362,16 @@ def test_td_prepare_duckdb_secret_and_extensions_home_env_var_set_but_empty(
 
     td = TIMDEXDataset(timdex_dataset_with_runs.location)
 
-    df = td.conn.query("""
+    df = (
+        td.conn.query("""
         select
             current_setting('secret_directory') as secret_directory,
             current_setting('extension_directory') as extension_directory
         ;
-        """).to_df().iloc[0]
+        """)
+        .to_df()
+        .iloc[0]
+    )
     assert df.secret_directory == "/tmp/.duckdb/secrets"
     assert df.extension_directory == "/tmp/.duckdb/extensions"
 

tests/utils.py

@@ -118,7 +118,7 @@ def generate_sample_embeddings_for_run(
     for _idx, record in records_metadata.iterrows():
         embedding_vector = [random.random() for _ in range(embedding_dimensions)]
         embedding_object = json.dumps(
-            {f"token{x+1}": (x + 1) / 10 for x in range(embedding_dimensions)}
+            {f"token{x + 1}": (x + 1) / 10 for x in range(embedding_dimensions)}
         ).encode()
 
         yield DatasetEmbedding(

timdex_dataset_api/dataset.py

@@ -111,7 +111,6 @@ class TIMDEXDatasetConfig:
 
 
 class TIMDEXDataset:
-
     def __init__(
         self,
         location: str,
@@ -211,7 +210,7 @@ def load_pyarrow_dataset(self, parquet_files: list[str] | None = None) -> ds.Dat
 
         logger.info(
             f"Dataset successfully loaded: '{self.data_records_root}', "
-            f"{round(time.perf_counter()-start_time, 2)}s"
+            f"{round(time.perf_counter() - start_time, 2)}s"
         )
 
         return dataset
@@ -496,12 +495,12 @@ def read_batches_iter(
 
             batch_rps = int(batch_yield_count / (time.perf_counter() - batch_time))
             logger.debug(
-                f"read_batches_iter batch {i+1}, yielded: {batch_yield_count} "
+                f"read_batches_iter batch {i + 1}, yielded: {batch_yield_count} "
                 f"@ {batch_rps} records/second, total yielded: {total_yield_count}"
             )
 
         logger.debug(
-            f"read_batches_iter() elapsed: {round(time.perf_counter()-start_time, 2)}s"
+            f"read_batches_iter() elapsed: {round(time.perf_counter() - start_time, 2)}s"
         )
 
     def _iter_meta_chunks(
@@ -525,7 +524,6 @@ def _iter_meta_chunks(
 
         total_yielded = 0
         while True:
-
             # enforce limit if passed
             if limit is not None:
                 remaining = limit - total_yielded
@@ -586,7 +584,8 @@ def _build_data_query_for_chunk(
         filenames = list(meta_chunk_df["filename"].unique())
         if self.location_scheme == "s3":
             filenames = [
-                f"s3://{f.removeprefix('s3://')}" for f in filenames  # type: ignore[union-attr]
+                f"s3://{f.removeprefix('s3://')}"
+                for f in filenames  # type: ignore[union-attr]
             ]
         parquet_list_sql = "[" + ",".join(f"'{f}'" for f in filenames) + "]"
 

timdex_dataset_api/embeddings.py

@@ -136,7 +136,6 @@ def to_dict(
 
 
 class TIMDEXEmbeddings:
-
     def __init__(self, timdex_dataset: "TIMDEXDataset"):
         """Init TIMDEXEmbeddings.
 
@@ -174,7 +173,7 @@ def _setup_embeddings_views(self) -> None:
 
         logger.debug(
             "Embeddings views setup for TIMDEXEmbeddings, "
-            f"{round(time.perf_counter()-start_time,2)}s"
+            f"{round(time.perf_counter() - start_time, 2)}s"
         )
 
     def _create_embeddings_view(self, conn: DuckDBPyConnection) -> None:
@@ -377,7 +376,7 @@ def read_batches_iter(
             rows_per_batch=self.timdex_dataset.config.read_batch_size
         )
 
-        logger.debug(f"read() elapsed: {round(time.perf_counter()-start_time, 2)}s")
+        logger.debug(f"read() elapsed: {round(time.perf_counter() - start_time, 2)}s")
 
     def _build_query(
         self,

timdex_dataset_api/metadata.py

@@ -37,7 +37,6 @@
 
 
 class TIMDEXDatasetMetadata:
-
     def __init__(self, timdex_dataset: "TIMDEXDataset") -> None:
         """Init TIMDEXDatasetMetadata.
 
@@ -174,7 +173,7 @@ def _create_full_dataset_table(self, conn: DuckDBPyConnection) -> None:
         query = f"""
             create or replace table records as (
                 select
-                    {','.join(ORDERED_METADATA_COLUMN_NAMES)}
+                    {",".join(ORDERED_METADATA_COLUMN_NAMES)}
                 from read_parquet(
                     '{self.location}/data/records/**/*.parquet',
                      hive_partitioning=true,
@@ -215,7 +214,7 @@ def _setup_metadata_schema(self) -> None:
 
         logger.debug(
             "Metadata schema setup for TIMDEXDatasetMetadata, "
-            f"{round(time.perf_counter()-start_time,2)}s"
+            f"{round(time.perf_counter() - start_time, 2)}s"
         )
 
     def _attach_database_file(self, conn: DuckDBPyConnection) -> None:
@@ -278,11 +277,11 @@ def _create_records_union_view(self, conn: DuckDBPyConnection) -> None:
             create or replace view metadata.records as
             (
                 select
-                    {','.join(ORDERED_METADATA_COLUMN_NAMES)}
+                    {",".join(ORDERED_METADATA_COLUMN_NAMES)}
                 from static_db.records
                 union all
                 select
-                    {','.join(ORDERED_METADATA_COLUMN_NAMES)}
+                    {",".join(ORDERED_METADATA_COLUMN_NAMES)}
                 from metadata.append_deltas
             );
             """)
@@ -380,10 +379,14 @@ def merge_append_deltas(self) -> None:
         s3_client = S3Client()
 
         # get filenames of append deltas
-        append_delta_filenames = self.conn.query("""
+        append_delta_filenames = (
+            self.conn.query("""
                 select distinct(append_delta_filename)
                 from metadata.append_deltas
-                """).to_df()["append_delta_filename"].to_list()
+                """)
+            .to_df()["append_delta_filename"]
+            .to_list()
+        )
 
         if len(append_delta_filenames) == 0:
             logger.info("no append deltas found")
@@ -408,7 +411,7 @@ def merge_append_deltas(self) -> None:
             self.conn.execute(f"""
                 insert into local_static_db.records
                 select
-                    {','.join(ORDERED_METADATA_COLUMN_NAMES)}
+                    {",".join(ORDERED_METADATA_COLUMN_NAMES)}
                 from metadata.append_deltas
                 """)
 
@@ -433,7 +436,7 @@ def merge_append_deltas(self) -> None:
 
         logger.debug(
             "append deltas merged into the static metadata database file: "
-            f"{self.metadata_database_path}, {time.perf_counter()-start_time}s"
+            f"{self.metadata_database_path}, {time.perf_counter() - start_time}s"
         )
 
     def write_append_delta_duckdb(self, filepath: str) -> None:
@@ -451,13 +454,13 @@ def write_append_delta_duckdb(self, filepath: str) -> None:
 
         # ensure s3:// schema prefix is present
         if self.location_scheme == "s3":
-            filepath = f"s3://{filepath.removeprefix("s3://")}"
+            filepath = f"s3://{filepath.removeprefix('s3://')}"
 
         # perform query + write as one SQL statement
         sql = f"""
         copy (
             select
-                {','.join(ORDERED_METADATA_COLUMN_NAMES)}
+                {",".join(ORDERED_METADATA_COLUMN_NAMES)}
             from read_parquet(
                 '{filepath}',
                 hive_partitioning=true,
@@ -469,7 +472,7 @@ def write_append_delta_duckdb(self, filepath: str) -> None:
         self.conn.execute(sql)
 
         logger.debug(
-            f"Append delta written: {output_path}, {time.perf_counter()-start_time}s"
+            f"Append delta written: {output_path}, {time.perf_counter() - start_time}s"
         )
 
     def build_keyset_paginated_metadata_query(

timdex_dataset_api/utils.py

@@ -138,7 +138,7 @@ def create_connection(self, path: str = ":memory:") -> DuckDBPyConnection:
         conn.execute("SET enable_progress_bar = false;")
         self.configure_connection(conn)
         logger.debug(
-            f"DuckDB connection created, {round(time.perf_counter()-start_time,2)}s"
+            f"DuckDB connection created, {round(time.perf_counter() - start_time, 2)}s"
         )
         return conn