diff --git a/bigquery/client.py b/bigquery/client.py index 1ac3cb3..1642bdb 100644 --- a/bigquery/client.py +++ b/bigquery/client.py @@ -442,6 +442,90 @@ def create_table(self, dataset, table, schema, expiration_time=None): else: return {} + def update_table(self, dataset, table, schema): + """Update an existing table in the dataset. + + Args: + dataset: the dataset to update the table in. + table: the name of table to update. + schema: table schema dict. + + Returns: + bool indicating if the table was successfully updated 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().update( + 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 update table {0}.{1}\n' + 'Http Error: {2}').format(dataset, table, + e.content)) + if self.swallow_results: + return False + 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 index d8869ee..a0f1440 100644 --- a/bigquery/tests/test_client.py +++ b/bigquery/tests/test_client.py @@ -1588,6 +1588,156 @@ def test_table_create_body_with_expiration_time(self): self.mock_tables.insert.return_value.execute.assert_called_with() +class TestUpdateTable(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_update_failed(self): + """Ensure that if updating the table fails, False is returned, + or if swallow_results is False an empty dict is returned.""" + + self.mock_tables.update.return_value.execute.side_effect = ( + HttpError(HttpResponse(404), 'There was an error'.encode('utf8'))) + + actual = self.client.update_table(self.dataset, self.table, + self.schema) + + self.assertFalse(actual) + + self.client.swallow_results = False + + actual = self.client.update_table(self.dataset, self.table, + self.schema) + + self.assertEqual(actual, {}) + + self.client.swallow_results = True + + self.mock_tables.update.assert_called_with( + projectId=self.project, datasetId=self.dataset, body=self.body) + + self.mock_tables.update.return_value.execute.assert_called_with() + + def test_table_update_success(self): + """Ensure that if updating the table succeeds, True is returned, + or if swallow_results is False the actual response is returned.""" + + self.mock_tables.update.return_value.execute.side_effect = [{ + 'status': 'foo'}, {'status': 'bar'}] + + actual = self.client.update_table(self.dataset, self.table, + self.schema) + + self.assertTrue(actual) + + self.client.swallow_results = False + + actual = self.client.update_table(self.dataset, self.table, + self.schema) + + self.assertEqual(actual, {'status': 'bar'}) + + self.client.swallow_results = True + + self.mock_tables.update.assert_called_with( + projectId=self.project, datasetId=self.dataset, body=self.body) + + 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):