Update bigtable.column_family to use V2 protos.

Note that {Create,Update,Delete}ColumnFamily messages all collapse to
ModifyColumnFamilies.

Author: tseaver

gcloud/bigtable/column_family.py

@@ -20,10 +20,10 @@
 from google.protobuf import duration_pb2
 
 from gcloud._helpers import _total_seconds
-from gcloud.bigtable._generated import (
-    bigtable_table_data_pb2 as data_v1_pb2)
-from gcloud.bigtable._generated import (
-    bigtable_table_service_messages_pb2 as messages_v1_pb2)
+from gcloud.bigtable._generated_v2 import (
+    table_pb2 as table_v2_pb2)
+from gcloud.bigtable._generated_v2 import (
+    bigtable_table_admin_pb2 as table_admin_v2_pb2)
 
 
 def _timedelta_to_duration_pb(timedelta_val):
@@ -111,10 +111,10 @@ def __eq__(self, other):
     def to_pb(self):
         """Converts the garbage collection rule to a protobuf.
 
-        :rtype: :class:`.data_v1_pb2.GcRule`
+        :rtype: :class:`.table_v2_pb2.GcRule`
         :returns: The converted current object.
         """
-        return data_v1_pb2.GcRule(max_num_versions=self.max_num_versions)
+        return table_v2_pb2.GcRule(max_num_versions=self.max_num_versions)
 
 
 class MaxAgeGCRule(GarbageCollectionRule):
@@ -135,11 +135,11 @@ def __eq__(self, other):
     def to_pb(self):
         """Converts the garbage collection rule to a protobuf.
 
-        :rtype: :class:`.data_v1_pb2.GcRule`
+        :rtype: :class:`.table_v2_pb2.GcRule`
         :returns: The converted current object.
         """
         max_age = _timedelta_to_duration_pb(self.max_age)
-        return data_v1_pb2.GcRule(max_age=max_age)
+        return table_v2_pb2.GcRule(max_age=max_age)
 
 
 class GCRuleUnion(GarbageCollectionRule):
@@ -160,12 +160,12 @@ def __eq__(self, other):
     def to_pb(self):
         """Converts the union into a single GC rule as a protobuf.
 
-        :rtype: :class:`.data_v1_pb2.GcRule`
+        :rtype: :class:`.table_v2_pb2.GcRule`
         :returns: The converted current object.
         """
