feat(session): refactor JobMetadata extraction and truncate query strings

Author: shuoweil

bigframes/session/metrics.py

@@ -52,6 +52,85 @@ class JobMetadata:
     output_rows: Optional[int] = None
     source_format: Optional[str] = None
 
+    @classmethod
+    def from_job(
+        cls, query_job: Union[QueryJob, LoadJob], exec_seconds: Optional[float] = None
+    ) -> "JobMetadata":
+        query_text = getattr(query_job, "query", None)
+        if query_text and len(query_text) > 1024:
+            query_text = query_text[:1021] + "..."
+
+        metadata = cls(
+            job_id=query_job.job_id,
+            location=query_job.location,
+            project=query_job.project,
+            creation_time=query_job.created,
+            start_time=query_job.started,
+            end_time=query_job.ended,
+            duration_seconds=exec_seconds,
+            status=query_job.state,
+            job_type=query_job.job_type,
+            error_result=query_job.error_result,
+            query=query_text,
+            job_url=f"https://console.cloud.google.com/bigquery?project={query_job.project}&j=bq:{query_job.location}:{query_job.job_id}&page=queryresults"
+            if getattr(query_job, "job_id", None)
+            else None,
+        )
+        if isinstance(query_job, QueryJob):
+            metadata.cached = getattr(query_job, "cache_hit", None)
+            metadata.destination_table = (
+                str(query_job.destination) if query_job.destination else None
+            )
+            metadata.total_bytes_processed = getattr(
+                query_job, "total_bytes_processed", None
+            )
+            metadata.total_slot_ms = getattr(query_job, "slot_millis", None)
+        elif isinstance(query_job, LoadJob):
+            metadata.output_rows = getattr(query_job, "output_rows", None)
+            metadata.input_files = getattr(query_job, "input_files", None)
+            metadata.input_bytes = getattr(query_job, "input_bytes", None)
+            metadata.destination_table = (
+                str(query_job.destination)
+                if getattr(query_job, "destination", None)
+                else None
+            )
+            if getattr(query_job, "source_uris", None):
+                metadata.source_uris = list(query_job.source_uris)
+            if query_job.configuration and hasattr(
+                query_job.configuration, "source_format"
+            ):
+                metadata.source_format = query_job.configuration.source_format
+
+        return metadata
+
+    @classmethod
+    def from_row_iterator(
+        cls, row_iterator: bq_table.RowIterator, exec_seconds: Optional[float] = None
+    ) -> "JobMetadata":
+        query_text = getattr(row_iterator, "query", None)
+        if query_text and len(query_text) > 1024:
+            query_text = query_text[:1021] + "..."
+
+        return cls(
+            job_id=getattr(row_iterator, "job_id", None),
+            query_id=getattr(row_iterator, "query_id", None),
+            location=getattr(row_iterator, "location", None),
+            project=getattr(row_iterator, "project", None),
+            creation_time=getattr(row_iterator, "created", None),
+            start_time=getattr(row_iterator, "started", None),
+            end_time=getattr(row_iterator, "ended", None),
+            duration_seconds=exec_seconds,
+            status="DONE",
+            total_bytes_processed=getattr(row_iterator, "total_bytes_processed", None),
+            total_slot_ms=getattr(row_iterator, "slot_millis", None),
+            job_type="query",
+            cached=getattr(row_iterator, "cache_hit", None),
+            query=query_text,
+            job_url=f"https://console.cloud.google.com/bigquery?project={getattr(row_iterator, 'project', '')}&j=bq:{getattr(row_iterator, 'location', '')}:{getattr(row_iterator, 'job_id', '')}&page=queryresults"
+            if getattr(row_iterator, "job_id", None)
+            else None,
+        )
+
 
 @dataclasses.dataclass
 class ExecutionMetrics:
