diff --git a/airflow/providers/google/cloud/transfers/bigquery_to_mysql.py b/airflow/providers/google/cloud/transfers/bigquery_to_mysql.py index 6adfd5e792c71..be0de9ab7d51d 100644 --- a/airflow/providers/google/cloud/transfers/bigquery_to_mysql.py +++ b/airflow/providers/google/cloud/transfers/bigquery_to_mysql.py @@ -129,21 +129,17 @@ def _bq_get_data(self): impersonation_chain=self.impersonation_chain, ) - conn = hook.get_conn() - cursor = conn.cursor() i = 0 while True: - response = cursor.get_tabledata( + response = hook.list_rows( dataset_id=self.dataset_id, table_id=self.table_id, max_results=self.batch_size, selected_fields=self.selected_fields, start_index=i * self.batch_size, ) - - if 'rows' in response: - rows = response['rows'] - else: + rows = [dict(r) for r in response] + if len(rows) == 0: self.log.info('Job Finished') return diff --git a/tests/providers/google/cloud/transfers/test_bigquery_to_mysql.py b/tests/providers/google/cloud/transfers/test_bigquery_to_mysql.py index acdf63dd0b103..a046020fd0568 100644 --- a/tests/providers/google/cloud/transfers/test_bigquery_to_mysql.py +++ b/tests/providers/google/cloud/transfers/test_bigquery_to_mysql.py @@ -39,12 +39,11 @@ def test_execute_good_request_to_bq(self, mock_hook): operator.execute(None) # fmt: off - mock_hook.return_value.get_conn.return_value.cursor.return_value.get_tabledata\ - .assert_called_once_with( - dataset_id=TEST_DATASET, - table_id=TEST_TABLE_ID, - max_results=1000, - selected_fields=None, - start_index=0, - ) + mock_hook.return_value.list_rows.assert_called_once_with( + dataset_id=TEST_DATASET, + table_id=TEST_TABLE_ID, + max_results=1000, + selected_fields=None, + start_index=0, + ) # fmt: on