Fix insert_all method of BigQueryHook to support tables without schema (

#13138) Fixes the insert_all method of BigQueryHook to support tables that does not have provided schema. Co-authored-by: Manuel Bordes <manuel.bordes@gamesys.co.uk>
apache · Jan 4, 2021 · 3a3e739 · 3a3e739
1 parent 35e4a3b
commit 3a3e739
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 9 deletions.
diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py
@@ -876,11 +876,11 @@ def insert_all(
         """
         self.log.info('Inserting %s row(s) into table %s:%s.%s', len(rows), project_id, dataset_id, table_id)
 
-        table = self._resolve_table_reference(
-            table_resource={}, project_id=project_id, dataset_id=dataset_id, table_id=table_id
-        )
-        errors = self.get_client().insert_rows(
-            table=Table.from_api_repr(table),
+        table_ref = TableReference(dataset_ref=DatasetReference(project_id, dataset_id), table_id=table_id)
+        bq_client = self.get_client(project_id=project_id)
+        table = bq_client.get_table(table_ref)
+        errors = bq_client.insert_rows(
+            table=table,
             rows=rows,
             ignore_unknown_values=ignore_unknown_values,
             skip_invalid_rows=skip_invalid_rows,

diff --git a/tests/providers/google/cloud/hooks/test_bigquery.py b/tests/providers/google/cloud/hooks/test_bigquery.py
@@ -677,9 +677,8 @@ def test_invalid_source_format(self, mock_get_service):
         ):
             self.hook.run_load("test.test", "test_schema.json", ["test_data.json"], source_format="json")
 
-    @mock.patch("airflow.providers.google.cloud.hooks.bigquery.Table")
     @mock.patch("airflow.providers.google.cloud.hooks.bigquery.Client")
-    def test_insert_all_succeed(self, mock_client, mock_table):
+    def test_insert_all_succeed(self, mock_client):
         rows = [{"json": {"a_key": "a_value_0"}}]
 
         self.hook.insert_all(
@@ -690,9 +689,9 @@ def test_insert_all_succeed(self, mock_client, mock_table):
             ignore_unknown_values=True,
             skip_invalid_rows=True,
         )
-        mock_table.from_api_repr.assert_called_once_with({"tableReference": TABLE_REFERENCE_REPR})
+        mock_client.return_value.get_table.assert_called_once_with(TABLE_REFERENCE)
         mock_client.return_value.insert_rows.assert_called_once_with(
-            table=mock_table.from_api_repr.return_value,
+            table=mock_client.return_value.get_table.return_value,
             rows=rows,
             ignore_unknown_values=True,
             skip_invalid_rows=True,