Add an patch table method to client

Antoine Cezar · Antoine Cezar · commit 5c6afe74bdcc · 2016-01-18T14:02:48.000+01:00
Fix tylertreat#66
diff --git a/bigquery/client.py b/bigquery/client.py
@@ -484,6 +484,48 @@ def update_table(self, dataset, table, schema):
             else:
                 return {}
 
+    def patch_table(self, dataset, table, schema):
+        """Patch an existing table in the dataset.
+
+        Args:
+            dataset: the dataset to patch the table in.
+            table: the name of table to patch.
+            schema: table schema dict.
+
+        Returns:
+            bool indicating if the table was successfully patched or not,
+            or response from BigQuery if swallow_results is set for False.
+        """
+
+        body = {
+            'schema': {'fields': schema},
+            'tableReference': {
+                'tableId': table,
+                'projectId': self.project_id,
+                'datasetId': dataset
+            }
+        }
+
+        try:
+            result = self.bigquery.tables().patch(
+                projectId=self.project_id,
+                datasetId=dataset,
+                body=body
+            ).execute()
+            if self.swallow_results:
+                return True
+            else:
+                return result
+
+        except HttpError as e:
+            logging.error(('Cannot patch table {0}.{1}\n'
+                           'Http Error: {2}').format(dataset, table,
+                                                     e.content))
+            if self.swallow_results:
+                return False
+            else:
+                return {}
+
     def create_view(self, dataset, view, query):
         """Create a new view in the dataset.
 
diff --git a/bigquery/tests/test_client.py b/bigquery/tests/test_client.py
@@ -1663,6 +1663,81 @@ def test_table_update_success(self):
         self.mock_tables.update.return_value.execute.assert_called_with()
 
 
+class TestPatchTable(unittest.TestCase):
+
+    def setUp(self):
+        self.mock_bq_service = mock.Mock()
+        self.mock_tables = mock.Mock()
+        self.mock_bq_service.tables.return_value = self.mock_tables
+        self.table = 'table'
+        self.schema = [
+            {'name': 'foo', 'type': 'STRING', 'mode': 'nullable'},
+            {'name': 'bar', 'type': 'FLOAT', 'mode': 'nullable'}
+        ]
+        self.project = 'project'
+        self.dataset = 'dataset'
+        self.client = client.BigQueryClient(self.mock_bq_service, self.project)
+        self.body = {
+            'schema': {'fields': self.schema},
+            'tableReference': {
+                'tableId': self.table, 'projectId': self.project,
+                'datasetId': self.dataset}
+        }
+        self.expiration_time = 1437513693000
+
+    def test_table_patch_failed(self):
+        """Ensure that if patching the table fails, False is returned,
+        or if swallow_results is False an empty dict is returned."""
+
+        self.mock_tables.patch.return_value.execute.side_effect = (
+            HttpError(HttpResponse(404), 'There was an error'.encode('utf8')))
+
+        actual = self.client.patch_table(self.dataset, self.table,
+                                         self.schema)
+
+        self.assertFalse(actual)
+
+        self.client.swallow_results = False
+
+        actual = self.client.patch_table(self.dataset, self.table,
+                                         self.schema)
+
+        self.assertEqual(actual, {})
+
+        self.client.swallow_results = True
+
+        self.mock_tables.patch.assert_called_with(
+            projectId=self.project, datasetId=self.dataset, body=self.body)
+
+        self.mock_tables.patch.return_value.execute.assert_called_with()
+
+    def test_table_patch_success(self):
+        """Ensure that if patching the table succeeds, True is returned,
+        or if swallow_results is False the actual response is returned."""
+
+        self.mock_tables.patch.return_value.execute.side_effect = [{
+            'status': 'foo'}, {'status': 'bar'}]
+
+        actual = self.client.patch_table(self.dataset, self.table,
+                                         self.schema)
+
+        self.assertTrue(actual)
+
+        self.client.swallow_results = False
+
+        actual = self.client.patch_table(self.dataset, self.table,
+                                         self.schema)
+
+        self.assertEqual(actual, {'status': 'bar'})
+
+        self.client.swallow_results = True
+
+        self.mock_tables.patch.assert_called_with(
+            projectId=self.project, datasetId=self.dataset, body=self.body)
+
+        self.mock_tables.patch.return_value.execute.assert_called_with()
+
+
 class TestCreateView(unittest.TestCase):
 
     def setUp(self):
