chore: log and labels update (#674)

* chore: log and labels update

* remove unused logic

* Update unit test.

* fixes for mypy

* lint update

Author: Genesis929

bigframes/core/log_adapter.py

@@ -21,6 +21,9 @@
 _api_methods: List = []
 _excluded_methods = ["__setattr__", "__getattr__"]
 
+# Stack to track method calls
+_call_stack: List = []
+
 
 def class_logger(decorated_cls):
     """Decorator that adds logging functionality to each method of the class."""
@@ -38,10 +41,17 @@ def wrapper(*args, **kwargs):
         class_name = decorated_cls.__name__  # Access decorated class name
         api_method_name = str(method.__name__)
         full_method_name = f"{class_name.lower()}-{api_method_name}"
-        # Track regular and "dunder" methods
-        if api_method_name.startswith("__") or not api_method_name.startswith("_"):
+
+        # Track directly called methods
+        if len(_call_stack) == 0:
             add_api_method(full_method_name)
-        return method(*args, **kwargs)
+
+        _call_stack.append(full_method_name)
+
+        try:
+            return method(*args, **kwargs)
+        finally:
+            _call_stack.pop()
 
     return wrapper
 

bigframes/session/_io/bigquery/__init__.py

@@ -23,7 +23,7 @@
 import textwrap
 import types
 import typing
-from typing import Dict, Iterable, Mapping, Optional, Sequence, Tuple, Union
+from typing import Dict, Iterable, Mapping, Optional, Tuple, Union
 
 import bigframes_vendored.pandas.io.gbq as third_party_pandas_gbq
 import google.api_core.exceptions
@@ -43,11 +43,15 @@
 
 def create_job_configs_labels(
     job_configs_labels: Optional[Dict[str, str]],
-    api_methods: Sequence[str],
+    api_methods: typing.List[str],
 ) -> Dict[str, str]:
     if job_configs_labels is None:
         job_configs_labels = {}
 
+    if api_methods:
+        job_configs_labels["bigframes-api"] = api_methods[0]
+        del api_methods[0]
+
     labels = list(
         itertools.chain(
             job_configs_labels.keys(),
@@ -198,10 +202,11 @@ def start_query_with_client(
     """
     Starts query job and waits for results.
     """
-    api_methods = log_adapter.get_and_reset_api_methods()
-    job_config.labels = create_job_configs_labels(
-        job_configs_labels=job_config.labels, api_methods=api_methods
-    )
+    if not job_config.dry_run:
+        api_methods = log_adapter.get_and_reset_api_methods()
+        job_config.labels = create_job_configs_labels(
+            job_configs_labels=job_config.labels, api_methods=api_methods
+        )
 
     try:
         query_job = bq_client.query(sql, job_config=job_config, timeout=timeout)

bigframes/session/_io/bigquery/read_gbq_table.py

@@ -112,8 +112,6 @@ def get_table_metadata(
     # atomically.
     table = bqclient.get_table(table_ref)
 
-    # TODO(b/336521938): Refactor to make sure we set the "bigframes-api"
-    # whereever we execute a query.
     job_config = bigquery.QueryJobConfig()
     job_config.labels["bigframes-api"] = api_name
     snapshot_timestamp = list(
@@ -344,8 +342,6 @@ def get_time_travel_datetime_and_table_metadata(
     # atomically.
     table = bqclient.get_table(table_ref)
 
-    # TODO(b/336521938): Refactor to make sure we set the "bigframes-api"
-    # whereever we execute a query.
     job_config = bigquery.QueryJobConfig()
     job_config.labels["bigframes-api"] = api_name
     snapshot_timestamp = list(

tests/unit/session/test_io_bigquery.py

@@ -30,38 +30,32 @@ def test_create_job_configs_labels_is_none():
     labels = io_bq.create_job_configs_labels(
         job_configs_labels=None, api_methods=api_methods
     )
-    expected_dict = {
-        "recent-bigframes-api-0": "agg",
-        "recent-bigframes-api-1": "series-mode",
-    }
+    expected_dict = {"bigframes-api": "agg", "recent-bigframes-api-0": "series-mode"}
     assert labels is not None
     assert labels == expected_dict
 
 
 def test_create_job_configs_labels_length_limit_not_met():
     cur_labels = {
-        "bigframes-api": "read_pandas",
         "source": "bigquery-dataframes-temp",
     }
     api_methods = ["agg", "series-mode"]
     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": "agg",
-        "recent-bigframes-api-1": "series-mode",
+        "bigframes-api": "agg",
+        "recent-bigframes-api-0": "series-mode",
     }
     assert labels is not None
-    assert len(labels) == 4
+    assert len(labels) == 3
     assert labels == expected_dict
 
 
 def test_create_job_configs_labels_log_adaptor_call_method_under_length_limit():
     log_adapter.get_and_reset_api_methods()
     cur_labels = {
-        "bigframes-api": "read_pandas",
         "source": "bigquery-dataframes-temp",
     }
     df = bpd.DataFrame(
@@ -76,14 +70,10 @@ def test_create_job_configs_labels_log_adaptor_call_method_under_length_limit():
         job_configs_labels=cur_labels, api_methods=api_methods
     )
     expected_dict = {
-        "bigframes-api": "read_pandas",
         "source": "bigquery-dataframes-temp",
-        "recent-bigframes-api-0": "series-__init__",
-        "recent-bigframes-api-1": "dataframe-max",
-        "recent-bigframes-api-2": "dataframe-__init__",
-        "recent-bigframes-api-3": "dataframe-head",
-        "recent-bigframes-api-4": "dataframe-__init__",
-        "recent-bigframes-api-5": "dataframe-__init__",
+        "bigframes-api": "dataframe-max",
+        "recent-bigframes-api-0": "dataframe-head",
+        "recent-bigframes-api-1": "dataframe-__init__",
     }
     assert labels == expected_dict
 
@@ -94,7 +84,7 @@ def test_create_job_configs_labels_length_limit_met_and_labels_is_none():
         {"col1": [1, 2], "col2": [3, 4]}, session=resources.create_bigquery_session()
     )
     # Test running methods more than the labels' length limit
-    for i in range(66):
+    for i in range(100):
         df.head()
     api_methods = log_adapter._api_methods
 
@@ -112,7 +102,7 @@ def test_create_job_configs_labels_length_limit_met():
         "bigframes-api": "read_pandas",
         "source": "bigquery-dataframes-temp",
     }
-    for i in range(60):
+    for i in range(100):
         key = f"bigframes-api-test-{i}"
         value = f"test{i}"
         cur_labels[key] = value