diff --git a/bigtable/google/cloud/bigtable/cluster.py b/bigtable/google/cloud/bigtable/cluster.py index 48b335c5196e..c2418576dde9 100644 --- a/bigtable/google/cloud/bigtable/cluster.py +++ b/bigtable/google/cloud/bigtable/cluster.py @@ -22,8 +22,7 @@ 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 _compute_type_url -from google.cloud.operation import register_type_url +from google.cloud.operation import register_type _CLUSTER_NAME_RE = re.compile(r'^projects/(?P[^/]+)/' @@ -34,10 +33,7 @@ """Default number of nodes to use when creating a cluster.""" -_UPDATE_CLUSTER_METADATA_URL = _compute_type_url( - messages_v2_pb2.UpdateClusterMetadata) -register_type_url( - _UPDATE_CLUSTER_METADATA_URL, messages_v2_pb2.UpdateClusterMetadata) +register_type(messages_v2_pb2.UpdateClusterMetadata) def _prepare_create_request(cluster): diff --git a/bigtable/google/cloud/bigtable/instance.py b/bigtable/google/cloud/bigtable/instance.py index afa4066a75b0..41dda563c843 100644 --- a/bigtable/google/cloud/bigtable/instance.py +++ b/bigtable/google/cloud/bigtable/instance.py @@ -27,8 +27,7 @@ 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 _compute_type_url -from google.cloud.operation import register_type_url +from google.cloud.operation import register_type _EXISTING_INSTANCE_LOCATION_ID = 'see-existing-cluster' @@ -36,12 +35,8 @@ r'instances/(?P[a-z][-a-z0-9]*)$') -_CREATE_INSTANCE_METADATA_URL = _compute_type_url( - messages_v2_pb2.CreateInstanceMetadata) -register_type_url( - _CREATE_INSTANCE_METADATA_URL, messages_v2_pb2.CreateInstanceMetadata) -_INSTANCE_METADATA_URL = _compute_type_url(data_v2_pb2.Instance) -register_type_url(_INSTANCE_METADATA_URL, data_v2_pb2.Instance) +register_type(messages_v2_pb2.CreateInstanceMetadata) +register_type(data_v2_pb2.Instance) def _prepare_create_request(instance): diff --git a/bigtable/unit_tests/test_cluster.py b/bigtable/unit_tests/test_cluster.py index 82185cef030f..d55510ec5147 100644 --- a/bigtable/unit_tests/test_cluster.py +++ b/bigtable/unit_tests/test_cluster.py @@ -283,7 +283,6 @@ def test_update(self): from google.cloud.bigtable._generated import ( bigtable_instance_admin_pb2 as messages_v2_pb2) from unit_tests._testing import _FakeStub - from google.cloud.bigtable.cluster import _UPDATE_CLUSTER_METADATA_URL NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -307,10 +306,12 @@ def test_update(self): 'operations/projects/%s/instances/%s/clusters/%s/operations/%d' % (self.PROJECT, self.INSTANCE_ID, self.CLUSTER_ID, OP_ID)) metadata = messages_v2_pb2.UpdateClusterMetadata(request_time=NOW_PB) + type_url = 'type.googleapis.com/%s' % ( + messages_v2_pb2.UpdateClusterMetadata.DESCRIPTOR.full_name,) response_pb = operations_pb2.Operation( name=OP_NAME, metadata=Any( - type_url=_UPDATE_CLUSTER_METADATA_URL, + type_url=type_url, value=metadata.SerializeToString() ) ) diff --git a/bigtable/unit_tests/test_instance.py b/bigtable/unit_tests/test_instance.py index 223cfd2033ff..73d1cce7ff82 100644 --- a/bigtable/unit_tests/test_instance.py +++ b/bigtable/unit_tests/test_instance.py @@ -236,8 +236,6 @@ def test_create(self): from unit_tests._testing import _FakeStub from google.cloud.operation import Operation from google.cloud.bigtable.cluster import DEFAULT_SERVE_NODES - from google.cloud.bigtable.instance import ( - _CREATE_INSTANCE_METADATA_URL) NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -247,10 +245,12 @@ def test_create(self): # Create response_pb metadata = messages_v2_pb2.CreateInstanceMetadata(request_time=NOW_PB) + type_url = 'type.googleapis.com/%s' % ( + messages_v2_pb2.CreateInstanceMetadata.DESCRIPTOR.full_name,) response_pb = operations_pb2.Operation( name=self.OP_NAME, metadata=Any( - type_url=_CREATE_INSTANCE_METADATA_URL, + type_url=type_url, value=metadata.SerializeToString(), ) ) diff --git a/core/google/cloud/operation.py b/core/google/cloud/operation.py index 839e5e3eaf30..b98222143b4c 100644 --- a/core/google/cloud/operation.py +++ b/core/google/cloud/operation.py @@ -39,17 +39,20 @@ def _compute_type_url(klass, prefix=_GOOGLE_APIS_PREFIX): return '%s/%s' % (prefix, name) -def register_type_url(type_url, klass): +def register_type(klass, type_url=None): """Register a klass as the factory for a given type URL. - :type type_url: str - :param type_url: URL naming the type - :type klass: type :param klass: class to be used as a factory for the given type + :type type_url: str + :param type_url: (Optional) URL naming the type. If not provided, + infers the URL from the type descriptor. + :raises: ValueError if a registration already exists for the URL. """ + if type_url is None: + type_url = _compute_type_url(klass) if type_url in _TYPE_URL_MAP: if _TYPE_URL_MAP[type_url] is not klass: raise ValueError("Conflict: %s" % (_TYPE_URL_MAP[type_url],)) @@ -123,7 +126,7 @@ class Operation(object): """Metadata about the current operation (as a protobuf). Code that uses operations must register the metadata types (via - :func:`register_type_url`) to ensure that the metadata fields can be + :func:`register_type`) to ensure that the metadata fields can be converted into the correct types. """ diff --git a/core/unit_tests/test_operation.py b/core/unit_tests/test_operation.py index e67bf21f6f23..a30ec97f19ae 100644 --- a/core/unit_tests/test_operation.py +++ b/core/unit_tests/test_operation.py @@ -44,48 +44,63 @@ def test_w_prefix(self): '%s/%s' % (PREFIX, Struct.DESCRIPTOR.full_name)) -class Test_register_type_url(unittest.TestCase): +class Test_register_type(unittest.TestCase): - def _callFUT(self, type_url, klass): - from google.cloud.operation import register_type_url - register_type_url(type_url, klass) + def _callFUT(self, klass, type_url=None): + from google.cloud.operation import register_type + register_type(klass, type_url=type_url) - def test_simple(self): + def test_explicit(self): from google.cloud import operation as MUT from google.cloud._testing import _Monkey - TYPE_URI = 'testing.google-cloud-python.com/testing' + + type_url = 'testing.google-cloud-python.com/testing' klass = object() type_url_map = {} with _Monkey(MUT, _TYPE_URL_MAP=type_url_map): - self._callFUT(TYPE_URI, klass) + self._callFUT(klass, type_url) + + self.assertEqual(type_url_map, {type_url: klass}) - self.assertEqual(type_url_map, {TYPE_URI: klass}) + def test_default(self): + from google.protobuf.struct_pb2 import Struct + from google.cloud._testing import _Monkey + from google.cloud import operation as MUT + + type_url_map = {} + with _Monkey(MUT, _TYPE_URL_MAP=type_url_map): + self._callFUT(Struct) + + type_url = MUT._compute_type_url(Struct) + self.assertEqual(type_url_map, {type_url: Struct}) def test_w_same_class(self): from google.cloud import operation as MUT from google.cloud._testing import _Monkey - TYPE_URI = 'testing.google-cloud-python.com/testing' + + type_url = 'testing.google-cloud-python.com/testing' klass = object() - type_url_map = {TYPE_URI: klass} + type_url_map = {type_url: klass} with _Monkey(MUT, _TYPE_URL_MAP=type_url_map): - self._callFUT(TYPE_URI, klass) + self._callFUT(klass, type_url) - self.assertEqual(type_url_map, {TYPE_URI: klass}) + self.assertEqual(type_url_map, {type_url: klass}) def test_w_conflict(self): from google.cloud import operation as MUT from google.cloud._testing import _Monkey - TYPE_URI = 'testing.google-cloud-python.com/testing' + + type_url = 'testing.google-cloud-python.com/testing' klass, other = object(), object() - type_url_map = {TYPE_URI: other} + type_url_map = {type_url: other} with _Monkey(MUT, _TYPE_URL_MAP=type_url_map): with self.assertRaises(ValueError): - self._callFUT(TYPE_URI, klass) + self._callFUT(klass, type_url) - self.assertEqual(type_url_map, {TYPE_URI: other}) + self.assertEqual(type_url_map, {type_url: other}) class OperationTests(unittest.TestCase):