feat: only show customer facing query

Author: shuoweil

bigframes/session/_io/bigquery/read_gbq_table.py

@@ -256,7 +256,6 @@ def check_if_index_columns_are_unique(
     index_cols: Sequence[str],
     *,
     publisher: bigframes.core.events.Publisher,
-    metrics: Optional[bigframes.session.metrics.ExecutionMetrics] = None,
 ) -> Tuple[str, ...]:
     import bigframes.core.sql
     import bigframes.session._io.bigquery
@@ -274,7 +273,7 @@ def check_if_index_columns_are_unique(
         timeout=None,
         location=None,
         project=None,
-        metrics=metrics,
+        metrics=None,
         query_with_job=False,
         publisher=publisher,
     )

bigframes/session/loader.py

@@ -926,7 +926,6 @@ def read_gbq_table(
                 table=table,
                 index_cols=index_cols,
                 publisher=self._publisher,
-                metrics=self._metrics,
             )
             if publish_execution:
                 self._publisher.publish(

bigframes/session/metrics.py

@@ -200,7 +200,6 @@ def count_job_stats(
                 self.jobs.append(metadata)
 
         else:
-            # Handle other job types (e.g. LoadJob)
             self.execution_count += 1
             duration = (
                 (query_job.ended - query_job.created).total_seconds()

tests/unit/session/test_execution_history.py

@@ -20,7 +20,6 @@
 
 import bigframes
 import bigframes.session
-import bigframes.session.metrics as metrics
 
 NOW = datetime.datetime.now(datetime.timezone.utc)
 
@@ -29,6 +28,7 @@ def test_execution_history_returns_dataframe_with_correct_job_id(monkeypatch):
     query_job = unittest.mock.create_autospec(bigquery.QueryJob, instance=True)
     query_job.job_id = "job_pandas"
     query_job.location = "US"
+    query_job.project = "test-project"
     query_job.created = NOW
     query_job.started = NOW
     query_job.ended = NOW + datetime.timedelta(seconds=1)
@@ -40,25 +40,43 @@ def test_execution_history_returns_dataframe_with_correct_job_id(monkeypatch):
     query_job.cache_hit = False
     query_job.query = "SELECT 1"
     query_job.configuration.dry_run = False
+    query_job.statement_type = "SELECT"
 
-    # Mock the clients provider to avoid actual BQ client creation
-    clients_provider = unittest.mock.create_autospec(
-        bigframes.session.clients.ClientsProvider
-    )
     bq_client = unittest.mock.create_autospec(bigquery.Client)
     bq_client.project = "test-project"
     bq_client.default_query_job_config = bigquery.QueryJobConfig()
     bq_client.query_and_wait.return_value = iter([[NOW]])
+    bq_client.get_job.return_value = query_job
+
+    clients_provider = unittest.mock.create_autospec(
+        bigframes.session.clients.ClientsProvider
+    )
     clients_provider.bqclient = bq_client
 
+    row_iterator = unittest.mock.create_autospec(
+        bigquery.table.RowIterator, instance=True
+    )
+    row_iterator.job_id = "job_pandas"
+    row_iterator.location = "US"
+    row_iterator.project = "test-project"
+
+    import bigframes.session._io.bigquery as bf_io_bigquery
+
+    monkeypatch.setattr(
+        bf_io_bigquery,
+        "start_query_with_client",
+        lambda *args, **kwargs: (row_iterator, None),
+    )
+
     session = bigframes.session.Session(clients_provider=clients_provider)
+    session._loader.read_gbq_table = unittest.mock.MagicMock()  # type: ignore[method-assign]
 
-    # Mock get_global_session to return our session
     monkeypatch.setattr(
         bigframes.core.global_session, "get_global_session", lambda: session
     )
 
-    session._metrics.count_job_stats(query_job=query_job)
+    # Act
+    session.read_gbq_query("SELECT 1")
 
     df = bigframes.execution_history()
 
@@ -67,10 +85,11 @@ def test_execution_history_returns_dataframe_with_correct_job_id(monkeypatch):
     assert df.iloc[0]["job_id"] == "job_pandas"
 
 
-def test_execution_history_includes_query_job_metadata():
+def test_execution_history_includes_query_job_metadata(monkeypatch):
     query_job = unittest.mock.create_autospec(bigquery.QueryJob, instance=True)
     query_job.job_id = "job1"
     query_job.location = "US"
+    query_job.project = "test-project"
     query_job.created = NOW
     query_job.started = NOW
     query_job.ended = NOW + datetime.timedelta(seconds=1)
@@ -82,21 +101,39 @@ def test_execution_history_includes_query_job_metadata():
     query_job.cache_hit = False
     query_job.query = "SELECT 1"
     query_job.configuration.dry_run = False
+    query_job.statement_type = "SELECT"
 
-    # Mock the clients provider to avoid actual BQ client creation
-    clients_provider = unittest.mock.create_autospec(
-        bigframes.session.clients.ClientsProvider
-    )
-    # We need to mock bqclient specifically as it's accessed during Session init
     bq_client = unittest.mock.create_autospec(bigquery.Client)
     bq_client.project = "test-project"
     bq_client.default_query_job_config = bigquery.QueryJobConfig()
-    # Mock clock sync query
     bq_client.query_and_wait.return_value = iter([[NOW]])
+    bq_client.get_job.return_value = query_job
+
+    clients_provider = unittest.mock.create_autospec(
+        bigframes.session.clients.ClientsProvider
+    )
     clients_provider.bqclient = bq_client
 
+    row_iterator = unittest.mock.create_autospec(
+        bigquery.table.RowIterator, instance=True
+    )
+    row_iterator.job_id = "job1"
+    row_iterator.location = "US"
+    row_iterator.project = "test-project"
+
+    import bigframes.session._io.bigquery as bf_io_bigquery
+
+    monkeypatch.setattr(
+        bf_io_bigquery,
+        "start_query_with_client",
+        lambda *args, **kwargs: (row_iterator, None),
+    )
+
     session = bigframes.session.Session(clients_provider=clients_provider)
-    session._metrics.count_job_stats(query_job=query_job)
+    session._loader.read_gbq_table = unittest.mock.MagicMock()  # type: ignore[method-assign]
+
+    # Act
+    session.read_gbq_query("SELECT 1")
 
     df = session.execution_history()
 
@@ -107,40 +144,78 @@ def test_execution_history_includes_query_job_metadata():
     assert "creation_time" in df.columns
 
 
-def test_execution_history_tracks_query_job_metrics():
+def test_execution_history_tracks_query_job_metrics(monkeypatch):
     query_job = unittest.mock.create_autospec(bigquery.QueryJob, instance=True)
     query_job.job_id = "job1"
     query_job.location = "US"
+    query_job.project = "test-project"
     query_job.created = NOW
-    query_job.started = NOW + datetime.timedelta(seconds=1)
-    query_job.ended = NOW + datetime.timedelta(seconds=3)
+    query_job.started = NOW
+    query_job.ended = NOW + datetime.timedelta(seconds=1)
     query_job.state = "DONE"
-    query_job.total_bytes_processed = 1024
-    query_job.slot_millis = 100
     query_job.job_type = "query"
     query_job.error_result = None
     query_job.cache_hit = False
     query_job.query = "SELECT 1"
     query_job.configuration.dry_run = False
+    query_job.statement_type = "SELECT"
+
+    bq_client = unittest.mock.create_autospec(bigquery.Client)
+    bq_client.project = "test-project"
+    bq_client.default_query_job_config = bigquery.QueryJobConfig()
+    bq_client.query_and_wait.return_value = iter([[NOW]])
+    bq_client.get_job.return_value = query_job
+
+    clients_provider = unittest.mock.create_autospec(
+        bigframes.session.clients.ClientsProvider
+    )
+    clients_provider.bqclient = bq_client
+
+    row_iterator = unittest.mock.create_autospec(
+        bigquery.table.RowIterator, instance=True
+    )
+    row_iterator.job_id = "job1"
+    row_iterator.location = "US"
+    row_iterator.project = "test-project"
+
+    import bigframes.session._io.bigquery as bf_io_bigquery
+
+    monkeypatch.setattr(
+        bf_io_bigquery,
+        "start_query_with_client",
+        lambda *args, **kwargs: (row_iterator, None),
+    )
+
+    import bigframes.session.metrics as bfm_metrics
+
+    monkeypatch.setattr(
+        bfm_metrics,
+        "get_performance_stats",
+        lambda job: (10, 1024, 100, 3.0),
+    )
+
+    session = bigframes.session.Session(clients_provider=clients_provider)
+    session._loader.read_gbq_table = unittest.mock.MagicMock()  # type: ignore[method-assign]
 
-    execution_metrics = metrics.ExecutionMetrics()
-    execution_metrics.count_job_stats(query_job=query_job)
+    # Act
+    session.read_gbq_query("SELECT 1")
 
-    assert len(execution_metrics.jobs) == 1
-    job = execution_metrics.jobs[0]
-    assert job.job_id == "job1"
-    assert job.status == "DONE"
-    assert job.total_bytes_processed == 1024
-    assert job.duration_seconds == 3.0
+    df = session.execution_history()
 
+    assert len(df) == 1
+    assert df.iloc[0]["job_id"] == "job1"
+    assert df.iloc[0]["total_bytes_processed"] == 1024
+    assert df.iloc[0]["duration_seconds"] == 3.0
 
-def test_execution_history_tracks_row_iterator_metrics():
+
+def test_execution_history_tracks_row_iterator_metrics(monkeypatch):
     row_iterator = unittest.mock.create_autospec(
         bigquery.table.RowIterator, instance=True
     )
-    row_iterator.job_id = "job2"
+    row_iterator.job_id = None  # Force fallback to RowIterator metrics in loader.py
     row_iterator.query_id = "query2"
     row_iterator.location = "US"
+    row_iterator.project = "test-project"
     row_iterator.created = NOW
     row_iterator.started = NOW + datetime.timedelta(seconds=1)
     row_iterator.ended = NOW + datetime.timedelta(seconds=2)
@@ -149,22 +224,43 @@ def test_execution_history_tracks_row_iterator_metrics():
     row_iterator.cache_hit = True
     row_iterator.query = "SELECT 2"
 
-    execution_metrics = metrics.ExecutionMetrics()
-    execution_metrics.count_job_stats(row_iterator=row_iterator)
+    bq_client = unittest.mock.create_autospec(bigquery.Client)
+    bq_client.project = "test-project"
+    bq_client.default_query_job_config = bigquery.QueryJobConfig()
+    bq_client.query_and_wait.return_value = iter([[NOW]])
+
+    clients_provider = unittest.mock.create_autospec(
+        bigframes.session.clients.ClientsProvider
+    )
+    clients_provider.bqclient = bq_client
+
+    import bigframes.session._io.bigquery as bf_io_bigquery
+
+    monkeypatch.setattr(
+        bf_io_bigquery,
+        "start_query_with_client",
+        lambda *args, **kwargs: (row_iterator, None),
+    )
+
+    session = bigframes.session.Session(clients_provider=clients_provider)
+
+    # Act
+    session.read_gbq_query("SELECT 2")
+
+    df = session.execution_history()
 
-    assert len(execution_metrics.jobs) == 1
-    job = execution_metrics.jobs[0]
-    assert job.job_id == "job2"
-    assert job.query_id == "query2"
-    assert job.status == "DONE"
-    assert job.cached is True
-    assert job.duration_seconds == 2.0
+    assert len(df) == 1
+    assert df.iloc[0]["query_id"] == "query2"
+    assert df.iloc[0]["total_bytes_processed"] == 512
+    # duration is 2 - 0 = 2 seconds
+    assert df.iloc[0]["duration_seconds"] == 2.0
 
 
-def test_execution_history_tracks_load_job_metrics():
+def test_execution_history_tracks_load_job_metrics(monkeypatch):
     load_job = unittest.mock.create_autospec(bigquery.LoadJob, instance=True)
     load_job.job_id = "job3"
     load_job.location = "US"
+    load_job.project = "test-project"
     load_job.created = NOW
     load_job.started = NOW
     load_job.ended = NOW + datetime.timedelta(seconds=5)
@@ -179,17 +275,43 @@ def test_execution_history_tracks_load_job_metrics():
     load_job.source_uris = ["gs://bucket/file.csv"]
     load_job.configuration.source_format = "CSV"
 
-    execution_metrics = metrics.ExecutionMetrics()
-    execution_metrics.count_job_stats(query_job=load_job)
-
-    assert len(execution_metrics.jobs) == 1
-    job = execution_metrics.jobs[0]
-    assert job.job_id == "job3"
-    assert job.job_type == "load"
-    assert job.duration_seconds == 5.0
-    assert job.output_rows == 100
-    assert job.input_files == 1
-    assert job.input_bytes == 1024
-    assert job.destination_table == "project.dataset.table"
-    assert job.source_uris == ["gs://bucket/file.csv"]
-    assert job.source_format == "CSV"
+    bq_client = unittest.mock.create_autospec(bigquery.Client)
+    bq_client.project = "test-project"
+    bq_client.default_query_job_config = bigquery.QueryJobConfig()
+    bq_client.query_and_wait.return_value = iter([[NOW]])
+    bq_client.load_table_from_file.return_value = load_job
+
+    storage_manager = unittest.mock.MagicMock()
+    storage_manager.create_temp_table.return_value = (
+        bigquery.TableReference.from_string("project.dataset.table")
+    )
+
+    clients_provider = unittest.mock.create_autospec(
+        bigframes.session.clients.ClientsProvider
+    )
+    clients_provider.bqclient = bq_client
+
+    import bigframes.core.bq_data as bq_data
+
+    monkeypatch.setattr(
+        bq_data.GbqNativeTable,
+        "from_table",
+        lambda *args, **kwargs: unittest.mock.MagicMock(),
+    )
+    monkeypatch.setattr(bq_data.BigqueryDataSource, "__post_init__", lambda self: None)
+
+    session = bigframes.session.Session(clients_provider=clients_provider)
+    session._storage_manager = storage_manager
+
+    import io
+
+    # Act
+    buffer = io.BytesIO(b"col1,col2\n1,2\n3,4")
+    session.read_csv(buffer, write_engine="bigquery_load")
+
+    df = session.execution_history()
+
+    assert len(df) == 1
+    assert df.iloc[0]["job_id"] == "job3"
+    assert df.iloc[0]["job_type"] == "load"
+    assert df.iloc[0]["duration_seconds"] == 5.0