Include star for ref columns (#193)

* initial commit

* added test for supporting "*" for reference columns

* modified custom sql statement default behavior

* removed reference columns from custom sql statement rule

* fixed unit tests and linting.

* fixed unit tests and linting.

* fixed unit tests and linting.

* incorporated PR comments

* incorporated PR comments

* incorporated PR comments and added test for incorrectly specified reference-columns

* fixed test pip install

Author: AmandeepSinghCS

clouddq/classes/dq_rule_binding.py

@@ -27,6 +27,7 @@
 from clouddq.classes.dq_reference_columns import DqReferenceColumns
 from clouddq.classes.dq_row_filter import DqRowFilter
 from clouddq.classes.dq_rule import DqRule
+from clouddq.integration.bigquery.bigquery_client import BigQueryClient
 from clouddq.utils import assert_not_none_or_empty
 from clouddq.utils import get_from_dict_and_assert
 from clouddq.utils import get_keys_from_dict_and_assert_oneof
@@ -277,6 +278,7 @@ def resolve_reference_columns_config(
     def resolve_all_configs_to_dict(
         self: DqRuleBinding,
         configs_cache: dq_configs_cache.DqConfigsCache,
+        bigquery_client: BigQueryClient,
     ) -> dict:
         """
 
@@ -334,11 +336,27 @@ def resolve_all_configs_to_dict(
             row_filter_config = self.resolve_row_filter_config(configs_cache)
             # resolve reference columns config
             if self.reference_columns_id:
+                include_all_reference_columns = False
                 include_reference_columns = self.resolve_reference_columns_config(
                     configs_cache
                 ).include_reference_columns
+                if "*" in include_reference_columns:
+                    if len(include_reference_columns) == 1:
+                        column_names = bigquery_client.get_table_columns(
+                            table=table_entity.get_table_name(),
+                            project_id=table_entity.instance_name,
+                        )
+                        include_reference_columns = sorted(column_names)
+                        include_all_reference_columns = True
+                    elif len(include_reference_columns) > 1:
+                        raise ValueError(
+                            f"The specified reference columns id {self.reference_columns_id} \n"
+                            f"must contain only '*' or column names but not both.\n "
+                            f"Current reference columns configs contains \n {include_reference_columns} columns"
+                        )
             else:
                 include_reference_columns = []
+                include_all_reference_columns = False
 
             return dict(
                 {
@@ -348,6 +366,7 @@ def resolve_all_configs_to_dict(
                     "column_id": self.column_id,
                     "reference_columns_id": self.reference_columns_id,
                     "include_reference_columns": include_reference_columns,
+                    "include_all_reference_columns": include_all_reference_columns,
                     "column_configs": dict(column_configs.dict_values()),
                     "rule_ids": list(self.rule_ids),
                     "rule_configs_dict": rule_configs_dict,

clouddq/integration/bigquery/bigquery_client.py

@@ -269,3 +269,16 @@ def get_table_schema(self, table: str, project_id: str = None) -> dict:
 
         logger.debug(f"Schema for table {table} is: {columns_dict}")
         return columns_dict
+
+    def get_table_columns(self, table: str, project_id: str = None) -> set:
+
+        client = self.get_connection(project_id=project_id)
+        try:
+            table_ref = client.get_table(table)
+        except KeyError as error:
+            raise KeyError(f"\n\nInput table `{table}` is not valid.\n{error}")
+
+        column_names = {column.name for column in table_ref.schema}
+        logger.debug(f"List of columns in table {table} is: {column_names}")
+
+        return column_names

clouddq/lib.py

@@ -207,6 +207,7 @@ def prepare_configs_from_rule_binding_id(
     )
     resolved_rule_binding_configs = rule_binding.resolve_all_configs_to_dict(
         configs_cache=configs_cache,
+        bigquery_client=bigquery_client,
     )
     configs: dict[typing.Any, typing.Any] = {
         "configs": dict(resolved_rule_binding_configs)
@@ -285,7 +286,7 @@ def get_high_watermark_value(
             AND rule_binding_id = '{rule_binding_id}'
             AND progress_watermark IS TRUE ;"""
     else:
-        query = f"""SELECT 
+        query = """SELECT
             TIMESTAMP("1970-01-01 00:00:00") as high_watermark,
             CURRENT_TIMESTAMP() as current_timestamp_value ;"""
     logger.info(f"High watermark query is \n {query}")

configs/reference_columns/reference-columns.yml

@@ -19,3 +19,7 @@ reference_columns:
       - row_id
       - contact_type
       - value
+
+  INCLUDE_ALL_REFERENCE_COLUMNS:
+    include_reference_columns:
+      - *

tests/BUILD

@@ -238,3 +238,11 @@ py_test(
     legacy_create_init = 0,
     deps = DEPS,
 )
+
+py_test(
+    name = "test_reference_columns",
+    srcs = SRCS,
+    data = DATA,
+    legacy_create_init = 0,
+    deps = DEPS,
+)

tests/conftest.py

@@ -273,6 +273,10 @@ def source_dq_rules_configs_file_path():
 def source_dq_advanced_rules_configs_path():
     return Path("docs").joinpath("examples").joinpath("advanced_rules").absolute()
 
+@pytest.fixture(scope="session")
+def source_dq_reference_configs_file_path():
+    return Path("tests").joinpath("resources").joinpath("dq_reference_configs.yml").absolute()
+
 @pytest.fixture(scope="session")
 def test_profiles_dir():
     return Path("tests").joinpath("resources", "test_dbt_profiles_dir").absolute()
@@ -396,6 +400,7 @@ def temp_configs_dir(
         lines = lines.replace("<my-gcp-dataplex-zone-id>", gcp_dataplex_zone_id)
         lines = lines.replace("<my_bigquery_dataset_id>", gcp_dataplex_bigquery_dataset_id)
         source_file.write(lines)
+
     yield configs_path.absolute()
     if os.path.exists(temp_clouddq_dir):
         shutil.rmtree(temp_clouddq_dir)
@@ -476,6 +481,33 @@ def temp_configs_from_dq_rules_config_file(
     if os.path.exists(temp_clouddq_dir):
         os.unlink(temp_clouddq_dir)
 
+@pytest.fixture(scope="function")
+def temp_configs_from_dq_reference_configs_file(
+        gcp_project_id,
+        gcp_dataplex_bigquery_dataset_id,
+        gcp_dataplex_region,
+        gcp_dataplex_lake_name,
+        gcp_dataplex_zone_id,
+        source_dq_reference_configs_file_path,
+        tmp_path):
+    # Create temp directory
+    temp_clouddq_dir = Path(tmp_path).joinpath("clouddq_test_dq_rules_configs")
+    # Copy over tests/resources/dq_reference_configs.yaml
+    registry_defaults = shutil.copyfile(source_dq_reference_configs_file_path, temp_clouddq_dir)
+    # Prepare entity_uri configs
+    with open(registry_defaults) as source_file:
+        lines = source_file.read()
+    with open(registry_defaults, "w") as source_file:
+        lines = lines.replace("<my-gcp-dataplex-lake-id>", gcp_dataplex_lake_name)
+        lines = lines.replace("<my-gcp-dataplex-region-id>", gcp_dataplex_region)
+        lines = lines.replace("<my-gcp-project-id>", gcp_project_id)
+        lines = lines.replace("<my-gcp-dataplex-zone-id>", gcp_dataplex_zone_id)
+        lines = lines.replace("<my_bigquery_dataset_id>", gcp_dataplex_bigquery_dataset_id)
+        source_file.write(lines)
+    yield temp_clouddq_dir.absolute()
+    if os.path.exists(temp_clouddq_dir):
+        os.unlink(temp_clouddq_dir)
+
 @pytest.fixture(scope="function")
 def temp_configs_from_dq_advanced_rules_configs(
         gcp_project_id,

tests/integration/test_bigquery_client.py

@@ -194,6 +194,31 @@ def test_get_table_schema_with_default_partitioning(self,
         else:
             print(f'`{table_name}` does not exists in the project `{project_id}`')
 
+    def test_get_table_columns(self,
+        gcp_project_id,
+        test_bigquery_client):
+        table_name = "bigquery-public-data.github_repos.commits"
+        table_ref = test_bigquery_client.table_from_string(
+            table_name
+        )
+        project_id = table_ref.project
+        is_table_exists = test_bigquery_client.is_table_exists(
+            table=table_name,
+            project_id=project_id
+        )
+        if is_table_exists:
+            table_columns = test_bigquery_client.get_table_columns(
+                table=table_name,
+                project_id=project_id
+            )
+            expected_table_columns = {
+                    'commit', 'tree', 'parent', 'author', 'committer', 'subject', 'message',
+                    'trailer', 'difference', 'difference_truncated', 'repo_name', 'encoding'
+            }
+            assert table_columns == expected_table_columns
+        else:
+            print(f'`{table_name}` does not exists in the project `{project_id}`')
+
 
 if __name__ == "__main__":
     raise SystemExit(pytest.main([__file__, '-vv', '-rP', '-n', 'auto']))

tests/integration/test_dataplex_metadata_uri_templates.py

@@ -85,6 +85,7 @@ def test_rule_bindings_class_resolve_configs(
         test_rule_bindings_collection_team_4,
         test_default_dataplex_configs_cache,
         test_dataplex_metadata_defaults_configs,
+        test_bigquery_client,
     ):
         for key, value in test_rule_bindings_collection_team_4.items():
             rule_binding = DqRuleBinding.from_dict(
@@ -94,7 +95,8 @@ def test_rule_bindings_class_resolve_configs(
             rule_binding.resolve_table_entity_config(configs_cache=test_default_dataplex_configs_cache)
             rule_binding.resolve_rule_config_list(configs_cache=test_default_dataplex_configs_cache)
             rule_binding.resolve_row_filter_config(configs_cache=test_default_dataplex_configs_cache)
-            rule_binding.resolve_all_configs_to_dict(configs_cache=test_default_dataplex_configs_cache)
+            rule_binding.resolve_all_configs_to_dict(configs_cache=test_default_dataplex_configs_cache,
+                                                     bigquery_client=test_bigquery_client)
 
     def test_prepare_configs_from_rule_binding(
         self,
@@ -199,6 +201,7 @@ def test_rule_bindings_class_resolve_configs_from_file(
             test_default_dataplex_configs_cache_from_file,
             test_dataplex_metadata_defaults_configs,
             test_rule_bindings_collection_from_configs_file,
+            test_bigquery_client,
     ):
         for rule_binding_id, rule_binding_configs in test_rule_bindings_collection_from_configs_file.items():
             rule_binding = DqRuleBinding.from_dict(
@@ -208,7 +211,8 @@ def test_rule_bindings_class_resolve_configs_from_file(
             rule_binding.resolve_table_entity_config(configs_cache=test_default_dataplex_configs_cache_from_file)
             rule_binding.resolve_rule_config_list(configs_cache=test_default_dataplex_configs_cache_from_file)
             rule_binding.resolve_row_filter_config(configs_cache=test_default_dataplex_configs_cache_from_file)
-            rule_binding.resolve_all_configs_to_dict(configs_cache=test_default_dataplex_configs_cache_from_file)
+            rule_binding.resolve_all_configs_to_dict(configs_cache=test_default_dataplex_configs_cache_from_file,
+                                                     bigquery_client=test_bigquery_client)
 
     def test_prepare_configs_from_configs_file(
             self,
@@ -313,6 +317,7 @@ def test_rule_bindings_class_resolve_gcs_configs(
         test_rule_bindings_collection_team_5,
         test_default_dataplex_configs_cache,
         test_dataplex_metadata_defaults_configs,
+        test_bigquery_client,
     ):
         for key, value in test_rule_bindings_collection_team_5.items():
             rule_binding = DqRuleBinding.from_dict(
@@ -322,7 +327,8 @@ def test_rule_bindings_class_resolve_gcs_configs(
             rule_binding.resolve_table_entity_config(configs_cache=test_default_dataplex_configs_cache)
             rule_binding.resolve_rule_config_list(configs_cache=test_default_dataplex_configs_cache)
             rule_binding.resolve_row_filter_config(configs_cache=test_default_dataplex_configs_cache)
-            rule_binding.resolve_all_configs_to_dict(configs_cache=test_default_dataplex_configs_cache)
+            rule_binding.resolve_all_configs_to_dict(configs_cache=test_default_dataplex_configs_cache,
+                                                     bigquery_client=test_bigquery_client)
 
     def test_prepare_configs_from_gcs_rule_binding(
         self,