BigTable: Modify system test for new GAPIC code (#5302)

* Provide new auto-generated layer for Bigtable.

* Change bigtable_pb2 imports to use from gapic library.

* Add retry for read rows

* Add parameter start_inclusive to _create_row_request

* Add retry for Deadline Exceeded on read rows

* Refactor yield_rows retry

* Add grpc google iam v1 on setup.py on bigtable

* Change routing_header to use to_grpc_metadata

Author: aneepct

packages/google-cloud-bigtable/google/cloud/bigtable/client.py

@@ -29,13 +29,21 @@
 """
 
 
-from google.cloud.bigtable.instance import Instance
-from google.cloud.bigtable.instance import _EXISTING_INSTANCE_LOCATION_ID
+from google.api_core.gapic_v1 import client_info
 
 from google.cloud import bigtable_v2
 from google.cloud import bigtable_admin_v2
 
+from google.cloud.bigtable import __version__
+from google.cloud.bigtable.instance import Instance
+from google.cloud.bigtable.instance import _EXISTING_INSTANCE_LOCATION_ID
+
+from google.cloud.client import ClientWithProject
 
+
+_CLIENT_INFO = client_info.ClientInfo(
+    client_library_version=__version__)
+SPANNER_ADMIN_SCOPE = 'https://www.googleapis.com/auth/spanner.admin'
 ADMIN_SCOPE = 'https://www.googleapis.com/auth/bigtable.admin'
 """Scope for interacting with the Cluster Admin and Table Admin APIs."""
 DATA_SCOPE = 'https://www.googleapis.com/auth/bigtable.data'
@@ -44,7 +52,7 @@
 """Scope for reading table data."""
 
 
-class Client(object):
+class Client(ClientWithProject):
     """Client for interacting with Google Cloud Bigtable API.
 
     .. note::
@@ -81,6 +89,9 @@ class Client(object):
     :raises: :class:`ValueError <exceptions.ValueError>` if both ``read_only``
              and ``admin`` are :data:`True`
     """
+    _table_data_client = None
+    _table_admin_client = None
+    _instance_admin_client = None
 
     def __init__(self, project=None, credentials=None,
                  read_only=False, admin=False, channel=None):
@@ -90,13 +101,11 @@ def __init__(self, project=None, credentials=None,
 
         # NOTE: We set the scopes **before** calling the parent constructor.
         #       It **may** use those scopes in ``with_scopes_if_required``.
-        self.project = project
         self._read_only = bool(read_only)
         self._admin = bool(admin)
         self._channel = channel
-        self._credentials = credentials
         self.SCOPE = self._get_scopes()
-        super(Client, self).__init__()
+        super(Client, self).__init__(project=project, credentials=credentials)
 
     def _get_scopes(self):
         """Get the scopes corresponding to admin / read-only state.
@@ -130,21 +139,27 @@ def project_path(self):
         :rtype: str
         :returns: Return a fully-qualified project string.
         """
-        instance_client = self._instance_admin_client
+        instance_client = self.instance_admin_client
         return instance_client.project_path(self.project)
 
     @property
-    def _table_data_client(self):
+    def table_data_client(self):
         """Getter for the gRPC stub used for the Table Admin API.
 
         :rtype: :class:`.bigtable_v2.BigtableClient`
         :returns: A BigtableClient object.
         """
-        return bigtable_v2.BigtableClient(channel=self._channel,
-                                          credentials=self._credentials)
+        if self._table_data_client is None:
+            if not self._admin:
+                raise ValueError('Client is not an admin client.')
+            self._table_data_client = (
+                bigtable_v2.BigtableClient(credentials=self._credentials,
+                                           client_info=_CLIENT_INFO))
+
+        return self._table_data_client
 
     @property
-    def _table_admin_client(self):
+    def table_admin_client(self):
         """Getter for the gRPC stub used for the Table Admin API.
 
         :rtype: :class:`.bigtable_admin_pb2.BigtableTableAdmin`
@@ -153,13 +168,17 @@ def _table_admin_client(self):
                  client is not an admin client or if it has not been
                  :meth:`start`-ed.
         """
-        if not self._admin:
-            raise ValueError('Client is not an admin client.')
-        return bigtable_admin_v2.BigtableTableAdminClient(
-            channel=self._channel, credentials=self._credentials)
+        if self._table_admin_client is None:
+            if not self._admin:
+                raise ValueError('Client is not an admin client.')
+            self._table_admin_client = (
+                bigtable_admin_v2.BigtableTableAdminClient(
+                    credentials=self._credentials, client_info=_CLIENT_INFO))
+
+        return self._table_admin_client
 
     @property
-    def _instance_admin_client(self):
+    def instance_admin_client(self):
         """Getter for the gRPC stub used for the Table Admin API.
 
         :rtype: :class:`.bigtable_admin_pb2.BigtableInstanceAdmin`
@@ -168,10 +187,14 @@ def _instance_admin_client(self):
                  client is not an admin client or if it has not been
                  :meth:`start`-ed.
         """
-        if not self._admin:
-            raise ValueError('Client is not an admin client.')
-        return bigtable_admin_v2.BigtableInstanceAdminClient(
-            channel=self._channel, credentials=self._credentials)
+        if self._instance_admin_client is None:
+            if not self._admin:
+                raise ValueError('Client is not an admin client.')
+            self._instance_admin_client = (
+                bigtable_admin_v2.BigtableInstanceAdminClient(
+                    credentials=self._credentials, client_info=_CLIENT_INFO))
+
+        return self._instance_admin_client
 
     def instance(self, instance_id, location=_EXISTING_INSTANCE_LOCATION_ID,
                  display_name=None):
@@ -202,4 +225,4 @@ def list_instances(self):
         :rtype: :class:`~google.api_core.page_iterator.Iterator`
         :returns: A list of Instance.
         """
-        return self._instance_admin_client.list_instances(self.project_path)
+        return self.instance_admin_client.list_instances(self.project_path)

packages/google-cloud-bigtable/google/cloud/bigtable/cluster.py

@@ -69,7 +69,7 @@ def name(self):
         :rtype: str
         :returns: The cluster name.
         """
-        return self._instance._client._instance_admin_client.cluster_path(
+        return self._instance._client.instance_admin_client.cluster_path(
             self._instance._client.project, self._instance.instance_id,
             self.cluster_id)
 
@@ -90,7 +90,7 @@ def __ne__(self, other):
 
     def reload(self):
         """Reload the metadata for this cluster."""
-        self._instance._client._instance_admin_client.get_cluster(self.name)
+        self._instance._client.instance_admin_client.get_cluster(self.name)
 
     def create(self):
         """Create this cluster.
@@ -113,7 +113,7 @@ def create(self):
                   create operation.
         """
         client = self._instance._client
-        return client._instance_admin_client.create_cluster(
+        return client.instance_admin_client.create_cluster(
             self._instance.name, self.cluster_id, {})
 
     def update(self, location='', serve_nodes=0):
@@ -147,7 +147,7 @@ def update(self, location='', serve_nodes=0):
                   update operation.
         """
         client = self._instance._client
-        return client._instance_admin_client.update_cluster(
+        return client.instance_admin_client.update_cluster(
             self.name, location, serve_nodes)
 
     def delete(self):
@@ -171,4 +171,4 @@ def delete(self):
           permanently deleted.
         """
         client = self._instance._client
-        client._instance_admin_client.delete_cluster(self.name)
+        client.instance_admin_client.delete_cluster(self.name)

packages/google-cloud-bigtable/google/cloud/bigtable/instance.py

@@ -18,13 +18,16 @@
 import re
 
 from google.cloud.bigtable.table import Table
+from google.cloud.bigtable.cluster import DEFAULT_SERVE_NODES
 
 from google.cloud.bigtable_admin_v2 import enums
+from google.cloud.bigtable_admin_v2.types import instance_pb2
 
 
 _EXISTING_INSTANCE_LOCATION_ID = 'see-existing-cluster'
 _INSTANCE_NAME_RE = re.compile(r'^projects/(?P<project>[^/]+)/'
                                r'instances/(?P<instance_id>[a-z][-a-z0-9]*)$')
+_STORAGE_TYPE_UNSPECIFIED = enums.StorageType.STORAGE_TYPE_UNSPECIFIED
 
 
 class Instance(object):
@@ -59,15 +62,31 @@ class Instance(object):
                          Cloud Console UI. (Must be between 4 and 30
                          characters.) If this value is not set in the
                          constructor, will fall back to the instance ID.
+
+    :type serve_nodes: int
+    :param serve_nodes: (Optional) The number of nodes in the instance's
+                        cluster; used to set up the instance's cluster.
+
+    :type default_storage_type: int
+    :param default_storage_type: (Optional) The default values are
+                                    STORAGE_TYPE_UNSPECIFIED = 0: The user did
+                                    not specify a storage type.
+                                    SSD = 1: Flash (SSD) storage should be
+                                    used.
+                                    HDD = 2: Magnetic drive (HDD) storage
+                                    should be used.
     """
 
     def __init__(self, instance_id, client,
                  location_id=_EXISTING_INSTANCE_LOCATION_ID,
-                 display_name=None):
+                 display_name=None, serve_nodes=DEFAULT_SERVE_NODES,
+                 default_storage_type=_STORAGE_TYPE_UNSPECIFIED):
         self.instance_id = instance_id
         self.display_name = display_name or instance_id
         self._cluster_location_id = location_id
+        self._cluster_serve_nodes = serve_nodes
         self._client = client
+        self._default_storage_type = default_storage_type
 
     @classmethod
     def from_pb(cls, instance_pb, client):
@@ -140,6 +159,15 @@ def __eq__(self, other):
     def __ne__(self, other):
         return not self == other
 
+    def reload(self):
+        """Reload the metadata for this instance."""
+        instance_pb = self._client._instance_admin_client.get_instance(
+            self.name)
+
+        # NOTE: _update_from_pb does not check that the project and
+        #       instance ID on the response match the request.
+        self._update_from_pb(instance_pb)
+
     def create(self):
         """Create this instance.
 
@@ -160,10 +188,25 @@ def create(self):
         :returns: The long-running operation corresponding to the create
                     operation.
         """
+        clusters = {}
+        cluster_id = '{}-cluster'.format(self.instance_id)
+        cluster_name = self._client._instance_admin_client.cluster_path(
+            self._client.project, self.instance_id, cluster_id)
+        location = self._client._instance_admin_client.location_path(
+            self._client.project, self._cluster_location_id)
+        cluster = instance_pb2.Cluster(
+            name=cluster_name, location=location,
+            serve_nodes=self._cluster_serve_nodes,
+            default_storage_type=self._default_storage_type)
+        instance = instance_pb2.Instance(
+            display_name=self.display_name
+        )
+        clusters[cluster_id] = cluster
         parent = self._client.project_path
+
         return self._client._instance_admin_client.create_instance(
-            parent=parent, instance_id=self.instance_id, instance={},
-            clusters={})
+            parent=parent, instance_id=self.instance_id, instance=instance,
+            clusters=clusters)
 
     def update(self):
         """Update this instance.

packages/google-cloud-bigtable/tests/system.py

@@ -86,12 +86,13 @@ def setUpModule():
     if not Config.IN_EMULATOR:
         retry = RetryErrors(GrpcRendezvous,
                             error_predicate=_retry_on_unavailable)
-        instances, failed_locations = retry(Config.CLIENT.list_instances)()
 
-        if len(failed_locations) != 0:
+        instances_response = retry(Config.CLIENT.list_instances)()
+
+        if len(instances_response.failed_locations) != 0:
             raise ValueError('List instances failed in module set up.')
 
-        EXISTING_INSTANCES[:] = instances
+        EXISTING_INSTANCES[:] = instances_response.instances
 
         # After listing, create the test instance.
         created_op = Config.INSTANCE.create()
@@ -116,11 +117,12 @@ def tearDown(self):
             instance.delete()
 
     def test_list_instances(self):
-        instances, failed_locations = Config.CLIENT.list_instances()
-        self.assertEqual(failed_locations, [])
+        instances_response = Config.CLIENT.list_instances()
+        self.assertEqual(instances_response.failed_locations, [])
         # We have added one new instance in `setUpModule`.
-        self.assertEqual(len(instances), len(EXISTING_INSTANCES) + 1)
-        for instance in instances:
+        self.assertEqual(len(instances_response.instances),
+                         len(EXISTING_INSTANCES) + 1)
+        for instance in instances_response.instances:
             instance_existence = (instance in EXISTING_INSTANCES or
                                   instance == Config.INSTANCE)
             self.assertTrue(instance_existence)

packages/google-cloud-bigtable/tests/unit/_testing.py

@@ -15,9 +15,23 @@
 """Mocks used to emulate gRPC generated objects."""
 
 
+import mock
+
+
 class _FakeStub(object):
     """Acts as a gPRC stub."""
 
     def __init__(self, *results):
         self.results = results
         self.method_calls = []
+
+
+def _make_credentials():
+    import google.auth.credentials
+
+    class _CredentialsWithScopes(
+            google.auth.credentials.Credentials,
+            google.auth.credentials.Scoped):
+        pass
+
+    return mock.Mock(spec=_CredentialsWithScopes)