Revert "fix(core): Make log_adapter state thread-local"

This reverts commit e2c38d3.

Author: shuoweil

bigframes/core/log_adapter.py

@@ -15,14 +15,12 @@
 import functools
 import inspect
 import threading
-from typing import Optional
+from typing import List, Optional
 
 from google.cloud import bigquery
 import pandas
 
-_thread_local_data = threading.local()
-_thread_local_data._api_methods = []
-_thread_local_data._call_stack = []
+_lock = threading.Lock()
 
 # The limit is 64 (https://cloud.google.com/bigquery/docs/labels-intro#requirements),
 # but leave a few spare for internal labels to be added.
@@ -32,8 +30,12 @@
 PANDAS_PARAM_TRACKING_TASK = "pandas_param_tracking"
 LOG_OVERRIDE_NAME = "__log_override_name__"
 
+_api_methods: List = []
 _excluded_methods = ["__setattr__", "__getattr__"]
 
+# Stack to track method calls
+_call_stack: List = []
+
 
 def submit_pandas_labels(
     bq_client: Optional[bigquery.Client],
@@ -170,14 +172,11 @@ def wrapper(*args, **kwargs):
                 base_name = custom_base_name
 
             full_method_name = f"{base_name.lower()}-{api_method_name}"
-            if not hasattr(_thread_local_data, "_call_stack"):
-                _thread_local_data._call_stack = []
-
             # Track directly called methods
-            if len(_thread_local_data._call_stack) == 0:
+            if len(_call_stack) == 0:
                 add_api_method(full_method_name)
 
-            _thread_local_data._call_stack.append(full_method_name)
+            _call_stack.append(full_method_name)
 
             try:
                 return method(*args, **kwargs)
@@ -186,7 +185,7 @@ def wrapper(*args, **kwargs):
                 # or not fully supported (NotImplementedError) in BigFrames.
                 # Logging is currently supported only when we can access the bqclient through
                 # _block.session.bqclient.
-                if len(_thread_local_data._call_stack) == 1:
+                if len(_call_stack) == 1:
                     submit_pandas_labels(
                         _get_bq_client(*args, **kwargs),
                         base_name,
@@ -197,7 +196,7 @@ def wrapper(*args, **kwargs):
                     )
                 raise e
             finally:
-                _thread_local_data._call_stack.pop()
+                _call_stack.pop()
 
         return wrapper
 
@@ -215,21 +214,19 @@ def property_logger(prop):
     def shared_wrapper(prop):
         @functools.wraps(prop)
         def wrapped(*args, **kwargs):
-            if not hasattr(_thread_local_data, "_call_stack"):
-                _thread_local_data._call_stack = []
             qualname_parts = getattr(prop, "__qualname__", prop.__name__).split(".")
             class_name = qualname_parts[-2] if len(qualname_parts) > 1 else ""
             property_name = prop.__name__
             full_property_name = f"{class_name.lower()}-{property_name.lower()}"
 
-            if len(_thread_local_data._call_stack) == 0:
+            if len(_call_stack) == 0:
                 add_api_method(full_property_name)
 
-            _thread_local_data._call_stack.append(full_property_name)
+            _call_stack.append(full_property_name)
             try:
                 return prop(*args, **kwargs)
             finally:
-                _thread_local_data._call_stack.pop()
+                _call_stack.pop()
 
         return wrapped
 
@@ -254,26 +251,23 @@ def wrapper(func):
 
 
 def add_api_method(api_method_name):
-    if not hasattr(_thread_local_data, "_api_methods"):
-        _thread_local_data._api_methods = []
-
-    # Push the method to the front of the _api_methods list
-    _thread_local_data._api_methods.insert(
-        0, api_method_name.replace("<", "").replace(">", "")
-    )
-    # Keep the list length within the maximum limit
-    _thread_local_data._api_methods = _thread_local_data._api_methods[:MAX_LABELS_COUNT]
+    global _lock
+    global _api_methods
+    with _lock:
+        # Push the method to the front of the _api_methods list
+        _api_methods.insert(0, api_method_name.replace("<", "").replace(">", ""))
+        # Keep the list length within the maximum limit (adjust MAX_LABELS_COUNT as needed)
+        _api_methods = _api_methods[:MAX_LABELS_COUNT]
 
 
 def get_and_reset_api_methods(dry_run: bool = False):
-    if not hasattr(_thread_local_data, "_api_methods"):
-        _thread_local_data._api_methods = []
-
-    previous_api_methods = list(_thread_local_data._api_methods)
+    global _lock
+    with _lock:
+        previous_api_methods = list(_api_methods)
 
-    # dry_run might not make a job resource, so only reset the log on real queries.
-    if not dry_run:
-        _thread_local_data._api_methods.clear()
+        # dry_run might not make a job resource, so only reset the log on real queries.
+        if not dry_run:
+            _api_methods.clear()
     return previous_api_methods
 
 

bigframes/testing/mocks.py

@@ -27,10 +27,7 @@
 import bigframes
 import bigframes.clients
 import bigframes.core.global_session
-import bigframes.core.log_adapter
 import bigframes.dataframe
-import bigframes.session._io.bigquery
-from bigframes.session._io.bigquery import create_job_configs_labels
 import bigframes.session.clients
 
 """Utilities for creating test resources."""
@@ -93,18 +90,6 @@ def query_mock(
         job_config: Optional[google.cloud.bigquery.QueryJobConfig] = None,
         **kwargs,
     ):
-        job_config = (
-            job_config
-            if job_config is not None
-            else google.cloud.bigquery.QueryJobConfig()
-        )
-        api_methods = bigframes.core.log_adapter.get_and_reset_api_methods(
-            dry_run=job_config.dry_run
-        )
-        job_config.labels = create_job_configs_labels(
-            job_configs_labels=job_config.labels,
-            api_methods=api_methods,
-        )
         queries.append(query)
         job_configs.append(copy.deepcopy(job_config))
         query_job = mock.create_autospec(google.cloud.bigquery.QueryJob, instance=True)