You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by tu...@apache.org on 2022/11/29 10:45:08 UTC
[arrow-rs] branch master updated: object_store: add support for using proxy_url for connection testing (#3109)
This is an automated email from the ASF dual-hosted git repository.
tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new b2bfe9c5f object_store: add support for using proxy_url for connection testing (#3109)
b2bfe9c5f is described below
commit b2bfe9c5fb4f09ef9f7b1fd67ec767f6d42cf8d8
Author: Sumit <su...@users.noreply.github.com>
AuthorDate: Tue Nov 29 11:45:02 2022 +0100
object_store: add support for using proxy_url for connection testing (#3109)
---
object_store/src/aws/client.rs | 23 +++++++++++-----
object_store/src/aws/mod.rs | 58 +++++++++++++++++++++++++++++++++++++---
object_store/src/azure/client.rs | 25 ++++++++++++-----
object_store/src/azure/mod.rs | 11 +++++++-
object_store/src/gcp/mod.rs | 55 ++++++++++++++++++++++++++++++++++++-
5 files changed, 153 insertions(+), 19 deletions(-)
diff --git a/object_store/src/aws/client.rs b/object_store/src/aws/client.rs
index a07cdb3c6..e51fe415c 100644
--- a/object_store/src/aws/client.rs
+++ b/object_store/src/aws/client.rs
@@ -88,6 +88,9 @@ pub(crate) enum Error {
#[snafu(display("Got invalid multipart response: {}", source))]
InvalidMultipartResponse { source: quick_xml::de::DeError },
+
+ #[snafu(display("Unable to use proxy url: {}", source))]
+ ProxyUrl { source: reqwest::Error },
}
impl From<Error> for crate::Error {
@@ -201,6 +204,7 @@ pub struct S3Config {
pub credentials: Box<dyn CredentialProvider>,
pub retry_config: RetryConfig,
pub allow_http: bool,
+ pub proxy_url: Option<String>,
}
impl S3Config {
@@ -216,13 +220,20 @@ pub(crate) struct S3Client {
}
impl S3Client {
- pub fn new(config: S3Config) -> Self {
- let client = reqwest::ClientBuilder::new()
- .https_only(!config.allow_http)
- .build()
- .unwrap();
+ pub fn new(config: S3Config) -> Result<Self> {
+ let builder = reqwest::ClientBuilder::new().https_only(!config.allow_http);
+ let client = match &config.proxy_url {
+ Some(ref url) => {
+ let pr = reqwest::Proxy::all(url)
+ .map_err(|source| Error::ProxyUrl { source })?;
+ builder.proxy(pr)
+ }
+ _ => builder,
+ }
+ .build()
+ .unwrap();
- Self { config, client }
+ Ok(Self { config, client })
}
/// Returns the config
diff --git a/object_store/src/aws/mod.rs b/object_store/src/aws/mod.rs
index 4a810658c..cf7a5542e 100644
--- a/object_store/src/aws/mod.rs
+++ b/object_store/src/aws/mod.rs
@@ -36,7 +36,7 @@ use bytes::Bytes;
use chrono::{DateTime, Utc};
use futures::stream::BoxStream;
use futures::TryStreamExt;
-use reqwest::Client;
+use reqwest::{Client, Proxy};
use snafu::{OptionExt, ResultExt, Snafu};
use std::collections::BTreeSet;
use std::ops::Range;
@@ -120,6 +120,9 @@ enum Error {
#[snafu(display("Error reading token file: {}", source))]
ReadTokenFile { source: std::io::Error },
+
+ #[snafu(display("Unable to use proxy url: {}", source))]
+ ProxyUrl { source: reqwest::Error },
}
impl From<Error> for super::Error {
@@ -363,6 +366,7 @@ pub struct AmazonS3Builder {
virtual_hosted_style_request: bool,
metadata_endpoint: Option<String>,
profile: Option<String>,
+ proxy_url: Option<String>,
}
impl AmazonS3Builder {
@@ -537,6 +541,12 @@ impl AmazonS3Builder {
self
}
+ /// Set the proxy_url to be used by the underlying client
+ pub fn with_proxy_url(mut self, proxy_url: impl Into<String>) -> Self {
+ self.proxy_url = Some(proxy_url.into());
+ self
+ }
+
/// Set the AWS profile name, see <https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-profiles.html>
///
/// This makes use of [aws-config] to provide credentials and therefore requires
@@ -561,6 +571,14 @@ impl AmazonS3Builder {
let bucket = self.bucket_name.context(MissingBucketNameSnafu)?;
let region = self.region.context(MissingRegionSnafu)?;
+ let clientbuilder = match self.proxy_url {
+ Some(ref url) => {
+ let pr: Proxy =
+ Proxy::all(url).map_err(|source| Error::ProxyUrl { source })?;
+ Client::builder().proxy(pr)
+ }
+ None => Client::builder(),
+ };
let credentials = match (self.access_key_id, self.secret_access_key, self.token) {
(Some(key_id), Some(secret_key), token) => {
info!("Using Static credential provider");
@@ -590,7 +608,7 @@ impl AmazonS3Builder {
let endpoint = format!("https://sts.{}.amazonaws.com", region);
// Disallow non-HTTPs requests
- let client = Client::builder().https_only(true).build().unwrap();
+ let client = clientbuilder.https_only(true).build().unwrap();
Box::new(WebIdentityProvider {
cache: Default::default(),
@@ -611,7 +629,7 @@ impl AmazonS3Builder {
info!("Using Instance credential provider");
// The instance metadata endpoint is access over HTTP
- let client = Client::builder().https_only(false).build().unwrap();
+ let client = clientbuilder.https_only(false).build().unwrap();
Box::new(InstanceCredentialProvider {
cache: Default::default(),
@@ -653,9 +671,10 @@ impl AmazonS3Builder {
credentials,
retry_config: self.retry_config,
allow_http: self.allow_http,
+ proxy_url: self.proxy_url,
};
- let client = Arc::new(S3Client::new(config));
+ let client = Arc::new(S3Client::new(config).unwrap());
Ok(AmazonS3 { client })
}
@@ -898,4 +917,35 @@ mod tests {
let err = integration.delete(&location).await.unwrap_err();
assert!(matches!(err, crate::Error::NotFound { .. }), "{}", err);
}
+
+ #[tokio::test]
+ async fn s3_test_proxy_url() {
+ let s3 = AmazonS3Builder::new()
+ .with_access_key_id("access_key_id")
+ .with_secret_access_key("secret_access_key")
+ .with_region("region")
+ .with_bucket_name("bucket_name")
+ .with_allow_http(true)
+ .with_proxy_url("https://example.com")
+ .build();
+
+ assert!(s3.is_ok());
+
+ let s3 = AmazonS3Builder::new()
+ .with_access_key_id("access_key_id")
+ .with_secret_access_key("secret_access_key")
+ .with_region("region")
+ .with_bucket_name("bucket_name")
+ .with_allow_http(true)
+ .with_proxy_url("asdf://example.com")
+ .build();
+
+ assert!(match s3 {
+ Err(crate::Error::Generic { source, .. }) => matches!(
+ source.downcast_ref(),
+ Some(crate::aws::Error::ProxyUrl { .. })
+ ),
+ _ => false,
+ })
+ }
}
diff --git a/object_store/src/azure/client.rs b/object_store/src/azure/client.rs
index ece07853a..d8cfdd1c7 100644
--- a/object_store/src/azure/client.rs
+++ b/object_store/src/azure/client.rs
@@ -27,7 +27,7 @@ use chrono::{DateTime, TimeZone, Utc};
use itertools::Itertools;
use reqwest::{
header::{HeaderValue, CONTENT_LENGTH, IF_NONE_MATCH, RANGE},
- Client as ReqwestClient, Method, Response, StatusCode,
+ Client as ReqwestClient, Method, Proxy, Response, StatusCode,
};
use serde::{Deserialize, Deserializer, Serialize};
use snafu::{ResultExt, Snafu};
@@ -82,6 +82,9 @@ pub(crate) enum Error {
Authorization {
source: crate::azure::credential::Error,
},
+
+ #[snafu(display("Unable to use proxy url: {}", source))]
+ ProxyUrl { source: reqwest::Error },
}
impl From<Error> for crate::Error {
@@ -124,6 +127,7 @@ pub struct AzureConfig {
pub allow_http: bool,
pub service: Url,
pub is_emulator: bool,
+ pub proxy_url: Option<String>,
}
impl AzureConfig {
@@ -148,13 +152,20 @@ pub(crate) struct AzureClient {
impl AzureClient {
/// create a new instance of [AzureClient]
- pub fn new(config: AzureConfig) -> Self {
- let client = reqwest::ClientBuilder::new()
- .https_only(!config.allow_http)
- .build()
- .unwrap();
+ pub fn new(config: AzureConfig) -> Result<Self> {
+ let builder = ReqwestClient::builder();
+
+ let client = if let Some(url) = config.proxy_url.as_ref() {
+ let pr = Proxy::all(url).map_err(|source| Error::ProxyUrl { source });
+ builder.proxy(pr.unwrap())
+ } else {
+ builder
+ }
+ .https_only(!config.allow_http)
+ .build()
+ .unwrap();
- Self { config, client }
+ Ok(Self { config, client })
}
/// Returns the config
diff --git a/object_store/src/azure/mod.rs b/object_store/src/azure/mod.rs
index f7ca4cf4e..060b4b2d2 100644
--- a/object_store/src/azure/mod.rs
+++ b/object_store/src/azure/mod.rs
@@ -360,6 +360,7 @@ pub struct MicrosoftAzureBuilder {
use_emulator: bool,
retry_config: RetryConfig,
allow_http: bool,
+ proxy_url: Option<String>,
}
impl Debug for MicrosoftAzureBuilder {
@@ -500,6 +501,12 @@ impl MicrosoftAzureBuilder {
self
}
+ /// Set the proxy_url to be used by the underlying client
+ pub fn with_proxy_url(mut self, proxy_url: impl Into<String>) -> Self {
+ self.proxy_url = Some(proxy_url.into());
+ self
+ }
+
/// Configure a connection to container with given name on Microsoft Azure
/// Blob store.
pub fn build(self) -> Result<MicrosoftAzure> {
@@ -516,6 +523,7 @@ impl MicrosoftAzureBuilder {
retry_config,
allow_http,
authority_host,
+ proxy_url,
} = self;
let container = container_name.ok_or(Error::MissingContainerName {})?;
@@ -567,9 +575,10 @@ impl MicrosoftAzureBuilder {
container,
credentials: auth,
is_emulator,
+ proxy_url,
};
- let client = Arc::new(client::AzureClient::new(config));
+ let client = Arc::new(client::AzureClient::new(config)?);
Ok(MicrosoftAzure { client })
}
diff --git a/object_store/src/gcp/mod.rs b/object_store/src/gcp/mod.rs
index 0ef4d3564..0da92fdbe 100644
--- a/object_store/src/gcp/mod.rs
+++ b/object_store/src/gcp/mod.rs
@@ -41,6 +41,7 @@ use chrono::{DateTime, Utc};
use futures::{stream::BoxStream, StreamExt, TryStreamExt};
use percent_encoding::{percent_encode, NON_ALPHANUMERIC};
use reqwest::header::RANGE;
+use reqwest::Proxy;
use reqwest::{header, Client, Method, Response, StatusCode};
use snafu::{ResultExt, Snafu};
use tokio::io::AsyncWrite;
@@ -122,6 +123,9 @@ enum Error {
#[snafu(display("GCP credential error: {}", source))]
Credential { source: credential::Error },
+
+ #[snafu(display("Unable to use proxy url: {}", source))]
+ ProxyUrl { source: reqwest::Error },
}
impl From<Error> for super::Error {
@@ -741,6 +745,7 @@ pub struct GoogleCloudStorageBuilder {
service_account_path: Option<String>,
client: Option<Client>,
retry_config: RetryConfig,
+ proxy_url: Option<String>,
}
impl GoogleCloudStorageBuilder {
@@ -782,6 +787,12 @@ impl GoogleCloudStorageBuilder {
self
}
+ /// Set proxy url used for connection
+ pub fn with_proxy_url(mut self, proxy_url: impl Into<String>) -> Self {
+ self.proxy_url = Some(proxy_url.into());
+ self
+ }
+
/// Configure a connection to Google Cloud Storage, returning a
/// new [`GoogleCloudStorage`] and consuming `self`
pub fn build(self) -> Result<GoogleCloudStorage> {
@@ -790,12 +801,24 @@ impl GoogleCloudStorageBuilder {
service_account_path,
client,
retry_config,
+ proxy_url,
} = self;
let bucket_name = bucket_name.ok_or(Error::MissingBucketName {})?;
let service_account_path =
service_account_path.ok_or(Error::MissingServiceAccountPath)?;
- let client = client.unwrap_or_else(Client::new);
+
+ let client = match (proxy_url, client) {
+ (_, Some(client)) => client,
+ (Some(url), None) => {
+ let pr = Proxy::all(&url).map_err(|source| Error::ProxyUrl { source })?;
+ Client::builder()
+ .proxy(pr)
+ .build()
+ .map_err(|source| Error::ProxyUrl { source })?
+ }
+ (None, None) => Client::new(),
+ };
let credentials = reader_credentials_file(service_account_path)?;
@@ -1015,4 +1038,34 @@ mod test {
err
)
}
+
+ #[tokio::test]
+ async fn gcs_test_proxy_url() {
+ use std::io::Write;
+ use tempfile::NamedTempFile;
+ let mut tfile = NamedTempFile::new().unwrap();
+ let creds = r#"{"private_key": "private_key", "client_email":"client_email", "disable_oauth":true}"#;
+ write!(tfile, "{}", creds).unwrap();
+ let service_account_path = tfile.path();
+ let gcs = GoogleCloudStorageBuilder::new()
+ .with_service_account_path(service_account_path.to_str().unwrap())
+ .with_bucket_name("foo")
+ .with_proxy_url("https://example.com")
+ .build();
+ assert!(dbg!(gcs).is_ok());
+
+ let gcs = GoogleCloudStorageBuilder::new()
+ .with_service_account_path(service_account_path.to_str().unwrap())
+ .with_bucket_name("foo")
+ .with_proxy_url("asdf://example.com")
+ .build();
+
+ assert!(match gcs {
+ Err(ObjectStoreError::Generic { source, .. }) => matches!(
+ source.downcast_ref(),
+ Some(crate::gcp::Error::ProxyUrl { .. })
+ ),
+ _ => false,
+ })
+ }
}