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

@@ -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

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
+               <https://cloud.google.com/iam/docs/overview#permissions>`_.
+               `Bigtable Permissions
+               <https://cloud.google.com/bigtable/docs/access-control>`_.
+
+        :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.
 

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)

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:[email protected]", "user:[email protected]"]
+        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:[email protected]", "user:[email protected]"]
+        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("[email protected]"),
+            Policy.service_account("[email protected]"),
+        ]
+
+        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