Skip to content

Commit eecd8f4

Browse files
tylertreattylertreat
committed
Merge pull request tylertreat#78 from izberg-marketplace/feat_table_update_patch
Feat table update patch
2 parents 117b916 + 5c6afe7 commit eecd8f4

File tree

2 files changed

+234
-0
lines changed

2 files changed

+234
-0
lines changed

‎bigquery/client.py‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -442,6 +442,90 @@ def create_table(self, dataset, table, schema, expiration_time=None):
442442
else:
443443
return {}
444444

445+
def update_table(self, dataset, table, schema):
446+
"""Update an existing table in the dataset.
447+
448+
Args:
449+
dataset: the dataset to update the table in.
450+
table: the name of table to update.
451+
schema: table schema dict.
452+
453+
Returns:
454+
bool indicating if the table was successfully updated or not,
455+
or response from BigQuery if swallow_results is set for False.
456+
"""
457+
458+
body = {
459+
'schema': {'fields': schema},
460+
'tableReference': {
461+
'tableId': table,
462+
'projectId': self.project_id,
463+
'datasetId': dataset
464+
}
465+
}
466+
467+
try:
468+
result = self.bigquery.tables().update(
469+
projectId=self.project_id,
470+
datasetId=dataset,
471+
body=body
472+
).execute()
473+
if self.swallow_results:
474+
return True
475+
else:
476+
return result
477+
478+
except HttpError as e:
479+
logging.error(('Cannot update table {0}.{1}\n'
480+
'Http Error: {2}').format(dataset, table,
481+
e.content))
482+
if self.swallow_results:
483+
return False
484+
else:
485+
return {}
486+
487+
def patch_table(self, dataset, table, schema):
488+
"""Patch an existing table in the dataset.
489+
490+
Args:
491+
dataset: the dataset to patch the table in.
492+
table: the name of table to patch.
493+
schema: table schema dict.
494+
495+
Returns:
496+
bool indicating if the table was successfully patched or not,
497+
or response from BigQuery if swallow_results is set for False.
498+
"""
499+
500+
body = {
501+
'schema': {'fields': schema},
502+
'tableReference': {
503+
'tableId': table,
504+
'projectId': self.project_id,
505+
'datasetId': dataset
506+
}
507+
}
508+
509+
try:
510+
result = self.bigquery.tables().patch(
511+
projectId=self.project_id,
512+
datasetId=dataset,
513+
body=body
514+
).execute()
515+
if self.swallow_results:
516+
return True
517+
else:
518+
return result
519+
520+
except HttpError as e:
521+
logging.error(('Cannot patch table {0}.{1}\n'
522+
'Http Error: {2}').format(dataset, table,
523+
e.content))
524+
if self.swallow_results:
525+
return False
526+
else:
527+
return {}
528+
445529
def create_view(self, dataset, view, query):
446530
"""Create a new view in the dataset.
447531

‎bigquery/tests/test_client.py‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1588,6 +1588,156 @@ def test_table_create_body_with_expiration_time(self):
15881588
self.mock_tables.insert.return_value.execute.assert_called_with()
15891589

15901590

1591+
class TestUpdateTable(unittest.TestCase):
1592+
1593+
def setUp(self):
1594+
self.mock_bq_service = mock.Mock()
1595+
self.mock_tables = mock.Mock()
1596+
self.mock_bq_service.tables.return_value = self.mock_tables
1597+
self.table = 'table'
1598+
self.schema = [
1599+
{'name': 'foo', 'type': 'STRING', 'mode': 'nullable'},
1600+
{'name': 'bar', 'type': 'FLOAT', 'mode': 'nullable'}
1601+
]
1602+
self.project = 'project'
1603+
self.dataset = 'dataset'
1604+
self.client = client.BigQueryClient(self.mock_bq_service, self.project)
1605+
self.body = {
1606+
'schema': {'fields': self.schema},
1607+
'tableReference': {
1608+
'tableId': self.table, 'projectId': self.project,
1609+
'datasetId': self.dataset}
1610+
}
1611+
self.expiration_time = 1437513693000
1612+
1613+
def test_table_update_failed(self):
1614+
"""Ensure that if updating the table fails, False is returned,
1615+
or if swallow_results is False an empty dict is returned."""
1616+
1617+
self.mock_tables.update.return_value.execute.side_effect = (
1618+
HttpError(HttpResponse(404), 'There was an error'.encode('utf8')))
1619+
1620+
actual = self.client.update_table(self.dataset, self.table,
1621+
self.schema)
1622+
1623+
self.assertFalse(actual)
1624+
1625+
self.client.swallow_results = False
1626+
1627+
actual = self.client.update_table(self.dataset, self.table,
1628+
self.schema)
1629+
1630+
self.assertEqual(actual, {})
1631+
1632+
self.client.swallow_results = True
1633+
1634+
self.mock_tables.update.assert_called_with(
1635+
projectId=self.project, datasetId=self.dataset, body=self.body)
1636+
1637+
self.mock_tables.update.return_value.execute.assert_called_with()
1638+
1639+
def test_table_update_success(self):
1640+
"""Ensure that if updating the table succeeds, True is returned,
1641+
or if swallow_results is False the actual response is returned."""
1642+
1643+
self.mock_tables.update.return_value.execute.side_effect = [{
1644+
'status': 'foo'}, {'status': 'bar'}]
1645+
1646+
actual = self.client.update_table(self.dataset, self.table,
1647+
self.schema)
1648+
1649+
self.assertTrue(actual)
1650+
1651+
self.client.swallow_results = False
1652+
1653+
actual = self.client.update_table(self.dataset, self.table,
1654+
self.schema)
1655+
1656+
self.assertEqual(actual, {'status': 'bar'})
1657+
1658+
self.client.swallow_results = True
1659+
1660+
self.mock_tables.update.assert_called_with(
1661+
projectId=self.project, datasetId=self.dataset, body=self.body)
1662+
1663+
self.mock_tables.update.return_value.execute.assert_called_with()
1664+
1665+
1666+
class TestPatchTable(unittest.TestCase):
1667+
1668+
def setUp(self):
1669+
self.mock_bq_service = mock.Mock()
1670+
self.mock_tables = mock.Mock()
1671+
self.mock_bq_service.tables.return_value = self.mock_tables
1672+
self.table = 'table'
1673+
self.schema = [
1674+
{'name': 'foo', 'type': 'STRING', 'mode': 'nullable'},
1675+
{'name': 'bar', 'type': 'FLOAT', 'mode': 'nullable'}
1676+
]
1677+
self.project = 'project'
1678+
self.dataset = 'dataset'
1679+
self.client = client.BigQueryClient(self.mock_bq_service, self.project)
1680+
self.body = {
1681+
'schema': {'fields': self.schema},
1682+
'tableReference': {
1683+
'tableId': self.table, 'projectId': self.project,
1684+
'datasetId': self.dataset}
1685+
}
1686+
self.expiration_time = 1437513693000
1687+
1688+
def test_table_patch_failed(self):
1689+
"""Ensure that if patching the table fails, False is returned,
1690+
or if swallow_results is False an empty dict is returned."""
1691+
1692+
self.mock_tables.patch.return_value.execute.side_effect = (
1693+
HttpError(HttpResponse(404), 'There was an error'.encode('utf8')))
1694+
1695+
actual = self.client.patch_table(self.dataset, self.table,
1696+
self.schema)
1697+
1698+
self.assertFalse(actual)
1699+
1700+
self.client.swallow_results = False
1701+
1702+
actual = self.client.patch_table(self.dataset, self.table,
1703+
self.schema)
1704+
1705+
self.assertEqual(actual, {})
1706+
1707+
self.client.swallow_results = True
1708+
1709+
self.mock_tables.patch.assert_called_with(
1710+
projectId=self.project, datasetId=self.dataset, body=self.body)
1711+
1712+
self.mock_tables.patch.return_value.execute.assert_called_with()
1713+
1714+
def test_table_patch_success(self):
1715+
"""Ensure that if patching the table succeeds, True is returned,
1716+
or if swallow_results is False the actual response is returned."""
1717+
1718+
self.mock_tables.patch.return_value.execute.side_effect = [{
1719+
'status': 'foo'}, {'status': 'bar'}]
1720+
1721+
actual = self.client.patch_table(self.dataset, self.table,
1722+
self.schema)
1723+
1724+
self.assertTrue(actual)
1725+
1726+
self.client.swallow_results = False
1727+
1728+
actual = self.client.patch_table(self.dataset, self.table,
1729+
self.schema)
1730+
1731+
self.assertEqual(actual, {'status': 'bar'})
1732+
1733+
self.client.swallow_results = True
1734+
1735+
self.mock_tables.patch.assert_called_with(
1736+
projectId=self.project, datasetId=self.dataset, body=self.body)
1737+
1738+
self.mock_tables.patch.return_value.execute.assert_called_with()
1739+
1740+
15911741
class TestCreateView(unittest.TestCase):
15921742

15931743
def setUp(self):

0 commit comments

Comments
 (0)