Skip to content

chore: log and labels update #674

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Genesis929 merged 5 commits into from
May 10, 2024
Merged

chore: log and labels update #674

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
ff38bd1
chore: log and labels update
Genesis929 May 9, 2024
f468815
remove unused logic
Genesis929 May 9, 2024
34ae2e7
Update unit test.
Genesis929 May 9, 2024
0927a85
fixes for mypy
Genesis929 May 9, 2024
83d47e7
lint update
Genesis929 May 9, 2024
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/log_adapter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand All @@ -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

Expand Down
17 changes: 11 additions & 6 deletions bigframes/session/_io/bigquery/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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(),
Expand Down Expand Up @@ -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)
Expand Down
4 changes: 0 additions & 4 deletions bigframes/session/_io/bigquery/read_gbq_table.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down
28 changes: 9 additions & 19 deletions tests/unit/session/test_io_bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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

Expand All @@ -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

Expand All @@ -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
Expand Down