@@ -134,16 +134,14 @@ def repr_query_job_html(query_job: Optional[bigquery.QueryJob]):
134134 return res
135135
136136
137- current_display : Optional [display .HTML ] = None
138137current_display_id : Optional [str ] = None
139- previous_display_html : str = ""
140138
141139
142140def progress_callback (
143141 event : bigframes .core .events .Event ,
144142):
145143 """Displays a progress bar while the query is running"""
146- global current_display , current_display_id , previous_display_html
144+ global current_display_id
147145
148146 try :
149147 import bigframes ._config
@@ -162,59 +160,43 @@ def progress_callback(
162160 if progress_bar == "notebook" :
163161 import IPython .display as display
164162
165- if (
166- isinstance (event , bigframes .core .events .ExecutionStarted )
167- or current_display is None
168- or current_display_id is None
169- ):
170- previous_display_html = ""
171- current_display_id = str (random .random ())
172- current_display = display .HTML ("Starting." )
173- display .display (
174- current_display ,
175- display_id = current_display_id ,
176- )
163+ display_html = None
164+
165+ if isinstance (event , bigframes .core .events .ExecutionStarted ):
166+ # Start a new context for progress output.
167+ current_display_id = None
168+
169+ elif isinstance (event , bigframes .core .events .BigQuerySentEvent ):
170+ display_html = render_bqquery_sent_event_html (event )
177171
178- if isinstance (event , bigframes .core .events .BigQuerySentEvent ):
179- previous_display_html = render_bqquery_sent_event_html (event )
180- display .update_display (
181- display .HTML (previous_display_html ),
182- display_id = current_display_id ,
183- )
184172 elif isinstance (event , bigframes .core .events .BigQueryRetryEvent ):
185- previous_display_html = render_bqquery_retry_event_html (event )
186- display .update_display (
187- display .HTML (previous_display_html ),
188- display_id = current_display_id ,
189- )
173+ display_html = render_bqquery_retry_event_html (event )
174+
190175 elif isinstance (event , bigframes .core .events .BigQueryReceivedEvent ):
191- previous_display_html = render_bqquery_received_event_html (event )
192- display .update_display (
193- display .HTML (previous_display_html ),
194- display_id = current_display_id ,
195- )
176+ display_html = render_bqquery_received_event_html (event )
177+
196178 elif isinstance (event , bigframes .core .events .BigQueryFinishedEvent ):
197- previous_display_html = render_bqquery_finished_event_html (event )
198- display . update_display (
199- display . HTML ( previous_display_html ),
200- display_id = current_display_id ,
201- )
202- elif isinstance ( event , bigframes . core . events . ExecutionFinished ) :
203- if previous_display_html :
179+ display_html = render_bqquery_finished_event_html (event )
180+
181+ elif isinstance ( event , bigframes . core . events . SessionClosed ):
182+ display_html = f"Session { event . session_id } closed." ) ,
183+
184+ if display_html :
185+ if current_display_id :
204186 display .update_display (
205- display .HTML (f"✅ Completed. { previous_display_html } " ),
187+ display .HTML (display_html ),
206188 display_id = current_display_id ,
207189 )
208-
209- elif isinstance ( event , bigframes . core . events . SessionClosed ):
210- display .update_display (
211- display . HTML ( f"Session { event . session_id } closed." ) ,
212- display_id = current_display_id ,
213- )
190+ else :
191+ display . display (
192+ display .HTML ( display_html ),
193+ display_id = current_display_id ,
194+ )
195+
214196 elif progress_bar == "terminal" :
215- if isinstance ( event , bigframes . core . events . ExecutionStarted ):
216- print ( "Starting execution." )
217- elif isinstance (event , bigframes .core .events .BigQuerySentEvent ):
197+ message = None
198+
199+ if isinstance (event , bigframes .core .events .BigQuerySentEvent ):
218200 message = render_bqquery_sent_event_plaintext (event )
219201 print (message )
220202 elif isinstance (event , bigframes .core .events .BigQueryRetryEvent ):
@@ -226,8 +208,6 @@ def progress_callback(
226208 elif isinstance (event , bigframes .core .events .BigQueryFinishedEvent ):
227209 message = render_bqquery_finished_event_plaintext (event )
228210 print (message )
229- elif isinstance (event , bigframes .core .events .ExecutionFinished ):
230- print ("Execution done." )
231211
232212
233213def wait_for_job (job : GenericJob , progress_bar : Optional [str ] = None ):
0 commit comments