enable saving run key for ticks that skip runs because of run key ide…

…mpotence (#7130)

* enable saving run key info for ticks that skip runs due to run key idempotence

* comments

python_modules/dagster/dagster/core/scheduler/instigation.py

@@ -298,8 +298,8 @@ def with_reason(self, skip_reason):
         check.opt_str_param(skip_reason, "skip_reason")
         return self._replace(tick_data=self.tick_data.with_reason(skip_reason))
 
-    def with_run(self, run_id, run_key=None):
-        return self._replace(tick_data=self.tick_data.with_run(run_id, run_key))
+    def with_run_info(self, run_id=None, run_key=None):
+        return self._replace(tick_data=self.tick_data.with_run_info(run_id, run_key))
 
     def with_cursor(self, cursor):
         return self._replace(tick_data=self.tick_data.with_cursor(cursor))
@@ -497,13 +497,14 @@ def with_status(self, status, error=None, timestamp=None, failure_count=None):
             )
         )
 
-    def with_run(self, run_id, run_key=None):
-        check.str_param(run_id, "run_id")
+    def with_run_info(self, run_id=None, run_key=None):
+        check.opt_str_param(run_id, "run_id")
+        check.opt_str_param(run_key, "run_key")
         return TickData(
             **merge_dicts(
                 self._asdict(),
                 {
-                    "run_ids": [*self.run_ids, run_id],
+                    "run_ids": [*self.run_ids, run_id] if run_id else self.run_ids,
                     "run_keys": [*self.run_keys, run_key] if run_key else self.run_keys,
                 },
             )

python_modules/dagster/dagster/daemon/sensor.py

@@ -89,8 +89,8 @@ def update_state(self, status, **kwargs):
         if origin_run_id:
             self._tick = self._tick.with_origin_run(origin_run_id)
 
-    def add_run(self, run_id, run_key=None):
-        self._tick = self._tick.with_run(run_id, run_key)
+    def add_run_info(self, run_id=None, run_key=None):
+        self._tick = self._tick.with_run_info(run_id, run_key)
 
     def set_should_update_cursor_on_failure(self, should_update_cursor_on_failure: bool):
         self._should_update_cursor_on_failure = should_update_cursor_on_failure
@@ -442,6 +442,7 @@ def _evaluate_sensor(
 
         if isinstance(run, SkippedSensorRun):
             skipped_runs.append(run)
+            context.add_run_info(run_id=None, run_key=run_request.run_key)
             yield
             continue
 
@@ -469,7 +470,7 @@ def _evaluate_sensor(
 
         _check_for_debug_crash(sensor_debug_crash_flags, "RUN_LAUNCHED")
 
-        context.add_run(run_id=run.run_id, run_key=run_request.run_key)
+        context.add_run_info(run_id=run.run_id, run_key=run_request.run_key)
 
     if skipped_runs:
         run_keys = [skipped.run_key for skipped in skipped_runs]
@@ -562,10 +563,8 @@ def _get_or_create_sensor_run(
 
     if run:
         if run.status != PipelineRunStatus.NOT_STARTED:
-            # A run already exists and was launched for this time period,
-            # but the daemon must have crashed before the tick could be put
-            # into a SUCCESS state
-            context.logger.info(f"Skipping run for {run_request.run_key}, found {run.run_id}.")
+            # A run already exists and was launched for this run key, but the daemon must have
+            # crashed before the tick could be updated
             return SkippedSensorRun(run_key=run_request.run_key, existing_run=run)
         else:
             context.logger.info(

python_modules/dagster/dagster/scheduler/scheduler.py

@@ -50,8 +50,8 @@ def update_state(self, status, error=None, **kwargs):
         if skip_reason:
             self._tick = self._tick.with_reason(skip_reason=skip_reason)
 
-    def add_run(self, run_id, run_key=None):
-        self._tick = self._tick.with_run(run_id, run_key)
+    def add_run_info(self, run_id=None, run_key=None):
+        self._tick = self._tick.with_run_info(run_id, run_key)
 
     def _write(self):
         self._instance.update_tick(self._tick)
@@ -459,7 +459,7 @@ def _schedule_runs_at_time(
                 logger.info(
                     f"Run {run.run_id} already completed for this execution of {external_schedule.name}"
                 )
-                tick_context.add_run(run_id=run.run_id, run_key=run_request.run_key)
+                tick_context.add_run_info(run_id=run.run_id, run_key=run_request.run_key)
                 yield
                 continue
             else:
@@ -490,7 +490,7 @@ def _schedule_runs_at_time(
                 yield error_info
 
         _check_for_debug_crash(debug_crash_flags, "RUN_LAUNCHED")
-        tick_context.add_run(run_id=run.run_id, run_key=run_request.run_key)
+        tick_context.add_run_info(run_id=run.run_id, run_key=run_request.run_key)
         _check_for_debug_crash(debug_crash_flags, "RUN_ADDED")
         yield
 

python_modules/dagster/dagster/utils/test/schedule_storage.py

@@ -250,7 +250,7 @@ def test_update_tick_to_success(self, storage):
         current_time = time.time()
         tick = storage.create_tick(self.build_schedule_tick(current_time))
 
-        updated_tick = tick.with_status(TickStatus.SUCCESS).with_run(run_id="1234")
+        updated_tick = tick.with_status(TickStatus.SUCCESS).with_run_info(run_id="1234")
         assert updated_tick.status == TickStatus.SUCCESS
 
         storage.update_tick(updated_tick)
@@ -507,7 +507,7 @@ def test_update_sensor_tick_to_success(self, storage):
         current_time = time.time()
         tick = storage.create_tick(self.build_sensor_tick(current_time))
 
-        updated_tick = tick.with_status(TickStatus.SUCCESS).with_run(run_id="1234")
+        updated_tick = tick.with_status(TickStatus.SUCCESS).with_run_info(run_id="1234")
         assert updated_tick.status == TickStatus.SUCCESS
 
         storage.update_tick(updated_tick)

python_modules/dagster/dagster_tests/daemon_tests/test_sensor_run.py

@@ -835,6 +835,10 @@ def test_launch_once(capfd):
                 freeze_datetime,
                 TickStatus.SKIPPED,
             )
+            assert ticks[0].run_keys
+            assert len(ticks[0].run_keys) == 1
+            assert not ticks[0].run_ids
+
             captured = capfd.readouterr()
             assert (
                 'Skipping 1 run for sensor run_key_sensor already completed with run keys: ["only_once"]'