Fixes the sending of an empty list to BigQuery list_rows (#12307)

* Fixes an issue that was causing an empty list being sent to the BigQuery client `list_rows` method resulting in no schema being returned. * Added a test to check that providing an empty list for `selected_fields` results in `list_rows` being called wth `None`.
apache · Nov 12, 2020 · 32b59f8 · 32b59f8
1 parent 571f831
commit 32b59f8
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 2 deletions.
diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py
@@ -1296,10 +1296,14 @@ def list_rows(
  :return: list of rows
  """
  location = location or self.location
- selected_fields = selected_fields or []
  if isinstance(selected_fields, str):
  selected_fields = selected_fields.split(",")
 
+ if selected_fields:
+ selected_fields = [SchemaField(n, "") for n in selected_fields]
+ else:
+ selected_fields = None
+
  table = self._resolve_table_reference(
  table_resource={},
  project_id=project_id,
@@ -1309,7 +1313,7 @@ def list_rows(
 
  result = self.get_client(project_id=project_id, location=location).list_rows(
  table=Table.from_api_repr(table),
- selected_fields=[SchemaField(n, "") for n in selected_fields],
+ selected_fields=selected_fields,
  max_results=max_results,
  page_token=page_token,
  start_index=start_index,

diff --git a/tests/providers/google/cloud/hooks/test_bigquery.py b/tests/providers/google/cloud/hooks/test_bigquery.py
@@ -462,6 +462,27 @@ def test_list_rows(self, mock_client, mock_schema, mock_table):
  start_index=5,
  )
 
+ @mock.patch("airflow.providers.google.cloud.hooks.bigquery.Table")
+ @mock.patch("airflow.providers.google.cloud.hooks.bigquery.Client")
+ def test_list_rows_with_empty_selected_fields(self, mock_client, mock_table):
+ self.hook.list_rows(
+ dataset_id=DATASET_ID,
+ table_id=TABLE_ID,
+ max_results=10,
+ page_token="page123",
+ selected_fields=[],
+ start_index=5,
+ location=LOCATION,
+ )
+ mock_table.from_api_repr.assert_called_once_with({"tableReference": TABLE_REFERENCE_REPR})
+ mock_client.return_value.list_rows.assert_called_once_with(
+ table=mock_table.from_api_repr.return_value,
+ max_results=10,
+ page_token='page123',
+ selected_fields=None,
+ start_index=5,
+ )
+
  @mock.patch("airflow.providers.google.cloud.hooks.bigquery.Table")
  @mock.patch("airflow.providers.google.cloud.hooks.bigquery.Client")
  def test_run_table_delete(self, mock_client, mock_table):