From b7ba9183a5384839e05afc9d26ddc4e5619ab3bf Mon Sep 17 00:00:00 2001 From: Leonid Emar-Kar <46078689+Emar-Kar@users.noreply.github.com> Date: Tue, 26 Nov 2019 22:17:38 +0300 Subject: [PATCH] feat(bigtable): add table level IAM policy controls (#9877) * feat(bigtable): add table level IAM policy controls * remove extra lines * system test * black * fix system test * update system-tests * chg comment lines * comment correction --- bigtable/docs/snippets_table.py | 49 +++++++++++ bigtable/google/cloud/bigtable/table.py | 69 ++++++++++++++++ bigtable/tests/system.py | 38 +++++++++ bigtable/tests/unit/test_table.py | 105 ++++++++++++++++++++++++ 4 files changed, 261 insertions(+) diff --git a/bigtable/docs/snippets_table.py b/bigtable/docs/snippets_table.py index 0fbb16bf74ad..702cf31b1447 100644 --- a/bigtable/docs/snippets_table.py +++ b/bigtable/docs/snippets_table.py @@ -356,6 +356,55 @@ def test_bigtable_get_cluster_states(): assert CLUSTER_ID in get_cluster_states +def test_bigtable_table_test_iam_permissions(): + table_policy = Config.INSTANCE.table("table_id_iam_policy") + table_policy.create() + assert table_policy.exists + + # [START bigtable_table_test_iam_permissions] + from google.cloud.bigtable import Client + + client = Client(admin=True) + instance = client.instance(INSTANCE_ID) + table = instance.table("table_id_iam_policy") + + permissions = ["bigtable.tables.mutateRows", "bigtable.tables.readRows"] + permissions_allowed = table.test_iam_permissions(permissions) + # [END bigtable_table_test_iam_permissions] + assert permissions_allowed == permissions + + +def test_bigtable_table_set_iam_policy_then_get_iam_policy(): + table_policy = Config.INSTANCE.table("table_id_iam_policy") + assert table_policy.exists + service_account_email = Config.CLIENT._credentials.service_account_email + + # [START bigtable_table_set_iam_policy] + from google.cloud.bigtable import Client + from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + + client = Client(admin=True) + instance = client.instance(INSTANCE_ID) + table = instance.table("table_id_iam_policy") + new_policy = Policy() + new_policy[BIGTABLE_ADMIN_ROLE] = [Policy.service_account(service_account_email)] + + policy_latest = table.set_iam_policy(new_policy) + # [END bigtable_table_set_iam_policy] + assert len(policy_latest.bigtable_admins) > 0 + + # [START bigtable_table_get_iam_policy] + from google.cloud.bigtable import Client + + client = Client(admin=True) + instance = client.instance(INSTANCE_ID) + table = instance.table("table_id_iam_policy") + policy = table.get_iam_policy() + # [END bigtable_table_get_iam_policy] + assert len(policy.bigtable_admins) > 0 + + def test_bigtable_table_exists(): # [START bigtable_check_table_exists] from google.cloud.bigtable import Client diff --git a/bigtable/google/cloud/bigtable/table.py b/bigtable/google/cloud/bigtable/table.py index 4ced9fbde0c2..69379b21d57e 100644 --- a/bigtable/google/cloud/bigtable/table.py +++ b/bigtable/google/cloud/bigtable/table.py @@ -28,6 +28,7 @@ from google.cloud.bigtable.column_family import ColumnFamily from google.cloud.bigtable.batcher import MutationsBatcher from google.cloud.bigtable.batcher import FLUSH_COUNT, MAX_ROW_BYTES +from google.cloud.bigtable.policy import Policy from google.cloud.bigtable.row import AppendRow from google.cloud.bigtable.row import ConditionalRow from google.cloud.bigtable.row import DirectRow @@ -138,6 +139,74 @@ def name(self): project=project, instance=instance_id, table=self.table_id ) + def get_iam_policy(self): + """Gets the IAM access control policy for this table. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_table_get_iam_policy] + :end-before: [END bigtable_table_get_iam_policy] + + :rtype: :class:`google.cloud.bigtable.policy.Policy` + :returns: The current IAM policy of this table. + """ + table_client = self._instance._client.table_admin_client + resp = table_client.get_iam_policy(resource=self.name) + return Policy.from_pb(resp) + + def set_iam_policy(self, policy): + """Sets the IAM access control policy for this table. Replaces any + existing policy. + + For more information about policy, please see documentation of + class `google.cloud.bigtable.policy.Policy` + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_table_set_iam_policy] + :end-before: [END bigtable_table_set_iam_policy] + + :type policy: :class:`google.cloud.bigtable.policy.Policy` + :param policy: A new IAM policy to replace the current IAM policy + of this table. + + :rtype: :class:`google.cloud.bigtable.policy.Policy` + :returns: The current IAM policy of this table. + """ + table_client = self._instance._client.table_admin_client + resp = table_client.set_iam_policy(resource=self.name, policy=policy.to_pb()) + return Policy.from_pb(resp) + + def test_iam_permissions(self, permissions): + """Tests whether the caller has the given permissions for this table. + Returns the permissions that the caller has. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_table_test_iam_permissions] + :end-before: [END bigtable_table_test_iam_permissions] + + :type permissions: list + :param permissions: The set of permissions to check for + the ``resource``. Permissions with wildcards (such as '*' + or 'storage.*') are not allowed. For more information see + `IAM Overview + `_. + `Bigtable Permissions + `_. + + :rtype: list + :returns: A List(string) of permissions allowed on the table. + """ + table_client = self._instance._client.table_admin_client + resp = table_client.test_iam_permissions( + resource=self.name, permissions=permissions + ) + return list(resp.permissions) + def column_family(self, column_family_id, gc_rule=None): """Factory to create a column family associated with this table. diff --git a/bigtable/tests/system.py b/bigtable/tests/system.py index ae43bb10ecdf..e9e3ab79179e 100644 --- a/bigtable/tests/system.py +++ b/bigtable/tests/system.py @@ -29,6 +29,8 @@ from google.cloud._helpers import UTC from google.cloud.bigtable.client import Client from google.cloud.bigtable.column_family import MaxVersionsGCRule +from google.cloud.bigtable.policy import Policy +from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE from google.cloud.bigtable.row_filters import ApplyLabelFilter from google.cloud.bigtable.row_filters import ColumnQualifierRegexFilter from google.cloud.bigtable.row_filters import RowFilterChain @@ -688,6 +690,42 @@ def test_create_table(self): sorted_tables = sorted(tables, key=name_attr) self.assertEqual(sorted_tables, expected_tables) + def test_test_iam_permissions(self): + temp_table_id = "test-test-iam-policy-table" + temp_table = Config.INSTANCE_DATA.table(temp_table_id) + temp_table.create() + self.tables_to_delete.append(temp_table) + + permissions = ["bigtable.tables.mutateRows", "bigtable.tables.readRows"] + permissions_allowed = temp_table.test_iam_permissions(permissions) + self.assertEqual(permissions, permissions_allowed) + + def test_get_iam_policy(self): + temp_table_id = "test-get-iam-policy-table" + temp_table = Config.INSTANCE_DATA.table(temp_table_id) + temp_table.create() + self.tables_to_delete.append(temp_table) + + policy = temp_table.get_iam_policy().to_api_repr() + self.assertEqual(policy["etag"], "ACAB") + self.assertEqual(policy["version"], 0) + + def test_set_iam_policy(self): + temp_table_id = "test-set-iam-policy-table" + temp_table = Config.INSTANCE_DATA.table(temp_table_id) + temp_table.create() + self.tables_to_delete.append(temp_table) + + new_policy = Policy() + service_account_email = Config.CLIENT._credentials.service_account_email + new_policy[BIGTABLE_ADMIN_ROLE] = [ + Policy.service_account(service_account_email) + ] + policy_latest = temp_table.set_iam_policy(new_policy).to_api_repr() + + self.assertEqual(policy_latest["bindings"][0]["role"], "roles/bigtable.admin") + self.assertIn(service_account_email, policy_latest["bindings"][0]["members"][0]) + def test_create_table_with_families(self): temp_table_id = "test-create-table-with-failies" temp_table = Config.INSTANCE_DATA.table(temp_table_id) diff --git a/bigtable/tests/unit/test_table.py b/bigtable/tests/unit/test_table.py index 495d8660d1f7..d4bb621c28c0 100644 --- a/bigtable/tests/unit/test_table.py +++ b/bigtable/tests/unit/test_table.py @@ -1048,6 +1048,111 @@ def test_mutations_batcher_factory(self): self.assertEqual(mutation_batcher.flush_count, flush_count) self.assertEqual(mutation_batcher.max_row_bytes, max_row_bytes) + def test_get_iam_policy(self): + from google.cloud.bigtable_admin_v2.gapic import bigtable_table_admin_client + from google.iam.v1 import policy_pb2 + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + + credentials = _make_credentials() + client = self._make_client( + project="project-id", credentials=credentials, admin=True + ) + instance = client.instance(instance_id=self.INSTANCE_ID) + table = self._make_one(self.TABLE_ID, instance) + + version = 1 + etag = b"etag_v1" + members = ["serviceAccount:service_acc1@test.com", "user:user1@test.com"] + bindings = [{"role": BIGTABLE_ADMIN_ROLE, "members": members}] + iam_policy = policy_pb2.Policy(version=version, etag=etag, bindings=bindings) + + table_api = mock.create_autospec( + bigtable_table_admin_client.BigtableTableAdminClient + ) + client._table_admin_client = table_api + table_api.get_iam_policy.return_value = iam_policy + + result = table.get_iam_policy() + + table_api.get_iam_policy.assert_called_once_with(resource=table.name) + self.assertEqual(result.version, version) + self.assertEqual(result.etag, etag) + admins = result.bigtable_admins + self.assertEqual(len(admins), len(members)) + for found, expected in zip(sorted(admins), sorted(members)): + self.assertEqual(found, expected) + + def test_set_iam_policy(self): + from google.cloud.bigtable_admin_v2.gapic import bigtable_table_admin_client + from google.iam.v1 import policy_pb2 + from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + + credentials = _make_credentials() + client = self._make_client( + project="project-id", credentials=credentials, admin=True + ) + instance = client.instance(instance_id=self.INSTANCE_ID) + table = self._make_one(self.TABLE_ID, instance) + + version = 1 + etag = b"etag_v1" + members = ["serviceAccount:service_acc1@test.com", "user:user1@test.com"] + bindings = [{"role": BIGTABLE_ADMIN_ROLE, "members": sorted(members)}] + iam_policy_pb = policy_pb2.Policy(version=version, etag=etag, bindings=bindings) + + table_api = mock.create_autospec( + bigtable_table_admin_client.BigtableTableAdminClient + ) + client._table_admin_client = table_api + table_api.set_iam_policy.return_value = iam_policy_pb + + iam_policy = Policy(etag=etag, version=version) + iam_policy[BIGTABLE_ADMIN_ROLE] = [ + Policy.user("user1@test.com"), + Policy.service_account("service_acc1@test.com"), + ] + + result = table.set_iam_policy(iam_policy) + + table_api.set_iam_policy.assert_called_once_with( + resource=table.name, policy=iam_policy_pb + ) + self.assertEqual(result.version, version) + self.assertEqual(result.etag, etag) + admins = result.bigtable_admins + self.assertEqual(len(admins), len(members)) + for found, expected in zip(sorted(admins), sorted(members)): + self.assertEqual(found, expected) + + def test_test_iam_permissions(self): + from google.cloud.bigtable_admin_v2.gapic import bigtable_table_admin_client + from google.iam.v1 import iam_policy_pb2 + + credentials = _make_credentials() + client = self._make_client( + project="project-id", credentials=credentials, admin=True + ) + instance = client.instance(instance_id=self.INSTANCE_ID) + table = self._make_one(self.TABLE_ID, instance) + + permissions = ["bigtable.tables.mutateRows", "bigtable.tables.readRows"] + + response = iam_policy_pb2.TestIamPermissionsResponse(permissions=permissions) + + table_api = mock.create_autospec( + bigtable_table_admin_client.BigtableTableAdminClient + ) + table_api.test_iam_permissions.return_value = response + client._table_admin_client = table_api + + result = table.test_iam_permissions(permissions) + + self.assertEqual(result, permissions) + table_api.test_iam_permissions.assert_called_once_with( + resource=table.name, permissions=permissions + ) + class Test__RetryableMutateRowsWorker(unittest.TestCase): from grpc import StatusCode