Fix parsing of optional mode field in BigQuery Result Schema (#26786)

apache · Oct 1, 2022 · cee610a · cee610a
1 parent b7203cd
commit cee610a
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py
@@ -2999,14 +2999,15 @@ def _format_schema_for_description(schema: dict) -> list:
     """
     description = []
     for field in schema["fields"]:
+        mode = field.get("mode", "NULLABLE")
         field_description = (
             field["name"],
             field["type"],
             None,
             None,
             None,
             None,
-            field["mode"] == "NULLABLE",
+            mode == "NULLABLE",
         )
         description.append(field_description)
     return description

diff --git a/tests/providers/google/cloud/hooks/test_bigquery.py b/tests/providers/google/cloud/hooks/test_bigquery.py
@@ -1255,11 +1255,17 @@ def test_format_schema_for_description(self):
             "schema": {
                 "fields": [
                     {"name": "field_1", "type": "STRING", "mode": "NULLABLE"},
+                    {"name": "field_2", "type": "STRING"},
+                    {"name": "field_3", "type": "STRING", "mode": "REPEATED"},
                 ]
             },
         }
         description = _format_schema_for_description(test_query_result["schema"])
-        assert description == [('field_1', 'STRING', None, None, None, None, True)]
+        assert description == [
+            ('field_1', 'STRING', None, None, None, None, True),
+            ('field_2', 'STRING', None, None, None, None, True),
+            ('field_3', 'STRING', None, None, None, None, False),
+        ]
 
     @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_service")
     @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.insert_job")