def bootstrap_bq_history(**kwargs):
"""Bootstraps BigQuery history tables for archival purposes.
Args:
**kwargs: keyword args including a user_email with which to run the
Directory API client methods (required for BigQuery streaming).
"""
del kwargs # Unused, but comes by default.
client = bigquery.BigQueryClient()
client.initialize_tables()
def stream(self):
"""Streams the row to BigQuery."""
logging.info('Streaming row to table %s', self.model_type)
bq_client = bigquery.BigQueryClient()
try:
bq_client.stream_row(self.model_type, self._to_bq_format())
except bigquery.InsertError:
logging.error('Unable to stream row, see logs.')
return
else:
self._set_streamed()
def test_initialize_tables(self, mock_schema):
with mock.patch.object(bigquery, 'bigquery'):
mock_client = bigquery.BigQueryClient()
mock_client._dataset = mock.Mock()
mock_client.initialize_tables()
mock_schema.assert_called()
mock_client._dataset.create.assert_called()
mock_client._dataset.table.called_with(constants.BIGQUERY_DEVICE_TABLE)
mock_client._dataset.table.called_with(constants.BIGQUERY_SHELF_TABLE)
def test_initialize_tables__dataset_exists(self, mock_schema, unused):
del unused
mock_client = bigquery.BigQueryClient()
mock_client._dataset = mock.Mock()
mock_client._dataset.create.side_effect = cloud.exceptions.Conflict(
'Already Exists: Dataset Loaner')
mock_client.initialize_tables()
mock_schema.assert_called()
mock_client._dataset.create.assert_called()
def stream_rows(cls):
"""Streams all unstreamed rows if a threshold has been reached."""
logging.info('Streaming rows to BigQuery.')
if not cls.threshold_reached():
logging.info('Not streaming rows, thresholds not met.')
return
bq_client = bigquery.BigQueryClient()
rows = cls._fetch_unstreamed_rows()
tables = _format_for_bq(rows)
try:
for table_name in tables:
bq_client.stream_table(table_name, tables[table_name])
except bigquery.InsertError:
logging.error('Unable to stream rows.')
return
_set_streamed(rows)
def get_history(self, request):
"""Gets historical data for a given device."""
self.check_xsrf_token(self.request_state)
client = bigquery.BigQueryClient()
device = _get_device(request.device)
serial = device.serial_number
info = client.get_device_info(serial)
if not info:
raise endpoints.NotFoundException(
'No history for the requested serial number.')
historical_device = device_messages.Device()
response = device_messages.HistoryResponse()
historical_device.asset_tag = info[0][5]['asset_tag']
for row in info:
response.devices.append(historical_device)
response.timestamp.append(row[1])
response.actor.append(row[2])
response.summary.append(row[4])
return response
def setUp(self):
super(BigQueryClientTest, self).setUp()
bq_patcher = mock.patch.object(gcloud_bq, 'Client', autospec=True)
self.addCleanup(bq_patcher.stop)
self.bq_mock = bq_patcher.start()
self.dataset_ref = mock.Mock(spec=gcloud_bq.DatasetReference)
self.table = mock.Mock(spec=gcloud_bq.Table)
self.table.schema = []
self.dataset_ref.table.return_value = self.table
with mock.patch.object(bigquery.BigQueryClient,
'__init__',
return_value=None):
self.client = bigquery.BigQueryClient()
self.client._client = self.bq_mock()
self.client._dataset_ref = self.dataset_ref
self.client._client.insert_rows.return_value = None
self.client._client.get_table.return_value = self.table
self.nested_schema = [
gcloud_bq.SchemaField('nested_string_attribute', 'STRING',
'NULLABLE')
]
self.entity_schema = [
gcloud_bq.SchemaField('string_attribute', 'STRING', 'NULLABLE'),
gcloud_bq.SchemaField('integer_attribute', 'INTEGER', 'NULLABLE'),
gcloud_bq.SchemaField('boolean_attribute', 'BOOLEAN', 'NULLABLE'),
gcloud_bq.SchemaField('nested_attribute',
'RECORD',
'NULLABLE',
fields=self.nested_schema)
]
test_device = device_model.Device(serial_number='abc123',
chrome_device_id='123123')
test_device.put()
test_row = bigquery_row_model.BigQueryRow.add(
test_device, datetime.datetime.utcnow(), loanertest.USER_EMAIL,
'Enroll', 'This is a test')
self.test_row_dict = test_row.to_json_dict()
self.test_table = [
(self.test_row_dict['ndb_key'], self.test_row_dict['timestamp'],
self.test_row_dict['actor'], self.test_row_dict['method'],
self.test_row_dict['summary'], self.test_row_dict['entity'])
]