Repository: incubator-airflow Updated Branches: refs/heads/master b1deb3318 -> f36ae3ac2
[AIRFLOW-2166] Restore BQ run_query dialect param Restores the use_legacy_sql parameter in the run_query method of BigQueryBaseCursor. This method was removed by commit d5d2c01f37f345458d9eeb8cdfbb0e77b55eb7ea, which introduced a backward-incompatible change for direct calls to the cursor methods. Closes #3087 from ji-han/AIRFLOW-2166 Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/f36ae3ac Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/f36ae3ac Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/f36ae3ac Branch: refs/heads/master Commit: f36ae3ac2da745eacf2c99ae4ee8aa8dc4c8594f Parents: b1deb33 Author: Winston Huang <wins...@quizlet.com> Authored: Fri Mar 2 14:40:34 2018 -0800 Committer: Chris Riccomini <criccom...@apache.org> Committed: Fri Mar 2 14:40:41 2018 -0800 ---------------------------------------------------------------------- airflow/contrib/hooks/bigquery_hook.py | 9 ++++++++- tests/contrib/hooks/test_bigquery_hook.py | 15 +++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/f36ae3ac/airflow/contrib/hooks/bigquery_hook.py ---------------------------------------------------------------------- diff --git a/airflow/contrib/hooks/bigquery_hook.py b/airflow/contrib/hooks/bigquery_hook.py index d937f1e..c6499d3 100644 --- a/airflow/contrib/hooks/bigquery_hook.py +++ b/airflow/contrib/hooks/bigquery_hook.py @@ -455,6 +455,7 @@ class BigQueryBaseCursor(LoggingMixin): allow_large_results=False, flatten_results=False, udf_config=False, + use_legacy_sql=None, maximum_billing_tier=None, maximum_bytes_billed=None, create_disposition='CREATE_IF_NEEDED', @@ -485,6 +486,9 @@ class BigQueryBaseCursor(LoggingMixin): :type flatten_results: boolean :param udf_config: The User Defined Function configuration for the query. See https://cloud.google.com/bigquery/user-defined-functions for details. + :param use_legacy_sql: Whether to use legacy SQL (true) or standard SQL (false). + If `None`, defaults to `self.use_legacy_sql`. + :type use_legacy_sql: boolean :type udf_config: list :param maximum_billing_tier: Positive integer that serves as a multiplier of the basic price. @@ -523,10 +527,13 @@ class BigQueryBaseCursor(LoggingMixin): "Please only use one or more of the following options: {1}" .format(schema_update_options, allowed_schema_update_options)) + if use_legacy_sql is None: + use_legacy_sql = self.use_legacy_sql + configuration = { 'query': { 'query': bql, - 'useLegacySql': self.use_legacy_sql, + 'useLegacySql': use_legacy_sql, 'maximumBillingTier': maximum_billing_tier, 'maximumBytesBilled': maximum_bytes_billed, 'priority': priority http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/f36ae3ac/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 a5dd595..6c6bed6 100644 --- a/tests/contrib/hooks/test_bigquery_hook.py +++ b/tests/contrib/hooks/test_bigquery_hook.py @@ -240,6 +240,21 @@ class TestBigQueryBaseCursor(unittest.TestCase): mock_jobs.cancel.assert_called_with(projectId=project_id, jobId=running_job_id) + @mock.patch.object(hook.BigQueryBaseCursor, 'run_with_configuration') + 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) + + @mock.patch.object(hook.BigQueryBaseCursor, 'run_with_configuration') + 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) + class TestTimePartitioningInRunJob(unittest.TestCase):
