DENG-877 Check backfills entries for tables with shredder_mitigation in metadata (#6588)

lucia-vargas-a · web-flow · commit bb6d9dadf260 · 2024-12-03T12:24:16.000Z
* Use attrs instead of attr for correct validators and attributes. Add constraint to stop backfills when table metadata has shredder_mitigation label and the backfill doesn't use mitigation.

* Use attrs instead of attr for correct validators and attributes. Add constraint to stop backfills when table metadata has shredder_mitigation label and the backfill doesn't use mitigation.

* Fix test.

* add validator and test for empty destination_table.

* add function docs
diff --git a/bigquery_etl/backfill/shredder_mitigation.py b/bigquery_etl/backfill/shredder_mitigation.py
@@ -11,7 +11,7 @@
 from types import NoneType
 from typing import Any, Optional, Tuple
 
-import attr
+import attrs
 import click
 from dateutil import parser
 from gcloud.exceptions import NotFound  # type: ignore
@@ -83,44 +83,46 @@ def from_schema_type(schema_type):
         return None
 
 
-@attr.define(eq=True)
+@attrs.define(eq=True)
 class Column:
     """Representation of a column in a query, with relevant details for shredder mitigation."""
 
     name: str
-    data_type: DataTypeGroup = attr.field(default=DataTypeGroup.UNDETERMINED)
-    column_type: ColumnType = attr.field(default=ColumnType.UNDETERMINED)
-    status: ColumnStatus = attr.field(default=ColumnStatus.UNDETERMINED)
-
-    @data_type.validator
-    def validate_data_type(self, _attribute, value):
-        """Check that the type of data_type is as expected."""
-        if not isinstance(value, DataTypeGroup):
-            raise ValueError(f"Invalid {value} with type: {type(value)}.")
-
-    @column_type.validator
-    def validate_column_type(self, _attribute, value):
-        """Check that the type of parameter column_type is as expected."""
-        if not isinstance(value, ColumnType):
-            raise ValueError(f"Invalid data type for: {value}.")
-
-    @status.validator
-    def validate_status(self, _attribute, value):
-        """Check that the type of parameter column_status is as expected."""
-        if not isinstance(value, ColumnStatus):
-            raise ValueError(f"Invalid data type for: {value}.")
-
-
-@attr.define(eq=True)
+    data_type: DataTypeGroup = attrs.field(
+        default=DataTypeGroup.UNDETERMINED,
+        validator=attrs.validators.instance_of(DataTypeGroup),
+    )
+    column_type: ColumnType = attrs.field(
+        default=ColumnType.UNDETERMINED,
+        validator=attrs.validators.instance_of(ColumnType),
+    )
+    status: ColumnStatus = attrs.field(
+        default=ColumnStatus.UNDETERMINED,
+        validator=attrs.validators.instance_of(ColumnStatus),
+    )
+
+
+@attrs.define(eq=True)
 class Subset:
     """Representation of a subset/CTEs in the query and the actions related to this subset."""
 
+    @staticmethod
+    def attr_not_null(instance, attribute, value):
+        """Raise an exception if the value is None or empty."""
+        if value is None or value == "":
+            raise click.ClickException(
+                f"{attribute.name} not given and it's required to continue."
+            )
+
     client: bigquery.Client
-    destination_table: str = attr.field(default="")
-    query_cte: str = attr.field(default="")
-    dataset: str = attr.field(default=TEMP_DATASET)
-    project_id: str = attr.field(default=DEFAULT_PROJECT_ID)
-    expiration_days: Optional[float] = attr.field(default=None)
+    destination_table: str = attrs.field(
+        default="",
+        validator=attr_not_null,
+    )
+    query_cte: str = attrs.field(default="")
+    dataset: str = attrs.field(default=TEMP_DATASET)
+    project_id: str = attrs.field(default=DEFAULT_PROJECT_ID)
+    expiration_days: Optional[float] = attrs.field(default=None)
 
     @property
     def version(self):
@@ -243,7 +245,7 @@ def get_query_path_results(
         partition_type = (
             "DATE" if self.partitioning["type"] == "DAY" else self.partitioning["type"]
         )
-        parameters = None
+        parameters = []
         if partition_field is not None:
             parameters = [
                 bigquery.ScalarQueryParameter(
@@ -546,9 +548,9 @@ def generate_query_with_shredder_mitigation(
         ) = classify_columns(
             new_table_row, previous_group_by, new_group_by, existing_schema, new_schema
         )
-    except TypeError as e:
+    except TypeError:
         raise click.ClickException(
-            f"Table {destination_table} did not return any rows for {backfill_date}.\n{e}"
+            f"Table {destination_table} did not return any rows for {backfill_date}."
         )
 
     if not common_dimensions or not added_dimensions or not metrics:
@@ -690,7 +692,8 @@ def generate_query_with_shredder_mitigation(
             if metric.data_type != DataTypeGroup.FLOAT
         ]
         + [
-            f"ROUND({previous_agg.query_cte}.{metric.name}, 10) - "  # Round FLOAT to avoid exponential numbers.
+            # Round FLOAT to avoid exponential numbers.
+            f"ROUND({previous_agg.query_cte}.{metric.name}, 10) - "
             f"ROUND(COALESCE({new_agg.query_cte}.{metric.name}, 0), 10) AS {metric.name}"
             for metric in metrics
             if metric.data_type == DataTypeGroup.FLOAT
diff --git a/bigquery_etl/cli/backfill.py b/bigquery_etl/cli/backfill.py
@@ -52,6 +52,7 @@
 from ..config import ConfigLoader
 from ..deploy import FailedDeployException, SkippedDeployException, deploy_table
 from ..metadata.parse_metadata import METADATA_FILE, Metadata
+from ..metadata.validate_metadata import SHREDDER_MITIGATION_LABEL
 from ..schema import SCHEMA_FILE, Schema
 
 logging.basicConfig(level=logging.INFO)
@@ -523,6 +524,25 @@ def _initiate_backfill(
     custom_query_path = None
     checks = None
     custom_checks_name = None
+
+    metadata = Metadata.from_file(
+        Path("sql") / project / dataset / table / METADATA_FILE
+    )
+
+    # Stop if the metadata contains shredder mitigation label and the backfill doesn't.
+    if (
+        SHREDDER_MITIGATION_LABEL in metadata.labels
+        and entry.shredder_mitigation is False
+    ):
+        click.echo(
+            click.style(
+                f"This backfill cannot continue.\nManaged backfills for tables with metadata label"
+                f" {SHREDDER_MITIGATION_LABEL} require using --shredder_mitigation.",
+                fg="yellow",
+            )
+        )
+        sys.exit(1)
+
     if entry.shredder_mitigation is True:
         click.echo(
             click.style(
diff --git a/tests/backfill/test_shredder_mitigation.py b/tests/backfill/test_shredder_mitigation.py
@@ -756,6 +756,24 @@ def test_version(self, mock_client):
                 == f"Invalid or missing table version in {test_subset.destination_table}."
             )
 
+    @patch("google.cloud.bigquery.Client")
+    def test_destination_table_not_given(self, mock_client):
+        """Test valid version in destination table."""
+        test_destination_table = ""
+
+        with pytest.raises(click.ClickException) as e:
+            Subset(
+                mock_client,
+                test_destination_table,
+                None,
+                self.dataset,
+                self.project_id,
+                None,
+            )
+        assert str(e.value.message) == (
+            "destination_table not given and it's required to continue."
+        )
+
     @patch("google.cloud.bigquery.Client")
     def test_partitioning(self, mock_client, runner):
         """Test that partitioning type and value associated to a subset are returned as expected."""
@@ -1153,10 +1171,10 @@ def test_generate_query_as_expected(
     @patch("google.cloud.bigquery.Client")
     @patch("bigquery_etl.backfill.shredder_mitigation.classify_columns")
     def test_generate_query_failed_for_missing_partitioning(
-        self, mock_classify_columns, mock_client, runner
+        self, mock_classify_columns, mock_client, runner, capfd
     ):
-        """Test that function raises exception for required query parameter missing
-        in metadata, instead of generating wrong query."""
+        """Test that function raises exception for required query parameter -partitioning-
+        missing in metadata, instead of generating wrong query."""
         existing_schema = {
             "fields": [
                 {"name": "column_1", "type": "DATE", "mode": "NULLABLE"},
@@ -1229,7 +1247,7 @@ def test_generate_query_failed_for_missing_partitioning(
                 [],
             )
 
-            with pytest.raises(TypeError) as e:
+            with pytest.raises((TypeError, IndexError)) as e:
                 generate_query_with_shredder_mitigation(
                     client=mock_client,
                     project_id=self.project_id,
@@ -1238,7 +1256,13 @@ def test_generate_query_failed_for_missing_partitioning(
                     staging_table_name=self.staging_table_name,
                     backfill_date=PREVIOUS_DATE,
                 )
-            assert str(e.value) == "'NoneType' object is not iterable"
+
+            assert isinstance(e.value, (TypeError, IndexError))
+            assert str(e.value) in {
+                "'NoneType' object is not iterable",
+                "Invalid type",
+                "list index out of range",
+            }
 
     @patch("google.cloud.bigquery.Client")
     @patch("bigquery_etl.backfill.shredder_mitigation.classify_columns")
diff --git a/tests/cli/test_cli_backfill.py b/tests/cli/test_cli_backfill.py
@@ -1029,7 +1029,6 @@ def test_validate_backfill_invalid_start_date_greater_than_end_date(self, runner
             assert result.exit_code == 1
             assert "Invalid start date" in str(result.exception)
 
-    #
     def test_validate_backfill_invalid_start_date_greater_than_entry_date(self, runner):
         with runner.isolated_filesystem():
             SQL_DIR = "sql/moz-fx-data-shared-prod/test/test_query_v1"
@@ -2657,3 +2656,136 @@ def test_initiate_partitioned_backfill_with_invalid_billing_project_from_entry_s
 
             assert result.exit_code == 1
             assert "Invalid billing project" in str(result.exception)
+
+    @patch("bigquery_etl.cli.backfill.deploy_table")
+    @patch("bigquery_etl.backfill.utils._should_initiate")
+    def test_dont_initiate_if_label_and_backfill_entry_dont_match(
+        self, mock_should_initiate, mock_deploy_table, runner
+    ):
+        """Test that process stops if the table has label shredder_mitigation and the backfill doesn't."""
+        path = Path("sql/moz-fx-data-shared-prod/test/test_query_v1")
+        mock_should_initiate.return_value = True
+        mock_deploy_table.return_value = None
+        expected_error_output = (
+            "This backfill cannot continue.\nManaged backfills for "
+            "tables with metadata label shredder_mitigation require "
+            "using --shredder_mitigation."
+        )
+
+        with runner.isolated_filesystem():
+            os.makedirs(path, exist_ok=True)
+
+            with open(
+                "sql/moz-fx-data-shared-prod/test/test_query_v1/query.sql", "w"
+            ) as f:
+                f.write("SELECT 1")
+
+            with open(path / "metadata.yaml", "w") as f:
+                f.write(
+                    "friendly_name: Test\ndescription: Test\nlabels:\n  incremental: true\n  shredder_mitigation: true"
+                )
+
+            with open(
+                "sql/moz-fx-data-shared-prod/test/dataset_metadata.yaml", "w"
+            ) as f:
+                f.write(yaml.dump(DATASET_METADATA_CONF))
+
+            backfill_entry_1 = Backfill(
+                date(2021, 5, 3),
+                date(2021, 1, 3),
+                date(2021, 5, 3),
+                [date(2021, 2, 3)],
+                VALID_REASON,
+                [VALID_WATCHER],
+                BackfillStatus.INITIATE,
+                shredder_mitigation=False,
+            )
+
+            backfill_file = (
+                Path("sql/moz-fx-data-shared-prod/test/test_query_v1") / BACKFILL_FILE
+            )
+            backfill_file.write_text(backfill_entry_1.to_yaml())
+            assert BACKFILL_FILE in os.listdir(
+                "sql/moz-fx-data-shared-prod/test/test_query_v1"
+            )
+            backfills = Backfill.entries_from_file(backfill_file)
+            assert backfills[0] == backfill_entry_1
+
+            result = runner.invoke(
+                initiate,
+                [
+                    "moz-fx-data-shared-prod.test.test_query_v1",
+                    "--parallelism=0",
+                ],
+            )
+
+            assert result.exit_code == 1
+            assert expected_error_output in result.output
+
+    @patch("bigquery_etl.cli.backfill.deploy_table")
+    @patch("bigquery_etl.backfill.utils._should_initiate")
+    def test_initiate_if_label_and_backfill_entry_match(
+        self, mock_should_initiate, mock_deploy_table, runner
+    ):
+        """Test that the process continues if both table & backfill use shredder_mitigation."""
+        path = Path("sql/moz-fx-data-shared-prod/test/test_query_v1")
+        mock_should_initiate.return_value = True
+        mock_deploy_table.return_value = None
+
+        with runner.isolated_filesystem():
+            os.makedirs(path, exist_ok=True)
+
+            with open(
+                "sql/moz-fx-data-shared-prod/test/test_query_v1/query.sql", "w"
+            ) as f:
+                f.write("SELECT 1")
+
+            with open(path / "metadata.yaml", "w") as f:
+                f.write(
+                    "friendly_name: Test\ndescription: Test\nlabels:\n  incremental: true\n  shredder_mitigation: true"
+                )
+
+            with open(
+                "sql/moz-fx-data-shared-prod/test/dataset_metadata.yaml", "w"
+            ) as f:
+                f.write(yaml.dump(DATASET_METADATA_CONF))
+
+            backfill_entry_1 = Backfill(
+                date(2021, 5, 3),
+                date(2021, 1, 3),
+                date(2021, 5, 3),
+                [date(2021, 2, 3)],
+                VALID_REASON,
+                [VALID_WATCHER],
+                BackfillStatus.INITIATE,
+                shredder_mitigation=True,
+            )
+
+            backfill_file = (
+                Path("sql/moz-fx-data-shared-prod/test/test_query_v1") / BACKFILL_FILE
+            )
+            backfill_file.write_text(backfill_entry_1.to_yaml())
+            assert BACKFILL_FILE in os.listdir(
+                "sql/moz-fx-data-shared-prod/test/test_query_v1"
+            )
+            backfills = Backfill.entries_from_file(backfill_file)
+            assert backfills[0] == backfill_entry_1
+
+            with patch(
+                "bigquery_etl.cli.backfill.query_backfill", return_value=None
+            ) as mock_backfill:
+                with patch(
+                    "bigquery_etl.cli.backfill.generate_query_with_shredder_mitigation",
+                    return_value=("a", "b"),
+                ) as mock_shredder_mitigation:
+                    result = runner.invoke(
+                        initiate,
+                        [
+                            "moz-fx-data-shared-prod.test.test_query_v1",
+                            "--parallelism=0",
+                        ],
+                    )
+
+                assert result.exit_code == 0
+                mock_shredder_mitigation.call_count == 2
+                mock_backfill.call_count == 2
