Fix calling get_client in BigQueryHook.table_exists (#9916)

Adding `project_id` argument to `get_client` method otherwise this call always falls back to the default connection id.
apache · Jul 22, 2020 · c4244e1 · c4244e1
1 parent c460529
commit c4244e1
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
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')