fix: preserve read_gbq_colab label in anywidget mode

Author: shuoweil

packages/bigframes/bigframes/session/_io/bigquery/__init__.py

@@ -68,8 +68,14 @@ def create_job_configs_labels(
         job_configs_labels = dict(job_configs_labels)
 
     if api_methods and "bigframes-api" not in job_configs_labels:
-        job_configs_labels["bigframes-api"] = api_methods[0]
-        del api_methods[0]
+        colab_idx = next(
+            (i for i, m in enumerate(api_methods) if "read_gbq_colab" in m), None
+        )
+        if colab_idx is not None:
+            job_configs_labels["bigframes-api"] = api_methods.pop(colab_idx)
+        else:
+            job_configs_labels["bigframes-api"] = api_methods[0]
+            del api_methods[0]
 
     # Make sure we always populate bigframes-api with _something_, even if we
     # have a code path which doesn't populate the list of api_methods. See

packages/bigframes/tests/unit/session/test_read_gbq_colab.py

@@ -47,6 +47,37 @@ 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 as the primary bigframes-api label in anywidget mode."""
+    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 = []
+    bigframes_api_labels = []
+    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())
+        if "bigframes-api" in job_config.labels:
+            bigframes_api_labels.append(job_config.labels["bigframes-api"])
+
+    assert "session-read_gbq_colab" in label_values
+    assert "session-read_gbq_colab" in bigframes_api_labels
+
+
 @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)