fix: support labels for schema validation (#260)

data_validation/__main__.py

@@ -187,7 +187,7 @@ def build_config_managers_from_args(args):
             filter_config = cli_tools.get_filters(args.filters)
         if args.threshold:
             threshold = args.threshold
-        labels = cli_tools.get_labels(args.labels)
+    labels = cli_tools.get_labels(args.labels)
 
     mgr = state_manager.StateManager()
     source_client = clients.get_data_client(mgr.get_connection_config(args.source_conn))

data_validation/cli_tools.py

@@ -409,9 +409,6 @@ def _configure_row_parser(row_parser):
         "-pk",
         help="Comma separated list of primary key columns 'col_a,col_b'",
     )
-    row_parser.add_argument(
-        "--labels", "-l", help="Key value pair labels for validation run"
-    )
     row_parser.add_argument(
         "--threshold",
         "-th",
@@ -499,9 +496,6 @@ def _configure_column_parser(column_parser):
         "-pk",
         help="Comma separated list of primary key columns 'col_a,col_b'",
     )
-    column_parser.add_argument(
-        "--labels", "-l", help="Key value pair labels for validation run"
-    )
     column_parser.add_argument(
         "--threshold",
         "-th",
@@ -542,6 +536,9 @@ def _add_common_arguments(parser):
     parser.add_argument(
         "--bq-result-handler", "-bqrh", help="BigQuery result handler config details"
     )
+    parser.add_argument(
+        "--labels", "-l", help="Key value pair labels for validation run"
+    )
     parser.add_argument(
         "--service-account",
         "-sa",

data_validation/schema_validation.py

@@ -66,20 +66,25 @@ def execute(self):
         df.insert(loc=1, column="validation_name", value="Schema")
         df.insert(loc=2, column="validation_type", value="Schema")
 
-        df.insert(loc=3, column="start_time", value=self.run_metadata.start_time)
-        df.insert(loc=4, column="end_time", value=self.run_metadata.end_time)
+        df.insert(
+            loc=3,
+            column="labels",
+            value=[self.run_metadata.labels for _ in range(len(df.index))],
+        )
+        df.insert(loc=4, column="start_time", value=self.run_metadata.start_time)
+        df.insert(loc=5, column="end_time", value=self.run_metadata.end_time)
 
         df.insert(
-            loc=5,
+            loc=6,
             column="source_table_name",
             value=self.config_manager.full_source_table,
         )
         df.insert(
-            loc=6,
+            loc=7,
             column="target_table_name",
             value=self.config_manager.full_target_table,
         )
-        df.insert(loc=9, column="aggregation_type", value="Schema")
+        df.insert(loc=10, column="aggregation_type", value="Schema")
 
         del df["error_result.details"]
         return df

tests/unit/test_schema_validation.py

@@ -51,6 +51,10 @@
     consts.CONFIG_AGGREGATES: [],
     consts.CONFIG_THRESHOLD: 0.0,
     consts.CONFIG_RESULT_HANDLER: None,
+    consts.CONFIG_LABELS: [
+        ("label_1_name", "label_1_value"),
+        ("label_2_name", "label_2_value"),
+    ],
     consts.CONFIG_FORMAT: "table",
 }
 
@@ -193,5 +197,6 @@ def test_execute(module_under_test, fs):
     assert len(result_df) == len(source_data[0]) + 1
     assert result_df["source_agg_value"].astype(float).sum() == 7
     assert result_df["target_agg_value"].astype(float).sum() == 7
+    assert result_df.labels[0] == SAMPLE_SCHEMA_CONFIG[consts.CONFIG_LABELS]
     assert failures["source_column_name"].to_list() == ["id", "N/A"]
     assert failures["target_column_name"].to_list() == ["N/A", "id_new"]