shahar1 commented on code in PR #39315:
URL: https://github.com/apache/airflow/pull/39315#discussion_r1585358900


##########
airflow/providers/google/cloud/hooks/bigquery.py:
##########
@@ -2390,6 +2397,50 @@ def var_print(var_name):
 
         return project_id, dataset_id, table_id
 
+    @GoogleBaseHook.fallback_to_default_project_id
+    def get_query_results(
+        self,
+        job_id: str,
+        location: str,
+        max_results: int | None = None,
+        selected_fields: list[str] | str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
+        retry: Retry = DEFAULT_RETRY,
+        job_retry: Retry = DEFAULT_JOB_RETRY,
+    ) -> list[dict[str, Any]]:
+        """
+        Get query results given a job_id.
+
+        :param job_id: The ID of the job.
+            The ID must contain only letters (a-z, A-Z), numbers (0-9), 
underscores (_), or
+            dashes (-). The maximum length is 1,024 characters.
+        :param location: The location used for the operation.
+        :param selected_fields: List of fields to return (comma-separated). If
+            unspecified, all fields are returned.
+        :param max_results: The maximum number of records (rows) to be fetched
+            from the table.
+        :param project_id: Google Cloud Project where the job ran.
+        :param retry: How to retry the RPC.
+        :param job_retry: How to retry failed jobs.
+
+        :return: List of rows where columns are filtered by selected fields, 
when given
+
+        :raises: AirflowException
+        """
+        if isinstance(selected_fields, str):
+            selected_fields = selected_fields.split(",")
+        job = self.get_job(job_id=job_id, project_id=project_id, 
location=location)
+        if not isinstance(job, QueryJob):
+            raise AirflowException(f"Job '{job_id}' is not a query job")
+
+        if job.state != "DONE":
+            raise AirflowException(f"Job '{job_id}' is not in DONE state")
+
+        rows = [
+            dict(row) for row in list(job.result(max_results=max_results, 
retry=retry, job_retry=job_retry))

Review Comment:
   Nope, it is already iterable - good call :)



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to