satisfy the mypy

Author: TrevorBergeron

packages/bigframes/bigframes/blob/_functions.py

@@ -103,15 +103,14 @@ def _create_udf(self):
 \"\"\"
         """
 
-        bf_io_bigquery.start_query_with_client(
+        bf_io_bigquery.start_query_with_job(
             self._session.bqclient,
             sql,
             job_config=bigquery.QueryJobConfig(),
             metrics=self._session._metrics,
             location=None,
             project=None,
             timeout=None,
-            query_with_job=True,
             publisher=self._session._publisher,
         )
 

packages/bigframes/bigframes/functions/_function_client.py

@@ -142,15 +142,14 @@ def _create_bq_function(self, create_function_ddl: str) -> None:
         # TODO(swast): plumb through the original, user-facing api_name.
         import bigframes.session._io.bigquery
 
-        _, query_job = bigframes.session._io.bigquery.start_query_with_client(
+        _, query_job = bigframes.session._io.bigquery.start_query_with_job(
             cast(bigquery.Client, self._session.bqclient),
             create_function_ddl,
             job_config=bigquery.QueryJobConfig(),
             location=None,
             project=None,
             timeout=None,
             metrics=None,
-            query_with_job=True,
             publisher=self._session._publisher,
         )
         logger.info(f"Created bigframes function {query_job.ddl_target_routine}")

packages/bigframes/bigframes/functions/function.py

@@ -183,10 +183,9 @@ def __call__(self, *args, **kwargs):
 
         args_string = ", ".join([sg_sql.to_sql(sg_sql.literal(v)) for v in args])
         sql = f"SELECT `{str(self._udf_def.routine_ref)}`({args_string})"
-        iter, job = bf_io_bigquery.start_query_with_client(
+        iter, job = bf_io_bigquery.start_query_with_job(
             self._session.bqclient,
             sql=sql,
-            query_with_job=True,
             job_config=bigquery.QueryJobConfig(),
             publisher=self._session._publisher,
         )  # type: ignore
@@ -257,10 +256,9 @@ def __call__(self, *args, **kwargs):
 
         args_string = ", ".join([sg_sql.to_sql(sg_sql.literal(v)) for v in args])
         sql = f"SELECT `{str(self._udf_def.routine_ref)}`({args_string})"
-        iter, job = bf_io_bigquery.start_query_with_client(
+        iter, job = bf_io_bigquery.start_query_with_job(
             self._session.bqclient,
             sql=sql,
-            query_with_job=True,
             job_config=bigquery.QueryJobConfig(),
             publisher=self._session._publisher,
         )  # type: ignore

packages/bigframes/bigframes/session/__init__.py

@@ -2310,15 +2310,14 @@ def _start_query_ml_ddl(
         # so we must reset any encryption set in the job config
         # https://cloud.google.com/bigquery/docs/customer-managed-encryption#encrypt-model
         job_config.destination_encryption_configuration = None
-        iterator, query_job = bf_io_bigquery.start_query_with_client(
+        iterator, query_job = bf_io_bigquery.start_query_with_job(
             self.bqclient,
             sql,
             job_config=job_config,
             metrics=self._metrics,
             location=None,
             project=None,
             timeout=None,
-            query_with_job=True,
             job_retry=third_party_gcb_retry.DEFAULT_ML_JOB_RETRY,
             publisher=self._publisher,
             session=self,
@@ -2340,15 +2339,14 @@ def _create_object_table(self, path: str, connection: str) -> str:
                 uris = ['{path}']);
             """
         )
