[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:
Giovanni Briggs 2016-11-11 09:57:47 -08:00 коммит произвёл Siddharth Anand
Родитель 868bc83137
Коммит 98f32184a6
2 изменённых файлов: 20 добавлений и 0 удалений

Просмотреть файл

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