diff --git a/packages/bigframes/bigframes/session/_io/bigquery/__init__.py b/packages/bigframes/bigframes/session/_io/bigquery/__init__.py index 58bd5bd61748..3d60bcc80744 100644 --- a/packages/bigframes/bigframes/session/_io/bigquery/__init__.py +++ b/packages/bigframes/bigframes/session/_io/bigquery/__init__.py @@ -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) job_configs_labels["bigframes-api"] = api_methods[0] del api_methods[0] diff --git a/packages/bigframes/tests/unit/session/test_io_bigquery.py b/packages/bigframes/tests/unit/session/test_io_bigquery.py index 9f44a02ef49b..e6fa7a901eca 100644 --- a/packages/bigframes/tests/unit/session/test_io_bigquery.py +++ b/packages/bigframes/tests/unit/session/test_io_bigquery.py @@ -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() @@ -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 diff --git a/packages/bigframes/tests/unit/session/test_read_gbq_colab.py b/packages/bigframes/tests/unit/session/test_read_gbq_colab.py index a168ccbad5e6..fc4181b6a2b4 100644 --- a/packages/bigframes/tests/unit/session/test_read_gbq_colab.py +++ b/packages/bigframes/tests/unit/session/test_read_gbq_colab.py @@ -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)