Add history import run method

Author: dpeng817

python_modules/dagster/dagster/_core/storage/legacy_storage.py

@@ -1,4 +1,5 @@
 from collections.abc import Iterable, Mapping, Sequence
+from datetime import datetime
 from typing import TYPE_CHECKING, AbstractSet, Optional, Union  # noqa: UP035
 
 from dagster import _check as check
@@ -198,6 +199,11 @@ def from_config_value(
     def add_run(self, dagster_run: "DagsterRun") -> "DagsterRun":
         return self._storage.run_storage.add_run(dagster_run)
 
+    def add_historical_run(
+        self, dagster_run: "DagsterRun", run_creation_time: datetime
+    ) -> "DagsterRun":
+        return self._storage.run_storage.add_historical_run(dagster_run, run_creation_time)
+
     def handle_run_event(self, run_id: str, event: "DagsterEvent") -> None:
         return self._storage.run_storage.handle_run_event(run_id, event)
 

python_modules/dagster/dagster/_core/storage/runs/base.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from collections.abc import Mapping, Sequence
+from datetime import datetime
 from typing import TYPE_CHECKING, Optional, Union
 
 from typing_extensions import TypedDict
@@ -54,6 +55,12 @@ def add_run(self, dagster_run: DagsterRun) -> DagsterRun:
             dagster_run (DagsterRun): The run to add.
         """
 
+    @abstractmethod
+    def add_historical_run(
+        self, dagster_run: DagsterRun, run_creation_time: datetime
+    ) -> DagsterRun:
+        """Add a historical run to storage."""
+
     @abstractmethod
     def handle_run_event(self, run_id: str, event: DagsterEvent) -> None:
         """Update run storage in accordance to a pipeline run related DagsterEvent.

python_modules/dagster/dagster/_core/storage/runs/sql_run_storage.py

@@ -110,7 +110,9 @@ def fetchone(self, query: SqlAlchemyQuery) -> Optional[Any]:
             else:
                 return conn.execute(query).fetchone()
 
-    def add_run(self, dagster_run: DagsterRun) -> DagsterRun:
+    def _get_run_insertion_values(
+        self, dagster_run: DagsterRun, run_creation_time: Optional[datetime] = None
+    ) -> dict[str, Any]:
         check.inst_param(dagster_run, "dagster_run", DagsterRun)
 
         if dagster_run.job_snapshot_id and not self.has_job_snapshot(dagster_run.job_snapshot_id):
@@ -132,24 +134,39 @@ def add_run(self, dagster_run: DagsterRun) -> DagsterRun:
         }
         if self.has_backfill_id_column():
             values["backfill_id"] = dagster_run.tags.get(BACKFILL_ID_TAG)
+        if run_creation_time:
+            values["create_timestamp"] = run_creation_time
+        return values
 
-        runs_insert = RunsTable.insert().values(**values)
+    def _core_add_run(self, col_values: dict[str, Any], tags: Mapping[str, str]) -> None:
+        run_id = col_values["run_id"]
+        runs_insert = RunsTable.insert().values(**col_values)
         with self.connect() as conn:
             try:
                 conn.execute(runs_insert)
             except db_exc.IntegrityError as exc:
                 raise DagsterRunAlreadyExists from exc
 
-            tags_to_insert = dagster_run.tags_for_storage()
-            if tags_to_insert:
+            if tags:
                 conn.execute(
                     RunTagsTable.insert(),
-                    [
-                        dict(run_id=dagster_run.run_id, key=k, value=v)
-                        for k, v in tags_to_insert.items()
-                    ],
+                    [dict(run_id=run_id, key=k, value=v) for k, v in tags.items()],
                 )
 
+    def add_run(self, dagster_run: DagsterRun) -> DagsterRun:
+        self._core_add_run(
+            self._get_run_insertion_values(dagster_run, get_current_datetime()),
+            dagster_run.tags_for_storage(),
+        )
+        return dagster_run
+
+    def add_historical_run(
+        self, dagster_run: DagsterRun, run_creation_time: datetime
+    ) -> DagsterRun:
+        self._core_add_run(
+            self._get_run_insertion_values(dagster_run, run_creation_time),
+            dagster_run.tags_for_storage(),
+        )
         return dagster_run
 
     def handle_run_event(self, run_id: str, event: DagsterEvent) -> None:

python_modules/dagster/dagster_tests/storage_tests/utils/run_storage.py

@@ -3,7 +3,7 @@
 import time
 import unittest
 from contextlib import contextmanager
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from typing import Optional
 from uuid import uuid4
 
@@ -405,6 +405,18 @@ def test_get_run_tags(self, storage: RunStorage):
         # empty tag_keys implies nothing instead of everything
         assert storage.get_run_tags(tag_keys=[]) == []
 
+    def test_add_historical_run(self, storage):
+        run = DagsterRun(run_id=make_new_run_id(), job_name="some_job", tags={"foo": "bar"})
+        storage.add_historical_run(run, datetime(2025, 1, 1, tzinfo=timezone.utc))
+        record = storage.get_run_records()[0]
+        assert record is not None
+        assert record.dagster_run.tags == run.tags
+        assert record.create_timestamp == datetime(2025, 1, 1, tzinfo=timezone.utc)
+        run_tags = storage.get_run_tags(tag_keys=["foo"])
+        assert len(run_tags) == 1
+        assert run_tags[0][0] == "foo"
+        assert run_tags[0][1] == {"bar"}
+
     def test_fetch_by_filter(self, storage):
         assert storage
         one = make_new_run_id()