Skip to content

test: verify read_gbq_colab label preservation in anywidget display mode #17887

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
shuoweil merged 5 commits into from
Jul 28, 2026
+33 −1
Merged

test: verify read_gbq_colab label preservation in anywidget display mode #17887

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
d0b590e
fix: preserve read_gbq_colab label in anywidget mode
shuoweil Jul 23, 2026
dc53247
Update packages/bigframes/bigframes/session/_io/bigquery/__init__.py
shuoweil Jul 23, 2026
a8a07ba
test: skip anywidget test when dependencies missing
shuoweil Jul 24, 2026
c563483
Merge branch 'main' into shuowei-sql-cell-track
shuoweil Jul 24, 2026
9185cdd
test(bigframes): verify read_gbq_colab label preservation in anywidge…
shuoweil Jul 24, 2026
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
1 change: 1 addition & 0 deletions packages/bigframes/bigframes/session/_io/bigquery/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@ def create_job_configs_labels(
job_configs_labels = dict(job_configs_labels)

if api_methods and "bigframes-api" not in job_configs_labels:
api_methods = list(api_methods)

@sycai sycai Jul 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Non-rhetorical question: do we need this line of change? It alters the behavior of the original code, which alters the original list that was passed in at line 73.

@shuoweil shuoweil Jul 28, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, this shallow copy is intentional as a safety guard. There are two reasons. (1) at line 73, del api_methods[0]
removes the primary API label. Creating a copy ensures this deletion stys local to the helper function and does not mutate the caller's original list if reused or logs elsewhere.
(2) api_methods may sometime pass in tuple. We have a unit test reate_job_configs_labels(labels, ("read_gbq_colab", "head")) will pass tuples to label. In this case, tuple is immutable. This copy is a guard to avoid errors for del api_methods[0].

job_configs_labels["bigframes-api"] = api_methods[0]
del api_methods[0]

Expand Down
3 changes: 2 additions & 1 deletion packages/bigframes/tests/unit/session/test_io_bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,7 +172,7 @@ def test_add_and_trim_labels_length_limit_met():
for _ in range(52):
df.head()

io_bq.add_and_trim_labels(job_config=job_config)
io_bq.add_and_trim_labels(job_config=job_config, session=df._session)
assert job_config.labels is not None
assert len(job_config.labels) == 56
assert "dataframe-max" not in job_config.labels.values()
Expand Down Expand Up @@ -218,6 +218,7 @@ def test_start_query_with_job_labels_length_limit_met(
timeout=timeout,
metrics=None,
publisher=bigframes.core.events.Publisher(),
session=df._session,
)

assert job_config.labels is not None
Expand Down
30 changes: 30 additions & 0 deletions packages/bigframes/tests/unit/session/test_read_gbq_colab.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,36 @@ def test_read_gbq_colab_includes_label():
assert "session-read_gbq_colab" in label_values


def test_read_gbq_colab_includes_label_in_anywidget_mode():
"""Make sure read_gbq_colab label is preserved in recent-bigframes-api labels in anywidget mode."""
pytest.importorskip("anywidget")
pytest.importorskip("traitlets")

import bigframes
import bigframes.display.html as bf_html

bqclient = mock.create_autospec(bigquery.Client, instance=True)
bqclient.project = "proj"
session = mocks.create_bigquery_session(bqclient=bqclient)
df = session._read_gbq_colab("SELECT 'read-gbq-colab-test'")

with bigframes.option_context("display.render_mode", "anywidget"):
_ = bf_html.get_anywidget_bundle(df)

label_values = []
for kall in itertools.chain(
bqclient.query_and_wait.call_args_list,
bqclient._query_and_wait_bigframes.call_args_list,
bqclient.query.call_args_list,
):
job_config = kall.kwargs.get("job_config")
if job_config is None:
continue
label_values.extend(job_config.labels.values())

assert "session-read_gbq_colab" in label_values


@pytest.mark.parametrize("dry_run", [True, False])
def test_read_gbq_colab_includes_formatted_values_in_dry_run(monkeypatch, dry_run):
bqclient = mock.create_autospec(bigquery.Client, instance=True)
Expand Down
Loading