BigTable: Add split keys on create table - v2 (googleapis#5513)

aneepct · tseaver · commit 941a18aad07b · 2018-07-11T16:47:33.000-04:00
diff --git a/bigtable/google/cloud/bigtable/table.py b/bigtable/google/cloud/bigtable/table.py
@@ -22,15 +22,17 @@
 from google.api_core.retry import if_exception_type
 from google.api_core.retry import Retry
 from google.cloud._helpers import _to_bytes
-from google.cloud.bigtable_v2.proto import (
-    bigtable_pb2 as data_messages_v2_pb2)
 from google.cloud.bigtable.column_family import _gc_rule_from_pb
 from google.cloud.bigtable.column_family import ColumnFamily
 from google.cloud.bigtable.row import AppendRow
 from google.cloud.bigtable.row import ConditionalRow
 from google.cloud.bigtable.row import DirectRow
 from google.cloud.bigtable.row_data import PartialRowsData
 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 (
+    bigtable_table_admin_pb2 as table_admin_messages_v2_pb2)
 
 
 # Maximum number of mutations in bulk (MutateRowsRequest message):
@@ -175,19 +177,35 @@ def __eq__(self, other):
     def __ne__(self, other):
         return not self == other
 
-    def create(self):
+    def create(self, initial_split_keys=None):
         """Creates this table.
 
         .. note::
 
             A create request returns a
             :class:`._generated.table_pb2.Table` but we don't use
             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.
         """
         table_client = self._instance._client.table_admin_client
         instance_name = self._instance.name
-        table_client.create_table(
-            parent=instance_name, table_id=self.table_id, table={})
+
+        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
+
+        table_client.create_table(parent=instance_name,
+                                  table_id=self.table_id, table={},
+                                  initial_splits=splits)
 
     def exists(self):
         """Check whether the table exists.
diff --git a/bigtable/tests/system.py b/bigtable/tests/system.py
@@ -242,6 +242,19 @@ def test_create_table(self):
         sorted_tables = sorted(tables, key=name_attr)
         self.assertEqual(sorted_tables, expected_tables)
 
+    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',
+                              b'split_key_20', b'']
+        temp_table = Config.INSTANCE.table(temp_table_id)
+        temp_table.create(initial_split_keys=initial_split_keys)
+        self.tables_to_delete.append(temp_table)
+
+        # Read Sample Row Keys for created splits
+        sample_row_keys = temp_table.sample_row_keys()
+        self.assertEqual(set([srk.row_key for srk in sample_row_keys]),
+                         set(initial_split_keys))
+
     def test_create_column_family(self):
         temp_table_id = 'test-create-column-family'
         temp_table = Config.INSTANCE.table(temp_table_id)
diff --git a/bigtable/tests/unit/test_table.py b/bigtable/tests/unit/test_table.py
@@ -282,6 +282,44 @@ def _create_test_helper(self):
     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)
+        from google.cloud.bigtable_admin_v2.proto import (
+            bigtable_table_admin_pb2 as table_admin_messages_v2_pb2)
+
+        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)
+
+        splits = []
+        for split_key in split_keys:
+            splits.append(
+                table_admin_messages_v2_pb2.CreateTableRequest.Split(
+                    key=split_key))
+
+        table_api.create_table.assert_called_once_with(
+            parent=self.INSTANCE_NAME,
+            table={},
+            table_id=self.TABLE_ID,
+            initial_splits=splits)
+
     def test_exists(self):
         from google.cloud.bigtable_admin_v2.proto import (
             table_pb2 as table_data_v2_pb2)
