Make the global _api_methods and lock threads

bigframes/core/log_adapter.py

@@ -14,41 +14,30 @@
 
 import functools
 import threading
+from typing import List
 
 _lock = threading.Lock()
 MAX_LABELS_COUNT = 64
+_api_methods: List = []
 
 
-def class_logger(api_methods=None):
+def class_logger(decorated_cls):
     """Decorator that adds logging functionality to each method of the class."""
+    for attr_name, attr_value in decorated_cls.__dict__.items():
+        if callable(attr_value):
+            setattr(decorated_cls, attr_name, method_logger(attr_value))
+    return decorated_cls
 
-    def decorator(decorated_cls):
-        for attr_name, attr_value in decorated_cls.__dict__.items():
 
-            if callable(attr_value):
-                setattr(
-                    decorated_cls, attr_name, method_logger(attr_value, decorated_cls)
-                )
-
-        # Initialize or extend _api_methods attribute
-        decorated_cls._api_methods = getattr(decorated_cls, "_api_methods", [])
-        if api_methods:
-            decorated_cls._api_methods.extend(api_methods)
-
-        return decorated_cls
-
-    return decorator
-
-
-def method_logger(method, cls):
+def method_logger(method):
     """Decorator that adds logging functionality to a method."""
 
     @functools.wraps(method)
     def wrapper(*args, **kwargs):
         api_method_name = str(method.__name__)
         # Track regular and "dunder" methods
         if api_method_name.startswith("__") or not api_method_name.startswith("_"):
-            add_api_method(api_method_name, cls)
+            add_api_method(api_method_name)
         try:
             result = method(*args, **kwargs)
             return result
@@ -58,18 +47,19 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-def add_api_method(api_method_name, cls):
+def add_api_method(api_method_name):
     global _lock
+    global _api_methods
     with _lock:
         # Push the method to the front of the _api_methods list
-        cls._api_methods.insert(0, api_method_name)
+        _api_methods.insert(0, api_method_name)
         # Keep the list length within the maximum limit (adjust MAX_LABELS_COUNT as needed)
-        cls._api_methods = cls._api_methods[:MAX_LABELS_COUNT]
+        _api_methods = _api_methods[:MAX_LABELS_COUNT]
 
 
-def get_and_reset_api_methods(cls):
+def get_and_reset_api_methods():
     global _lock
     with _lock:
-        previous_api_methods = list(cls._api_methods)
-        cls._api_methods.clear()
+        previous_api_methods = list(_api_methods)
+        _api_methods.clear()
     return previous_api_methods

bigframes/session/__init__.py

@@ -113,7 +113,6 @@ def _is_query(query_or_table: str) -> bool:
     return re.search(r"\s", query_or_table.strip(), re.MULTILINE) is not None
 
 
