chore: add callback argument to _read_gbq_colab

Author: tswast

bigframes/session/__init__.py

@@ -24,7 +24,6 @@
 import typing
 from typing import (
     Any,
-    Callable,
     Dict,
     IO,
     Iterable,
@@ -504,6 +503,9 @@ def _read_gbq_colab(
         *,
         pyformat_args: Optional[Dict[str, Any]] = None,
         dry_run: bool = False,
+        callback: abc.Callable[
+            [Any], None
+        ] = lambda _: None,  # TODO: use an Event class not Any.
     ) -> Union[dataframe.DataFrame, pandas.Series]:
         """A version of read_gbq that has the necessary default values for use in colab integrations.
 
@@ -535,6 +537,7 @@ def _read_gbq_colab(
             index_col=bigframes.enums.DefaultIndexKind.NULL,
             force_total_order=False,
             dry_run=typing.cast(Union[Literal[False], Literal[True]], dry_run),
+            callback=callback,
         )
 
     @overload
@@ -992,7 +995,7 @@ def read_csv(
                 pandas.Series,
                 pandas.Index,
                 np.ndarray[Any, Any],
-                Callable[[Any], bool],
+                abc.Callable[[Any], bool],
             ]
         ] = None,
         dtype: Optional[Dict] = None,

bigframes/session/loader.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+from collections import abc
 import copy
 import dataclasses
 import datetime
@@ -22,6 +23,7 @@
 import os
 import typing
 from typing import (
+    Any,
     Dict,
     Generator,
     Hashable,
@@ -743,6 +745,9 @@ def read_gbq_query(
         filters: third_party_pandas_gbq.FiltersType = (),
         dry_run: bool = False,
         force_total_order: Optional[bool] = None,
+        callback: abc.Callable[
+            [Any], None
+        ] = lambda _: None,  # TODO: use an Event class not Any.
     ) -> dataframe.DataFrame | pandas.Series:
         import bigframes.dataframe as dataframe