BigTable: Add data app profile id

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,7 +390,7 @@ 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
 
 
@@ -403,10 +410,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 +476,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 +505,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 +538,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 +564,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 +578,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 +587,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

@@ -362,7 +362,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
 
@@ -502,6 +503,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)])
 
@@ -1146,12 +1148,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'
@@ -1234,6 +1238,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 (