From d01b812f738d2ce60e24dec630f8d3ff8fa6c9a1 Mon Sep 17 00:00:00 2001 From: Eric Johnson 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