fix: fix mypy error

Author: shuoweil

bigframes/session/loader.py

@@ -47,6 +47,8 @@
 from google.cloud import bigquery_storage_v1
 import google.cloud.bigquery
 import google.cloud.bigquery as bigquery
+from google.cloud.bigquery.job.load import LoadJob
+from google.cloud.bigquery.job.query import QueryJob
 import google.cloud.bigquery.table
 from google.cloud.bigquery_storage_v1 import types as bq_storage_types
 import pandas
@@ -562,7 +564,7 @@ def _start_generic_job(self, job: formatting_helpers.GenericJob):
         else:
             job.result()
 
-        if self._metrics is not None and isinstance(job, google.cloud.bigquery.job.Job):
+        if self._metrics is not None and isinstance(job, (QueryJob, LoadJob)):
             self._metrics.count_job_stats(query_job=job)
 
     @overload

bigframes/session/metrics.py

@@ -17,9 +17,11 @@
 import dataclasses
 import datetime
 import os
-from typing import Any, Mapping, Optional, Tuple
+from typing import Any, Mapping, Optional, Tuple, Union
 
 import google.cloud.bigquery as bigquery
+from google.cloud.bigquery.job.load import LoadJob
+from google.cloud.bigquery.job.query import QueryJob
 import google.cloud.bigquery.table as bq_table
 
 LOGGING_NAME_ENV_VAR = "BIGFRAMES_PERFORMANCE_LOG_NAME"
@@ -62,7 +64,7 @@ class ExecutionMetrics:
 
     def count_job_stats(
         self,
-        query_job: Optional[bigquery.job.Job] = None,
+        query_job: Optional[Union[QueryJob, LoadJob]] = None,
         row_iterator: Optional[bq_table.RowIterator] = None,
     ):
         if query_job is None:
@@ -106,7 +108,7 @@ def count_job_stats(
                 )
             )
 
-        elif query_job.configuration.dry_run:
+        elif isinstance(query_job, QueryJob) and query_job.configuration.dry_run:
             query_char_count = len(getattr(query_job, "query", ""))
 
             # TODO(tswast): Pass None after making benchmark publishing robust to missing data.

notebooks/dataframes/job_history.ipynb

@@ -4,8 +4,8 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Job History Manual Test\\n",
-    "\\n",
+    "# Job History Manual Test\n",
+    "\n",
     "This notebook demonstrates and manually tests the `bigframes.pandas.job_history()` functionality."
    ]
   },
@@ -15,11 +15,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import pandas as pd\\n",
-    "import bigframes.pandas as bpd\\n",
-    "\\n",
-    "# Set options if needed, e.g. project/location\\n",
-    "# bpd.options.bigquery.project = \"YOUR_PROJECT\"\\n",
+    "import pandas as pd\n",
+    "import bigframes.pandas as bpd\n",
+    "\n",
+    "# Set options if needed, e.g. project/location\\n\n",
+    "# bpd.options.bigquery.project = \"YOUR_PROJECT\"\\n\n",
     "# bpd.options.bigquery.location = \"US\""
    ]
   },
@@ -36,7 +36,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "df = bpd.read_gbq(\"SELECT 1 as a, 2 as b\")\\n",
+    "df = bpd.read_gbq(\"SELECT 1 as a, 2 as b\")\n",
     "df.head()"
    ]
   },
@@ -53,8 +53,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "local_df = pd.DataFrame({'col1': [1, 2, 3], 'col2': ['a', 'b', 'c']})\\n",
-    "bf_df = bpd.read_pandas(local_df)\\n",
+    "local_df = pd.DataFrame({'col1': [1, 2, 3], 'col2': ['a', 'b', 'c']})\n",
+    "bf_df = bpd.read_pandas(local_df)\n",
     "bf_df.head()"
    ]
   },
@@ -71,8 +71,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Perform a simple aggregation to trigger a computation\\n",
-    "agg_df = bf_df.groupby('col2').sum()\\n",
+    "# Perform a simple aggregation to trigger a computation\\n\n",
+    "agg_df = bf_df.groupby('col2').sum()\n",
     "agg_df.head()"
    ]
   },
@@ -89,7 +89,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "history = bpd.job_history()\\n",
+    "history = bpd.job_history()\n",
     "history"
    ]
   },
@@ -106,20 +106,20 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Display key columns to verify data population\\n",
-    "cols_to_check = [\\n",
-    "    'job_id', \\n",
-    "    'job_type', \\n",
-    "    'creation_time', \\n",
-    "    'duration_seconds', \\n",
-    "    'total_bytes_processed', \\n",
-    "    'query', \\n",
-    "    'input_files',  # Should be populated for Load Job\\n",
-    "    'destination_table'\\n",
-    "]\\n",
-    "\\n",
-    "# Filter columns that exist in the history DataFrame\\n",
-    "existing_cols = [col for col in cols_to_check if col in history.columns]\\n",
+    "# Display key columns to verify data population\\n\n",
+    "cols_to_check = [\n",
+    "    'job_id',\n",
+    "    'job_type',\n",
+    "    'creation_time',\n",
+    "    'duration_seconds',\n",
+    "    'total_bytes_processed',\n",
+    "    'query',\n",
+    "    'input_files',  # Should be populated for Load Job\n",
+    "    'destination_table'\n",
+    "]\n",
+    "\n",
+    "# Filter columns that exist in the history DataFrame\n",
+    "existing_cols = [col for col in cols_to_check if col in history.columns]\n",
     "history[existing_cols]"
    ]
   }