From c74c71201ee406b2f76c0edfe45f2782e1bc6dd6 Mon Sep 17 00:00:00 2001 From: Peter Lamut Date: Thu, 6 May 2021 14:18:22 +0200 Subject: [PATCH 1/5] test: add column ACLs test with real policy tag --- noxfile.py | 3 ++ tests/system/test_client.py | 59 +++++++++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+) diff --git a/noxfile.py b/noxfile.py index a52025635..60e2591d9 100644 --- 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) diff --git a/tests/system/test_client.py b/tests/system/test_client.py index b4b0c053d..ac077ddf9 100644 --- a/tests/system/test_client.py +++ b/tests/system/test_client.py @@ -381,6 +381,65 @@ 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 + from google.cloud.datacatalog_v1beta1 import PolicyTagManagerClient + from google.cloud.datacatalog_v1beta1.types import PolicyTag + from google.cloud.datacatalog_v1beta1.types import Taxonomy + + policy_tag_client = PolicyTagManagerClient() + + new_taxonomy = Taxonomy( + display_name="Custom test taxonomy", + description="This taxonomy is ony used for a test.", + activated_policy_types=[Taxonomy.PolicyType.FINE_GRAINED_ACCESS_CONTROL], + ) + taxonomy = policy_tag_client.create_taxonomy( + parent=Config.CLIENT.project, taxonomy=new_taxonomy + ) + self.to_delete.insert(0, taxonomy) + + parent_policy_tag = policy_tag_client.create_policy_tag( + parent=taxonomy.name, + policy_tag=PolicyTag( + display_name="Parent policy tag", parent_policy_tag=None + ), + ) + child_policy_tag = policy_tag_client.create_policy_tag( + parent=taxonomy.name, + policy_tag=PolicyTag( + display_name="Child policy tag", + parent_policy_tag=parent_policy_tag.name, + ), + ) + self.to_delete.insert(0, parent_policy_tag) + self.to_delete.insert(0, child_policy_tag) + + 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.assertEqual(table.schema[0].policy_tags, [parent_policy_tag.name]) + self.assertEqual(table.schema[1].policy_tags, [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 From a232afcc133e07b30dd57e29373f8140939ce5af Mon Sep 17 00:00:00 2001 From: Peter Lamut Date: Mon, 24 May 2021 14:47:53 +0200 Subject: [PATCH 2/5] Use v1 version of the datacatalog client --- tests/system/test_client.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/system/test_client.py b/tests/system/test_client.py index ac077ddf9..3b3940647 100644 --- a/tests/system/test_client.py +++ b/tests/system/test_client.py @@ -383,9 +383,9 @@ def test_create_table_with_policy(self): def test_create_table_with_real_custom_policy(self): from google.cloud.bigquery.schema import PolicyTagList - from google.cloud.datacatalog_v1beta1 import PolicyTagManagerClient - from google.cloud.datacatalog_v1beta1.types import PolicyTag - from google.cloud.datacatalog_v1beta1.types import Taxonomy + from google.cloud.datacatalog_v1 import PolicyTagManagerClient + from google.cloud.datacatalog_v1.types import PolicyTag + from google.cloud.datacatalog_v1.types import Taxonomy policy_tag_client = PolicyTagManagerClient() From 1fa67439dc7b70b3af4dbf471ff804ccd7a19a91 Mon Sep 17 00:00:00 2001 From: Peter Lamut Date: Mon, 24 May 2021 15:40:43 +0200 Subject: [PATCH 3/5] Install datacatalog in pre-releease tests --- noxfile.py | 1 + 1 file changed, 1 insertion(+) diff --git a/noxfile.py b/noxfile.py index 60e2591d9..10e6f6ddb 100644 --- a/noxfile.py +++ b/noxfile.py @@ -214,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", From c9ebd3f0a43e76e25f26136eb8e2a42adde88ad3 Mon Sep 17 00:00:00 2001 From: Peter Lamut Date: Wed, 16 Jun 2021 18:25:02 +0200 Subject: [PATCH 4/5] Adjust test to actually make it work --- tests/system/test_client.py | 28 +++++++++++++++++++--------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/tests/system/test_client.py b/tests/system/test_client.py index 3b3940647..a088b524d 100644 --- a/tests/system/test_client.py +++ b/tests/system/test_client.py @@ -68,6 +68,7 @@ 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 test_utils.retry import RetryErrors from test_utils.retry import RetryInstanceState @@ -183,6 +184,10 @@ 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.PolicyTag, datacatalog_types.Taxonomy) + ): + pass # TODO: add logic else: doomed.delete() @@ -384,30 +389,31 @@ def test_create_table_with_policy(self): def test_create_table_with_real_custom_policy(self): from google.cloud.bigquery.schema import PolicyTagList from google.cloud.datacatalog_v1 import PolicyTagManagerClient - from google.cloud.datacatalog_v1.types import PolicyTag - from google.cloud.datacatalog_v1.types import Taxonomy policy_tag_client = PolicyTagManagerClient() + taxonomy_parent = f"projects/{Config.CLIENT.project}/locations/us" - new_taxonomy = Taxonomy( + new_taxonomy = datacatalog_types.Taxonomy( display_name="Custom test taxonomy", description="This taxonomy is ony used for a test.", - activated_policy_types=[Taxonomy.PolicyType.FINE_GRAINED_ACCESS_CONTROL], + activated_policy_types=[ + datacatalog_types.Taxonomy.PolicyType.FINE_GRAINED_ACCESS_CONTROL + ], ) taxonomy = policy_tag_client.create_taxonomy( - parent=Config.CLIENT.project, taxonomy=new_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=PolicyTag( + 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=PolicyTag( + policy_tag=datacatalog_types.PolicyTag( display_name="Child policy tag", parent_policy_tag=parent_policy_tag.name, ), @@ -437,8 +443,12 @@ def test_create_table_with_real_custom_policy(self): self.to_delete.insert(0, table) self.assertTrue(_table_exists(table)) - self.assertEqual(table.schema[0].policy_tags, [parent_policy_tag.name]) - self.assertEqual(table.schema[1].policy_tags, [child_policy_tag.name]) + 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 045ab70c9d51e227ef77985bf9f1dd8304c60877 Mon Sep 17 00:00:00 2001 From: Peter Lamut Date: Thu, 17 Jun 2021 10:00:50 +0200 Subject: [PATCH 5/5] Make sure taxonomy is properly cleaned up --- tests/system/test_client.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/tests/system/test_client.py b/tests/system/test_client.py index a088b524d..f91004eac 100644 --- a/tests/system/test_client.py +++ b/tests/system/test_client.py @@ -69,6 +69,7 @@ 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 @@ -168,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 @@ -184,10 +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.PolicyTag, datacatalog_types.Taxonomy) - ): - pass # TODO: add logic + elif isinstance(doomed, datacatalog_types.Taxonomy): + policy_tag_client.delete_taxonomy(name=doomed.name) else: doomed.delete() @@ -388,7 +389,6 @@ def test_create_table_with_policy(self): def test_create_table_with_real_custom_policy(self): from google.cloud.bigquery.schema import PolicyTagList - from google.cloud.datacatalog_v1 import PolicyTagManagerClient policy_tag_client = PolicyTagManagerClient() taxonomy_parent = f"projects/{Config.CLIENT.project}/locations/us" @@ -400,6 +400,7 @@ def test_create_table_with_real_custom_policy(self): datacatalog_types.Taxonomy.PolicyType.FINE_GRAINED_ACCESS_CONTROL ], ) + taxonomy = policy_tag_client.create_taxonomy( parent=taxonomy_parent, taxonomy=new_taxonomy ) @@ -418,8 +419,6 @@ def test_create_table_with_real_custom_policy(self): parent_policy_tag=parent_policy_tag.name, ), ) - self.to_delete.insert(0, parent_policy_tag) - self.to_delete.insert(0, child_policy_tag) dataset = self.temp_dataset( _make_dataset_id("create_table_with_real_custom_policy")