-        union = data_v1_pb2.GcRule.Union(
+        union = table_v2_pb2.GcRule.Union(
             rules=[rule.to_pb() for rule in self.rules])
-        return data_v1_pb2.GcRule(union=union)
+        return table_v2_pb2.GcRule(union=union)
 
 
 class GCRuleIntersection(GarbageCollectionRule):
@@ -186,12 +186,12 @@ def __eq__(self, other):
     def to_pb(self):
         """Converts the intersection into a single GC rule as a protobuf.
 
-        :rtype: :class:`.data_v1_pb2.GcRule`
+        :rtype: :class:`.table_v2_pb2.GcRule`
         :returns: The converted current object.
         """
-        intersection = data_v1_pb2.GcRule.Intersection(
+        intersection = table_v2_pb2.GcRule.Intersection(
             rules=[rule.to_pb() for rule in self.rules])
-        return data_v1_pb2.GcRule(intersection=intersection)
+        return table_v2_pb2.GcRule(intersection=intersection)
 
 
 class ColumnFamily(object):
@@ -251,21 +251,22 @@ def __ne__(self, other):
     def create(self):
         """Create this column family."""
         if self.gc_rule is None:
-            column_family = data_v1_pb2.ColumnFamily()
+            column_family = table_v2_pb2.ColumnFamily()
         else:
-            column_family = data_v1_pb2.ColumnFamily(
+            column_family = table_v2_pb2.ColumnFamily(
                 gc_rule=self.gc_rule.to_pb())
-        request_pb = messages_v1_pb2.CreateColumnFamilyRequest(
-            name=self._table.name,
-            column_family_id=self.column_family_id,
-            column_family=column_family,
+        request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest(
+            name=self._table.name)
+        request_pb.modifications.add(
+            id=self.column_family_id,
+            create=column_family,
         )
         client = self._table._cluster._client
-        # We expect a `.data_v1_pb2.ColumnFamily`. We ignore it since the only
+        # 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.
-        client._table_stub.CreateColumnFamily(request_pb,
-                                              client.timeout_seconds)
+        client._table_stub.ModifyColumnFamilies(request_pb,
+                                                client.timeout_seconds)
 
     def update(self):
         """Update this column family.
@@ -275,30 +276,40 @@ def update(self):
             Only the GC rule can be updated. By changing the column family ID,
             you will simply be referring to a different column family.
         """
-        request_kwargs = {'name': self.name}
-        if self.gc_rule is not None:
-            request_kwargs['gc_rule'] = self.gc_rule.to_pb()
-        request_pb = data_v1_pb2.ColumnFamily(**request_kwargs)
+        if self.gc_rule is None:
+            column_family = table_v2_pb2.ColumnFamily()
+        else:
+            column_family = table_v2_pb2.ColumnFamily(
+                gc_rule=self.gc_rule.to_pb())
+        request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest(
+            name=self._table.name)
+        request_pb.modifications.add(
+            id=self.column_family_id,
+            update=column_family)
         client = self._table._cluster._client
-        # We expect a `.data_v1_pb2.ColumnFamily`. We ignore it since the only
+        # 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.
-        client._table_stub.UpdateColumnFamily(request_pb,
-                                              client.timeout_seconds)
+        client._table_stub.ModifyColumnFamilies(request_pb,
+                                                client.timeout_seconds)
 
     def delete(self):
         """Delete this column family."""
-        request_pb = messages_v1_pb2.DeleteColumnFamilyRequest(name=self.name)
+        request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest(
+            name=self._table.name)
+        request_pb.modifications.add(
+            id=self.column_family_id,
+            drop=True)
         client = self._table._cluster._client
         # We expect a `google.protobuf.empty_pb2.Empty`
-        client._table_stub.DeleteColumnFamily(request_pb,
-                                              client.timeout_seconds)
+        client._table_stub.ModifyColumnFamilies(request_pb,
+                                                client.timeout_seconds)
 
 
 def _gc_rule_from_pb(gc_rule_pb):
     """Convert a protobuf GC rule to a native object.
 
-    :type gc_rule_pb: :class:`.data_v1_pb2.GcRule`
+    :type gc_rule_pb: :class:`.table_v2_pb2.GcRule`
     :param gc_rule_pb: The GC rule to convert.
 
     :rtype: :class:`GarbageCollectionRule` or :data:`NoneType <types.NoneType>`

gcloud/bigtable/test_column_family.py

@@ -389,8 +389,8 @@ def test___ne__(self):
         self.assertNotEqual(column_family1, column_family2)
 
     def _create_test_helper(self, gc_rule=None):
-        from gcloud.bigtable._generated import (
-            bigtable_table_service_messages_pb2 as messages_v1_pb2)
+        from gcloud.bigtable._generated_v2 import (
+            bigtable_table_admin_pb2 as table_admin_v2_pb2)
         from gcloud.bigtable._testing import _FakeStub
 
         project_id = 'project-id'
@@ -411,12 +411,12 @@ def _create_test_helper(self, gc_rule=None):
         if gc_rule is None:
             column_family_pb = _ColumnFamilyPB()
         else:
-            column_family_pb = _ColumnFamilyPB(
-                gc_rule=gc_rule.to_pb())
-        request_pb = messages_v1_pb2.CreateColumnFamilyRequest(
-            name=table_name,
-            column_family_id=column_family_id,
-            column_family=column_family_pb,
+            column_family_pb = _ColumnFamilyPB(gc_rule=gc_rule.to_pb())
+        request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest(
+            name=table_name)
+        request_pb.modifications.add(
+            id=column_family_id,
+            create=column_family_pb,
         )
 
         # Create response_pb
@@ -434,7 +434,7 @@ def _create_test_helper(self, gc_rule=None):
         self.assertEqual(stub.results, ())
         self.assertEqual(result, expected_result)
         self.assertEqual(stub.method_calls, [(
-            'CreateColumnFamily',
+            'ModifyColumnFamilies',
             (request_pb, timeout_seconds),
             {},
         )])
@@ -449,6 +449,8 @@ def test_create_with_gc_rule(self):
 
     def _update_test_helper(self, gc_rule=None):
         from gcloud.bigtable._testing import _FakeStub
+        from gcloud.bigtable._generated_v2 import (
+            bigtable_table_admin_pb2 as table_admin_v2_pb2)
 
         project_id = 'project-id'
         zone = 'zone'
@@ -458,8 +460,6 @@ def _update_test_helper(self, gc_rule=None):
         timeout_seconds = 28
         table_name = ('projects/' + project_id + '/zones/' + zone +
                       '/clusters/' + cluster_id + '/tables/' + table_id)
-        column_family_name = (
-            table_name + '/columnFamilies/' + column_family_id)
 
         client = _Client(timeout_seconds=timeout_seconds)
         table = _Table(table_name, client=client)
@@ -468,12 +468,15 @@ def _update_test_helper(self, gc_rule=None):
 
         # Create request_pb
         if gc_rule is None:
-            request_pb = _ColumnFamilyPB(name=column_family_name)
+            column_family_pb = _ColumnFamilyPB()
         else:
-            request_pb = _ColumnFamilyPB(
-                name=column_family_name,
-                gc_rule=gc_rule.to_pb(),
-            )
+            column_family_pb = _ColumnFamilyPB(gc_rule=gc_rule.to_pb())
+        request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest(
+            name=table_name)
+        request_pb.modifications.add(
+            id=column_family_id,
+            update=column_family_pb,
+        )
 
         # Create response_pb
         response_pb = _ColumnFamilyPB()
@@ -490,7 +493,7 @@ def _update_test_helper(self, gc_rule=None):
         self.assertEqual(stub.results, ())
         self.assertEqual(result, expected_result)
         self.assertEqual(stub.method_calls, [(
-            'UpdateColumnFamily',
+            'ModifyColumnFamilies',
             (request_pb, timeout_seconds),
             {},
         )])
@@ -505,8 +508,8 @@ def test_update_with_gc_rule(self):
 
     def test_delete(self):
         from google.protobuf import empty_pb2
-        from gcloud.bigtable._generated import (
-            bigtable_table_service_messages_pb2 as messages_v1_pb2)
+        from gcloud.bigtable._generated_v2 import (
+            bigtable_table_admin_pb2 as table_admin_v2_pb2)
         from gcloud.bigtable._testing import _FakeStub
 
         project_id = 'project-id'
@@ -517,16 +520,17 @@ def test_delete(self):
         timeout_seconds = 7
         table_name = ('projects/' + project_id + '/zones/' + zone +
                       '/clusters/' + cluster_id + '/tables/' + table_id)
-        column_family_name = (
-            table_name + '/columnFamilies/' + column_family_id)
 
         client = _Client(timeout_seconds=timeout_seconds)
         table = _Table(table_name, client=client)
         column_family = self._makeOne(column_family_id, table)
 
         # Create request_pb
-        request_pb = messages_v1_pb2.DeleteColumnFamilyRequest(
-            name=column_family_name)
+        request_pb = table_admin_v2_pb2.ModifyColumnFamiliesRequest(
+            name=table_name)
+        request_pb.modifications.add(
+            id=column_family_id,
+            drop=True)
 
         # Create response_pb
         response_pb = empty_pb2.Empty()
@@ -543,7 +547,7 @@ def test_delete(self):
         self.assertEqual(stub.results, ())
         self.assertEqual(result, expected_result)
         self.assertEqual(stub.method_calls, [(
-            'DeleteColumnFamily',
+            'ModifyColumnFamilies',
             (request_pb, timeout_seconds),
             {},
         )])
@@ -623,27 +627,27 @@ def WhichOneof(cls, name):
 
 
 def _GcRulePB(*args, **kw):
-    from gcloud.bigtable._generated import (
-        bigtable_table_data_pb2 as data_v1_pb2)
-    return data_v1_pb2.GcRule(*args, **kw)
+    from gcloud.bigtable._generated_v2 import (
+        table_pb2 as table_v2_pb2)
+    return table_v2_pb2.GcRule(*args, **kw)
 
 
 def _GcRuleIntersectionPB(*args, **kw):
-    from gcloud.bigtable._generated import (
-        bigtable_table_data_pb2 as data_v1_pb2)
-    return data_v1_pb2.GcRule.Intersection(*args, **kw)
+    from gcloud.bigtable._generated_v2 import (
+        table_pb2 as table_v2_pb2)
+    return table_v2_pb2.GcRule.Intersection(*args, **kw)
 
 
 def _GcRuleUnionPB(*args, **kw):
-    from gcloud.bigtable._generated import (
-        bigtable_table_data_pb2 as data_v1_pb2)
-    return data_v1_pb2.GcRule.Union(*args, **kw)
+    from gcloud.bigtable._generated_v2 import (
+        table_pb2 as table_v2_pb2)
+    return table_v2_pb2.GcRule.Union(*args, **kw)
 
 
 def _ColumnFamilyPB(*args, **kw):
-    from gcloud.bigtable._generated import (
-        bigtable_table_data_pb2 as data_v1_pb2)
-    return data_v1_pb2.ColumnFamily(*args, **kw)
+    from gcloud.bigtable._generated_v2 import (
+        table_pb2 as table_v2_pb2)
+    return table_v2_pb2.ColumnFamily(*args, **kw)
 
 
 class _Cluster(object):