@@ -87,25 +166,7 @@ def count_job_stats(
             self.execution_secs += exec_seconds
 
             self.jobs.append(
-                JobMetadata(
-                    job_id=getattr(row_iterator, "job_id", None),
-                    query_id=getattr(row_iterator, "query_id", None),
-                    location=getattr(row_iterator, "location", None),
-                    project=getattr(row_iterator, "project", None),
-                    creation_time=created,
-                    start_time=getattr(row_iterator, "started", None),
-                    end_time=ended,
-                    duration_seconds=exec_seconds,
-                    status="DONE",
-                    total_bytes_processed=bytes_processed,
-                    total_slot_ms=slot_millis,
-                    job_type="query",
-                    cached=getattr(row_iterator, "cache_hit", None),
-                    query=getattr(row_iterator, "query", None),
-                    job_url=f"https://console.cloud.google.com/bigquery?project={getattr(row_iterator, 'project', '')}&j=bq:{getattr(row_iterator, 'location', '')}:{getattr(row_iterator, 'job_id', '')}&page=queryresults"
-                    if getattr(row_iterator, "job_id", None)
-                    else None,
-                )
+                JobMetadata.from_row_iterator(row_iterator, exec_seconds=exec_seconds)
             )
 
         elif isinstance(query_job, QueryJob) and query_job.configuration.dry_run:
@@ -125,28 +186,10 @@ def count_job_stats(
                 self.slot_millis += slot_millis or 0
                 self.execution_secs += exec_seconds or 0
 
-                self.jobs.append(
-                    JobMetadata(
-                        job_id=query_job.job_id,
-                        location=query_job.location,
-                        project=query_job.project,
-                        creation_time=query_job.created,
-                        start_time=query_job.started,
-                        end_time=query_job.ended,
-                        duration_seconds=exec_seconds,
-                        status=query_job.state,
-                        total_bytes_processed=bytes_processed,
-                        total_slot_ms=slot_millis,
-                        job_type=query_job.job_type,
-                        error_result=query_job.error_result,
-                        cached=query_job.cache_hit,
-                        query=query_job.query,
-                        destination_table=str(query_job.destination)
-                        if query_job.destination
-                        else None,
-                        job_url=f"https://console.cloud.google.com/bigquery?project={query_job.project}&j=bq:{query_job.location}:{query_job.job_id}&page=queryresults",
-                    )
-                )
+                metadata = JobMetadata.from_job(query_job, exec_seconds=exec_seconds)
+                metadata.total_bytes_processed = bytes_processed
+                metadata.total_slot_ms = slot_millis
+                self.jobs.append(metadata)
 
         else:
             # Handle other job types (e.g. LoadJob)
@@ -156,36 +199,7 @@ def count_job_stats(
                 if query_job.ended and query_job.created
                 else None
             )
-
-            job_metadata = JobMetadata(
-                job_id=query_job.job_id,
-                location=query_job.location,
-                project=query_job.project,
-                creation_time=query_job.created,
-                start_time=query_job.started,
-                end_time=query_job.ended,
-                duration_seconds=duration,
-                status=query_job.state,
-                job_type=query_job.job_type,
-                error_result=query_job.error_result,
-                job_url=f"https://console.cloud.google.com/bigquery?project={query_job.project}&j=bq:{query_job.location}:{query_job.job_id}&page=queryresults",
-            )
-
-            if isinstance(query_job, bigquery.LoadJob):
-                job_metadata.output_rows = getattr(query_job, "output_rows", None)
-                job_metadata.input_files = getattr(query_job, "input_files", None)
-                job_metadata.input_bytes = getattr(query_job, "input_bytes", None)
-                job_metadata.destination_table = (
-                    str(query_job.destination) if query_job.destination else None
-                )
-                if query_job.source_uris:
-                    job_metadata.source_uris = list(query_job.source_uris)
-                if query_job.configuration and hasattr(
-                    query_job.configuration, "source_format"
-                ):
-                    job_metadata.source_format = query_job.configuration.source_format
-
-            self.jobs.append(job_metadata)
+            self.jobs.append(JobMetadata.from_job(query_job, exec_seconds=duration))
 
         # For pytest runs only, log information about the query job
         # to a file in order to create a performance report.