Add request middleware to BQ

bigquery/Cargo.toml

@@ -27,6 +27,8 @@ base64 = "0.21"
 bigdecimal = { version="0.3", features=["serde"] }
 num-bigint = "0.4"
 backon = "0.4"
+reqwest-middleware = "0.2"
+anyhow = "1.0"
 
 google-cloud-auth = { optional = true, version = "0.13", path="../foundation/auth", default-features=false }
 

bigquery/src/client.rs

@@ -33,7 +33,7 @@ const JOB_RETRY_REASONS: [&str; 3] = ["backendError", "rateLimitExceeded", "inte
 
 #[derive(Debug)]
 pub struct ClientConfig {
-    http: reqwest::Client,
+    http: reqwest_middleware::ClientWithMiddleware,
     bigquery_endpoint: Cow<'static, str>,
     token_source_provider: Box<dyn TokenSourceProvider>,
     environment: Environment,
@@ -65,7 +65,7 @@ impl ClientConfig {
         grpc_token_source_provider: Box<dyn TokenSourceProvider>,
     ) -> Self {
         Self {
-            http: reqwest::Client::default(),
+            http: reqwest_middleware::ClientBuilder::new(reqwest::Client::default()).build(),
             bigquery_endpoint: "https://bigquery.googleapis.com".into(),
             token_source_provider: http_token_source_provider,
             environment: Environment::GoogleCloud(grpc_token_source_provider),
@@ -81,7 +81,7 @@ impl ClientConfig {
         self.streaming_read_config = value;
         self
     }
-    pub fn with_http_client(mut self, value: reqwest::Client) -> Self {
+    pub fn with_http_client(mut self, value: reqwest_middleware::ClientWithMiddleware) -> Self {
         self.http = value;
         self
     }

bigquery/src/http/bigquery_client.rs

@@ -1,6 +1,7 @@
 use std::sync::Arc;
 
-use reqwest::{Client, RequestBuilder, Response};
+use reqwest::Response;
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use google_cloud_token::TokenSource;
 
@@ -139,7 +140,12 @@ pub(crate) mod test {
         .unwrap();
         let cred = tsp.source_credentials.clone();
         let ts = tsp.token_source();
-        let client = BigqueryClient::new(ts, "https://bigquery.googleapis.com", reqwest::Client::new(), false);
+        let client = BigqueryClient::new(
+            ts,
+            "https://bigquery.googleapis.com",
+            reqwest_middleware::ClientBuilder::new(reqwest::Client::new()).build(),
+            false,
+        );
         (client, cred.unwrap().project_id.unwrap())
     }
 

bigquery/src/http/dataset/delete.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 pub fn build(base_url: &str, client: &Client, project_id: &str, dataset_id: &str) -> RequestBuilder {
     let url = format!("{}/projects/{}/datasets/{}", base_url, project_id, dataset_id);

bigquery/src/http/dataset/get.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 pub fn build(base_url: &str, client: &Client, project_id: &str, dataset_id: &str) -> RequestBuilder {
     let url = format!("{}/projects/{}/datasets/{}", base_url, project_id, dataset_id);

bigquery/src/http/dataset/insert.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::dataset::Dataset;
 

bigquery/src/http/dataset/list.rs

@@ -1,6 +1,6 @@
 use std::collections::HashMap;
 
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::dataset::DatasetReference;
 

bigquery/src/http/dataset/patch.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::dataset::Dataset;
 

bigquery/src/http/error.rs

@@ -10,11 +10,24 @@ pub enum Error {
     #[error(transparent)]
     HttpClient(#[from] reqwest::Error),
 
+    /// An error from the HTTP client.
+    #[error(transparent)]
+    HttpMiddleware(anyhow::Error),
+
     /// An error from a token source.
     #[error("token source failed: {0}")]
     TokenSource(Box<dyn std::error::Error + Send + Sync>),
 }
 
+impl From<reqwest_middleware::Error> for Error {
+    fn from(error: reqwest_middleware::Error) -> Self {
+        match error {
+            reqwest_middleware::Error::Reqwest(err) => Error::HttpClient(err),
+            reqwest_middleware::Error::Middleware(err) => Error::HttpMiddleware(err),
+        }
+    }
+}
+
 #[derive(Debug, serde::Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct ErrorResponse {

bigquery/src/http/job/cancel.rs

@@ -1,5 +1,5 @@
 use reqwest::header::CONTENT_LENGTH;
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::job::Job;
 

bigquery/src/http/job/delete.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 pub fn build(base_url: &str, client: &Client, project_id: &str, job_id: &str) -> RequestBuilder {
     let url = format!("{}/projects/{}/jobs/{}", base_url, project_id, job_id);

bigquery/src/http/job/get.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 #[derive(Clone, PartialEq, serde::Deserialize, serde::Serialize, Debug, Default)]
 #[serde(rename_all = "camelCase")]

bigquery/src/http/job/get_query_results.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::job::JobReference;
 use crate::http::table::TableSchema;

bigquery/src/http/job/insert.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::job::Job;
 

bigquery/src/http/job/list.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::job::{JobConfiguration, JobReference, JobState, JobStatistics, JobStatus};
 use crate::http::types::ErrorProto;

bigquery/src/http/job/query.rs

@@ -1,6 +1,6 @@
 use std::collections::HashMap;
 
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::dataset::DatasetReference;
 use crate::http::job::{DmlStats, JobReference, SessionInfo};

bigquery/src/http/model/delete.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 pub fn build(base_url: &str, client: &Client, project_id: &str, dataset_id: &str, model_id: &str) -> RequestBuilder {
     let url = format!(

bigquery/src/http/model/get.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 pub fn build(base_url: &str, client: &Client, project_id: &str, dataset_id: &str, model_id: &str) -> RequestBuilder {
     let url = format!(

bigquery/src/http/model/list.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::model::{ModelReference, ModelType};
 

bigquery/src/http/model/patch.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::model::Model;
 

bigquery/src/http/routine/delete.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 pub fn build(base_url: &str, client: &Client, project_id: &str, dataset_id: &str, routine_id: &str) -> RequestBuilder {
     let url = format!(

bigquery/src/http/routine/get.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 pub fn build(base_url: &str, client: &Client, project_id: &str, dataset_id: &str, routine_id: &str) -> RequestBuilder {
     let url = format!(

bigquery/src/http/routine/insert.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::routine::Routine;
 

bigquery/src/http/routine/list.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::routine::{Language, RemoteFunctionOptions, RoutineReference, RoutineType};
 

bigquery/src/http/routine/update.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::routine::Routine;
 

bigquery/src/http/row_access_policy/get_iam_policy.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::table::get_iam_policy::GetIamPolicyRequest;
 

bigquery/src/http/row_access_policy/list.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 use time::OffsetDateTime;
 
 use crate::http::row_access_policy::RowAccessPolicyReference;

bigquery/src/http/row_access_policy/test_iam_permissions.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::table::test_iam_permissions::TestIamPermissionsRequest;
 

bigquery/src/http/table/delete.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 pub fn build(base_url: &str, client: &Client, project_id: &str, dataset_id: &str, table_id: &str) -> RequestBuilder {
     let url = format!(

bigquery/src/http/table/get.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 pub fn build(base_url: &str, client: &Client, project_id: &str, dataset_id: &str, table_id: &str) -> RequestBuilder {
     let url = format!(

bigquery/src/http/table/get_iam_policy.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::types::GetPolicyOptions;
 

bigquery/src/http/table/insert.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::table::Table;
 

bigquery/src/http/table/list.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::table::{Clustering, RangePartitioning, TableReference, TimePartitioning};
 

bigquery/src/http/table/patch.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::table::Table;
 

bigquery/src/http/table/set_iam_policy.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 use crate::http::types::Policy;
 

bigquery/src/http/table/test_iam_permissions.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 #[derive(Clone, PartialEq, Eq, serde::Deserialize, serde::Serialize, Default, Debug)]
 #[serde(rename_all = "camelCase")]

bigquery/src/http/tabledata/insert_all.rs

@@ -1,4 +1,4 @@
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 use serde::Serialize;
 
 #[derive(Clone, PartialEq, Eq, serde::Deserialize, serde::Serialize, Debug)]

bigquery/src/http/tabledata/list.rs

@@ -1,6 +1,6 @@
 use std::fmt::Debug;
 
-use reqwest::{Client, RequestBuilder};
+use reqwest_middleware::{ClientWithMiddleware as Client, RequestBuilder};
 
 #[derive(Clone, PartialEq, Eq, serde::Deserialize, serde::Serialize, Debug)]
 #[serde(untagged)]

storage/src/http/mod.rs

@@ -32,7 +32,7 @@ pub enum Error {
 
     /// An error from one of the middleware used.
     #[error(transparent)]
-    HttpMiddleware(#[from] anyhow::Error),
+    HttpMiddleware(anyhow::Error),
 
     /// An error from a token source.
     #[error("token source failed: {0}")]