Bigtable system tests passing against V2 endpoints

gcloud/bigtable/_generated_v2/operations_grpc_pb2.py

@@ -1,4 +1,12 @@
-
+from google.longrunning.operations_pb2 import (
+    CancelOperationRequest,
+    DeleteOperationRequest,
+    GetOperationRequest,
+    ListOperationsRequest,
+    ListOperationsResponse,
+    Operation,
+    google_dot_protobuf_dot_empty__pb2,
+)
 from grpc.beta import implementations as beta_implementations
 from grpc.beta import interfaces as beta_interfaces
 from grpc.framework.common import cardinality

gcloud/bigtable/client.py

@@ -43,7 +43,9 @@
 from gcloud.bigtable._generated_v2 import (
     operations_grpc_pb2 as operations_grpc_v2_pb2)
 
+from gcloud.bigtable.cluster import DEFAULT_SERVE_NODES
 from gcloud.bigtable.instance import Instance
+from gcloud.bigtable.instance import _EXISTING_INSTANCE_LOCATION_ID
 from gcloud.client import _ClientFactoryMixin
 from gcloud.client import _ClientProjectMixin
 from gcloud.credentials import get_credentials
@@ -375,22 +377,33 @@ def __exit__(self, exc_type, exc_val, exc_t):
         """Stops the client as a context manager."""
         self.stop()
 
-    def instance(self, instance_id, display_name=None):
+    def instance(self, instance_id, location=_EXISTING_INSTANCE_LOCATION_ID,
+                 display_name=None, serve_nodes=DEFAULT_SERVE_NODES):
         """Factory to create a instance associated with this client.
 
         :type instance_id: str
         :param instance_id: The ID of the instance.
 
+        :type location: string
+        :param location: location name, in form
+                         ``projects/<project>/locations/<location>``; used to
+                         set up the instance's cluster.
+
         :type display_name: str
         :param display_name: (Optional) The display name for the instance in
                              the 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.
+
         :rtype: :class:`.Instance`
         :returns: an instance owned by this client.
         """
-        return Instance(instance_id, self, display_name=display_name)
+        return Instance(instance_id, self, location,
+                        display_name=display_name, serve_nodes=serve_nodes)
 
     def list_instances(self):
         """List instances owned by the project.

gcloud/bigtable/column_family.py

@@ -261,7 +261,7 @@ def create(self):
             id=self.column_family_id,
             create=column_family,
         )
-        client = self._table._cluster._client
+        client = self._table._instance._client
         # We expect a `.table_v2_pb2.ColumnFamily`. We ignore it since the only
         # data it contains are the GC rule and the column family ID already
         # stored on this instance.
@@ -286,7 +286,7 @@ def update(self):
         request_pb.modifications.add(
             id=self.column_family_id,
             update=column_family)
-        client = self._table._cluster._client
+        client = self._table._instance._client
         # We expect a `.table_v2_pb2.ColumnFamily`. We ignore it since the only
         # data it contains are the GC rule and the column family ID already
         # stored on this instance.
@@ -300,7 +300,7 @@ def delete(self):
         request_pb.modifications.add(
             id=self.column_family_id,
             drop=True)
-        client = self._table._cluster._client
+        client = self._table._instance._client
         # We expect a `google.protobuf.empty_pb2.Empty`
         client._table_stub.ModifyColumnFamilies(request_pb,
                                                 client.timeout_seconds)

gcloud/bigtable/instance.py

@@ -27,14 +27,17 @@
 from gcloud.bigtable._generated_v2 import (
     bigtable_table_admin_pb2 as table_messages_v2_pb2)
 from gcloud.bigtable.cluster import Cluster
+from gcloud.bigtable.cluster import DEFAULT_SERVE_NODES
 from gcloud.bigtable.table import Table
 
 
+_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]*)$')
 _OPERATION_NAME_RE = re.compile(r'^operations/projects/([^/]+)/'
-                                r'instances/([a-z][-a-z0-9]*)/operations/'
-                                r'(?P<operation_id>\d+)$')
+                                r'instances/([a-z][-a-z0-9]*)/'
+                                r'locations/(?P<location_id>[a-z][-a-z0-9]*)/'
+                                r'operations/(?P<operation_id>\d+)$')
 _TYPE_URL_BASE = 'type.googleapis.com/google.bigtable.'
 _ADMIN_TYPE_URL_BASE = _TYPE_URL_BASE + 'admin.v2.'
 _INSTANCE_CREATE_METADATA = _ADMIN_TYPE_URL_BASE + 'CreateInstanceMetadata'
@@ -53,13 +56,19 @@ def _prepare_create_request(instance):
     :returns: The CreateInstance request object containing the instance info.
     """
     parent_name = ('projects/' + instance._client.project)
