fix execution metadata

Author: TrevorBergeron

bigframes/session/bq_caching_executor.py

@@ -320,7 +320,7 @@ def _export_gbq(
 
         # TODO(swast): plumb through the api_name of the user-facing api that
         # caused this query.
-        iterator, _ = self._run_execute_query(
+        iterator, job = self._run_execute_query(
             sql=sql,
             job_config=job_config,
         )
@@ -338,7 +338,9 @@ def _export_gbq(
 
         return executor.EmptyExecuteResult(
             bf_schema=array_value.schema,
-            execution_metadata=executor.ExecutionMetadata.from_iterator(iterator),
+            execution_metadata=executor.ExecutionMetadata.from_iterator_and_job(
+                iterator, job
+            ),
         )
 
     def dry_run(
@@ -681,7 +683,9 @@ def _execute_plan_gbq(
             assert compiled.row_order is not None
             self.cache.cache_results_table(og_plan, result_bq_data)
 
-        execution_metadata = executor.ExecutionMetadata.from_iterator(iterator)
+        execution_metadata = executor.ExecutionMetadata.from_iterator_and_job(
+            iterator, query_job
+        )
         if result_bq_data is not None:
             return executor.BQTableExecuteResult(
                 data=result_bq_data,

bigframes/session/direct_gbq_execution.py

@@ -68,7 +68,9 @@ def execute(
         return executor.LocalExecuteResult(
             data=iterator.to_arrow(),
             bf_schema=plan.schema,
-            execution_metadata=executor.ExecutionMetadata.from_iterator(iterator),
+            execution_metadata=executor.ExecutionMetadata.from_iterator_and_job(
+                iterator, query_job
+            ),
         )
 
     def _run_execute_query(

bigframes/session/executor.py

@@ -171,10 +171,10 @@ class ExecutionMetadata:
     bytes_processed: Optional[int] = None
 
     @classmethod
-    def from_iterator(cls, iterator: bq_table.RowIterator) -> ExecutionMetadata:
-        return cls(
-            query_job=iterator.query, bytes_processed=iterator.total_bytes_processed
-        )
+    def from_iterator_and_job(
+        cls, iterator: bq_table.RowIterator, job: Optional[bigquery.QueryJob]
+    ) -> ExecutionMetadata:
+        return cls(query_job=job, bytes_processed=iterator.total_bytes_processed)
 
 
 class LocalExecuteResult(ExecuteResult):