Support project_id argument in BigQueryGetDataOperator (#25782)

sudohainguyen · web-flow · commit fc6dfa338a76 · 2022-08-20T23:14:09.000+02:00
diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
@@ -361,13 +361,16 @@ class BigQueryGetDataOperator(BaseOperator):
             task_id='get_data_from_bq',
             dataset_id='test_dataset',
             table_id='Transaction_partitions',
+            project_id='internal-gcp-project',
             max_results=100,
             selected_fields='DATE',
             gcp_conn_id='airflow-conn-id'
         )
 
     :param dataset_id: The dataset ID of the requested table. (templated)
     :param table_id: The table ID of the requested table. (templated)
+    :param project_id: (Optional) The name of the project where the data
+        will be returned from. (templated)
     :param max_results: The maximum number of records (rows) to be fetched
         from the table. (templated)
     :param selected_fields: List of fields to return (comma-separated). If
@@ -390,6 +393,7 @@ class BigQueryGetDataOperator(BaseOperator):
     template_fields: Sequence[str] = (
         'dataset_id',
         'table_id',
+        'project_id',
         'max_results',
         'selected_fields',
         'impersonation_chain',
@@ -401,6 +405,7 @@ def __init__(
         *,
         dataset_id: str,
         table_id: str,
+        project_id: Optional[str] = None,
         max_results: int = 100,
         selected_fields: Optional[str] = None,
         gcp_conn_id: str = 'google_cloud_default',
@@ -419,6 +424,7 @@ def __init__(
         self.delegate_to = delegate_to
         self.location = location
         self.impersonation_chain = impersonation_chain
+        self.project_id = project_id
 
     def execute(self, context: 'Context') -> list:
         self.log.info(
@@ -445,6 +451,7 @@ def execute(self, context: 'Context') -> list:
             max_results=self.max_results,
             selected_fields=self.selected_fields,
             location=self.location,
+            project_id=self.project_id,
         )
 
         self.log.info('Total extracted rows: %s', len(rows))
diff --git a/tests/providers/google/cloud/operators/test_bigquery.py b/tests/providers/google/cloud/operators/test_bigquery.py
@@ -717,6 +717,7 @@ def test_execute(self, mock_hook):
             task_id=TASK_ID,
             dataset_id=TEST_DATASET,
             table_id=TEST_TABLE_ID,
+            project_id=TEST_GCP_PROJECT_ID,
             max_results=max_results,
             selected_fields=selected_fields,
             location=TEST_DATASET_LOCATION,
@@ -725,6 +726,7 @@ def test_execute(self, mock_hook):
         mock_hook.return_value.list_rows.assert_called_once_with(
             dataset_id=TEST_DATASET,
             table_id=TEST_TABLE_ID,
+            project_id=TEST_GCP_PROJECT_ID,
             max_results=max_results,
             selected_fields=selected_fields,
             location=TEST_DATASET_LOCATION,
