BigTable: Add data app profile id

bigtable/google/cloud/bigtable/client.py

@@ -170,6 +170,7 @@ def _instance_admin_client(self):
         """
         if not self._admin:
             raise ValueError('Client is not an admin client.')
+
         return bigtable_admin_v2.BigtableInstanceAdminClient(
             channel=self._channel, credentials=self._credentials)
 

bigtable/google/cloud/bigtable/table.py

@@ -84,11 +84,15 @@ class Table(object):
 
     :type instance: :class:`~google.cloud.bigtable.instance.Instance`
     :param instance: The instance that owns the table.
+
+    :type: app_profile_id: str
+    :param app_profile_id: (Optional) The unique name of the AppProfile.
     """
 
-    def __init__(self, table_id, instance):
+    def __init__(self, table_id, instance, app_profile_id=None):
         self.table_id = table_id
         self._instance = instance
+        self._app_profile_id = app_profile_id
 
     @property
     def name(self):
@@ -227,8 +231,9 @@ def read_row(self, row_key, filter_=None):
         :raises: :class:`ValueError <exceptions.ValueError>` if a commit row
                  chunk is never encountered.
         """
-        request_pb = _create_row_request(self.name, row_key=row_key,
-                                         filter_=filter_)
+        request_pb = _create_row_request(
+            self.name, row_key=row_key, filter_=filter_,
+            app_profile_id=self._app_profile_id)
         client = self._instance._client
         rows_data = PartialRowsData(client._table_data_client._read_rows,
                                     request_pb)
@@ -276,7 +281,8 @@ def read_rows(self, start_key=None, end_key=None, limit=None,
         """
         request_pb = _create_row_request(
             self.name, start_key=start_key, end_key=end_key, filter_=filter_,
-            limit=limit, end_inclusive=end_inclusive)
+            limit=limit, end_inclusive=end_inclusive,
+            app_profile_id=self._app_profile_id)
         client = self._instance._client
         return PartialRowsData(client._table_data_client._read_rows,
                                request_pb)
@@ -310,7 +316,7 @@ def yield_rows(self, start_key=None, end_key=None, limit=None,
         """
         request_pb = _create_row_request(
             self.name, start_key=start_key, end_key=end_key, filter_=filter_,
-            limit=limit)
+            limit=limit, app_profile_id=self._app_profile_id)
         client = self._instance._client
         generator = YieldRowsData(client._table_data_client._read_rows,
                                   request_pb)
@@ -347,7 +353,8 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY):
                   sent. These will be in the same order as the `rows`.
         """
         retryable_mutate_rows = _RetryableMutateRowsWorker(
-            self._instance._client, self.name, rows)
+            self._instance._client, self.name, rows,
+            app_profile_id=self._app_profile_id)
         return retryable_mutate_rows(retry=retry)
 
     def sample_row_keys(self):
@@ -383,9 +390,50 @@ def sample_row_keys(self):
         """
         client = self._instance._client
         response_iterator = client._table_data_client.sample_row_keys(
-            self.name)
+            self.name, app_profile_id=self._app_profile_id)
         return response_iterator
 
+    def truncate(self, timeout=60):
+        """Truncate the table
+
+        :type timeout: float
+        :param timeout: (Optional) The amount of time, in seconds, to wait
+                        for the request to complete.
+
+        :raise: google.api_core.exceptions.GoogleAPICallError: If the
+                request failed for any reason.
+                google.api_core.exceptions.RetryError: If the request failed
+                due to a retryable error and retry attempts failed.
+                ValueError: If the parameters are invalid.
+        """
+        client = self._instance._client
+        table_admin_client = client._table_admin_client
+        table_admin_client.drop_row_range(self.name,
+                                          delete_all_data_from_table=True,
+                                          timeout=timeout)
+
+    def drop_by_prefix(self, row_key_prefix, timeout=60):
+        """
+        :type row_prefix: bytes
+        :param row_prefix: Delete all rows that start with this row key
+                            prefix. Prefix cannot be zero length.
+
+        :type timeout: float
+        :param timeout: (Optional) The amount of time, in seconds, to wait
+                        for the request to complete.
+
+        :raise: google.api_core.exceptions.GoogleAPICallError: If the
+                request failed for any reason.
+                google.api_core.exceptions.RetryError: If the request failed
+                due to a retryable error and retry attempts failed.
+                ValueError: If the parameters are invalid.
+        """
+        client = self._instance._client
+        table_admin_client = client._table_admin_client
+        table_admin_client.drop_row_range(
+            self.name, row_key_prefix=row_key_prefix.encode('utf-8'),
+            timeout=timeout)
+
 
 class _RetryableMutateRowsWorker(object):
     """A callable worker that can retry to mutate rows with transient errors.
