Fix BigQuery data extraction in BigQueryToMySqlOperator (#18073)

closes: #17198
apache · Sep 9, 2021 · cfb602a · cfb602a
1 parent 2767781
commit cfb602a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 15 deletions.
diff --git 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
@@ -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