Apply the same fix as #54 to query_all (#64)

Fixes #61.
lquerel · Sep 27, 2023 · b0bfa86 · b0bfa86
1 parent 0bba73f
commit b0bfa86
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/src/job.rs b/src/job.rs
@@ -101,15 +101,22 @@ impl JobApi {
                             project_id,
                             job_id,
                             GetQueryResultsParameters {
-                                page_token,
+                                page_token: page_token.clone(),
                                 max_results: page_size,
                                 ..Default::default()
                             },
                         )
                         .await?;
 
+                    // Waiting for the job to be completed.
+                    if !qr.job_complete.unwrap_or(false) {
+                        tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;
+                        continue;
+                    }
+
                     // Rows is present when the query finishes successfully.
-                    yield Ok(qr.rows.expect("Rows are not present"));
+                    // Rows is empty when query result is empty.
+                    yield Ok(qr.rows.unwrap_or_default());
 
                     page_token = match qr.page_token {
                         None => break,