@@ -403,10 +451,11 @@ class _RetryableMutateRowsWorker(object):
     )
     # pylint: enable=unsubscriptable-object
 
-    def __init__(self, client, table_name, rows):
+    def __init__(self, client, table_name, rows, app_profile_id=None):
         self.client = client
         self.table_name = table_name
         self.rows = rows
+        self.app_profile_id = app_profile_id
         self.responses_statuses = [None] * len(self.rows)
 
     def __call__(self, retry=DEFAULT_RETRY):
@@ -468,7 +517,8 @@ def _do_mutate_retryable_rows(self):
             return self.responses_statuses
 
         mutate_rows_request = _mutate_rows_request(
-            self.table_name, retryable_rows)
+            self.table_name, retryable_rows,
+            app_profile_id=self.app_profile_id)
         responses = self.client._table_data_client._mutate_rows(
             mutate_rows_request, retry=None)
 
@@ -496,7 +546,8 @@ def _do_mutate_retryable_rows(self):
 
 
 def _create_row_request(table_name, row_key=None, start_key=None, end_key=None,
-                        filter_=None, limit=None, end_inclusive=False):
+                        filter_=None, limit=None, end_inclusive=False,
+                        app_profile_id=None):
     """Creates a request to read rows in a table.
 
     :type table_name: str
@@ -528,6 +579,9 @@ def _create_row_request(table_name, row_key=None, start_key=None, end_key=None,
     :param end_inclusive: (Optional) Whether the ``end_key`` should be
                   considered inclusive. The default is False (exclusive).
 
+    :type: app_profile_id: str
+    :param app_profile_id: (Optional) The unique name of the AppProfile.
+
     :rtype: :class:`data_messages_v2_pb2.ReadRowsRequest`
     :returns: The ``ReadRowsRequest`` protobuf corresponding to the inputs.
     :raises: :class:`ValueError <exceptions.ValueError>` if both
@@ -551,6 +605,8 @@ def _create_row_request(table_name, row_key=None, start_key=None, end_key=None,
         request_kwargs['filter'] = filter_.to_pb()
     if limit is not None:
         request_kwargs['rows_limit'] = limit
+    if app_profile_id is not None:
+        request_kwargs['app_profile_id'] = app_profile_id
 
     message = data_messages_v2_pb2.ReadRowsRequest(**request_kwargs)
 
@@ -563,7 +619,7 @@ def _create_row_request(table_name, row_key=None, start_key=None, end_key=None,
     return message
 
 
-def _mutate_rows_request(table_name, rows):
+def _mutate_rows_request(table_name, rows, app_profile_id=None):
     """Creates a request to mutate rows in a table.
 
     :type table_name: str
@@ -572,12 +628,16 @@ def _mutate_rows_request(table_name, rows):
     :type rows: list
     :param rows: List or other iterable of :class:`.DirectRow` instances.
 
