lquerel / gcp-bigquery-client

GCP BigQuery Client (Rust)
Apache License 2.0
92 stars 60 forks source link

make ResultSet more general #97

Open imor opened 2 days ago

imor commented 2 days ago

NOTE: review and merge the following PR first: https://github.com/lquerel/gcp-bigquery-client/pull/96. This is because this PR's branch is cut from the previous PR's branch.

This PR makes ResultSet more general by only holding onto a Vec<TableRow> instead of a QueryResponse. This change makes it possible to create a ResultSet not only from a QueryResponse but also from a GetQueryResultsResponse. Need for this arose when I wanted to create a single function which will first call JobApi::query and if it returns job_complete=false then fall back to polling JobApi::get_query_results to get the results when the query job finishes. Such a function needed the ability to create a ResultSet from the results of either of these APIs. This is the function wrote that uses the refactored code:

async fn query(&self, query: String) -> Result<ResultSet, BQError> {
        let query_response = self
            .client
            .job()
            .query(&self.project_id, QueryRequest::new(query))
            .await?;

        if query_response.job_complete.unwrap_or(false) {
            info!("job complete. returning result set");
            Ok(ResultSet::new_from_query_response(query_response))
        } else {
            let job_id = query_response
                .job_reference
                .as_ref()
                .expect("missing job reference")
                .job_id
                .as_ref()
                .expect("missing job id");
            loop {
                info!("job incomplete. waiting for 5 seconds");
                sleep(Duration::from_secs(5)).await;
                let result = self
                    .client
                    .job()
                    .get_query_results(&self.project_id, job_id, Default::default())
                    .await?;
                if result.job_complete.unwrap_or(false) {
                    info!("job complete. returning result set");
                    let result_set = ResultSet::new_from_get_query_results_response(result);
                    break Ok(result_set);
                }
            }
        }
    }

Note: This is a breaking API change, so should be included only in a major version bumped release.