Repository: googleapis/google-cloud-python

Issue:
Bigtable: createTable should allow creation of column families
Table.create() should allow the user to create column families when the table itself is created.  Users should be able to supply a Column Family Id => GcRule policy mapping to accomplish this feature.


Discussion:

Resolving patch:
diff --git a/bigtable/google/cloud/bigtable/table.py b/bigtable/google/cloud/bigtable/table.py
--- a/bigtable/google/cloud/bigtable/table.py
+++ b/bigtable/google/cloud/bigtable/table.py
@@ -31,6 +31,8 @@
 from google.cloud.bigtable.row_data import YieldRowsData
 from google.cloud.bigtable_v2.proto import (
     bigtable_pb2 as data_messages_v2_pb2)
+from google.cloud.bigtable_admin_v2.proto import (
+    table_pb2 as admin_messages_v2_pb2)
 from google.cloud.bigtable_admin_v2.proto import (
     bigtable_table_admin_pb2 as table_admin_messages_v2_pb2)
 
@@ -88,7 +90,7 @@ class Table(object):
     :type instance: :class:`~google.cloud.bigtable.instance.Instance`
     :param instance: The instance that owns the table.
 
-    :type: app_profile_id: str
+    :type app_profile_id: str
     :param app_profile_id: (Optional) The unique name of the AppProfile.
     """
 
@@ -177,7 +179,7 @@ def __eq__(self, other):
     def __ne__(self, other):
         return not self == other
 
-    def create(self, initial_split_keys=None):
+    def create(self, initial_split_keys=[], column_families={}):
         """Creates this table.
 
         .. note::
@@ -187,25 +189,27 @@ def create(self, initial_split_keys=None):
             this response.
 
         :type initial_split_keys: list
-        :param initial_split_keys: The optional list of row keys in bytes that
-                                    will be used to initially split the table
-                                    into several tablets.
+        :param initial_split_keys: (Optional) list of row keys in bytes that
+                                   will be used to initially split the table
+                                   into several tablets.
+
+        :type column_families: dict
+        :param column_failies: (Optional) A map columns to create.  The key is
+                               the column_id str and the value is a
+                               :class:`GarbageCollectionRule`
         """
         table_client = self._instance._client.table_admin_client
         instance_name = self._instance.name
 
-        if initial_split_keys is not None:
-            splits = []
-            for initial_split_key in initial_split_keys:
-                splits.append(
-                    table_admin_messages_v2_pb2.CreateTableRequest.Split(
-                        key=initial_split_key))
-        else:
-            splits = None
+        families = {id: ColumnFamily(id, self, rule).to_pb()
+                    for (id, rule) in column_families.items()}
+        table = admin_messages_v2_pb2.Table(column_families=families)
+
+        split = table_admin_messages_v2_pb2.CreateTableRequest.Split
+        splits = [split(key=key) for key in initial_split_keys]
 
-        table_client.create_table(parent=instance_name,
-                                  table_id=self.table_id, table={},
-                                  initial_splits=splits)
+        table_client.create_table(parent=instance_name, table_id=self.table_id,
+                                  table=table, initial_splits=splits)
 
     def exists(self):
         """Check whether the table exists.


Test patch:
diff --git a/bigtable/tests/system.py b/bigtable/tests/system.py
--- a/bigtable/tests/system.py
+++ b/bigtable/tests/system.py
@@ -242,6 +242,22 @@ def test_create_table(self):
         sorted_tables = sorted(tables, key=name_attr)
         self.assertEqual(sorted_tables, expected_tables)
 
