Merge remote-tracking branch 'origin/main' into tswast-doctest-boilerplate

Author: tswast

bigframes/blob/_functions.py

@@ -99,6 +99,7 @@ def _create_udf(self):
             project=None,
             timeout=None,
             query_with_job=True,
+            publisher=self._session._publisher,
         )
 
         return udf_name

bigframes/core/events.py

@@ -0,0 +1,237 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import dataclasses
+import datetime
+import threading
+from typing import Any, Callable, Optional, Set
+import uuid
+
+import google.cloud.bigquery._job_helpers
+import google.cloud.bigquery.job.query
+import google.cloud.bigquery.table
+
+import bigframes.session.executor
+
+
+class Subscriber:
+    def __init__(self, callback: Callable[[Event], None], *, publisher: Publisher):
+        self._publisher = publisher
+        self._callback = callback
+        self._subscriber_id = uuid.uuid4()
+
+    def __call__(self, *args, **kwargs):
+        return self._callback(*args, **kwargs)
+
+    def __hash__(self) -> int:
+        return hash(self._subscriber_id)
+
+    def __eq__(self, value: object):
+        if not isinstance(value, Subscriber):
+            return NotImplemented
+        return value._subscriber_id == self._subscriber_id
+
+    def close(self):
+        self._publisher.unsubscribe(self)
+        del self._publisher
+        del self._callback
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_value is not None:
+            self(
+                UnknownErrorEvent(
+                    exc_type=exc_type,
+                    exc_value=exc_value,
+                    traceback=traceback,
+                )
+            )
+        self.close()
+
+
+class Publisher:
+    def __init__(self):
+        self._subscribers_lock = threading.Lock()
+        self._subscribers: Set[Subscriber] = set()
+
+    def subscribe(self, callback: Callable[[Event], None]) -> Subscriber:
+        # TODO(b/448176657): figure out how to handle subscribers/publishers in
+        # a background thread. Maybe subscribers should be thread-local?
+        subscriber = Subscriber(callback, publisher=self)
+        with self._subscribers_lock:
+            self._subscribers.add(subscriber)
+        return subscriber
+
+    def unsubscribe(self, subscriber: Subscriber):
+        with self._subscribers_lock:
+            self._subscribers.remove(subscriber)
+
+    def publish(self, event: Event):
+        with self._subscribers_lock:
+            for subscriber in self._subscribers:
+                subscriber(event)
+
+
+class Event:
+    pass
+
+
+@dataclasses.dataclass(frozen=True)
+class SessionClosed(Event):
+    session_id: str
+
+
+class ExecutionStarted(Event):
+    pass
+
+
+class ExecutionRunning(Event):
+    pass
+
+
+@dataclasses.dataclass(frozen=True)
+class ExecutionFinished(Event):
+    result: Optional[bigframes.session.executor.ExecuteResult] = None
+
+
+@dataclasses.dataclass(frozen=True)
+class UnknownErrorEvent(Event):
+    exc_type: Any
+    exc_value: Any
+    traceback: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class BigQuerySentEvent(ExecutionRunning):
+    """Query sent to BigQuery."""
+
+    query: str
+    billing_project: Optional[str] = None
+    location: Optional[str] = None
+    job_id: Optional[str] = None
+    request_id: Optional[str] = None
+
+    @classmethod
+    def from_bqclient(cls, event: google.cloud.bigquery._job_helpers.QuerySentEvent):
+        return cls(
+            query=event.query,
+            billing_project=event.billing_project,
+            location=event.location,
+            job_id=event.job_id,
+            request_id=event.request_id,
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class BigQueryRetryEvent(ExecutionRunning):
+    """Query sent another time because the previous attempt failed."""
+
+    query: str
+    billing_project: Optional[str] = None
+    location: Optional[str] = None
+    job_id: Optional[str] = None
+    request_id: Optional[str] = None
+
+    @classmethod
+    def from_bqclient(cls, event: google.cloud.bigquery._job_helpers.QueryRetryEvent):
+        return cls(
+            query=event.query,
+            billing_project=event.billing_project,
+            location=event.location,
+            job_id=event.job_id,
+            request_id=event.request_id,
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class BigQueryReceivedEvent(ExecutionRunning):
+    """Query received and acknowledged by the BigQuery API."""
+
+    billing_project: Optional[str] = None
+    location: Optional[str] = None
+    job_id: Optional[str] = None
+    statement_type: Optional[str] = None
+    state: Optional[str] = None
+    query_plan: Optional[list[google.cloud.bigquery.job.query.QueryPlanEntry]] = None
+    created: Optional[datetime.datetime] = None
+    started: Optional[datetime.datetime] = None
+    ended: Optional[datetime.datetime] = None
+
+    @classmethod
+    def from_bqclient(
+        cls, event: google.cloud.bigquery._job_helpers.QueryReceivedEvent
+    ):
+        return cls(
+            billing_project=event.billing_project,
+            location=event.location,
+            job_id=event.job_id,
+            statement_type=event.statement_type,
+            state=event.state,
+            query_plan=event.query_plan,
+            created=event.created,
+            started=event.started,
+            ended=event.ended,
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class BigQueryFinishedEvent(ExecutionRunning):
+    """Query finished successfully."""
+
+    billing_project: Optional[str] = None
+    location: Optional[str] = None
+    query_id: Optional[str] = None
+    job_id: Optional[str] = None
+    destination: Optional[google.cloud.bigquery.table.TableReference] = None
+    total_rows: Optional[int] = None
+    total_bytes_processed: Optional[int] = None
+    slot_millis: Optional[int] = None
+    created: Optional[datetime.datetime] = None
+    started: Optional[datetime.datetime] = None
+    ended: Optional[datetime.datetime] = None
+
+    @classmethod
+    def from_bqclient(
+        cls, event: google.cloud.bigquery._job_helpers.QueryFinishedEvent
+    ):
+        return cls(
+            billing_project=event.billing_project,
+            location=event.location,
+            query_id=event.query_id,
+            job_id=event.job_id,
+            destination=event.destination,
+            total_rows=event.total_rows,
+            total_bytes_processed=event.total_bytes_processed,
+            slot_millis=event.slot_millis,
+            created=event.created,
+            started=event.started,
+            ended=event.ended,
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class BigQueryUnknownEvent(ExecutionRunning):
+    """Got unknown event from the BigQuery client library."""
+
+    # TODO: should we just skip sending unknown events?
+
+    event: object
+
+    @classmethod
+    def from_bqclient(cls, event):
+        return cls(event)

bigframes/core/global_session.py

@@ -14,16 +14,19 @@
 
 """Utilities for managing a default, globally available Session object."""
 
+from __future__ import annotations
+
 import threading
 import traceback
-from typing import Callable, Optional, TypeVar
+from typing import Callable, Optional, TYPE_CHECKING, TypeVar
 import warnings
 
 import google.auth.exceptions
 
-import bigframes._config
 import bigframes.exceptions as bfe
-import bigframes.session
+
+if TYPE_CHECKING:
+    import bigframes.session
 
 _global_session: Optional[bigframes.session.Session] = None
 _global_session_lock = threading.Lock()
@@ -56,6 +59,9 @@ def close_session() -> None:
     Returns:
         None
     """
+    # Avoid troubles with circular imports.
+    import bigframes._config
+
     global _global_session, _global_session_lock, _global_session_state
 
     if bigframes._config.options.is_bigquery_thread_local:
@@ -88,6 +94,10 @@ def get_global_session():
 
     Creates the global session if it does not exist.
     """
+    # Avoid troubles with circular imports.
+    import bigframes._config
+    import bigframes.session
+
     global _global_session, _global_session_lock, _global_session_state
 
     if bigframes._config.options.is_bigquery_thread_local:

bigframes/dataframe.py

@@ -4666,24 +4666,24 @@ def to_string(
     ) -> str | None:
         return self.to_pandas(allow_large_results=allow_large_results).to_string(
             buf,
-            columns,  # type: ignore
-            col_space,
-            header,  # type: ignore
-            index,
-            na_rep,
-            formatters,
-            float_format,
-            sparsify,
-            index_names,
-            justify,
-            max_rows,
-            max_cols,
-            show_dimensions,
-            decimal,
-            line_width,
-            min_rows,
-            max_colwidth,
-            encoding,
+            columns=columns,  # type: ignore
+            col_space=col_space,
+            header=header,  # type: ignore
+            index=index,
+            na_rep=na_rep,
+            formatters=formatters,
+            float_format=float_format,
+            sparsify=sparsify,
+            index_names=index_names,
+            justify=justify,
+            max_rows=max_rows,
+            max_cols=max_cols,
+            show_dimensions=show_dimensions,
+            decimal=decimal,
+            line_width=line_width,
+            min_rows=min_rows,
+            max_colwidth=max_colwidth,
+            encoding=encoding,
         )
 
     def to_html(
@@ -4716,28 +4716,28 @@ def to_html(
     ) -> str:
         return self.to_pandas(allow_large_results=allow_large_results).to_html(
             buf,
-            columns,  # type: ignore
-            col_space,
-            header,
-            index,
-            na_rep,
-            formatters,
-            float_format,
-            sparsify,
-            index_names,
-            justify,  # type: ignore
-            max_rows,
-            max_cols,
-            show_dimensions,
-            decimal,
-            bold_rows,
-            classes,
-            escape,
-            notebook,
-            border,
-            table_id,
-            render_links,
-            encoding,
+            columns=columns,  # type: ignore
+            col_space=col_space,
+            header=header,
+            index=index,
+            na_rep=na_rep,
+            formatters=formatters,
+            float_format=float_format,
+            sparsify=sparsify,
+            index_names=index_names,
+            justify=justify,  # type: ignore
+            max_rows=max_rows,
+            max_cols=max_cols,
+            show_dimensions=show_dimensions,
+            decimal=decimal,
+            bold_rows=bold_rows,
+            classes=classes,
+            escape=escape,
+            notebook=notebook,
+            border=border,
+            table_id=table_id,
+            render_links=render_links,
+            encoding=encoding,
         )
 
     def to_markdown(
@@ -4749,7 +4749,7 @@ def to_markdown(
         allow_large_results: Optional[bool] = None,
         **kwargs,
     ) -> str | None:
-        return self.to_pandas(allow_large_results=allow_large_results).to_markdown(buf, mode, index, **kwargs)  # type: ignore
+        return self.to_pandas(allow_large_results=allow_large_results).to_markdown(buf, mode=mode, index=index, **kwargs)  # type: ignore
 
     def to_pickle(self, path, *, allow_large_results=None, **kwargs) -> None:
         return self.to_pandas(allow_large_results=allow_large_results).to_pickle(