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