Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions bigframes/core/global_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
24 changes: 16 additions & 8 deletions third_party/bigframes_vendored/pandas/_config/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,6 @@
import contextlib
import operator

import bigframes


class option_context(contextlib.ContextDecorator):
"""
Expand Down Expand Up @@ -35,8 +33,11 @@ def __init__(self, *args) -> None:
self.ops = list(zip(args[::2], args[1::2]))

def __enter__(self) -> None:
# Avoid problems with circular imports.
import bigframes._config

self.undo = [
(pat, operator.attrgetter(pat)(bigframes.options))
(pat, operator.attrgetter(pat)(bigframes._config.options))
for pat, _ in self.ops
# Don't try to undo changes to bigquery options. We're starting and
# closing a new thread-local session if those are set.
Expand All @@ -47,25 +48,32 @@ def __enter__(self) -> None:
self._set_option(pat, val)

def __exit__(self, *args) -> None:
# Avoid problems with circular imports.
import bigframes._config
import bigframes.core.global_session

if self.undo:
for pat, val in self.undo:
self._set_option(pat, val)

# TODO(tswast): What to do if someone nests several context managers
# with separate "bigquery" options? We might need a "stack" of
# sessions if we allow that.
if bigframes.options.is_bigquery_thread_local:
bigframes.close_session()
if bigframes._config.options.is_bigquery_thread_local:
bigframes.core.global_session.close_session()

# Reset bigquery_options so that we're no longer thread-local.
bigframes.options._local.bigquery_options = None
bigframes._config.options._local.bigquery_options = None

def _set_option(self, pat, val):
# Avoid problems with circular imports.
import bigframes._config

root, attr = pat.rsplit(".", 1)

# We are now using a thread-specific session.
if root == "bigquery":
bigframes.options._init_bigquery_thread_local()
bigframes._config.options._init_bigquery_thread_local()

parent = operator.attrgetter(root)(bigframes.options)
parent = operator.attrgetter(root)(bigframes._config.options)
setattr(parent, attr, val)