def test_run_query_with_arg(self, mocked_rwc): project_id = 12345 def run_with_config(config): self.assertEqual( config['query']['timePartitioning'], { 'field': 'test_field', 'type': 'DAY', 'expirationMs': 1000 } ) mocked_rwc.side_effect = run_with_config bq_hook = hook.BigQueryBaseCursor(mock.Mock(), project_id) bq_hook.run_query( sql='select 1', destination_dataset_table='my_dataset.my_table', time_partitioning={'type': 'DAY', 'field': 'test_field', 'expirationMs': 1000} ) assert mocked_rwc.call_count == 1
def test_insert_all_succeed(self, run_with_config): project_id = 'bq-project' dataset_id = 'bq_dataset' table_id = 'bq_table' rows = [{"json": {"a_key": "a_value_0"}}] body = { "rows": rows, "ignoreUnknownValues": False, "kind": "bigquery#tableDataInsertAllRequest", "skipInvalidRows": False, } mock_service = mock.Mock() method = (mock_service.tabledata.return_value.insertAll) method.return_value.execute.return_value = { "kind": "bigquery#tableDataInsertAllResponse" } cursor = hook.BigQueryBaseCursor(mock_service, 'project_id') cursor.insert_all(project_id, dataset_id, table_id, rows) method.assert_called_with(projectId=project_id, datasetId=dataset_id, tableId=table_id, body=body)
def test_run_load_with_arg(self, mocked_rwc): project_id = 12345 def run_with_config(config): self.assertEqual( config['load']['timePartitioning'], { 'field': 'test_field', 'type': 'DAY', 'expirationMs': 1000 } ) mocked_rwc.side_effect = run_with_config bq_hook = hook.BigQueryBaseCursor(mock.Mock(), project_id) bq_hook.run_load( destination_project_dataset_table='my_dataset.my_table', schema_fields=[], source_uris=[], time_partitioning={'type': 'DAY', 'field': 'test_field', 'expirationMs': 1000} ) assert mocked_rwc.call_count == 1
def test_create_empty_table_succeed(self): project_id = 'bq-project' dataset_id = 'bq_dataset' table_id = 'bq_table' mock_service = mock.Mock() method = mock_service.tables.return_value.insert cursor = hook.BigQueryBaseCursor(mock_service, project_id) cursor.create_empty_table( project_id=project_id, dataset_id=dataset_id, table_id=table_id) body = { 'tableReference': { 'tableId': table_id } } method.assert_called_once_with( projectId=project_id, datasetId=dataset_id, body=body )
def test_patch_view(self): project_id = 'bq-project' dataset_id = 'bq_dataset' view_id = 'bq_view' view_patched = { 'query': "SELECT * FROM `test-project-id.test_dataset_id.test_table_prefix*` LIMIT 500", 'useLegacySql': False } mock_service = mock.Mock() method = (mock_service.tables.return_value.patch) cursor = hook.BigQueryBaseCursor(mock_service, project_id) cursor.patch_table(dataset_id, view_id, project_id, view=view_patched) body = { 'view': view_patched } method.assert_called_once_with( projectId=project_id, datasetId=dataset_id, tableId=view_id, body=body )
def test_insert_all_fail(self): project_id = 'bq-project' dataset_id = 'bq_dataset' table_id = 'bq_table' rows = [ {"json": {"a_key": "a_value_0"}} ] mock_service = mock.Mock() method = mock_service.tabledata.return_value.insertAll method.return_value.execute.return_value = { "kind": "bigquery#tableDataInsertAllResponse", "insertErrors": [ { "index": 1, "errors": [] } ] } cursor = hook.BigQueryBaseCursor(mock_service, 'project_id') with self.assertRaises(Exception): cursor.insert_all(project_id, dataset_id, table_id, rows, fail_on_error=True)
def test_update_dataset(self): dataset_resource = { "kind": "bigquery#dataset", "location": "US", "id": "your-project:dataset_2_test", "datasetReference": { "projectId": "your-project", "datasetId": "dataset_2_test" } } dataset_id = "test_dataset" project_id = "project_test" mock_service = mock.Mock() method = (mock_service.datasets.return_value.update) cursor = hook.BigQueryBaseCursor(mock_service, project_id) cursor.update_dataset(dataset_id=dataset_id, project_id=project_id, dataset_resource=dataset_resource) method.assert_called_once_with(projectId=project_id, datasetId=dataset_id, body=dataset_resource)
def test_run_query_sql_dialect_override(self, run_with_config): for bool_val in [True, False]: cursor = hook.BigQueryBaseCursor(mock.Mock(), "project_id") cursor.run_query('query', use_legacy_sql=bool_val) args, kwargs = run_with_config.call_args self.assertIs(args[0]['query']['useLegacySql'], bool_val)
def test_run_query_sql_dialect_default(self, run_with_config): cursor = hook.BigQueryBaseCursor(mock.Mock(), "project_id") cursor.run_query('query') args, kwargs = run_with_config.call_args self.assertIs(args[0]['query']['useLegacySql'], True)
def test_nobql_nosql_param_error(self): with self.assertRaises(TypeError) as context: hook.BigQueryBaseCursor("test", "test").run_query(sql=None, bql=None) self.assertIn('missing 1 required positional', str(context.exception))
def test_bql_deprecation_warning(self): with warnings.catch_warnings(record=True) as w: hook.BigQueryBaseCursor( "test", "test").run_query(bql='select * from test_table') yield self.assertIn('Deprecated parameter `bql`', w[0].message.args[0])
def test_run_with_auto_detect(self, run_with_config): destination_project_dataset_table = "autodetect.table" cursor = hook.BigQueryBaseCursor(mock.Mock(), "project_id") cursor.run_load(destination_project_dataset_table, [], [], autodetect=True) args, kwargs = run_with_config.call_args self.assertIs(args[0]['load']['autodetect'], True)
def test_get_dataset_without_dataset_id(self): with mock.patch.object(hook.BigQueryHook, 'get_service'): with self.assertRaises(ValueError): hook.BigQueryBaseCursor( mock.Mock(), "test_create_empty_dataset").get_dataset( dataset_id="", project_id="project_test")
def test_create_empty_dataset_no_dataset_id_err(self): with self.assertRaises(ValueError): hook.BigQueryBaseCursor( mock.Mock(), "test_create_empty_dataset").create_empty_dataset( dataset_id="", project_id="")
def test_create_external_table_with_kms(self): project_id = "bq-project" dataset_id = "bq_dataset" table_id = "bq_table" external_project_dataset_table = "{}.{}.{}".format( project_id, dataset_id, table_id ) source_uris = ['test_data.csv'] source_format = 'CSV' autodetect = False compression = 'NONE' ignore_unknown_values = False max_bad_records = 10 skip_leading_rows = 1 field_delimiter = ',' quote_character = None allow_quoted_newlines = False allow_jagged_rows = False labels = {'label1': 'test1', 'label2': 'test2'} schema_fields = [ {"name": "id", "type": "STRING", "mode": "REQUIRED"} ] encryption_configuration = { "kms_key_name": "projects/p/locations/l/keyRings/k/cryptoKeys/c" } mock_service = mock.Mock() method = mock_service.tables.return_value.insert cursor = hook.BigQueryBaseCursor(mock_service, project_id) cursor.create_external_table( external_project_dataset_table=external_project_dataset_table, source_uris=source_uris, source_format=source_format, autodetect=autodetect, compression=compression, ignore_unknown_values=ignore_unknown_values, max_bad_records=max_bad_records, skip_leading_rows=skip_leading_rows, field_delimiter=field_delimiter, quote_character=quote_character, allow_jagged_rows=allow_jagged_rows, allow_quoted_newlines=allow_quoted_newlines, labels=labels, schema_fields=schema_fields, encryption_configuration=encryption_configuration ) body = { 'externalDataConfiguration': { 'autodetect': autodetect, 'sourceFormat': source_format, 'sourceUris': source_uris, 'compression': compression, 'ignoreUnknownValues': ignore_unknown_values, 'schema': {'fields': schema_fields}, 'maxBadRecords': max_bad_records, 'csvOptions': { 'skipLeadingRows': skip_leading_rows, 'fieldDelimiter': field_delimiter, 'quote': quote_character, 'allowQuotedNewlines': allow_quoted_newlines, 'allowJaggedRows': allow_jagged_rows } }, 'tableReference': { 'projectId': project_id, 'datasetId': dataset_id, 'tableId': table_id, }, 'labels': labels, "encryptionConfiguration": encryption_configuration, } method.assert_called_once_with( projectId=project_id, datasetId=dataset_id, body=body )
def test_invalid_source_format(self): with self.assertRaises(Exception) as context: hook.BigQueryBaseCursor("test", "test").run_load("test.test", "test_schema.json", ["test_data.json"], source_format="json") # since we passed 'json' in, and it's not valid, make sure it's present in the error string. self.assertIn("JSON", str(context.exception))