WIP: use callbacks for progress bars

Author: tswast

bigframes/display/anywidget.py

@@ -149,7 +149,7 @@ def _validate_page_size(self, proposal: Dict[str, Any]) -> int:
 
     def _update_progress(self, event):
         # TODO: use formatting helpers here.
-        self.progress_html = f"<code>{repr(event)}"
+        self.progress_html = f"<code>{repr(event)}</code>"
 
     def _get_next_batch(self) -> bool:
         """

bigframes/display/table_widget.js

@@ -60,6 +60,7 @@ function render({ model, el }) {
 	const pageSizeSelect = document.createElement("select");
 
 	// Add CSS classes
+	progressContainer.classList.add("progress-container");
 	tableContainer.classList.add("table-container");
 	footer.classList.add("footer");
 	paginationContainer.classList.add("pagination");
@@ -123,7 +124,7 @@ function render({ model, el }) {
 	}
 
 	/** Updates the HTML in the progress container. */
-	function handleTableHTMLChange() {
+	function handleProgressHTMLChange() {
 		// Note: Using innerHTML is safe here because the content is generated
 		// by a trusted backend (formatting_helpers).
 		progressContainer.innerHTML = model.get(ModelProperty.PROGRESS_HTML);
@@ -147,6 +148,7 @@ function render({ model, el }) {
 		}
 	});
 	model.on(Event.CHANGE_TABLE_HTML, handleTableHTMLChange);
+	model.on(Event.CHANGE_PROGRESS_HTML, handleProgressHTMLChange);
 
 	// Assemble the DOM
 	paginationContainer.appendChild(prevPage);

bigframes/formatting_helpers.py

@@ -42,6 +42,45 @@
 }
 
 
+def create_progress_bar_callback(
+    *,
+    progress_bar: Optional[str] = None,
+    callback: Callable = lambda _: None,
+) -> Callable:
+    if progress_bar == "auto":
+        progress_bar = "notebook" if in_ipython() else "terminal"
+
+    if progress_bar == "notebook":
+        loading_bar = display.HTML("")
+        display_id = str(random.random())
+        display.display(loading_bar, display_id=display_id)
+
+        def outer_callback(event):
+            callback(event)
+            display.update_display(
+                display.HTML(get_query_job_loading_html(event)),
+                display_id=display_id,
+            )
+
+    elif progress_bar == "terminal":
+        previous_bar_text = ""
+
+        def outer_callback(event):
+            nonlocal previous_bar_text
+
+            callback(event)
+
+            bar_text = get_query_job_loading_string(event)
+            if bar_text != previous_bar_text:
+                print(bar_text)
+                previous_bar_text = bar_text
+
+    else:
+        outer_callback = callback
+
+    return outer_callback
+
+
 def add_feedback_link(
     exception: Union[
         api_core_exceptions.RetryError, api_core_exceptions.GoogleAPICallError
@@ -125,7 +164,6 @@ def wait_for_query_job(
     query_job: bigquery.QueryJob,
     max_results: Optional[int] = None,
     page_size: Optional[int] = None,
-    progress_bar: Optional[str] = None,
     callback: Callable = lambda _: None,
 ) -> bigquery.table.RowIterator:
     """Return query results. Displays a progress bar while the query is running
@@ -141,36 +179,8 @@ def wait_for_query_job(
     Returns:
         A row iterator over the query results.
     """
-    if progress_bar == "auto":
-        progress_bar = "notebook" if in_ipython() else "terminal"
-
-    if progress_bar == "notebook":
-        loading_bar = display.HTML(get_query_job_loading_html(query_job))
-        display_id = str(random.random())
-        display.display(loading_bar, display_id=display_id)
-
-        def extended_callback(event):
-            callback(event)
-            display.update_display(
-                display.HTML(get_query_job_loading_html(query_job)),
-                display_id=display_id,
-            )
-
-    elif progress_bar == "terminal":
-        initial_loading_bar = get_query_job_loading_string(query_job)
-        print(initial_loading_bar)
-
-        def extended_callback(event):
-            callback(event)
-
-            if initial_loading_bar != get_query_job_loading_string(query_job):
-                print(get_query_job_loading_string(query_job))
-
-    else:
-        extended_callback = callback
-
     try:
-        extended_callback(
+        callback(
             # DONOTSUBMIT: we should create our own events.
             google.cloud.bigquery._job_helpers.QueryReceivedEvent(
                 billing_project=query_job.project,
@@ -184,11 +194,15 @@ def extended_callback(event):
                 ended=query_job.ended,
             )
         )
+        # TODO(tswast): Add a timeout so that progress bars can make updates as
+        # the query stats come int.
+        # TODO(tswast): Listen for cancellation on the callback (or maybe
+        # callbacks should just raise KeyboardInterrupt like IPython does?).
         query_results = query_job.result(
             page_size=page_size,
             max_results=max_results,
         )
-        extended_callback(
+        callback(
             # DONOTSUBMIT: we should create our own events.
             google.cloud.bigquery._job_helpers.QueryFinishedEvent(
                 billing_project=query_job.project,
@@ -206,13 +220,16 @@ def extended_callback(event):
         )
         return query_results
     except api_core_exceptions.RetryError as exc:
+        # TODO: turn this into a callback event, too.
         add_feedback_link(exc)
         raise
     except api_core_exceptions.GoogleAPICallError as exc:
+        # TODO: turn this into a callback event, too.
         add_feedback_link(exc)
         raise
     except KeyboardInterrupt:
         query_job.cancel()
+        # TODO: turn this into a callback event, too.
         print(
             f"Requested cancellation for {query_job.job_type.capitalize()}"
             f" job {query_job.job_id} in location {query_job.location}..."

bigframes/session/_io/bigquery/__init__.py

@@ -362,6 +362,12 @@ def start_query_with_client(
     """
     Starts query job and waits for results.
     """
+    opts = bigframes.options.display
+    progress_callback = formatting_helpers.create_progress_bar_callback(
+        progress_bar=opts.progress_bar,
+        callback=callback,
+    )
+
     try:
         # Note: Ensure no additional labels are added to job_config after this
         # point, as `add_and_trim_labels` ensures the label count does not
@@ -376,7 +382,7 @@ def start_query_with_client(
                 project=project,
                 api_timeout=timeout,
                 job_retry=job_retry,
-                callback=callback,
+                callback=progress_callback,
             )
             if metrics is not None:
                 metrics.count_job_stats(row_iterator=results_iterator)
@@ -395,7 +401,6 @@ def start_query_with_client(
             ex.message += CHECK_DRIVE_PERMISSIONS
         raise
 
-    opts = bigframes.options.display
     if opts.progress_bar is not None and not query_job.configuration.dry_run:
         results_iterator = formatting_helpers.wait_for_query_job(
             query_job,

bigframes/session/bq_caching_executor.py

@@ -321,7 +321,6 @@ def _export_gbq(
 
         # TODO(swast): plumb through the api_name of the user-facing api that
         # caused this query.
-        breakpoint()
         row_iter, query_job = self._run_execute_query(
             sql=sql,
             job_config=job_config,