+    def test_create_table_with_families(self):
+        temp_table_id = 'test-create-table-with-failies'
+        temp_table = Config.INSTANCE.table(temp_table_id)
+        gc_rule = MaxVersionsGCRule(1)
+        temp_table.create(column_families={COLUMN_FAMILY_ID1: gc_rule})
+        self.tables_to_delete.append(temp_table)
+
+        col_fams = temp_table.list_column_families()
+
+        self.assertEqual(len(col_fams), 1)
+        retrieved_col_fam = col_fams[COLUMN_FAMILY_ID1]
+        self.assertIs(retrieved_col_fam._table, temp_table)
+        self.assertEqual(retrieved_col_fam.column_family_id,
+                         COLUMN_FAMILY_ID1)
+        self.assertEqual(retrieved_col_fam.gc_rule, gc_rule)
+
     def test_create_table_with_split_keys(self):
         temp_table_id = 'foo-bar-baz-split-table'
         initial_split_keys = [b'split_key_1', b'split_key_10',
diff --git a/bigtable/tests/unit/test_table.py b/bigtable/tests/unit/test_table.py
--- a/bigtable/tests/unit/test_table.py
+++ b/bigtable/tests/unit/test_table.py
@@ -254,72 +254,53 @@ def test___ne__(self):
         table2 = self._make_one('table_id2', None)
         self.assertNotEqual(table1, table2)
 
-    def _create_test_helper(self):
+    def _create_test_helper(self, split_keys=[], column_families={}):
         from google.cloud.bigtable_admin_v2.gapic import (
-            bigtable_instance_admin_client, bigtable_table_admin_client)
-
-        table_api = mock.create_autospec(
-            bigtable_table_admin_client.BigtableTableAdminClient)
-        instance_api = mock.create_autospec(
-            bigtable_instance_admin_client.BigtableInstanceAdminClient)
-        credentials = _make_credentials()
-        client = self._make_client(project='project-id',
-                                   credentials=credentials, admin=True)
-        instance = client.instance(instance_id=self.INSTANCE_ID)
-        table = self._make_one(self.TABLE_ID, instance)
-
-        # Patch API calls
-        client._table_admin_client = table_api
-        client._instance_admin_client = instance_api
-
-        # Create expected_result.
-        expected_result = None  # create() has no return value.
-
-        # Perform the method and check the result.
-        result = table.create()
-        self.assertEqual(result, expected_result)
-
-    def test_create(self):
-        self._create_test_helper()
-
-    def test_create_with_split_keys(self):
-        from google.cloud.bigtable_admin_v2.gapic import (
-            bigtable_instance_admin_client, bigtable_table_admin_client)
+            bigtable_table_admin_client)
+        from google.cloud.bigtable_admin_v2.proto import table_pb2
         from google.cloud.bigtable_admin_v2.proto import (
             bigtable_table_admin_pb2 as table_admin_messages_v2_pb2)
+        from google.cloud.bigtable.column_family import ColumnFamily
 
         table_api = mock.create_autospec(
             bigtable_table_admin_client.BigtableTableAdminClient)
-        instance_api = (
-            bigtable_instance_admin_client.BigtableInstanceAdminClient(
-                mock.Mock()))
         credentials = _make_credentials()
         client = self._make_client(project='project-id',
                                    credentials=credentials, admin=True)
         instance = client.instance(instance_id=self.INSTANCE_ID)
         table = self._make_one(self.TABLE_ID, instance)
 
-        split_keys = [b'split1', b'split2', b'split3']
-
         # Patch API calls
         client._table_admin_client = table_api
-        client._instance_admin_client = instance_api
 
         # Perform the method and check the result.
-        table.create(split_keys)
+        table.create(column_families=column_families,
+                     initial_split_keys=split_keys)
+
+        families = {id: ColumnFamily(id, self, rule).to_pb()
+                    for (id, rule) in column_families.items()}
 
-        splits = []
-        for split_key in split_keys:
-            splits.append(
-                table_admin_messages_v2_pb2.CreateTableRequest.Split(
-                    key=split_key))
+        split = table_admin_messages_v2_pb2.CreateTableRequest.Split
+        splits = [split(key=split_key) for split_key in split_keys]
 
         table_api.create_table.assert_called_once_with(
             parent=self.INSTANCE_NAME,
-            table={},
+            table=table_pb2.Table(column_families=families),
             table_id=self.TABLE_ID,
             initial_splits=splits)
 
+    def test_create(self):
+        self._create_test_helper()
+
+    def test_create_with_families(self):
+        from google.cloud.bigtable.column_family import MaxVersionsGCRule
+
+        families = {"family": MaxVersionsGCRule(5)}
+        self._create_test_helper(column_families=families)
+
+    def test_create_with_split_keys(self):
+        self._create_test_helper(split_keys=[b'split1', b'split2', b'split3'])
+
     def test_exists(self):
         from google.cloud.bigtable_admin_v2.proto import (
             table_pb2 as table_data_v2_pb2)