Fix calling get_client in BigQueryHook.table_exists (#9916)

nathadfield · web-flow · commit c4244e18bb89 · 2020-07-22T08:10:56.000+02:00
Adding `project_id` argument to `get_client` method 
otherwise this call always falls back to the default connection id.
diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py
@@ -215,7 +215,7 @@ def table_exists(self, dataset_id: str, table_id: str, project_id: str) -> bool:
         """
         table_reference = TableReference(DatasetReference(project_id, dataset_id), table_id)
         try:
-            self.get_client().get_table(table_reference)
+            self.get_client(project_id=project_id).get_table(table_reference)
             return True
         except NotFound:
             return False
diff --git a/tests/providers/google/cloud/hooks/test_bigquery.py b/tests/providers/google/cloud/hooks/test_bigquery.py
@@ -104,18 +104,19 @@ def test_bigquery_insert_rows_not_implemented(self):
         with self.assertRaises(NotImplementedError):
             self.hook.insert_rows(table="table", rows=[1, 2])
 
-    @mock.patch("airflow.providers.google.cloud.hooks.bigquery.Client")
+    @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_client")
     def test_bigquery_table_exists_true(self, mock_client):
         result = self.hook.table_exists(project_id=PROJECT_ID, dataset_id=DATASET_ID, table_id=TABLE_ID)
         mock_client.return_value.get_table.assert_called_once_with(TABLE_REFERENCE)
+        mock_client.assert_called_once_with(project_id=PROJECT_ID)
         assert result is True
 
-    @mock.patch("airflow.providers.google.cloud.hooks.bigquery.Client")
+    @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_client")
     def test_bigquery_table_exists_false(self, mock_client):
         mock_client.return_value.get_table.side_effect = NotFound("Dataset not found")
-
         result = self.hook.table_exists(project_id=PROJECT_ID, dataset_id=DATASET_ID, table_id=TABLE_ID)
         mock_client.return_value.get_table.assert_called_once_with(TABLE_REFERENCE)
+        mock_client.assert_called_once_with(project_id=PROJECT_ID)
         assert result is False
 
     @mock.patch('airflow.providers.google.cloud.hooks.bigquery.read_gbq')
