Handle corrupt backfill rows in tag counts

Author: reachsrinivaschennupati-tech

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

@@ -10,7 +10,6 @@
 
 import sqlalchemy as db
 import sqlalchemy.exc as db_exc
-from dagster_shared.serdes import deserialize_values
 from dagster_shared.seven import JSONDecodeError
 from sqlalchemy.engine import Connection
 
@@ -1094,9 +1093,7 @@ def get_backfills_count(self, filters: BulkActionsFilter | None = None) -> int:
             # also have the requested tags. This requires fetching the backfills from the db and filtering them
             query = self._backfills_query(filters=filters)
             rows = self.fetchall(query)
-            backfill_candidates = deserialize_values(
-                (row["body"] for row in rows), PartitionBackfill
-            )
+            backfill_candidates = self._deserialize_backfill_rows(rows)
             return len(
                 self._apply_backfill_tags_filter_to_results(backfill_candidates, filters.tags)
             )

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

@@ -24,7 +24,7 @@
 from dagster._core.storage.root import LocalArtifactStorage
 from dagster._core.storage.runs.base import RunStorage
 from dagster._core.storage.runs.migration import REQUIRED_DATA_MIGRATIONS
-from dagster._core.storage.runs.schema import BulkActionsTable
+from dagster._core.storage.runs.schema import BackfillTagsTable, BulkActionsTable
 from dagster._core.storage.runs.sql_run_storage import SqlRunStorage
 from dagster._core.storage.tags import (
     BACKFILL_ID_TAG,
@@ -1450,6 +1450,14 @@ def test_get_backfills_skips_corrupt_backfill_rows(
             backfill_timestamp=time.time(),
         )
         storage.add_backfill(valid_backfill)
+        storage.add_run(
+            create_dagster_run(
+                run_id=make_new_run_id(),
+                job_name="some_pipeline",
+                status=DagsterRunStatus.SUCCESS,
+                tags={"foo": "bar", BACKFILL_ID_TAG: valid_backfill.backfill_id},
+            )
+        )
 
         with storage.connect() as conn:
             conn.execute(
@@ -1510,6 +1518,66 @@ def test_get_backfills_limit_skips_corrupt_rows_without_shrinking_page(
 
         assert "Skipping backfill corrupt because it could not be deserialized." in caplog.text
 
+    def test_get_backfills_count_with_tags_skips_corrupt_rows(
+        self, storage: RunStorage, caplog: pytest.LogCaptureFixture
+    ):
+        if not isinstance(storage, SqlRunStorage):
+            pytest.skip("storage is not SQL-backed")
+        if not self.supports_backfill_tags_filtering_queries():
+            pytest.skip("storage does not support filtering backfills by tag")
+        if not self.supports_backfills_count():
+            pytest.skip("storage does not support backfill count")
+
+        origin = self.fake_partition_set_origin("fake_partition_set")
+        valid_backfill = PartitionBackfill(
+            "valid",
+            partition_set_origin=origin,
+            status=BulkActionStatus.REQUESTED,
+            partition_names=["a", "b", "c"],
+            from_failure=False,
+            tags={"foo": "bar"},
+            backfill_timestamp=time.time(),
+        )
+        storage.add_backfill(valid_backfill)
+        storage.add_run(
+            create_dagster_run(
+                run_id=make_new_run_id(),
+                job_name="some_pipeline",
+                status=DagsterRunStatus.SUCCESS,
+                tags={"foo": "bar", BACKFILL_ID_TAG: valid_backfill.backfill_id},
+            )
+        )
+
+        with storage.connect() as conn:
+            conn.execute(
+                BulkActionsTable.insert().values(
+                    key="corrupt",
+                    status=BulkActionStatus.REQUESTED.value,
+                    timestamp=datetime_from_timestamp(time.time()),
+                    body='{"truncated"',
+                )
+            )
+            conn.execute(
+                BackfillTagsTable.insert().values(
+                    backfill_id="corrupt",
+                    key="foo",
+                    value="bar",
+                )
+            )
+        storage.add_run(
+            create_dagster_run(
+                run_id=make_new_run_id(),
+                job_name="some_pipeline",
+                status=DagsterRunStatus.SUCCESS,
+                tags={"foo": "bar", BACKFILL_ID_TAG: "corrupt"},
+            )
+        )
+
+        with caplog.at_level("WARNING", logger="dagster"):
+            assert storage.get_backfills_count(BulkActionsFilter(tags={"foo": "bar"})) == 1
+
+        assert "Skipping backfill corrupt because it could not be deserialized." in caplog.text
+
     def test_backfill_status_filtering(self, storage: RunStorage):
         origin = self.fake_partition_set_origin("fake_partition_set")
         backfills = storage.get_backfills()