Merge pull request #1913 from tseaver/bigtable-v2-column_family_crud

Update bigtable.column_family to use V2 protos

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>`