-        bf_io_bigquery.start_query_with_client(
+        bf_io_bigquery.start_query_with_job(
             self.bqclient,
             sql,
             job_config=bigquery.QueryJobConfig(),
             metrics=self._metrics,
             location=None,
             project=None,
             timeout=None,
-            query_with_job=True,
             publisher=self._publisher,
             session=self,
         )

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

@@ -262,73 +262,7 @@ def publish_bq_event(event):
     return publish_bq_event
 
 
-@overload
-def start_query_with_client(
-    bq_client: bigquery.Client,
-    sql: str,
-    *,
-    job_config: bigquery.QueryJobConfig,
-    location: Optional[str],
-    project: Optional[str],
-    timeout: Optional[float],
-    metrics: Optional[bigframes.session.metrics.ExecutionMetrics],
-    query_with_job: Literal[True],
-    publisher: bigframes.core.events.Publisher,
-    session=None,
-) -> Tuple[google.cloud.bigquery.table.RowIterator, bigquery.QueryJob]: ...
-
-
-@overload
-def start_query_with_client(
-    bq_client: bigquery.Client,
-    sql: str,
-    *,
-    job_config: bigquery.QueryJobConfig,
-    location: Optional[str],
-    project: Optional[str],
-    timeout: Optional[float],
-    metrics: Optional[bigframes.session.metrics.ExecutionMetrics],
-    query_with_job: Literal[False],
-    publisher: bigframes.core.events.Publisher,
-    session=None,
-) -> Tuple[google.cloud.bigquery.table.RowIterator, Optional[bigquery.QueryJob]]: ...
-
-
-@overload
-def start_query_with_client(
-    bq_client: bigquery.Client,
-    sql: str,
-    *,
-    job_config: bigquery.QueryJobConfig,
-    location: Optional[str],
-    project: Optional[str],
-    timeout: Optional[float],
-    metrics: Optional[bigframes.session.metrics.ExecutionMetrics],
-    query_with_job: Literal[True],
-    job_retry: google.api_core.retry.Retry,
-    publisher: bigframes.core.events.Publisher,
-    session=None,
-) -> Tuple[google.cloud.bigquery.table.RowIterator, bigquery.QueryJob]: ...
-
-
-@overload
-def start_query_with_client(
-    bq_client: bigquery.Client,
-    sql: str,
-    *,
-    job_config: bigquery.QueryJobConfig,
-    location: Optional[str],
-    project: Optional[str],
-    timeout: Optional[float],
-    metrics: Optional[bigframes.session.metrics.ExecutionMetrics],
-    query_with_job: Literal[False],
-    job_retry: google.api_core.retry.Retry,
-    publisher: bigframes.core.events.Publisher,
-    session=None,
-) -> Tuple[google.cloud.bigquery.table.RowIterator, Optional[bigquery.QueryJob]]: ...
-
-
-def start_query_with_client(
+def start_query_with_job(
     bq_client: bigquery.Client,
     sql: str,
     *,
@@ -337,7 +271,6 @@ def start_query_with_client(
     project: Optional[str] = None,
     timeout: Optional[float] = None,
     metrics: Optional[bigframes.session.metrics.ExecutionMetrics] = None,
-    query_with_job: bool = True,
     # TODO(tswast): We can stop providing our own default once we use a
     # google-cloud-bigquery version with
     # https://github.com/googleapis/python-bigquery/pull/2256 merged, likely
@@ -355,20 +288,6 @@ def start_query_with_client(
     add_and_trim_labels(job_config, session=session)
 
     try:
-        if not query_with_job:
-            results_iterator = bq_client._query_and_wait_bigframes(
-                sql,
-                job_config=job_config,
-                location=location,
-                project=project,
-                api_timeout=timeout,
-                job_retry=job_retry,
-                callback=create_bq_event_callback(publisher),
-            )
-            if metrics is not None:
-                metrics.count_job_stats(row_iterator=results_iterator)
-            return results_iterator, None
-
         query_job = bq_client.query(
             sql,
             job_config=job_config,
@@ -382,6 +301,61 @@ def start_query_with_client(
             ex.message += CHECK_DRIVE_PERMISSIONS
         raise
 
+    results_iterator = query_job.result()
+    _publish_events(
+        query_job=query_job,
+        total_rows=results_iterator.total_rows,
+        sql=sql,
+        publisher=publisher,
+        metrics=metrics,
+    )
+    return results_iterator, query_job
+
+
+def start_query_job_optional(
+    bq_client: bigquery.Client,
+    sql: str,
+    *,
+    job_config: bigquery.QueryJobConfig,
+    location: Optional[str] = None,
+    project: Optional[str] = None,
+    timeout: Optional[float] = None,
+    metrics: Optional[bigframes.session.metrics.ExecutionMetrics] = None,
+    # TODO(tswast): We can stop providing our own default once we use a
+    # google-cloud-bigquery version with
+    # https://github.com/googleapis/python-bigquery/pull/2256 merged, likely
+    # version 3.36.0 or later.
+    job_retry: google.api_core.retry.Retry = third_party_gcb_retry.DEFAULT_JOB_RETRY,
+    publisher: bigframes.core.events.Publisher,
+    session=None,
+):
+    add_and_trim_labels(job_config, session=session)
+    try:
+        results_iterator = bq_client._query_and_wait_bigframes(
+            sql,
+            job_config=job_config,
+            location=location,
+            project=project,
+            api_timeout=timeout,
+            job_retry=job_retry,
+            callback=create_bq_event_callback(publisher),
+        )
+        if metrics is not None:
+            metrics.count_job_stats(row_iterator=results_iterator)
+        return results_iterator, None
+    except google.api_core.exceptions.Forbidden as ex:
+        if "Drive credentials" in ex.message:
+            ex.message += CHECK_DRIVE_PERMISSIONS
+        raise
+
+
+def _publish_events(
+    query_job: bigquery.QueryJob,
+    sql: str,
+    total_rows: Optional[int],
+    publisher: bigframes.core.events.Publisher,
+    metrics: Optional[bigframes.session.metrics.ExecutionMetrics] = None,
+):
     if not query_job.configuration.dry_run:
         publisher.publish(
             bigframes.core.events.BigQuerySentEvent(
@@ -392,15 +366,14 @@ def start_query_with_client(
                 request_id=None,
             )
         )
-    results_iterator = query_job.result()
     if not query_job.configuration.dry_run:
         publisher.publish(
             bigframes.core.events.BigQueryFinishedEvent(
                 billing_project=query_job.project,
                 location=query_job.location,
                 job_id=query_job.job_id,
                 destination=query_job.destination,
-                total_rows=results_iterator.total_rows,
+                total_rows=total_rows,
                 total_bytes_processed=query_job.total_bytes_processed,
                 slot_millis=query_job.slot_millis,
                 created=query_job.created,
@@ -411,7 +384,6 @@ def start_query_with_client(
 
     if metrics is not None:
         metrics.count_job_stats(query_job=query_job)
-    return results_iterator, query_job
 
 
 def delete_tables_matching_session_id(
@@ -473,15 +445,14 @@ def create_bq_dataset_reference(
     """
     job_config = google.cloud.bigquery.QueryJobConfig()
 
-    _, query_job = start_query_with_client(
+    _, query_job = start_query_with_job(
         bq_client,
         "SELECT 1",
         location=location,
         job_config=job_config,
         project=project,
         timeout=None,
         metrics=None,
-        query_with_job=True,
         publisher=publisher,
     )
 

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

@@ -193,15 +193,14 @@ 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_client(
+            bigframes.session._io.bigquery.start_query_with_job_optional(
                 bq_client=bqclient,
                 sql=snapshot_sql,
                 job_config=bigquery.QueryJobConfig(dry_run=True),
                 location=None,
                 project=None,
                 timeout=None,
                 metrics=None,
-                query_with_job=False,
                 publisher=publisher,
             )
             return True
@@ -266,15 +265,14 @@ 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_client(
+    results, _ = bigframes.session._io.bigquery.start_query_with_job_optional(
         bq_client=bqclient,
         sql=is_unique_sql,
         job_config=job_config,
         timeout=None,
         location=None,
         project=None,
         metrics=None,
-        query_with_job=False,
         publisher=publisher,
     )
     row = next(iter(results))