From 8ef544f1157ccf14759e8feda7560766d2f283ce Mon Sep 17 00:00:00 2001 From: Peter Wicks Date: Wed, 31 Aug 2022 16:47:29 -0600 Subject: [PATCH] Don't throw an exception when a BQ cusor job has no schema --- airflow/providers/google/cloud/hooks/bigquery.py | 6 ++++-- tests/providers/google/cloud/hooks/test_bigquery.py | 11 +++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py index 6776942fc5c0b..3ec149b577a53 100644 --- a/airflow/providers/google/cloud/hooks/bigquery.py +++ b/airflow/providers/google/cloud/hooks/bigquery.py @@ -2695,8 +2695,10 @@ def execute(self, operation: str, parameters: Optional[dict] = None) -> None: self.job_id = self.hook.run_query(sql) query_results = self._get_query_result() - description = _format_schema_for_description(query_results["schema"]) - self.description = description + if "schema" in query_results: + self.description = _format_schema_for_description(query_results["schema"]) + else: + self.description = [] def executemany(self, operation: str, seq_of_parameters: list) -> None: """ diff --git a/tests/providers/google/cloud/hooks/test_bigquery.py b/tests/providers/google/cloud/hooks/test_bigquery.py index 9342e9dc7675c..6d69431f14d51 100644 --- a/tests/providers/google/cloud/hooks/test_bigquery.py +++ b/tests/providers/google/cloud/hooks/test_bigquery.py @@ -1268,6 +1268,17 @@ def test_description(self, mock_insert, mock_get_service): bq_cursor.execute("SELECT CURRENT_TIMESTAMP() as ts") assert bq_cursor.description == [("ts", "TIMESTAMP", None, None, None, None, True)] + @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_service") + @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.insert_job") + def test_description_no_schema(self, mock_insert, mock_get_service): + mock_get_query_results = mock_get_service.return_value.jobs.return_value.getQueryResults + mock_execute = mock_get_query_results.return_value.execute + mock_execute.return_value = {} + + bq_cursor = self.hook.get_cursor() + bq_cursor.execute("UPDATE airflow.test_table SET foo = 'bar'") + assert bq_cursor.description == [] + @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_service") def test_close(self, mock_get_service): bq_cursor = self.hook.get_cursor()