-    return messages_v2_pb2.CreateInstanceRequest(
+    message = messages_v2_pb2.CreateInstanceRequest(
         parent=parent_name,
         instance_id=instance.instance_id,
         instance=data_v2_pb2.Instance(
             display_name=instance.display_name,
         ),
     )
+    cluster = message.clusters[instance.instance_id]
+    cluster.name = instance.name + '/clusters/' + instance.instance_id
+    cluster.location = (
+        parent_name + '/locations/' + instance._cluster_location_id)
+    cluster.serve_nodes = instance._cluster_serve_nodes
+    return message
 
 
 def _parse_pb_any_to_native(any_val, expected_type=None):
@@ -91,25 +100,24 @@ def _process_operation(operation_pb):
     :param operation_pb: The long-running operation response from a
                          Create/Update/Undelete instance request.
 
-    :rtype: tuple
-    :returns: A pair of an integer and datetime stamp. The integer is the ID
-              of the operation (``operation_id``) and the timestamp when
-              the create operation began (``operation_begin``).
+    :rtype: (int, str, datetime)
+    :returns: (operation_id, location_id, operation_begin).
     :raises: :class:`ValueError <exceptions.ValueError>` if the operation name
              doesn't match the :data:`_OPERATION_NAME_RE` regex.
     """
     match = _OPERATION_NAME_RE.match(operation_pb.name)
     if match is None:
         raise ValueError('Operation name was not in the expected '
-                         'format after a instance modification.',
+                         'format after instance creation.',
                          operation_pb.name)
+    location_id = match.group('location_id')
     operation_id = int(match.group('operation_id'))
 
     request_metadata = _parse_pb_any_to_native(operation_pb.metadata)
     operation_begin = _pb_timestamp_to_datetime(
         request_metadata.request_time)
 
-    return operation_id, operation_begin
+    return operation_id, location_id, operation_begin
 
 
 class Operation(object):
@@ -128,14 +136,18 @@ class Operation(object):
     :type begin: :class:`datetime.datetime`
     :param begin: The time when the operation was started.
 
+    :type location_id: str
+    :param location_id: ID of the location in which the operation is running
+
     :type instance: :class:`Instance`
     :param instance: The instance that created the operation.
     """
 
-    def __init__(self, op_type, op_id, begin, instance=None):
+    def __init__(self, op_type, op_id, begin, location_id, instance=None):
         self.op_type = op_type
         self.op_id = op_id
         self.begin = begin
+        self.location_id = location_id
         self._instance = instance
         self._complete = False
 
@@ -145,6 +157,7 @@ def __eq__(self, other):
         return (other.op_type == self.op_type and
                 other.op_id == self.op_id and
                 other.begin == self.begin and
+                other.location_id == self.location_id and
                 other._instance == self._instance and
                 other._complete == self._complete)
 
@@ -162,8 +175,9 @@ def finished(self):
         if self._complete:
             raise ValueError('The operation has completed.')
 
-        operation_name = ('operations/' + self._instance.name +
-                          '/operations/%d' % (self.op_id,))
+        operation_name = (
+            'operations/%s/locations/%s/operations/%d' %
+            (self._instance.name, self.location_id, self.op_id))
         request_pb = operations_pb2.GetOperationRequest(name=operation_name)
         # We expect a `google.longrunning.operations_pb2.Operation`.
         operation_pb = self._instance._client._operations_stub.GetOperation(
@@ -199,17 +213,30 @@ class Instance(object):
     :param client: The client that owns the instance. Provides
                    authorization and a project ID.
 
+    :type location_id: str
+    :param location_id: ID of the location in which the instance will be
+                        created.  Required for instances which do not yet
+                        exist.
+
     :type display_name: str
     :param display_name: (Optional) The display name for the instance in the
                          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.
     """
 
     def __init__(self, instance_id, client,
-                 display_name=None):
+                 location_id=_EXISTING_INSTANCE_LOCATION_ID,
+                 display_name=None,
+                 serve_nodes=DEFAULT_SERVE_NODES):
         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
 
     def _update_from_pb(self, instance_pb):
@@ -246,8 +273,9 @@ def from_pb(cls, instance_pb, client):
         if match.group('project') != client.project:
             raise ValueError('Project ID on instance does not match the '
                              'project ID on the client')
+        instance_id = match.group('instance_id')
 
-        result = cls(match.group('instance_id'), client)
+        result = cls(instance_id, client, _EXISTING_INSTANCE_LOCATION_ID)
         result._update_from_pb(instance_pb)
         return result
 
@@ -262,6 +290,7 @@ def copy(self):
         """
         new_client = self._client.copy()
         return self.__class__(self.instance_id, new_client,
+                              self._cluster_location_id,
                               display_name=self.display_name)
 
     @property
@@ -332,8 +361,8 @@ def create(self):
         operation_pb = self._client._instance_stub.CreateInstance(
             request_pb, self._client.timeout_seconds)
 
-        op_id, op_begin = _process_operation(operation_pb)
-        return Operation('create', op_id, op_begin, instance=self)
+        op_id, loc_id, op_begin = _process_operation(operation_pb)
+        return Operation('create', op_id, op_begin, loc_id, instance=self)
 
     def update(self):
         """Update this instance.

gcloud/bigtable/row.py

@@ -395,7 +395,7 @@ def commit(self):
             mutations=mutations_list,
         )
         # We expect a `google.protobuf.empty_pb2.Empty`
-        client = self._table._cluster._client
+        client = self._table._instance._client
         client._data_stub.MutateRow(request_pb, client.timeout_seconds)
         self.clear()
 
@@ -512,7 +512,7 @@ def commit(self):
             false_mutations=false_mutations,
         )
         # We expect a `.messages_v2_pb2.CheckAndMutateRowResponse`
-        client = self._table._cluster._client
+        client = self._table._instance._client
         resp = client._data_stub.CheckAndMutateRow(
             request_pb, client.timeout_seconds)
         self.clear()
@@ -800,7 +800,7 @@ def commit(self):
             rules=self._rule_pb_list,
         )
         # We expect a `.data_v2_pb2.Row`
-        client = self._table._cluster._client
+        client = self._table._instance._client
         row_response = client._data_stub.ReadModifyWriteRow(
             request_pb, client.timeout_seconds)
 

gcloud/bigtable/test_client.py

@@ -526,20 +526,47 @@ def test_stop_while_stopped(self):
         # Make sure the cluster stub did not change.
         self.assertEqual(client._instance_stub_internal, instance_stub)
 
-    def test_instance_factory(self):
+    def test_instance_factory_defaults(self):
+        from gcloud.bigtable.cluster import DEFAULT_SERVE_NODES
         from gcloud.bigtable.instance import Instance
+        from gcloud.bigtable.instance import _EXISTING_INSTANCE_LOCATION_ID
 
         PROJECT = 'PROJECT'
         INSTANCE_ID = 'instance-id'
         DISPLAY_NAME = 'display-name'
-
         credentials = _Credentials()
         client = self._makeOne(project=PROJECT, credentials=credentials)
 
         instance = client.instance(INSTANCE_ID, display_name=DISPLAY_NAME)
+
+        self.assertTrue(isinstance(instance, Instance))
+        self.assertEqual(instance.instance_id, INSTANCE_ID)
+        self.assertEqual(instance.display_name, DISPLAY_NAME)
+        self.assertEqual(instance._cluster_location_id,
+                         _EXISTING_INSTANCE_LOCATION_ID)
+        self.assertEqual(instance._cluster_serve_nodes, DEFAULT_SERVE_NODES)
+        self.assertTrue(instance._client is client)
+
+    def test_instance_factory_w_explicit_serve_nodes(self):
+        from gcloud.bigtable.instance import Instance
+
+        PROJECT = 'PROJECT'
+        INSTANCE_ID = 'instance-id'
+        DISPLAY_NAME = 'display-name'
+        LOCATION_ID = 'locname'
+        SERVE_NODES = 5
+        credentials = _Credentials()
+        client = self._makeOne(project=PROJECT, credentials=credentials)
+
+        instance = client.instance(
+            INSTANCE_ID, display_name=DISPLAY_NAME,
+            location=LOCATION_ID, serve_nodes=SERVE_NODES)
+
         self.assertTrue(isinstance(instance, Instance))
         self.assertEqual(instance.instance_id, INSTANCE_ID)
         self.assertEqual(instance.display_name, DISPLAY_NAME)
+        self.assertEqual(instance._cluster_location_id, LOCATION_ID)
+        self.assertEqual(instance._cluster_serve_nodes, SERVE_NODES)
         self.assertTrue(instance._client is client)
 
     def test_list_instances(self):
@@ -549,6 +576,7 @@ def test_list_instances(self):
             bigtable_instance_admin_pb2 as messages_v2_pb2)
         from gcloud.bigtable._testing import _FakeStub
 
+        LOCATION = 'projects/' + self.PROJECT + '/locations/locname'
         FAILED_LOCATION = 'FAILED'
         INSTANCE_ID1 = 'instance-id1'
         INSTANCE_ID2 = 'instance-id2'
@@ -593,8 +621,8 @@ def test_list_instances(self):
         # Create expected_result.
         failed_locations = [FAILED_LOCATION]
         instances = [
-            client.instance(INSTANCE_ID1),
-            client.instance(INSTANCE_ID2),
+            client.instance(INSTANCE_ID1, LOCATION),
+            client.instance(INSTANCE_ID2, LOCATION),
         ]
         expected_result = (instances, failed_locations)
 

gcloud/bigtable/test_column_family.py

@@ -650,7 +650,7 @@ def _ColumnFamilyPB(*args, **kw):
     return table_v2_pb2.ColumnFamily(*args, **kw)
 
 
-class _Cluster(object):
+class _Instance(object):
 
     def __init__(self, client=None):
         self._client = client
@@ -666,4 +666,4 @@ class _Table(object):
 
     def __init__(self, name, client=None):
         self.name = name
-        self._cluster = _Cluster(client)
+        self._instance = _Instance(client)