-@log_adapter.class_logger
 class Session(
     third_party_pandas_gbq.GBQIOMixin,
     third_party_pandas_parquet.ParquetIOMixin,
@@ -1387,7 +1386,7 @@ def _start_query(
         """
         Starts query job and waits for results.
         """
-        api_methods = log_adapter.get_and_reset_api_methods(self)
+        api_methods = log_adapter.get_and_reset_api_methods()
         job_config = self._prepare_job_config(job_config)
         job_config.labels = bigframes_io.create_job_configs_labels(
             job_configs_labels=job_config.labels, api_methods=api_methods

tests/unit/core/test_log_adapter.py

@@ -8,7 +8,7 @@
 @pytest.fixture
 def test_instance():
     # Create a simple class for testing
-    @log_adapter.class_logger()
+    @log_adapter.class_logger
     class TestClass:
         def method1(self):
             pass
@@ -19,29 +19,26 @@ def method2(self):
     return TestClass()
 
 
-def test_class_logger_decorator(test_instance):
-    # Ensure that the class logger decorator adds _api_methods attribute
-    assert hasattr(test_instance, "_api_methods")
-    assert test_instance._api_methods == []
-
+def test_method_logging(test_instance):
+    test_instance.method1()
+    test_instance.method2()
 
-def test_add_api_method(test_instance):
-    # Ensure that add_api_method correctly adds a method to _api_methods
-    log_adapter.add_api_method("method3", test_instance)
-    assert test_instance._api_methods == ["method3"]
+    # Check if the methods were added to the _api_methods list
+    api_methods = log_adapter.get_and_reset_api_methods()
+    assert api_methods == ["method2", "method1"]
 
 
 def test_add_api_method_limit(test_instance):
     # Ensure that add_api_method correctly adds a method to _api_methods
     for i in range(70):
-        log_adapter.add_api_method("method3", test_instance)
-    assert len(test_instance._api_methods) == MAX_LABELS_COUNT
+        test_instance.method2()
+    assert len(log_adapter._api_methods) == MAX_LABELS_COUNT
 
 
 def test_get_and_reset_api_methods(test_instance):
     # Ensure that get_and_reset_api_methods returns a copy and resets the list
     test_instance.method1()
     test_instance.method2()
-    previous_methods = log_adapter.get_and_reset_api_methods(test_instance)
+    previous_methods = log_adapter.get_and_reset_api_methods()
     assert previous_methods == ["method2", "method1"]
-    assert test_instance._api_methods == []
+    assert log_adapter._api_methods == []

tests/unit/session/test_io_bigquery.py

@@ -20,6 +20,8 @@
 import pytest
 
 import bigframes
+from bigframes.core import log_adapter
+import bigframes.pandas as bpd
 import bigframes.session._io.bigquery as io_bq
 
 
@@ -56,6 +58,50 @@ def test_create_job_configs_labels_length_limit_not_met():
     assert labels == expected_dict
 
 
+def test_create_job_configs_labels_log_adaptor_call_method_under_length_limit():
+    cur_labels = {
+        "bigframes-api": "read_pandas",
+        "source": "bigquery-dataframes-temp",
+    }
+    df = bpd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
+    # Test running two methods
+    df.head()
+    df.max()
+    api_methods = log_adapter._api_methods
+
+    labels = io_bq.create_job_configs_labels(
+        job_configs_labels=cur_labels, api_methods=api_methods
+    )
+    expected_dict = {
+        "bigframes-api": "read_pandas",
+        "source": "bigquery-dataframes-temp",
+        "recent-bigframes-api-0": "__init__",
+        "recent-bigframes-api-1": "max",
+        "recent-bigframes-api-2": "__init__",
+        "recent-bigframes-api-3": "head",
+        "recent-bigframes-api-4": "__init__",
+    }
+    assert labels is not None
+    assert len(labels) == 7
+    assert labels == expected_dict
+
+
+def test_create_job_configs_labels_length_limit_met_and_labels_is_none():
+    df = bpd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
+    # Test running methods more than the labels' length limit
+    for i in range(66):
+        df.head()
+    api_methods = log_adapter._api_methods
+
+    labels = io_bq.create_job_configs_labels(
+        job_configs_labels=None, api_methods=api_methods
+    )
+    assert labels is not None
+    assert len(labels) == 64
+    assert "head" in labels.values()
+    assert "__init__" in labels.values()
+
+
 def test_create_job_configs_labels_length_limit_met():
     cur_labels = {
         "bigframes-api": "read_pandas",
@@ -66,14 +112,19 @@ def test_create_job_configs_labels_length_limit_met():
         value = f"test{i}"
         cur_labels[key] = value
     # If cur_labels length is 62, we can only add one label from api_methods
-    api_methods = ["agg", "series-mode", "head"]
+    df = bpd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
+    # Test running two methods
+    df.head()
+    df.max()
+    api_methods = log_adapter._api_methods
 
     labels = io_bq.create_job_configs_labels(
         job_configs_labels=cur_labels, api_methods=api_methods
     )
     assert labels is not None
     assert len(labels) == 64
-    assert "agg" in labels.values()
+    assert "max" in labels.values()
+    assert "__init__" in labels.values()
     assert "head" not in labels.values()
     assert "bigframes-api" in labels.keys()
     assert "source" in labels.keys()