def test_listing_tables(self): # type: () -> None dataset_ref1 = DatasetReference('my_project', 'dataset1') self.bq_client.create_dataset(Dataset(dataset_ref1)) self.bq_client.create_table(Table(TableReference(dataset_ref1, 'table1'), _TEST_SCHEMA)) self.bq_client.create_table(Table(TableReference(dataset_ref1, 'table2'), [])) six.assertCountEqual( self, [table_ref.table_id for table_ref in self.bq_client.list_tables(dataset_ref1)], ['table1', 'table2'])
def setUp(self): super(ClientWriteFromQueryTest, self).setUp() dataset_ref = DatasetReference(self.bq_client.project, 'my_dataset') schema = [SchemaField(name="a", field_type='INT64'), SchemaField(name="b", field_type='FLOAT64'), ] self.source_table = Table(TableReference(dataset_ref, 'source_table'), schema) self.destination_table = Table(TableReference(dataset_ref, 'destination_table'), schema) self.bq_client.create_dataset(Dataset(dataset_ref)) self.bq_client.create_table(self.source_table) # We don't create the destination table here; some tests do not want it created. # Stick two rows into source_table self.assertFalse(self.bq_client.insert_rows(self.source_table, [{'a': 1, 'b': 2.5}, {'a': 3, 'b': 4.25}]))
def load(project_id, bq_client, src_dataset_id, dst_dataset_id): """ Transform safely loaded tables and store results in target dataset. :param project_id: Identifies the BQ project :param bq_client: a BigQuery client object :param src_dataset_id: reference to source dataset object :param dst_dataset_id: reference to destination dataset object :return: List of BQ job_ids """ dst_dataset = Dataset(f'{bq_client.project}.{dst_dataset_id}') dst_dataset.description = f'Vocabulary cleaned and loaded from {src_dataset_id}' dst_dataset.labels = {'type': 'vocabulary'} dst_dataset.location = "US" bq_client.create_dataset(dst_dataset, exists_ok=True) src_tables = list(bq_client.list_tables(dataset=src_dataset_id)) job_config = QueryJobConfig() query_jobs = [] for src_table in src_tables: schema = bq.get_table_schema(src_table.table_id) destination = f'{project_id}.{dst_dataset_id}.{src_table.table_id}' table = bq_client.create_table(Table(destination, schema=schema), exists_ok=True) job_config.destination = table query = SELECT_TPL.render(project_id=project_id, dataset_id=src_dataset_id, table=src_table.table_id, fields=schema) query_job = bq_client.query(query, job_config=job_config) LOGGER.info(f'table:{destination} job_id:{query_job.job_id}') query_jobs.append(query_job) query_job.result() return query_jobs
def copy_fitbit_tables_from_views(client, from_dataset, to_dataset, table_prefix): """ Copies tables from views with prefix :param client: bq client :param from_dataset: dataset containing views :param to_dataset: dataset to create tables :param table_prefix: prefix added to table_ids :return: """ for table in FITBIT_TABLES: schema_list = bq.get_table_schema(table) fq_dest_table = f'{client.project}.{to_dataset}.{table}' dest_table = Table(fq_dest_table, schema=schema_list) dest_table = client.create_table(dest_table) LOGGER.info(f'Created empty table {fq_dest_table}') fields_name_str = ',\n'.join([item.name for item in schema_list]) fields_casted_str = ',\n'.join([ cast_to_schema_type(item.name, item.field_type) for item in schema_list ]) content_query = INSERT_QUERY.render(fq_dest_table=fq_dest_table, fields=fields_name_str, fields_casted=fields_casted_str, client=client, from_dataset=from_dataset, table_prefix=table_prefix, table=table) job = client.query(content_query) job.result() LOGGER.info(f'Copied fitbit tables from `{from_dataset}` to `{to_dataset}`')
def test_should_insert_records_to_non_partitioned_table(self): # given table_id = f'{self.dataset_manager.project_id}.{self.dataset_manager.dataset_name}.example_test_table' table = Table(table_id, schema=[ { "mode": "NULLABLE", "name": "example_field", "type": "STRING" }, ]) self.dataset_manager.create_table_from_schema('example_test_table', table=table) self.dataset_manager.insert('example_test_table', [{ "example_field": "example_field_value" }], partitioned=False) # then expected_result = [{"example_field": "example_field_value"}] actual_result = self.dataset_manager.collect(''' SELECT * FROM `{example_test_table}` ''').to_dict(orient='records') self.assertEqual(expected_result, actual_result)
def insert_rows_from_uri(client, job_location, dataset, table_name, table_schema, table_uri, load_job_id, timeout=10): """ Insert rows from remote :param client: BQ Client :param job_location : BQ job location :param dataset: BQ Dataset :param table_name: The table string name :param table_schema: The table schema :param table_uri: The GCS blob uri to insert :param load_job_id: Load job ID :param timeout : optional timeout, default to 10 :return: The created Table """ table_ref = dataset.table(table_name) table = Table(table_ref, table_schema) table_load_job_conf = bigquery.LoadJobConfig( schema=table_schema, source_format=bigquery.SourceFormat.NEWLINE_DELIMITED_JSON ) load_job = client.load_table_from_uri( table_uri, table, job_id=load_job_id, location=job_location, job_config=table_load_job_conf ) try: load_job.result(timeout=timeout) rows = client.list_rows(table, timeout=timeout) print_rows(rows) except Exception as error: logging.error(error) else: return table
def test_table_lookup(self): # type: () -> None dataset_ref1 = DatasetReference('my_project', 'dataset1') table_ref1 = TableReference(dataset_ref1, 'table1') table1 = Table(table_ref1, _TEST_SCHEMA) # Trying to get the same dataset/table in another project doesn't work. with self.assertRaisesRegexp(NotFound, 'other_project'): self.bq_client.get_table( TableReference(DatasetReference('other_project', 'dataset1'), 'table1')) # Trying to get the table before the dataset exists doesn't work with self.assertRaisesRegexp(NotFound, 'dataset1'): self.bq_client.get_table(table_ref1) self.bq_client.create_dataset(Dataset(dataset_ref1)) # Trying to get the table before the table exists doesn't work with self.assertRaises(NotFound): self.bq_client.get_table(table_ref1) self.bq_client.create_table(table1) # Assert the created table has the expected properties. table_found = self.bq_client.get_table(table_ref1) self.assertEqual(table1.project, "my_project") self.assertEqual(table1.dataset_id, "dataset1") self.assertEqual(table1.table_id, "table1") six.assertCountEqual(self, table_found.schema, _TEST_SCHEMA)
def make_table(project, dataset_id, table_id, friendly_name=None, description=None, expires=None, partitioning_type=None, view_use_legacy_sql=None, view_query=None, schema=None, labels=None): dataset_ref = DatasetReference(project, dataset_id) table_ref = TableReference(dataset_ref, table_id) table = Table(table_ref) table.friendly_name = friendly_name table.description = description table.expires = expires table.partitioning_type = partitioning_type if view_use_legacy_sql is not None: table.view_use_legacy_sql = view_use_legacy_sql if view_query is not None: table.view_query = view_query table.schema = schema if labels is not None: table.labels = labels return table
def create_view_task(): client = Client() dest_table_name = '{task}'.format(task=task) dest_table_ref = client.dataset( dataset_name, project=destination_dataset_project_id).table(dest_table_name) table = Table(dest_table_ref) sql_path = os.path.join( dags_folder, 'resources/stages/enrich/sqls/{task}.sql'.format(task=task)) sql = read_file(sql_path, environment) table.view_query = sql description_path = os.path.join( dags_folder, 'resources/stages/enrich/descriptions/{task}.txt'.format( task=task)) table.description = read_file(description_path) logging.info('Creating view: ' + json.dumps(table.to_api_repr())) try: table = client.create_table(table) except Conflict: # https://cloud.google.com/bigquery/docs/managing-views table = client.update_table(table, ['view_query']) assert table.table_id == dest_table_name
def create_view_task(ds, **kwargs): template_context = kwargs.copy() template_context['ds'] = ds template_context['params'] = environment client = Client() dest_table_name = '{task}'.format(task=task) dest_table_ref = client.dataset( dataset_name, project=destination_dataset_project_id).table(dest_table_name) table = Table(dest_table_ref) sql_path = os.path.join( dags_folder, 'resources/stages/enrich/sqls/{task}.sql'.format(task=task)) sql_template = read_file(sql_path) sql = kwargs['task'].render_template('', sql_template, template_context) table.view_query = sql description_path = os.path.join( dags_folder, 'resources/stages/enrich/descriptions/{task}.txt'.format( task=task)) table.description = read_file(description_path) logging.info('Creating view: ' + json.dumps(table.to_api_repr())) try: table = client.create_table(table) except Conflict: # https://cloud.google.com/bigquery/docs/managing-views table = client.update_table(table, ['view_query']) assert table.table_id == dest_table_name
def test_dataflow_event_big_query_success(sdc_builder, sdc_executor, gcp): """ Create data using Google BigQuery client and then check if Google BigQuery origin receives them using wiretap, including a query successfully read event. The pipeline looks like: google_bigquery >> wiretap """ pipeline_builder = sdc_builder.get_pipeline_builder() dataset_name = get_random_string(string.ascii_letters, 5) table_name = get_random_string(string.ascii_letters, 5) google_bigquery = pipeline_builder.add_stage('Google BigQuery', type='origin') query_str = f'SELECT * FROM {dataset_name}.{table_name} ORDER BY id' google_bigquery.set_attributes(query=query_str) wiretap = pipeline_builder.add_wiretap() events_wiretap = pipeline_builder.add_wiretap() google_bigquery >> wiretap.destination google_bigquery >= events_wiretap.destination pipeline = pipeline_builder.build().configure_for_environment(gcp) sdc_executor.add_pipeline(pipeline) bigquery_client = gcp.bigquery_client try: dataset_ref = Dataset(bigquery_client.dataset(dataset_name)) # Using Google bigquery client, create dataset, table and data inside table logger.info('Creating dataset %s using Google bigquery client ...', dataset_name) bigquery_client.create_dataset(dataset_ref) table = bigquery_client.create_table( Table(dataset_ref.table(table_name), schema=SCHEMA)) errors = bigquery_client.insert_rows(table, ROWS_TO_INSERT) assert not errors, 'Errors found when creating data using bigquery client' # Start pipeline and verify correct rows are received. logger.info('Starting pipeline') sdc_executor.start_pipeline(pipeline).wait_for_finished() assert len(wiretap.output_records) == len(ROWS_TO_INSERT),\ f'Expected {len(ROWS_TO_INSERT)} records, received {len(wiretap.output_records)}' rows_from_wiretap = get_rows_from_wiretap(wiretap) assert rows_from_wiretap == ROWS_TO_INSERT # We have exactly one output record, check that it is a big-query-success event assert len( events_wiretap.output_records ) == 1, f'Expected 1 records, received {len(events_wiretap.output_records)}' event_record = events_wiretap.output_records[0] event_type = event_record.header.values['sdc.event.type'] assert event_type == 'big-query-success', 'Received %s as event type (expected new-file)' % event_type finally: if dataset_ref: bigquery_client.delete_dataset(dataset_ref, delete_contents=True)
def test_multiple_batch(sdc_builder, sdc_executor, gcp, number_batches): """ Create data using Google BigQuery destination using different batch sizes and then check if it has been correctly created using BigQuery client. The pipeline looks like: dev_raw_data_source >> google_bigquery """ number_records = 1000 batch_size = number_records//number_batches pipeline_builder = sdc_builder.get_pipeline_builder() dev_data_generator = pipeline_builder.add_stage('Dev Data Generator') dev_data_generator.set_attributes(batch_size=batch_size, fields_to_generate=[ {"type": "INTEGER", "field": "age"}, {"type": "STRING", "field": "full_name"} ]) dataset_name = get_random_string(string.ascii_letters, 5) table_name = get_random_string(string.ascii_letters, 5) google_bigquery = pipeline_builder.add_stage(name=DESTINATION_STAGE_NAME, type='destination') google_bigquery.set_attributes(dataset=dataset_name, table_name=table_name) dev_data_generator >> google_bigquery pipeline = pipeline_builder.build().configure_for_environment(gcp) sdc_executor.add_pipeline(pipeline) bigquery_client = gcp.bigquery_client try: dataset_ref = Dataset(bigquery_client.dataset(dataset_name)) # Using Google bigquery client, create dataset, table and data inside table logger.info('Creating dataset %s using Google bigquery client ...', dataset_name) bigquery_client.create_dataset(dataset_ref) table = bigquery_client.create_table(Table(dataset_ref.table(table_name), schema=SCHEMA)) # Start pipeline and verify correct rows are received. logger.info('Starting pipeline') sdc_executor.start_pipeline(pipeline).wait_for_pipeline_output_records_count(number_records) sdc_executor.stop_pipeline(pipeline) data_from_bigquery = [tuple(row.values()) for row in bigquery_client.list_rows(table)] data_from_bigquery.sort() assert data_from_bigquery == data_from_bigquery history = sdc_executor.get_pipeline_history(pipeline) records = history.latest.metrics.counter('pipeline.batchInputRecords.counter').count assert len(data_from_bigquery) == records finally: if dataset_ref: bigquery_client.delete_dataset(dataset_ref, delete_contents=True)
def test_data_type(sdc_builder, sdc_executor, gcp, data_type, data, expected_data): """ Create data using Google BigQuery client and then check if Google BigQuery origin receives them using wiretap with one data format each. The pipeline looks like: google_bigquery >> wiretap """ pipeline_builder = sdc_builder.get_pipeline_builder() dataset_name = get_random_string(string.ascii_letters, 5) table_name = get_random_string(string.ascii_letters, 5) google_bigquery = pipeline_builder.add_stage('Google BigQuery', type='origin') query_str = f'SELECT * FROM {dataset_name}.{table_name}' google_bigquery.set_attributes(query=query_str) wiretap = pipeline_builder.add_wiretap() google_bigquery >> wiretap.destination pipeline = pipeline_builder.build().configure_for_environment(gcp) sdc_executor.add_pipeline(pipeline) bigquery_client = gcp.bigquery_client schema = [ SchemaField('id', 'INTEGER', mode='NULLABLE'), SchemaField('data', data_type, mode='NULLABLE') ] try: dataset_ref = Dataset(bigquery_client.dataset(dataset_name)) # Using Google bigquery client, create dataset, table and data inside table logger.info('Creating dataset %s using Google bigquery client ...', dataset_name) bigquery_client.create_dataset(dataset_ref) table = bigquery_client.create_table( Table(dataset_ref.table(table_name), schema=schema)) errors = bigquery_client.insert_rows(table, [(None, data)]) assert not errors, 'Errors found when creating data using bigquery client' # Start pipeline and verify correct rows are received. logger.info('Starting pipeline') sdc_executor.start_pipeline(pipeline).wait_for_finished() assert len( wiretap.output_records ) == 1, f'Expected 1 record, received {len(wiretap.output_records)}' received_data = wiretap.output_records[0].field['data'] assert received_data == expected_data finally: if dataset_ref: bigquery_client.delete_dataset(dataset_ref, delete_contents=True)
def test_google_bigquery_origin_batch_handling(sdc_builder, sdc_executor, gcp): """Verify proper batch handling by the BigQuery origin. In this test, we write 8 records to BigQuery with a batch size of 3, verifying that each batch is of size 3 and that 8 total records are captured. The small numbers are used because of the limitations described in SDC-8765. """ MAX_BATCH_SIZE = 3 TOTAL_RECORDS = 8 dataset_name = get_random_string(ascii_letters, 5) table_name = get_random_string(ascii_letters, 5) query_str = f'SELECT * FROM {dataset_name}.{table_name} ORDER BY number' pipeline_builder = sdc_builder.get_pipeline_builder() google_bigquery = pipeline_builder.add_stage('Google BigQuery', type='origin') google_bigquery.set_attributes(query=query_str, max_batch_size_in_records=MAX_BATCH_SIZE) trash = pipeline_builder.add_stage('Trash') google_bigquery >> trash pipeline = pipeline_builder.build().configure_for_environment(gcp) sdc_executor.add_pipeline(pipeline) bigquery_client = gcp.bigquery_client schema = [SchemaField('number', 'STRING', mode='required'), SchemaField('value', 'STRING', mode='required')] data = [dict(number=str(i), value=get_random_string()) for i in range(TOTAL_RECORDS)] dataset = Dataset(bigquery_client.dataset(dataset_name)) try: # Using Google bigquery client, create dataset, table and data inside table. logger.info('Creating dataset %s using Google bigquery client ...', dataset_name) bigquery_client.create_dataset(dataset) table = bigquery_client.create_table(Table(dataset.table(table_name), schema=schema)) errors = bigquery_client.insert_rows(table, data) assert not errors # Ceiling division is needed to capture all the complete batches along with the last partial one. snapshot = sdc_executor.capture_snapshot(pipeline, start_pipeline=True, batches=math.ceil(TOTAL_RECORDS / MAX_BATCH_SIZE)).snapshot # Assert that the batch size is being respected in each batch (including the last). In this case, # we'd expect batch sizes of 3, 3, and 2. for i, batch in enumerate(snapshot.snapshot_batches, start=1): # for 8 records, we'd expect batch sizes of 3, 3, and 2. assert (len(batch.stage_outputs[google_bigquery.instance_name].output) == MAX_BATCH_SIZE if i * MAX_BATCH_SIZE <= TOTAL_RECORDS else TOTAL_RECORDS % MAX_BATCH_SIZE) all_records = [record.field for batch in snapshot.snapshot_batches for record in batch.stage_outputs[google_bigquery.instance_name].output] assert all_records == data finally: bigquery_client.delete_dataset(dataset, delete_contents=True)
def test_listing_tables_with_max(self): # type: () -> None dataset_ref1 = DatasetReference('my_project', 'dataset1') self.bq_client.create_dataset(Dataset(dataset_ref1)) for i in range(10): self.bq_client.create_table(Table(TableReference(dataset_ref1, 'table{}'.format(i)), _TEST_SCHEMA)) self.assertEqual(5, len(self.bq_client.list_tables(dataset_ref1, max_results=5))) self.assertEqual(10, len(self.bq_client.list_tables(dataset_ref1, max_results=20))) self.assertEqual(10, len(self.bq_client.list_tables(dataset_ref1)))
def create_table(client, dataset_ref, table_name, is_partitioned=False): table_ref = dataset_ref.table(table_name) table_obj = Table(table_ref, schema=TABLE_SCHEMA) if is_partitioned: time_partitioning = TimePartitioning() time_partitioning.field = 'partition_value' table_obj.time_partitioning = time_partitioning return client.create_table(table_obj)
def test_google_bigquery_destination(sdc_builder, sdc_executor, gcp): """ Send data to Google BigQuery from Dev Raw Data Source and confirm that Google BigQuery destination successfully recieves them using Google BigQuery client. This is achieved by using a deduplicator which assures that there is only one ingest to Google BigQuery. The pipeline looks like: dev_raw_data_source >> record_deduplicator >> google_bigquery record_deduplicator >> trash """ pipeline_builder = sdc_builder.get_pipeline_builder() dev_raw_data_source = pipeline_builder.add_stage('Dev Raw Data Source') dev_raw_data_source.set_attributes(data_format='DELIMITED', header_line='WITH_HEADER', raw_data='\n'.join(CSV_DATA_TO_INSERT)) dataset_name = get_random_string(ascii_letters, 5) table_name = get_random_string(ascii_letters, 5) google_bigquery = pipeline_builder.add_stage('Google BigQuery', type='destination') google_bigquery.set_attributes(dataset=dataset_name, table_name=table_name) record_deduplicator = pipeline_builder.add_stage('Record Deduplicator') trash = pipeline_builder.add_stage('Trash') dev_raw_data_source >> record_deduplicator >> google_bigquery record_deduplicator >> trash pipeline = pipeline_builder.build(title='Google BigQuery Destination').configure_for_environment(gcp) sdc_executor.add_pipeline(pipeline) bigquery_client = gcp.bigquery_client schema = [SchemaField('full_name', 'STRING', mode='required'), SchemaField('age', 'INTEGER', mode='required')] dataset_ref = Dataset(bigquery_client.dataset(dataset_name)) try: logger.info('Creating dataset %s using Google BigQuery client ...', dataset_name) dataset = bigquery_client.create_dataset(dataset_ref) table = bigquery_client.create_table(Table(dataset_ref.table(table_name), schema=schema)) logger.info('Starting BigQuery Destination pipeline and waiting for it to produce records ...') sdc_executor.start_pipeline(pipeline).wait_for_pipeline_batch_count(1) logger.info('Stopping BigQuery Destination pipeline and getting the count of records produced in total ...') sdc_executor.stop_pipeline(pipeline) # Verify by reading records using Google BigQuery client data_from_bigquery = [tuple(row.values()) for row in bigquery_client.list_rows(table)] data_from_bigquery.sort() logger.debug('read_data = {}'.format(data_from_bigquery)) assert ROWS_EXPECTED == data_from_bigquery finally: bigquery_client.delete_dataset(dataset_ref, delete_contents=True)
def test_object_names_columns(sdc_builder, sdc_executor, gcp, column_name): """ Create data using Google BigQuery destination with different column names and then check if it has been correctly created using BigQuery client. The pipeline looks like: dev_raw_data_source >> google_bigquery """ data_to_insert = [(1, 'data')] rows_to_insert = [f'id,{column_name}'] + [','.join(str(element) for element in row) for row in data_to_insert] data_to_expect = [(1, 'data')] pipeline_builder = sdc_builder.get_pipeline_builder() dev_raw_data_source = pipeline_builder.add_stage('Dev Raw Data Source') dev_raw_data_source.set_attributes(data_format='DELIMITED', header_line='WITH_HEADER', stop_after_first_batch=True, raw_data='\n'.join(rows_to_insert)) dataset_name = get_random_string(string.ascii_letters, 5) table_name = get_random_string(string.ascii_letters, 5) google_bigquery = pipeline_builder.add_stage(name=DESTINATION_STAGE_NAME, type='destination') google_bigquery.set_attributes(dataset=dataset_name, table_name=table_name) dev_raw_data_source >> google_bigquery pipeline = pipeline_builder.build().configure_for_environment(gcp) sdc_executor.add_pipeline(pipeline) bigquery_client = gcp.bigquery_client schema = [SchemaField('id', 'INTEGER', mode='NULLABLE'), SchemaField(column_name, 'STRING', mode='NULLABLE')] try: dataset_ref = Dataset(bigquery_client.dataset(dataset_name)) # Using Google bigquery client, create dataset, table and data inside table logger.info('Creating dataset %s using Google bigquery client ...', dataset_name) bigquery_client.create_dataset(dataset_ref) table = bigquery_client.create_table(Table(dataset_ref.table(table_name), schema=schema)) # Start pipeline and verify correct rows are received. logger.info('Starting pipeline') sdc_executor.start_pipeline(pipeline).wait_for_finished() data_from_bigquery = [tuple(row.values()) for row in bigquery_client.list_rows(table)] data_from_bigquery.sort() assert data_to_expect == data_from_bigquery finally: if dataset_ref: bigquery_client.delete_dataset(dataset_ref, delete_contents=True)
def __create_test_table(self, table_name, dataset_id): table_schema = [ SchemaField("int_data", "INT64"), SchemaField("str_data", "STRING") ] table_reference = TableReference(dataset_id, table_name) test_table = Table(table_reference, table_schema) test_table.time_partitioning = TimePartitioning("DAY") self.__delete_if_exists(test_table) self.GCP_BIGQUERY_CLIENT.create_table(test_table) return test_table
def test_google_bigquery_origin(sdc_builder, sdc_executor, gcp): """ Create data using Google BigQuery client and then check if Google BigQuery origin receives them using snapshot. The pipeline looks like: google_bigquery >> trash """ pipeline_builder = sdc_builder.get_pipeline_builder() dataset_name = get_random_string(ascii_letters, 5) table_name = get_random_string(ascii_letters, 5) google_bigquery = pipeline_builder.add_stage('Google BigQuery', type='origin') query_str = f'SELECT * FROM {dataset_name}.{table_name} ORDER BY full_name' google_bigquery.set_attributes(query=query_str) trash = pipeline_builder.add_stage('Trash') google_bigquery >> trash pipeline = pipeline_builder.build( title='Google BigQuery').configure_for_environment(gcp) sdc_executor.add_pipeline(pipeline) bigquery_client = gcp.bigquery_client schema = [ SchemaField('full_name', 'STRING', mode='required'), SchemaField('age', 'INTEGER', mode='required') ] dataset_ref = Dataset(bigquery_client.dataset(dataset_name)) try: # Using Google bigquery client, create dataset, table and data inside table logger.info('Creating dataset %s using Google bigquery client ...', dataset_name) dataset = bigquery_client.create_dataset(dataset_ref) table = bigquery_client.create_table( Table(dataset_ref.table(table_name), schema=schema)) errors = bigquery_client.insert_rows(table, ROWS_TO_INSERT) assert errors == [] # Start pipeline and verify correct rows are received using snaphot. logger.info('Starting pipeline and snapshot') snapshot = sdc_executor.capture_snapshot(pipeline, start_pipeline=True).snapshot if sdc_executor.get_pipeline_status(pipeline) == 'RUNNING': sdc_executor.stop_pipeline(pipeline) rows_from_snapshot = [(record.value['value'][0]['value'], int(record.value['value'][1]['value'])) for record in snapshot[google_bigquery].output] assert rows_from_snapshot == ROWS_TO_INSERT finally: bigquery_client.delete_dataset(dataset_ref, delete_contents=True)
def get_table(self, dataset_name, table_name, schema): dataset = self.client.dataset(dataset_name) dataset_table_names = [ t.table_id for t in self.client.list_tables(dataset) ] table_ref = dataset.table(table_name) if table_name not in dataset_table_names: assert schema table = Table(table_ref, self.schema_object(schema)) table = self.client.create_table(table) else: table = self.client.get_table(table_ref) return table
def test_object_names_datasets(sdc_builder, sdc_executor, gcp, dataset_name): """ Create data using Google BigQuery destination with different dataset names and then check if it has been correctly created using BigQuery client. The pipeline looks like: dev_raw_data_source >> google_bigquery """ pipeline_builder = sdc_builder.get_pipeline_builder() dev_raw_data_source = pipeline_builder.add_stage('Dev Raw Data Source') dev_raw_data_source.set_attributes(data_format='DELIMITED', header_line='WITH_HEADER', stop_after_first_batch=True, raw_data='\n'.join(CSV_DATA_TO_INSERT)) table_name = get_random_string(string.ascii_letters, 5) # If tests fail for any reason, leftovers with equal names might lead to more errors dataset_name = f'{dataset_name}_{get_random_string(string.ascii_letters, 5)}' google_bigquery = pipeline_builder.add_stage(name=DESTINATION_STAGE_NAME, type='destination') google_bigquery.set_attributes(dataset=dataset_name, table_name=table_name) dev_raw_data_source >> google_bigquery pipeline = pipeline_builder.build().configure_for_environment(gcp) sdc_executor.add_pipeline(pipeline) bigquery_client = gcp.bigquery_client try: dataset_ref = Dataset(bigquery_client.dataset(dataset_name)) # Using Google bigquery client, create dataset, table and data inside table logger.info('Creating dataset %s using Google bigquery client ...', dataset_name) bigquery_client.create_dataset(dataset_ref) table = bigquery_client.create_table(Table(dataset_ref.table(table_name), schema=SCHEMA)) # Start pipeline and verify correct rows are received. logger.info('Starting pipeline') sdc_executor.start_pipeline(pipeline).wait_for_finished() data_from_bigquery = [tuple(row.values()) for row in bigquery_client.list_rows(table)] data_from_bigquery.sort() assert ROWS_EXPECTED == data_from_bigquery finally: if dataset_ref: bigquery_client.delete_dataset(dataset_ref, delete_contents=True)
def store_digital_health_status_data(project_id, json_data, destination_table, schema=None): """ Stores the fetched digital_health_sharing_status data in a BigQuery dataset. If the table doesn't exist, it will create that table. If the table does exist, it will create a partition in the designated table or append to the same partition. This is necessary for storing data has "RECORD" type fields which do not conform to a dataframe. The data is stored using a JSON file object since it is one of the ways BigQuery expects it. :param project_id: identifies the project :param json_data: list of json objects retrieved from process_digital_health_data_to_json :param destination_table: fully qualified destination table name as 'project.dataset.table' :param schema: a list of SchemaField objects corresponding to the destination table :return: returns the bq job_id for the loading of digital health data """ # Parameter check if not isinstance(project_id, str): raise RuntimeError( f'Please specify the project in which to create the table') client = get_client(project_id) if not schema: schema = get_table_schema(DIGITAL_HEALTH_SHARING_STATUS) try: table = client.get_table(destination_table) except NotFound: table = Table(destination_table, schema=schema) table.time_partitioning = TimePartitioning( type_=TimePartitioningType.DAY) table = client.create_table(table) file_obj = StringIO() for json_obj in json_data: json.dump(json_obj, file_obj) file_obj.write('\n') job_config = LoadJobConfig( source_format=SourceFormat.NEWLINE_DELIMITED_JSON, schema=schema) job = client.load_table_from_file(file_obj, table, rewind=True, job_config=job_config, job_id_prefix='ps_digital_health_load_') job.result() return job.job_id
def test_object_names_datasets(sdc_builder, sdc_executor, gcp, dataset_name): """ Create data using Google BigQuery client with specific table names and then check if Google BigQuery origin receives them using wiretap. The pipeline looks like: google_bigquery >> wiretap """ pipeline_builder = sdc_builder.get_pipeline_builder() table_name = get_random_string(string.ascii_letters, 5) dataset_name = f'{dataset_name}_{get_random_string(string.ascii_letters, 5)}' google_bigquery = pipeline_builder.add_stage('Google BigQuery', type='origin') query_str = f"SELECT * FROM `{dataset_name}`.`{table_name}` ORDER BY id" google_bigquery.set_attributes(query=query_str) wiretap = pipeline_builder.add_wiretap() google_bigquery >> wiretap.destination pipeline = pipeline_builder.build().configure_for_environment(gcp) sdc_executor.add_pipeline(pipeline) bigquery_client = gcp.bigquery_client try: dataset_ref = Dataset(bigquery_client.dataset(dataset_name)) # Using Google bigquery client, create dataset, table and data inside table logger.info('Creating dataset %s using Google bigquery client ...', dataset_name) bigquery_client.create_dataset(dataset_ref) table = bigquery_client.create_table( Table(dataset_ref.table(table_name), schema=SCHEMA)) errors = bigquery_client.insert_rows(table, ROWS_TO_INSERT) assert not errors, 'Errors found when creating data using bigquery client' # Start pipeline and verify correct rows are received. logger.info('Starting pipeline') sdc_executor.start_pipeline(pipeline).wait_for_finished() assert len(wiretap.output_records) == len(ROWS_TO_INSERT),\ f'Expected {len(ROWS_TO_INSERT)} records, received {len(wiretap.output_records)}' rows_from_wiretap = get_rows_from_wiretap(wiretap) assert rows_from_wiretap == ROWS_TO_INSERT finally: if dataset_ref: bigquery_client.delete_dataset(dataset_ref, delete_contents=True)
def copy_vocab_tables(cls, vocabulary_id): """ A function for copying the vocab tables to the test dataset_id :param vocabulary_id: :return: """ # Copy vocab tables over to the test dataset vocabulary_dataset = cls.client.get_dataset(vocabulary_id) for src_table in bq.list_tables(cls.client, vocabulary_dataset): schema = bq.get_table_schema(src_table.table_id) destination = f'{cls.project_id}.{cls.dataset_id}.{src_table.table_id}' dst_table = cls.client.create_table(Table(destination, schema=schema), exists_ok=True) cls.client.copy_table(src_table, dst_table)
def get_table(self, table_ref, retry=None): # type: (TableReference, Optional[Retry]) -> Table """Looks up a table. Args: table_ref: Table reference to find. retry: If provided, what retry strategy to use (unused in this implementation). Returns: Table found. """ del retry # Unused in this implementation. typed_dataframe = self._safe_lookup(table_ref.project, table_ref.dataset_id, table_ref.table_id) return Table(table_ref, typed_dataframe.to_bq_schema())
def createTable(tablename, global_dataset_ref): schema = [ #Enter Schema here. # SchemaField('url', 'STRING', mode='required'), # SchemaField('views', 'INTEGER', mode='required') ] table_ref = global_dataset_ref.table(tablename) table = Table(table_ref, schema=schema) table = bigquery_client.create_table(table) assert table.table_id == tablename return
def setUp(self): self.project_id = os.environ.get(PROJECT_ID) self.dataset_id = os.environ.get('COMBINED_DATASET_ID') self.dataset_ref = DatasetReference(self.project_id, self.dataset_id) self.client = bq.get_client(self.project_id) self.schema = [ SchemaField("person_id", "INT64"), SchemaField("first_name", "STRING"), SchemaField("last_name", "STRING"), SchemaField("algorithm", "STRING") ] self.ps_api_fields = [ dict(name='person_id', type='integer', mode='nullable'), dict(name='first_name', type='string', mode='nullable'), dict(name='last_name', type='string', mode='nullable') ] self.id_match_fields = [ dict(name='person_id', type='integer', mode='nullable'), dict(name='first_name', type='string', mode='nullable'), dict(name='last_name', type='string', mode='nullable'), dict(name='algorithm', type='string', mode='nullable') ] self.hpo_id = 'fake_site' self.id_match_table_id = f'{IDENTITY_MATCH_TABLE}_{self.hpo_id}' self.ps_values_table_id = f'ps_api_values_{self.hpo_id}' # Create and populate the ps_values site table schema = bq.get_table_schema(PS_API_VALUES) tablename = self.ps_values_table_id table = Table(f'{self.project_id}.{self.dataset_id}.{tablename}', schema=schema) table.time_partitioning = TimePartitioning( type_=TimePartitioningType.HOUR) table = self.client.create_table(table) populate_query = POPULATE_PS_VALUES.render( project_id=self.project_id, drc_dataset_id=self.dataset_id, ps_values_table_id=self.ps_values_table_id) job = self.client.query(populate_query) job.result()
def create(self, table_id, schema): """ Create a table in Google BigQuery given a table and schema Parameters ---------- table : str Name of table to be written schema : str Use the generate_bq_schema to generate your table schema from a dataframe. """ from google.cloud.bigquery import SchemaField from google.cloud.bigquery import Table if self.exists(table_id): raise TableCreationError( "Table {0} already " "exists".format(table_id) ) if not _Dataset(self.project_id, credentials=self.credentials).exists( self.dataset_id ): _Dataset( self.project_id, credentials=self.credentials, location=self.location, ).create(self.dataset_id) table_ref = self.client.dataset(self.dataset_id).table(table_id) table = Table(table_ref) # Manually create the schema objects, adding NULLABLE mode # as a workaround for # https://github.com/GoogleCloudPlatform/google-cloud-python/issues/4456 for field in schema["fields"]: if "mode" not in field: field["mode"] = "NULLABLE" table.schema = [ SchemaField.from_api_repr(field) for field in schema["fields"] ] try: self.client.create_table(table) except self.http_error as ex: self.process_http_error(ex)
def send_to_bq(dataset, table, row): bigquery_client = bigquery.Client(project='theo-home') table_ref = TableReference( dataset_ref=DatasetReference(dataset_id=dataset, project='theo-home'), table_id=table, ) schema = [SchemaField(name=field, field_type=bq_types[type(data)]) for field, data in row.items()] table = bigquery_client.create_table( Table(table_ref, schema=schema), exists_ok=True ) errors = bigquery_client.insert_rows(table, [row]) if errors: print(errors, file=sys.stderr)