DatasetRecord supports explicit run_timestamp

Why these changes are being introduced:

After it was discovered that a single ETL run (run_id) has multiple, different
run_timestamps in the dataset, it was clear we needed a way to pass an explicit
run_timestamp for writing instead of relying on a run_timestamp minted as part
of the TIMDEXDataset.write() method.

How this addresses that need:

DatasetRecord was given an optional run_timestamp property, that defaults
to using run_date and producing a full ISO timestamp from that.

Side effects of this change:
* Applications like Transmogrifier can pass an explicit run_timestamp
when writing to the dataset, ensuring that even multiple invocations
of Transmogrifier + TIMDEXDataset.write() can share the same timestamp
for an ETL run.

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

Author: ghukill

Pipfile

@@ -25,7 +25,6 @@ pytest = "*"
 ruff = "*"
 setuptools = "*"
 pip-audit = "*"
-pytest-freezegun = "*"
 
 [requires]
 python_version = "3.12"

Pipfile.lock

@@ -93,7 +93,6 @@ ignore = [
     "PLR0912",
     "PLR0913", 
     "PLR0915",
-    "S320",
     "S321",
     "S608",
     "TRY003"

pyproject.toml

@@ -182,16 +182,16 @@ def dataset_with_same_day_runs(tmp_path) -> TIMDEXDataset:
     # be represented.
     run_params.extend(
         [
-            (100, "alma", "2025-01-01", "full", "index", "run-1"),
-            (75, "alma", "2025-01-01", "full", "index", "run-2"),
-            (10, "alma", "2025-01-01", "daily", "index", "run-3"),
-            (20, "alma", "2025-01-02", "daily", "index", "run-4"),
-            (5, "alma", "2025-01-02", "daily", "delete", "run-5"),
+            (100, "alma", "2025-01-01", "full", "index", "run-1", "2025-01-01T01:00:00"),
+            (75, "alma", "2025-01-01", "full", "index", "run-2", "2025-01-01T02:00:00"),
+            (10, "alma", "2025-01-01", "daily", "index", "run-3", "2025-01-01T03:00:00"),
+            (20, "alma", "2025-01-02", "daily", "index", "run-4", "2025-01-01T04:00:00"),
+            (5, "alma", "2025-01-02", "daily", "delete", "run-5", "2025-01-01T05:00:00"),
         ]
     )
 
     for params in run_params:
-        num_records, source, run_date, run_type, action, run_id = params
+        num_records, source, run_date, run_type, action, run_id, run_timestamp = params
         records = generate_sample_records(
             num_records,
             timdex_record_id_prefix=source,
@@ -200,6 +200,7 @@ def dataset_with_same_day_runs(tmp_path) -> TIMDEXDataset:
             run_type=run_type,
             action=action,
             run_id=run_id,
+            run_timestamp=run_timestamp,
         )
         timdex_dataset.write(records)
 

tests/conftest.py

@@ -1,14 +1,13 @@
-# ruff: noqa: D205, D209, S105, S106, SLF001, PD901, PLR2004
+# ruff: noqa: D205, D209, SLF001, PLR2004
 
 import os
-from datetime import UTC, date, datetime
+from datetime import date
 from unittest.mock import MagicMock, patch
 
 import pyarrow as pa
 import pytest
 from pyarrow import fs
 
-from tests.utils import generate_sample_records
 from timdex_dataset_api.dataset import (
     DatasetNotLoadedError,
     TIMDEXDataset,
@@ -466,47 +465,6 @@ def test_dataset_current_records_index_filtering_accurate_records_yielded(
     ]
 
 
-@pytest.mark.freeze_time("2025-05-22 01:23:45.567890")
-def test_dataset_write_includes_minted_run_timestamp(tmp_path):
-    # create dataset
-    location = str(tmp_path / "one_run_at_frozen_time")
-    os.mkdir(location)
-    timdex_dataset = TIMDEXDataset(location)
-
-    run_id = "abc123"
-
-    # perform a single ETL run that should pickup the frozen time for run_timestamp
-    records = generate_sample_records(
-        10,
-        timdex_record_id_prefix="alma",
-        source="alma",
-        run_date="2025-05-22",
-        run_type="full",
-        action="index",
-        run_id=run_id,
-    )
-    timdex_dataset.write(records)
-    timdex_dataset.load()
-
-    # assert TIMDEXDataset.write() applies current time as run_timestamp
-    run_row_dict = next(timdex_dataset.read_dicts_iter())
-    assert "run_timestamp" in run_row_dict
-    assert run_row_dict["run_timestamp"] == datetime(
-        2025,
-        5,
-        22,
-        1,
-        23,
-        45,
-        567890,
-        tzinfo=UTC,
-    )
-
-    # assert the same run_timestamp is applied to all rows in the run
-    df = timdex_dataset.read_dataframe(run_id=run_id)
-    assert len(list(df.run_timestamp.unique())) == 1
-
-
 def test_dataset_load_current_records_gets_correct_same_day_full_run(
     dataset_with_same_day_runs,
 ):

tests/test_dataset.py

@@ -1,4 +1,4 @@
-# ruff: noqa: PLR2004, PD901
+# ruff: noqa: PLR2004
 
 import pandas as pd
 import pyarrow as pa

tests/test_read.py

@@ -1,5 +1,5 @@
 import re
-from datetime import date
+from datetime import UTC, date, datetime
 
 import pytest
 
@@ -71,7 +71,70 @@ def test_dataset_record_serialization():
         "action": "index",
         "run_id": "abc123",
         "run_record_offset": 0,
+        "run_timestamp": datetime(2024, 12, 1, 0, 0, tzinfo=UTC),
         "year": "2024",
         "month": "12",
         "day": "01",
     }
+
+
+@pytest.mark.parametrize(
+    ("run_timestamp_input", "expected_run_timestamp", "expected_exception"),
+    [
+        (
+            None,
+            None,
+            TypeError,  # expecting string, not None
+        ),
+        (
+            date(2025, 1, 1),
+            None,
+            TypeError,  # expecting string, not datetime object
+        ),
+        (
+            "2024-12-01T10:00:00Z",
+            datetime(2024, 12, 1, 10, 0, tzinfo=UTC),
+            None,
+        ),
+        (
+            "2024-12-01T23:59:59.999999+00:00",
+            datetime(2024, 12, 1, 23, 59, 59, 999999, tzinfo=UTC),
+            None,
+        ),
+    ],
+)
+def test_dataset_record_run_timestamp_parsing(
+    run_timestamp_input, expected_run_timestamp, expected_exception
+):
+    values = {
+        "timdex_record_id": "alma:123",
+        "source_record": b"<record><title>Hello World.</title></record>",
+        "transformed_record": b"""{"title":["Hello World."]}""",
+        "source": "libguides",
+        "run_date": "2024-12-01",
+        "run_type": "full",
+        "action": "index",
+        "run_id": "abc123",
+        "run_record_offset": 0,
+        "run_timestamp": run_timestamp_input,
+    }
+    if not expected_exception:
+        dataset_record = DatasetRecord(**values)
+        assert dataset_record.to_dict() == {
+            "timdex_record_id": "alma:123",
+            "source_record": b"<record><title>Hello World.</title></record>",
+            "transformed_record": b"""{"title":["Hello World."]}""",
+            "source": "libguides",
+            "run_date": date(2024, 12, 1),
+            "run_type": "full",
+            "action": "index",
+            "run_id": "abc123",
+            "run_record_offset": 0,
+            "run_timestamp": expected_run_timestamp,
+            "year": "2024",
+            "month": "12",
+            "day": "01",
+        }
+    else:
+        with pytest.raises(expected_exception):
+            DatasetRecord(**values)

tests/test_records.py

@@ -1,4 +1,4 @@
-# ruff: noqa: S105, S106, SLF001, PLR2004, PD901, D209, D205
+# ruff: noqa: PLR2004, D209, D205
 import math
 import os
 from unittest.mock import patch

tests/test_write.py

@@ -17,6 +17,7 @@ def generate_sample_records(
     run_type: str | None = "daily",
     action: str | None = "index",
     run_id: str | None = None,
+    run_timestamp: str | None = None,
 ) -> Iterator[DatasetRecord]:
     """Generate sample DatasetRecords."""
     if not run_id:
@@ -33,6 +34,7 @@ def generate_sample_records(
             action=action,
             run_id=run_id,
             run_record_offset=x,
+            run_timestamp=run_timestamp or run_date,
         )
 
 

tests/utils.py

@@ -423,17 +423,10 @@ def create_record_batches(
         Args:
             - records_iter: Iterator of DatasetRecord instances
         """
-        run_timestamp = datetime.now(UTC)
         for i, record_batch in enumerate(
             itertools.batched(records_iter, self.config.write_batch_size)
         ):
-            record_dicts = [
-                {
-                    **record.to_dict(),
-                    "run_timestamp": run_timestamp,
-                }
-                for record in record_batch
-            ]
+            record_dicts = [record.to_dict() for record in record_batch]
             batch = pa.RecordBatch.from_pylist(record_dicts)
             logger.debug(f"Yielding batch {i + 1} for dataset writing.")
             yield batch