+    :type: app_profile_id: str
+    :param app_profile_id: (Optional) The unique name of the AppProfile.
+
     :rtype: :class:`data_messages_v2_pb2.MutateRowsRequest`
     :returns: The ``MutateRowsRequest`` protobuf corresponding to the inputs.
     :raises: :exc:`~.table.TooManyMutationsError` if the number of mutations is
              greater than 100,000
     """
-    request_pb = data_messages_v2_pb2.MutateRowsRequest(table_name=table_name)
+    request_pb = data_messages_v2_pb2.MutateRowsRequest(
+        table_name=table_name, app_profile_id=app_profile_id)
     mutations_count = 0
     for row in rows:
         _check_row_table_name(table_name, row)

bigtable/tests/unit/test_table.py

@@ -337,7 +337,8 @@ def _read_row_helper(self, chunks, expected_result):
         request_pb = object()  # Returned by our mock.
         mock_created = []
 
-        def mock_create_row_request(table_name, row_key, filter_):
+        def mock_create_row_request(table_name, row_key, filter_,
+                                    app_profile_id=None):
             mock_created.append((table_name, row_key, filter_))
             return request_pb
 
@@ -462,6 +463,7 @@ def mock_create_row_request(table_name, **kwargs):
             'filter_': filter_obj,
             'limit': limit,
             'end_inclusive': False,
+            'app_profile_id': None
         }
         self.assertEqual(mock_created, [(table.name, created_kwargs)])
 
@@ -531,6 +533,34 @@ def test_sample_row_keys(self):
         result = table.sample_row_keys()
         self.assertEqual(result[0], expected_result)
 
+    def test_truncate(self):
+        channel = self._make_channel()
+        client = self._make_client(project='project-id', channel=channel,
+                                   admin=True)
+        instance = client.instance(instance_id=self.INSTANCE_ID)
+        table = self._make_one(self.TABLE_ID, instance)
+
+        expected_result = None  # truncate() has no return value.
+
+        result = table.truncate()
+
+        self.assertEqual(result, expected_result)
+
+    def test_drop_by_prefix(self):
+        channel = self._make_channel()
+        client = self._make_client(project='project-id', channel=channel,
+                                   admin=True)
+        instance = client.instance(instance_id=self.INSTANCE_ID)
+        table = self._make_one(self.TABLE_ID, instance)
+
+        expected_result = None  # drop_by_prefix() has no return value.
+
+        row_key_prefix = 'row-key-prefix'
+
+        result = table.drop_by_prefix(row_key_prefix=row_key_prefix)
+
+        self.assertEqual(result, expected_result)
+
 
 class Test__RetryableMutateRowsWorker(unittest.TestCase):
     from grpc import StatusCode
@@ -1014,12 +1044,14 @@ def test_do_mutate_retryable_rows_mismatch_num_responses(self):
 class Test__create_row_request(unittest.TestCase):
 
     def _call_fut(self, table_name, row_key=None, start_key=None, end_key=None,
-                  filter_=None, limit=None, end_inclusive=False):
+                  filter_=None, limit=None, end_inclusive=False,
+                  app_profile_id=None):
         from google.cloud.bigtable.table import _create_row_request
 
         return _create_row_request(
             table_name, row_key=row_key, start_key=start_key, end_key=end_key,
-            filter_=filter_, limit=limit, end_inclusive=end_inclusive)
+            filter_=filter_, limit=limit, end_inclusive=end_inclusive,
+            app_profile_id=app_profile_id)
 
     def test_table_name_only(self):
         table_name = 'table_name'
@@ -1102,6 +1134,19 @@ def test_with_limit(self):
         )
         self.assertEqual(result, expected_result)
 
+    def test_with_app_profile_id(self):
+        table_name = 'table_name'
+        limit = 1337
+        app_profile_id = 'app-profile-id'
+        result = self._call_fut(table_name, limit=limit,
+                                app_profile_id=app_profile_id)
+        expected_result = _ReadRowsRequestPB(
+            table_name=table_name,
+            rows_limit=limit,
+            app_profile_id=app_profile_id
+        )
+        self.assertEqual(result, expected_result)
+
 
 def _ReadRowsRequestPB(*args, **kw):
     from google.cloud.bigtable_v2.proto import (