Add BigQuery row to dictionary conversion in list

prefect_gcp/bigquery.py

@@ -3,7 +3,7 @@
 import os
 from functools import partial
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 from anyio import to_thread
 from prefect import get_run_logger, task
@@ -52,8 +52,9 @@ async def bigquery_query(
     to_dataframe: bool = False,
     job_config: Optional[dict] = None,
     project: Optional[str] = None,
+    result_transformer: Optional[Callable[[List["Row"]], Any]] = None,
     location: str = "US",
-) -> List["Row"]:
+) -> Any:
     """
     Runs a BigQuery query.
 
@@ -78,6 +79,7 @@ async def bigquery_query(
             (e.g., dataset references will be rejected).
         project: The project to initialize the BigQuery Client with; if not
             provided, will default to the one inferred from your credentials.
+        result_transformer: Function that can be passed to transform the result of a query before returning. The function will be passed the list of rows returned by BigQuery for the given query.
         location: Location of the dataset that will be queried.
 
     Returns:
@@ -156,10 +158,14 @@ def example_bigquery_query_flow():
         job_config=job_config,
     )
     result = await to_thread.run_sync(partial_query)
+
     if to_dataframe:
         return result.to_dataframe()
     else:
-        return list(result)
+        if result_transformer:
+            return result_transformer(result)
+        else:
+            return list(result)
 
 
 @task

tests/test_bigquery.py

@@ -17,8 +17,11 @@
 
 
 @pytest.mark.parametrize("to_dataframe", [False, True])
+@pytest.mark.parametrize("result_transformer", [None, lambda _: ("test_transformer",)])
 @pytest.mark.parametrize("dry_run_max_bytes", [None, 5, 15])
-def test_bigquery_query(to_dataframe, dry_run_max_bytes, gcp_credentials):
+def test_bigquery_query(
+    to_dataframe, result_transformer, dry_run_max_bytes, gcp_credentials
+):
     @flow
     def test_flow():
         return bigquery_query(
@@ -32,6 +35,7 @@ def test_flow():
             job_config={},
             project="project",
             location="US",
+            result_transformer=result_transformer,
         )
 
     if dry_run_max_bytes is not None and dry_run_max_bytes < 10:
@@ -42,7 +46,10 @@ def test_flow():
         if to_dataframe:
             assert result == "dataframe_query"
         else:
-            assert result == ["query"]
+            if result_transformer:
+                assert result == ("test_transformer",)
+            else:
+                assert result == ["query"]
 
 
 def test_bigquery_create_table(gcp_credentials):