test: add column ACLs test with real policy tag (googleapis#678)

plamut · web-flow · commit b2a689b05b4f · 2021-06-17T22:59:26.000+02:00
* test: add column ACLs test with real policy tag

* Use v1 version of the datacatalog client

* Install datacatalog in pre-releease tests

* Adjust test to actually make it work

* Make sure taxonomy is properly cleaned up
diff --git a/noxfile.py b/noxfile.py
@@ -142,6 +142,9 @@ def system(session):
     else:
         session.install("google-cloud-storage", "-c", constraints_path)
 
+    # Data Catalog needed for the column ACL test with a real Policy Tag.
+    session.install("google-cloud-datacatalog", "-c", constraints_path)
+
     session.install("-e", ".[all]", "-c", constraints_path)
     session.install("ipython", "-c", constraints_path)
 
@@ -211,6 +214,7 @@ def prerelease_deps(session):
     session.install("--pre", "grpcio", "pandas")
     session.install(
         "freezegun",
+        "google-cloud-datacatalog",
         "google-cloud-storage",
         "google-cloud-testutils",
         "IPython",
diff --git a/tests/system/test_client.py b/tests/system/test_client.py
@@ -68,6 +68,8 @@
 from google.cloud._helpers import UTC
 from google.cloud.bigquery import dbapi, enums
 from google.cloud import storage
+from google.cloud.datacatalog_v1 import types as datacatalog_types
+from google.cloud.datacatalog_v1 import PolicyTagManagerClient
 
 from test_utils.retry import RetryErrors
 from test_utils.retry import RetryInstanceState
@@ -167,6 +169,8 @@ def setUp(self):
         self.to_delete = [dataset]
 
     def tearDown(self):
+        policy_tag_client = PolicyTagManagerClient()
+
         def _still_in_use(bad_request):
             return any(
                 error["reason"] == "resourceInUse" for error in bad_request._errors
@@ -183,6 +187,8 @@ def _still_in_use(bad_request):
                 retry_in_use(Config.CLIENT.delete_dataset)(doomed, delete_contents=True)
             elif isinstance(doomed, (Table, bigquery.TableReference)):
                 retry_in_use(Config.CLIENT.delete_table)(doomed)
+            elif isinstance(doomed, datacatalog_types.Taxonomy):
+                policy_tag_client.delete_taxonomy(name=doomed.name)
             else:
                 doomed.delete()
 
@@ -381,6 +387,68 @@ def test_create_table_with_policy(self):
         table2 = Config.CLIENT.update_table(table, ["schema"])
         self.assertEqual(policy_2, table2.schema[1].policy_tags)
 
+    def test_create_table_with_real_custom_policy(self):
+        from google.cloud.bigquery.schema import PolicyTagList
+
+        policy_tag_client = PolicyTagManagerClient()
+        taxonomy_parent = f"projects/{Config.CLIENT.project}/locations/us"
+
+        new_taxonomy = datacatalog_types.Taxonomy(
+            display_name="Custom test taxonomy",
+            description="This taxonomy is ony used for a test.",
+            activated_policy_types=[
+                datacatalog_types.Taxonomy.PolicyType.FINE_GRAINED_ACCESS_CONTROL
+            ],
+        )
+
+        taxonomy = policy_tag_client.create_taxonomy(
+            parent=taxonomy_parent, taxonomy=new_taxonomy
+        )
+        self.to_delete.insert(0, taxonomy)
+
+        parent_policy_tag = policy_tag_client.create_policy_tag(
+            parent=taxonomy.name,
+            policy_tag=datacatalog_types.PolicyTag(
+                display_name="Parent policy tag", parent_policy_tag=None
+            ),
+        )
+        child_policy_tag = policy_tag_client.create_policy_tag(
+            parent=taxonomy.name,
+            policy_tag=datacatalog_types.PolicyTag(
+                display_name="Child policy tag",
+                parent_policy_tag=parent_policy_tag.name,
+            ),
+        )
+
+        dataset = self.temp_dataset(
+            _make_dataset_id("create_table_with_real_custom_policy")
+        )
+        table_id = "test_table"
+        policy_1 = PolicyTagList(names=[parent_policy_tag.name])
+        policy_2 = PolicyTagList(names=[child_policy_tag.name])
+
+        schema = [
+            bigquery.SchemaField(
+                "first_name", "STRING", mode="REQUIRED", policy_tags=policy_1
+            ),
+            bigquery.SchemaField(
+                "age", "INTEGER", mode="REQUIRED", policy_tags=policy_2
+            ),
+        ]
+        table_arg = Table(dataset.table(table_id), schema=schema)
+        self.assertFalse(_table_exists(table_arg))
+
+        table = helpers.retry_403(Config.CLIENT.create_table)(table_arg)
+        self.to_delete.insert(0, table)
+
+        self.assertTrue(_table_exists(table))
+        self.assertCountEqual(
+            list(table.schema[0].policy_tags.names), [parent_policy_tag.name]
+        )
+        self.assertCountEqual(
+            list(table.schema[1].policy_tags.names), [child_policy_tag.name]
+        )
+
     def test_create_table_w_time_partitioning_w_clustering_fields(self):
         from google.cloud.bigquery.table import TimePartitioning
         from google.cloud.bigquery.table import TimePartitioningType
