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']) ]