From d01b812f738d2ce60e24dec630f8d3ff8fa6c9a1 Mon Sep 17 00:00:00 2001
From: Eric Johnson <ericdjohnson@google.com>
Date: Fri, 12 Sep 2025 09:10:10 -0700
Subject: [PATCH] No public description

PiperOrigin-RevId: 806311764
---
 google/colab/sql/bigquery.py      | 11 ++++++++++-
 google/colab/sql/bigquery_test.py | 10 ++++++++--
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/google/colab/sql/bigquery.py b/google/colab/sql/bigquery.py
index 077d2e27..2763a870 100644
--- a/google/colab/sql/bigquery.py
+++ b/google/colab/sql/bigquery.py
@@ -12,6 +12,15 @@
 # Extracts the line and column number from a dry-run error message.
 _LINE_COLUMN_REGEX = re.compile(r'\[(\d+):(\d+)\]', re.MULTILINE)
 
+# Option pair to disable the progress bar when running queries.
+_NO_PROGRESS_BAR = ('display.progress_bar', None)
+
+# Option pair to add additional query labels for telemetry purposes.
+_QUERY_LABELS = (
+    'compute.extra_query_labels',
+    {'bigframes-connector': 'sql-cell'},
+)
+
 
 class TableReference(TypedDict):
   """A reference to a BigQuery table."""
@@ -140,7 +149,7 @@ def run(sql: str) -> bpd.DataFrame:
   """Executes the SQL and returns the BigQuery DataFrame."""
   # repr_mode cannot be set in context manager.
   bpd.options.display.repr_mode = 'anywidget'
-  with bpd.option_context('display.progress_bar', None):
+  with bpd.option_context(*_NO_PROGRESS_BAR, *_QUERY_LABELS):
     # pylint: disable=protected-access
     return bpd._read_gbq_colab(sql, pyformat_args=_get_ipython_locals())
 
diff --git a/google/colab/sql/bigquery_test.py b/google/colab/sql/bigquery_test.py
index 806acace..4389e63e 100644
--- a/google/colab/sql/bigquery_test.py
+++ b/google/colab/sql/bigquery_test.py
@@ -401,7 +401,10 @@ def test_run(self):
     self.assertEqual(bigquery.run(_QUERY), expected)
 
     self.mock_bpd_option_context.assert_called_once_with(
-        'display.progress_bar', None
+        'display.progress_bar',
+        None,
+        'compute.extra_query_labels',
+        {'bigframes-connector': 'sql-cell'},
     )
     self.mock_read_gbq_colab.assert_called_once_with(_QUERY, pyformat_args={})
     self.assertEqual(bpd.options.display.repr_mode, 'anywidget')
@@ -420,7 +423,10 @@ def test_run_with_pyformat_args(self, mock_get_ipython):
     self.assertEqual(bigquery.run(query), expected)
 
     self.mock_bpd_option_context.assert_called_once_with(
-        'display.progress_bar', None
+        'display.progress_bar',
+        None,
+        'compute.extra_query_labels',
+        {'bigframes-connector': 'sql-cell'},
     )
     self.mock_read_gbq_colab.assert_called_once_with(
         query, pyformat_args=user_ns