[AIRFLOW-611] source_format in BigQueryBaseCursor
Check source_format in BigQueryBaseCursor The edits to `bigquery_hook.py` are made to `BigQueryBaseCursor`. Closes #1873 from Jalepeno112/bug/AIRFLOW-611
This commit is contained in:
Родитель
868bc83137
Коммит
98f32184a6
|
@ -373,6 +373,19 @@ class BigQueryBaseCursor(object):
|
|||
:param field_delimiter: The delimiter to use when loading from a CSV.
|
||||
:type field_delimiter: string
|
||||
"""
|
||||
|
||||
# bigquery only allows certain source formats
|
||||
# we check to make sure the passed source format is valid
|
||||
# if it's not, we raise a ValueError
|
||||
# Refer to this link for more details:
|
||||
# https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs#configuration.query.tableDefinitions.(key).sourceFormat
|
||||
source_format = source_format.upper()
|
||||
allowed_formats = ["CSV", "NEWLINE_DELIMITED_JSON", "AVRO", "GOOGLE_SHEETS"]
|
||||
if source_format not in allowed_formats:
|
||||
raise ValueError("{0} is not a valid source format. "
|
||||
"Please use one of the following types: {1}"
|
||||
.format(source_format, allowed_formats))
|
||||
|
||||
destination_project, destination_dataset, destination_table = \
|
||||
_split_tablename(table_input=destination_project_dataset_table,
|
||||
default_project_id=self.project_id,
|
||||
|
|
|
@ -104,6 +104,13 @@ class TestBigQueryTableSplitter(unittest.TestCase):
|
|||
self.assertIn('Format exception for var_x:',
|
||||
str(context.exception), "")
|
||||
|
||||
class TestBigQueryHookSourceFormat(unittest.TestCase):
|
||||
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))
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
|
Загрузка…
Ссылка в новой задаче