Don't throw an exception when a BQ cusor job has no schema (#26096)

apache · Sep 7, 2022 · 12cbc0f · 12cbc0f
1 parent b7969d4
commit 12cbc0f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git 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
@@ -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()