Repository: incubator-airflow Updated Branches: refs/heads/master 49556917d -> 5747f5849
[AIRFLOW-2107] add time_partitioning to run_query on BigQueryBaseCursor Closes #3043 from marengaz/query_time_part Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/5747f584 Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/5747f584 Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/5747f584 Branch: refs/heads/master Commit: 5747f584992a6dc4758b37d80a8797f248fcc889 Parents: 4955691 Author: ben.marengo <[email protected]> Authored: Thu May 24 21:04:24 2018 +0100 Committer: Kaxil Naik <[email protected]> Committed: Thu May 24 21:04:33 2018 +0100 ---------------------------------------------------------------------- airflow/contrib/hooks/bigquery_hook.py | 71 ++++++---- airflow/contrib/operators/bigquery_operator.py | 17 ++- tests/contrib/hooks/test_bigquery_hook.py | 147 ++++++++++++-------- 3 files changed, 147 insertions(+), 88 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/5747f584/airflow/contrib/hooks/bigquery_hook.py ---------------------------------------------------------------------- diff --git a/airflow/contrib/hooks/bigquery_hook.py b/airflow/contrib/hooks/bigquery_hook.py index c0a70ab..f4a0a58 100644 --- a/airflow/contrib/hooks/bigquery_hook.py +++ b/airflow/contrib/hooks/bigquery_hook.py @@ -466,7 +466,8 @@ class BigQueryBaseCursor(LoggingMixin): create_disposition='CREATE_IF_NEEDED', query_params=None, schema_update_options=(), - priority='INTERACTIVE'): + priority='INTERACTIVE', + time_partitioning={}): """ Executes a BigQuery SQL query. Optionally persists results in a BigQuery table. See here: @@ -479,6 +480,7 @@ class BigQueryBaseCursor(LoggingMixin): :type bql: string :param destination_dataset_table: The dotted <dataset>.<table> BigQuery table to save the query results. + :type destination_dataset_table: string :param write_disposition: What to do if the table already exists in BigQuery. :type write_disposition: string @@ -516,6 +518,12 @@ class BigQueryBaseCursor(LoggingMixin): Possible values include INTERACTIVE and BATCH. The default value is INTERACTIVE. :type priority: string + :param time_partitioning: configure optional time partitioning fields i.e. + partition by field, type and + expiration as per API specifications. Note that 'field' is not available in + conjunction with dataset.table$partition. + :type time_partitioning: dict + """ # BigQuery also allows you to define how you want a table's schema to change @@ -553,14 +561,10 @@ class BigQueryBaseCursor(LoggingMixin): _split_tablename(table_input=destination_dataset_table, default_project_id=self.project_id) configuration['query'].update({ - 'allowLargeResults': - allow_large_results, - 'flattenResults': - flatten_results, - 'writeDisposition': - write_disposition, - 'createDisposition': - create_disposition, + 'allowLargeResults': allow_large_results, + 'flattenResults': flatten_results, + 'writeDisposition': write_disposition, + 'createDisposition': create_disposition, 'destinationTable': { 'projectId': destination_project, 'datasetId': destination_dataset, @@ -570,8 +574,7 @@ class BigQueryBaseCursor(LoggingMixin): if udf_config: assert isinstance(udf_config, list) configuration['query'].update({ - 'userDefinedFunctionResources': - udf_config + 'userDefinedFunctionResources': udf_config }) if query_params: @@ -581,6 +584,15 @@ class BigQueryBaseCursor(LoggingMixin): else: configuration['query']['queryParameters'] = query_params + time_partitioning = _cleanse_time_partitioning( + destination_dataset_table, + time_partitioning + ) + if time_partitioning: + configuration['query'].update({ + 'timePartitioning': time_partitioning + }) + if schema_update_options: if write_disposition not in ["WRITE_APPEND", "WRITE_TRUNCATE"]: raise ValueError("schema_update_options is only " @@ -800,7 +812,7 @@ class BigQueryBaseCursor(LoggingMixin): :param time_partitioning: configure optional time partitioning fields i.e. partition by field, type and expiration as per API specifications. Note that 'field' is not available in - concurrency with dataset.table$partition. + conjunction with dataset.table$partition. :type time_partitioning: dict """ @@ -853,21 +865,14 @@ class BigQueryBaseCursor(LoggingMixin): } } - # if it is a partitioned table ($ is in the table name) add partition load option - if '$' in destination_project_dataset_table: - if time_partitioning.get('field'): - raise AirflowException( - "Cannot specify field partition and partition name " - "(dataset.table$partition) at the same time" - ) - configuration['load']['timePartitioning'] = dict(type='DAY') - - # can specify custom time partitioning options based on a field, or adding - # expiration + time_partitioning = _cleanse_time_partitioning( + destination_project_dataset_table, + time_partitioning + ) if time_partitioning: - if not configuration.get('load', {}).get('timePartitioning'): - configuration['load']['timePartitioning'] = {} - configuration['load']['timePartitioning'].update(time_partitioning) + configuration['load'].update({ + 'timePartitioning': time_partitioning + }) if schema_fields: configuration['load']['schema'] = {'fields': schema_fields} @@ -1491,3 +1496,17 @@ def _split_tablename(table_input, default_project_id, var_name=None): project_id = default_project_id return project_id, dataset_id, table_id + + +def _cleanse_time_partitioning(destination_dataset_table, time_partitioning_in): + # if it is a partitioned table ($ is in the table name) add partition load option + time_partitioning_out = {} + if destination_dataset_table and '$' in destination_dataset_table: + assert not time_partitioning_in.get('field'), ( + "Cannot specify field partition and partition name " + "(dataset.table$partition) at the same time" + ) + time_partitioning_out['type'] = 'DAY' + + time_partitioning_out.update(time_partitioning_in) + return time_partitioning_out http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/5747f584/airflow/contrib/operators/bigquery_operator.py ---------------------------------------------------------------------- diff --git a/airflow/contrib/operators/bigquery_operator.py b/airflow/contrib/operators/bigquery_operator.py index dc5edb3..fb35b03 100644 --- a/airflow/contrib/operators/bigquery_operator.py +++ b/airflow/contrib/operators/bigquery_operator.py @@ -76,8 +76,17 @@ class BigQueryOperator(BaseOperator): :param query_params: a dictionary containing query parameter types and values, passed to BigQuery. :type query_params: dict - + :param priority: Specifies a priority for the query. + Possible values include INTERACTIVE and BATCH. + The default value is INTERACTIVE. + :type priority: string + :param time_partitioning: configure optional time partitioning fields i.e. + partition by field, type and + expiration as per API specifications. Note that 'field' is not available in + conjunction with dataset.table$partition. + :type time_partitioning: dict """ + template_fields = ('bql', 'destination_dataset_table') template_ext = ('.sql', ) ui_color = '#e4f0e8' @@ -99,6 +108,7 @@ class BigQueryOperator(BaseOperator): schema_update_options=(), query_params=None, priority='INTERACTIVE', + time_partitioning={}, *args, **kwargs): super(BigQueryOperator, self).__init__(*args, **kwargs) @@ -118,6 +128,7 @@ class BigQueryOperator(BaseOperator): self.query_params = query_params self.bq_cursor = None self.priority = priority + self.time_partitioning = time_partitioning def execute(self, context): if self.bq_cursor is None: @@ -140,7 +151,9 @@ class BigQueryOperator(BaseOperator): create_disposition=self.create_disposition, query_params=self.query_params, schema_update_options=self.schema_update_options, - priority=self.priority) + priority=self.priority, + time_partitioning=self.time_partitioning + ) def on_kill(self): super(BigQueryOperator, self).on_kill() http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/5747f584/tests/contrib/hooks/test_bigquery_hook.py ---------------------------------------------------------------------- diff --git a/tests/contrib/hooks/test_bigquery_hook.py b/tests/contrib/hooks/test_bigquery_hook.py index 46778b8..8841598 100644 --- a/tests/contrib/hooks/test_bigquery_hook.py +++ b/tests/contrib/hooks/test_bigquery_hook.py @@ -21,10 +21,11 @@ import unittest import mock -from airflow import AirflowException from airflow.contrib.hooks import bigquery_hook as hook from oauth2client.contrib.gce import HttpAccessTokenRefreshError +from airflow.contrib.hooks.bigquery_hook import _cleanse_time_partitioning + bq_available = True try: @@ -228,7 +229,7 @@ class TestBigQueryBaseCursor(unittest.TestCase): @mock.patch("airflow.contrib.hooks.bigquery_hook.LoggingMixin") @mock.patch("airflow.contrib.hooks.bigquery_hook.time") - def test_cancel_queries(self, mocked_logging, mocked_time): + def test_cancel_queries(self, mocked_time, mocked_logging): project_id = 12345 running_job_id = 3 @@ -263,71 +264,104 @@ class TestBigQueryBaseCursor(unittest.TestCase): class TestTimePartitioningInRunJob(unittest.TestCase): - class BigQueryBaseCursorTest(hook.BigQueryBaseCursor): - """Use this class to verify the load configuration""" - def run_with_configuration(self, configuration): - return configuration + @mock.patch("airflow.contrib.hooks.bigquery_hook.LoggingMixin") + @mock.patch("airflow.contrib.hooks.bigquery_hook.time") + @mock.patch.object(hook.BigQueryBaseCursor, 'run_with_configuration') + def test_run_load_default(self, mocked_rwc, mocked_time, mocked_logging): + project_id = 12345 - class Serv(object): - """mocks the behaviour of a succezsfull Job""" + def run_with_config(config): + self.assertIsNone(config['load'].get('timePartitioning')) + mocked_rwc.side_effect = run_with_config - class Job(object): - """mocks the behaviour of a succezsfull Job""" - def __getitem__(self, item=None): - return self + bq_hook = hook.BigQueryBaseCursor(mock.Mock(), project_id) + bq_hook.run_load( + destination_project_dataset_table='my_dataset.my_table', + schema_fields=[], + source_uris=[], + ) - def get(self, projectId, jobId=None): - return self.__getitem__(projectId) + mocked_rwc.assert_called_once() - def insert(self, projectId, body=None): - return self.get(projectId, body) + @mock.patch("airflow.contrib.hooks.bigquery_hook.LoggingMixin") + @mock.patch("airflow.contrib.hooks.bigquery_hook.time") + @mock.patch.object(hook.BigQueryBaseCursor, 'run_with_configuration') + def test_run_load_with_arg(self, mocked_rwc, mocked_time, mocked_logging): + project_id = 12345 - def execute(self): - return { - 'status': {'state': 'DONE'}, - 'jobReference': {'jobId': 0} + def run_with_config(config): + self.assertEqual( + config['load']['timePartitioning'], + { + 'field': 'test_field', + 'type': 'DAY', + 'expirationMs': 1000 } + ) + mocked_rwc.side_effect = run_with_config - def __int__(self, job='mock_load'): - self.job = job - - def jobs(self): - return self.Job() - - def test_the_job_execution_wont_break(self): - s = self.Serv() - bqc = hook.BigQueryBaseCursor(s, 'str') - job = bqc.run_load( - destination_project_dataset_table='test.teast', + 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} ) - self.assertEquals(job, 0) + mocked_rwc.assert_called_once() - def test_dollar_makes_partition(self): - s = self.Serv() - bqc = self.BigQueryBaseCursorTest(s, 'str') - cnfg = bqc.run_load( - destination_project_dataset_table='test.teast$20170101', - schema_fields=[], - source_uris=[], - src_fmt_configs={} + @mock.patch("airflow.contrib.hooks.bigquery_hook.LoggingMixin") + @mock.patch("airflow.contrib.hooks.bigquery_hook.time") + @mock.patch.object(hook.BigQueryBaseCursor, 'run_with_configuration') + def test_run_query_default(self, mocked_rwc, mocked_time, mocked_logging): + project_id = 12345 + + def run_with_config(config): + self.assertIsNone(config['query'].get('timePartitioning')) + mocked_rwc.side_effect = run_with_config + + bq_hook = hook.BigQueryBaseCursor(mock.Mock(), project_id) + bq_hook.run_query(bql='select 1') + + mocked_rwc.assert_called_once() + + @mock.patch("airflow.contrib.hooks.bigquery_hook.LoggingMixin") + @mock.patch("airflow.contrib.hooks.bigquery_hook.time") + @mock.patch.object(hook.BigQueryBaseCursor, 'run_with_configuration') + def test_run_query_with_arg(self, mocked_rwc, mocked_time, mocked_logging): + 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( + bql='select 1', + destination_dataset_table='my_dataset.my_table', + time_partitioning={'type': 'DAY', 'field': 'test_field', 'expirationMs': 1000} ) + + mocked_rwc.assert_called_once() + + def test_dollar_makes_partition(self): + tp_out = _cleanse_time_partitioning('test.teast$20170101', {}) expect = { 'type': 'DAY' } - self.assertEqual(cnfg['load'].get('timePartitioning'), expect) + self.assertEqual(tp_out, expect) def test_extra_time_partitioning_options(self): - s = self.Serv() - bqc = self.BigQueryBaseCursorTest(s, 'str') - cnfg = bqc.run_load( - destination_project_dataset_table='test.teast', - schema_fields=[], - source_uris=[], - time_partitioning={'type': 'DAY', 'field': 'test_field', 'expirationMs': 1000} + tp_out = _cleanse_time_partitioning( + 'test.teast', + {'type': 'DAY', 'field': 'test_field', 'expirationMs': 1000} ) expect = { @@ -335,20 +369,13 @@ class TestTimePartitioningInRunJob(unittest.TestCase): 'field': 'test_field', 'expirationMs': 1000 } - - self.assertEqual(cnfg['load'].get('timePartitioning'), expect) + self.assertEqual(tp_out, expect) def test_cant_add_dollar_and_field_name(self): - s = self.Serv() - bqc = self.BigQueryBaseCursorTest(s, 'str') - - with self.assertRaises(AirflowException): - tp_dict = {'type': 'DAY', 'field': 'test_field', 'expirationMs': 1000} - bqc.run_load( - destination_project_dataset_table='test.teast$20170101', - schema_fields=[], - source_uris=[], - time_partitioning=tp_dict + with self.assertRaises(AssertionError): + _cleanse_time_partitioning( + 'test.teast$20170101', + {'type': 'DAY', 'field': 'test_field', 'expirationMs': 1000} )
