more mypy fixes

Author: TrevorBergeron

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

@@ -278,7 +278,7 @@ def start_query_with_job(
     job_retry: google.api_core.retry.Retry = third_party_gcb_retry.DEFAULT_JOB_RETRY,
     publisher: bigframes.core.events.Publisher,
     session=None,
-) -> Tuple[google.cloud.bigquery.table.RowIterator, Optional[bigquery.QueryJob]]:
+) -> Tuple[google.cloud.bigquery.table.RowIterator, bigquery.QueryJob]:
     """
     Starts query job and waits for results.
     """
@@ -328,7 +328,7 @@ def start_query_job_optional(
     job_retry: google.api_core.retry.Retry = third_party_gcb_retry.DEFAULT_JOB_RETRY,
     publisher: bigframes.core.events.Publisher,
     session=None,
-):
+) -> google.cloud.bigquery.table.RowIterator:
     add_and_trim_labels(job_config, session=session)
     try:
         results_iterator = bq_client._query_and_wait_bigframes(
@@ -342,7 +342,7 @@ def start_query_job_optional(
         )
         if metrics is not None:
             metrics.count_job_stats(row_iterator=results_iterator)
-        return results_iterator, None
+        return results_iterator
     except google.api_core.exceptions.Forbidden as ex:
         if "Drive credentials" in ex.message:
             ex.message += CHECK_DRIVE_PERMISSIONS

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

@@ -193,7 +193,7 @@ def is_time_travel_eligible(
         )
         try:
             # If this succeeds, we know that time travel will for sure work.
-            bigframes.session._io.bigquery.start_query_with_job_optional(
+            bigframes.session._io.bigquery.start_query_job_optional(
                 bq_client=bqclient,
                 sql=snapshot_sql,
                 job_config=bigquery.QueryJobConfig(dry_run=True),
@@ -265,7 +265,7 @@ def check_if_index_columns_are_unique(
         index_cols, table.get_table_ref()
     )
     job_config = bigquery.QueryJobConfig()
-    results, _ = bigframes.session._io.bigquery.start_query_with_job_optional(
+    results = bigframes.session._io.bigquery.start_query_job_optional(
         bq_client=bqclient,
         sql=is_unique_sql,
         job_config=job_config,

packages/bigframes/bigframes/session/bigquery_session.py

@@ -116,7 +116,7 @@ def close(self):
             self._sessiondaemon.stop()
 
         if self._session_id is not None and self.bqclient is not None:
-            bfbqio.start_query_with_job_optional(
+            bfbqio.start_query_job_optional(
                 self.bqclient,
                 f"CALL BQ.ABORT_SESSION('{self._session_id}')",
                 job_config=bigquery.QueryJobConfig(),
@@ -166,7 +166,7 @@ def _keep_session_alive(self):
                 ]
             )
             try:
-                bfbqio.start_query_with_job_optional(
+                bfbqio.start_query_job_optional(
                     self.bqclient,
                     "SELECT 1",
                     job_config=job_config,

packages/bigframes/bigframes/session/bq_caching_executor.py

@@ -116,7 +116,7 @@ def __init__(
             bqstoragereadclient=bqstoragereadclient,
             metrics=self.metrics,
             publisher=self._publisher,
-            labels=labels,
+            labels=dict(labels),
         )
 
     def to_sql(
@@ -225,7 +225,7 @@ def _execute_gbq_table_export(
             results = self._execute_bigquery(array_value, execution_spec)
             assert isinstance(results, executor.BQTableExecuteResult)
             self._export_gbq_with_dml(results, dest_spec)
-            result = results
+            result: executor.ExecuteResult = results
         else:
             result = self._execute_gbq_query_only(array_value, execution_spec)
 

packages/bigframes/bigframes/session/direct_gbq_execution.py

@@ -46,9 +46,9 @@ def __init__(
         bqclient: bigquery.Client,
         bqstoragereadclient: google.cloud.bigquery_storage_v1.BigQueryReadClient,
         *,
+        publisher: bigframes.core.events.Publisher,
         compiler: Literal["ibis", "sqlglot"] = "ibis",
         metrics: Optional[bigframes.session.metrics.ExecutionMetrics] = None,
-        publisher: Optional[bigframes.core.events.Publisher] = None,
         labels: Mapping[str, str] = {},
     ):
         self.bqclient = bqclient
@@ -171,13 +171,16 @@ def _run_execute_query(
                     session=session,
                 )
             else:
-                return bq_io.start_query_with_job_optional(
-                    self.bqclient,
-                    sql,
-                    job_config=job_config,
-                    metrics=self._metrics,
-                    publisher=self._publisher,
-                    session=session,
+                return (
+                    bq_io.start_query_job_optional(
+                        self.bqclient,
+                        sql,
+                        job_config=job_config,
+                        metrics=self._metrics,
+                        publisher=self._publisher,
+                        session=session,
+                    ),
+                    None,
                 )
         except google.api_core.exceptions.BadRequest as e:
             # Unfortunately, this error type does not have a separate error code or exception type

packages/bigframes/bigframes/session/loader.py

@@ -1462,7 +1462,7 @@ def _start_query_with_job_optional(
         Do not execute dataframe through this API, instead use the executor.
         """
         job_config = self._prepare_job_config(job_config)
-        rows, _ = bf_io_bigquery.start_query_with_job_optional(
+        rows = bf_io_bigquery.start_query_job_optional(
             self._bqclient,
             sql,
             job_config=job_config,