Update bigtable to use operation.future

bigtable/google/cloud/bigtable/cluster.py

@@ -21,9 +21,7 @@
     instance_pb2 as data_v2_pb2)
 from google.cloud.bigtable._generated import (
     bigtable_instance_admin_pb2 as messages_v2_pb2)
-from google.cloud.operation import Operation
-from google.cloud.operation import register_type
-
+from google.cloud.future import operation
 
 _CLUSTER_NAME_RE = re.compile(r'^projects/(?P<project>[^/]+)/'
                               r'instances/(?P<instance>[^/]+)/clusters/'
@@ -33,9 +31,6 @@
 """Default number of nodes to use when creating a cluster."""
 
 
-register_type(messages_v2_pb2.UpdateClusterMetadata)
-
-
 def _prepare_create_request(cluster):
     """Creates a protobuf request for a CreateCluster request.
 
@@ -208,15 +203,18 @@ def create(self):
         :returns: The long-running operation corresponding to the
                   create operation.
         """
-        request_pb = _prepare_create_request(self)
-        # We expect a `google.longrunning.operations_pb2.Operation`.
         client = self._instance._client
+
+        # We expect a `google.longrunning.operations_pb2.Operation`.
+        request_pb = _prepare_create_request(self)
         operation_pb = client._instance_stub.CreateCluster(request_pb)
 
-        operation = Operation.from_pb(operation_pb, client)
-        operation.target = self
-        operation.caller_metadata['request_type'] = 'CreateCluster'
-        return operation
+        operation_future = operation.from_grpc(
+            operation_pb,
+            client._operations_stub,
+            data_v2_pb2.Cluster,
+            metadata_type=messages_v2_pb2.UpdateClusterMetadata)
+        return operation_future
 
     def update(self):
         """Update this cluster.
@@ -236,18 +234,21 @@ def update(self):
         :returns: The long-running operation corresponding to the
                   update operation.
         """
+        client = self._instance._client
+
+        # We expect a `google.longrunning.operations_pb2.Operation`.
         request_pb = data_v2_pb2.Cluster(
             name=self.name,
             serve_nodes=self.serve_nodes,
         )
-        # We expect a `google.longrunning.operations_pb2.Operation`.
-        client = self._instance._client
         operation_pb = client._instance_stub.UpdateCluster(request_pb)
 
-        operation = Operation.from_pb(operation_pb, client)
-        operation.target = self
-        operation.caller_metadata['request_type'] = 'UpdateCluster'
-        return operation
+        operation_future = operation.from_grpc(
+            operation_pb,
+            client._operations_stub,
+            data_v2_pb2.Cluster,
+            metadata_type=messages_v2_pb2.UpdateClusterMetadata)
+        return operation_future
 
     def delete(self):
         """Delete this cluster.

bigtable/google/cloud/bigtable/instance.py

@@ -26,19 +26,14 @@
 from google.cloud.bigtable.cluster import Cluster
 from google.cloud.bigtable.cluster import DEFAULT_SERVE_NODES
 from google.cloud.bigtable.table import Table
-from google.cloud.operation import Operation
-from google.cloud.operation import register_type
+from google.cloud.future import operation
 
 
 _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]*)$')
 
 
-register_type(messages_v2_pb2.CreateInstanceMetadata)
-register_type(data_v2_pb2.Instance)
-
-
 def _prepare_create_request(instance):
     """Creates a protobuf request for a CreateInstance request.
 
@@ -232,10 +227,12 @@ def create(self):
         # We expect a `google.longrunning.operations_pb2.Operation`.
         operation_pb = self._client._instance_stub.CreateInstance(request_pb)
 
-        operation = Operation.from_pb(operation_pb, self._client)
-        operation.target = self
-        operation.caller_metadata['request_type'] = 'CreateInstance'
-        return operation
+        operation_future = operation.from_grpc(
+            operation_pb,
+            self._client._operations_stub,
+            data_v2_pb2.Instance,
+            metadata_type=messages_v2_pb2.CreateInstanceMetadata)
+        return operation_future
 
     def update(self):
         """Update this instance.

bigtable/tests/system.py

@@ -32,7 +32,6 @@
 from google.cloud.environment_vars import BIGTABLE_EMULATOR
 
 from test_utils.retry import RetryErrors
-from test_utils.retry import RetryResult
 from test_utils.system import EmulatorCreds
 from test_utils.system import unique_resource_id
 
@@ -65,27 +64,6 @@ class Config(object):
     IN_EMULATOR = False
 
 
-def _wait_until_complete(operation, max_attempts=5):
-    """Wait until an operation has completed.
-
-    :type operation: :class:`google.cloud.operation.Operation`
-    :param operation: Operation that has not completed.
-
-    :type max_attempts: int
-    :param max_attempts: (Optional) The maximum number of times to check if
-                         the operation has completed. Defaults to 5.
-
-    :rtype: bool
-    :returns: Boolean indicating if the operation is complete.
-    """
-
-    def _operation_complete(result):
-        return result
-
-    retry = RetryResult(_operation_complete, max_tries=max_attempts)
-    return retry(operation.poll)()
-
-
 def _retry_on_unavailable(exc):
     """Retry only errors whose status code is 'UNAVAILABLE'."""
     from grpc import StatusCode
@@ -117,8 +95,7 @@ def setUpModule():
 
         # After listing, create the test instance.
         created_op = Config.INSTANCE.create()
-        if not _wait_until_complete(created_op):
-            raise RuntimeError('Instance creation exceed 5 seconds.')
+        created_op.result(timeout=10)
 
 
 def tearDownModule():
@@ -166,7 +143,7 @@ def test_create_instance(self):
         self.instances_to_delete.append(instance)
 
         # We want to make sure the operation completes.
-        self.assertTrue(_wait_until_complete(operation))
+        operation.result(timeout=10)
 
         # Create a new instance instance and make sure it is the same.
         instance_alt = Config.CLIENT.instance(ALT_INSTANCE_ID, LOCATION_ID)

bigtable/tests/unit/test_cluster.py

@@ -15,6 +15,8 @@
 
 import unittest
 
+import mock
+
 
 class TestCluster(unittest.TestCase):
 
@@ -232,7 +234,7 @@ def test_reload(self):
 
     def test_create(self):
         from google.longrunning import operations_pb2
-        from google.cloud.operation import Operation
+        from google.cloud.future import operation
         from google.cloud.bigtable._generated import (
             bigtable_instance_admin_pb2 as messages_v2_pb2)
         from tests.unit._testing import _FakeStub
@@ -256,13 +258,9 @@ def test_create(self):
         # Perform the method and check the result.
         result = cluster.create()
 
-        self.assertIsInstance(result, Operation)
-        self.assertEqual(result.name, OP_NAME)
-        self.assertIs(result.target, cluster)
-        self.assertIs(result.client, client)
+        self.assertIsInstance(result, operation.Operation)
+        self.assertEqual(result.operation.name, OP_NAME)
         self.assertIsNone(result.metadata)
-        self.assertEqual(result.caller_metadata,
-                         {'request_type': 'CreateCluster'})
 
         self.assertEqual(len(stub.method_calls), 1)
         api_name, args, kwargs = stub.method_calls[0]
@@ -278,7 +276,7 @@ def test_create(self):
     def test_update(self):
         import datetime
         from google.longrunning import operations_pb2
-        from google.cloud.operation import Operation
+        from google.cloud.future import operation
         from google.protobuf.any_pb2 import Any
         from google.cloud._helpers import _datetime_to_pb_timestamp
         from google.cloud.bigtable._generated import (
@@ -324,15 +322,11 @@ def test_update(self):
 
         result = cluster.update()
 
-        self.assertIsInstance(result, Operation)
-        self.assertEqual(result.name, OP_NAME)
-        self.assertIs(result.target, cluster)
-        self.assertIs(result.client, client)
+        self.assertIsInstance(result, operation.Operation)
+        self.assertEqual(result.operation.name, OP_NAME)
         self.assertIsInstance(result.metadata,
                               messages_v2_pb2.UpdateClusterMetadata)
         self.assertEqual(result.metadata.request_time, NOW_PB)
-        self.assertEqual(result.caller_metadata,
-                         {'request_type': 'UpdateCluster'})
 
         self.assertEqual(len(stub.method_calls), 1)
         api_name, args, kwargs = stub.method_calls[0]
@@ -448,6 +442,7 @@ class _Client(object):
     def __init__(self, project):
         self.project = project
         self.project_name = 'projects/' + self.project
+        self._operations_stub = mock.sentinel.operations_stub
 
     def __eq__(self, other):
         return (other.project == self.project and

bigtable/tests/unit/test_instance.py

@@ -15,6 +15,8 @@
 
 import unittest
 
+import mock
+
 
 class TestInstance(unittest.TestCase):
 
@@ -236,7 +238,7 @@ def test_create(self):
             bigtable_instance_admin_pb2 as messages_v2_pb2)
         from google.cloud._helpers import _datetime_to_pb_timestamp
         from tests.unit._testing import _FakeStub
-        from google.cloud.operation import Operation
+        from google.cloud.future import operation
         from google.cloud.bigtable.cluster import DEFAULT_SERVE_NODES
 
         NOW = datetime.datetime.utcnow()
@@ -263,15 +265,11 @@ def test_create(self):
         # Perform the method and check the result.
         result = instance.create()
 
-        self.assertIsInstance(result, Operation)
-        self.assertEqual(result.name, self.OP_NAME)
-        self.assertIs(result.target, instance)
-        self.assertIs(result.client, client)
+        self.assertIsInstance(result, operation.Operation)
+        self.assertEqual(result.operation.name, self.OP_NAME)
         self.assertIsInstance(result.metadata,
                               messages_v2_pb2.CreateInstanceMetadata)
         self.assertEqual(result.metadata.request_time, NOW_PB)
-        self.assertEqual(result.caller_metadata,
-                         {'request_type': 'CreateInstance'})
 
         self.assertEqual(len(stub.method_calls), 1)
         api_name, args, kwargs = stub.method_calls[0]
@@ -291,7 +289,7 @@ def test_create_w_explicit_serve_nodes(self):
         from google.cloud.bigtable._generated import (
             bigtable_instance_admin_pb2 as messages_v2_pb2)
         from tests.unit._testing import _FakeStub
-        from google.cloud.operation import Operation
+        from google.cloud.future import operation
 
         SERVE_NODES = 5
 
@@ -308,10 +306,8 @@ def test_create_w_explicit_serve_nodes(self):
         # Perform the method and check the result.
         result = instance.create()
 
-        self.assertIsInstance(result, Operation)
-        self.assertEqual(result.name, self.OP_NAME)
-        self.assertIs(result.target, instance)
-        self.assertIs(result.client, client)
+        self.assertIsInstance(result, operation.Operation)
+        self.assertEqual(result.operation.name, self.OP_NAME)
 
         self.assertEqual(len(stub.method_calls), 1)
         api_name, args, kwargs = stub.method_calls[0]
@@ -582,6 +578,7 @@ class _Client(object):
     def __init__(self, project):
         self.project = project
         self.project_name = 'projects/' + self.project
+        self._operations_stub = mock.sentinel.operations_stub
 
     def copy(self):
         from copy import deepcopy