working s3 cloud support

This commit is contained in:
Jack Wright 2024-12-05 16:58:21 -08:00
parent 34a0025aea
commit 6d51d23681
3 changed files with 92 additions and 8 deletions

60
Cargo.lock generated
View File

@ -4262,6 +4262,7 @@ dependencies = [
"nu-protocol",
"nu-utils",
"num",
"object_store 0.11.1",
"polars",
"polars-arrow",
"polars-io",
@ -4547,7 +4548,37 @@ dependencies = [
"ring",
"serde",
"serde_json",
"snafu",
"snafu 0.7.5",
"tokio",
"tracing",
"url",
"walkdir",
]
[[package]]
name = "object_store"
version = "0.11.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6eb4c22c6154a1e759d7099f9ffad7cc5ef8245f9efbab4a41b92623079c82f3"
dependencies = [
"async-trait",
"base64 0.22.1",
"bytes",
"chrono",
"futures",
"humantime",
"hyper 1.5.0",
"itertools 0.13.0",
"md-5",
"parking_lot",
"percent-encoding",
"quick-xml 0.36.2",
"rand",
"reqwest",
"ring",
"serde",
"serde_json",
"snafu 0.8.5",
"tokio",
"tracing",
"url",
@ -5080,7 +5111,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4643898a644f30c83737db85f942f8c8956b0c11190b39afec745218eae1746b"
dependencies = [
"avro-schema",
"object_store",
"object_store 0.10.2",
"polars-arrow-format",
"regex",
"simdutf8",
@ -5134,7 +5165,7 @@ dependencies = [
"memchr",
"memmap2",
"num-traits",
"object_store",
"object_store 0.10.2",
"once_cell",
"percent-encoding",
"polars-arrow",
@ -6939,7 +6970,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e4de37ad025c587a29e8f3f5605c00f70b98715ef90b9061a815b9e59e9042d6"
dependencies = [
"doc-comment",
"snafu-derive",
"snafu-derive 0.7.5",
]
[[package]]
name = "snafu"
version = "0.8.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "223891c85e2a29c3fe8fb900c1fae5e69c2e42415e3177752e8718475efa5019"
dependencies = [
"snafu-derive 0.8.5",
]
[[package]]
@ -6954,6 +6994,18 @@ dependencies = [
"syn 1.0.109",
]
[[package]]
name = "snafu-derive"
version = "0.8.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "03c3c6b7927ffe7ecaa769ee0e3994da3b8cafc8f444578982c83ecb161af917"
dependencies = [
"heck 0.5.0",
"proc-macro2",
"quote",
"syn 2.0.87",
]
[[package]]
name = "snap"
version = "1.1.1"

View File

@ -31,7 +31,7 @@ mimalloc = { version = "0.1.42" }
num = {version = "0.4"}
serde = { version = "1.0", features = ["derive"] }
sqlparser = { version = "0.49"}
polars-io = { version = "0.44", features = ["avro"]}
polars-io = { version = "0.44", features = ["avro", "cloud", "aws"]}
polars-arrow = { version = "0.44"}
polars-ops = { version = "0.44", features = ["pivot"]}
polars-plan = { version = "0.44", features = ["regex"]}
@ -48,6 +48,7 @@ hashbrown = { version = "0.14", features = ["rayon", "ahash", "serde", "raw"] }
aws-config = { version = "1.5", features = ["sso"] }
aws-credential-types = "1.2"
tokio = { version = "1.41.1", features = ["full"] }
object_store = { version = "0.11.1", features = ["aws"] }
[dependencies.polars]
features = [

View File

@ -3,9 +3,12 @@ use std::error::Error;
use aws_config::{BehaviorVersion, SdkConfig};
use aws_credential_types::{provider::ProvideCredentials, Credentials};
use nu_protocol::ShellError;
use object_store::aws::AmazonS3ConfigKey;
use polars_io::cloud::CloudOptions;
async fn aws_load_config() -> SdkConfig {
use crate::PolarsPlugin;
async fn load_aws_config() -> SdkConfig {
aws_config::load_defaults(BehaviorVersion::latest()).await
}
@ -31,6 +34,34 @@ async fn aws_creds(aws_config: &SdkConfig) -> Result<Option<Credentials>, ShellE
}
}
async fn aws_cloud_options() -> CloudOptions {
CloudOptions::
async fn build_aws_cloud_configs() -> Result<Vec<(AmazonS3ConfigKey, String)>, ShellError> {
let sdk_config = load_aws_config().await;
let creds = aws_creds(&sdk_config)
.await?
.ok_or(ShellError::GenericError {
error: "Could not determine AWS credentials".into(),
msg: "".into(),
span: None,
help: None,
inner: vec![],
})?;
let mut configs = vec![
(AmazonS3ConfigKey::AccessKeyId, creds.access_key_id().into()),
(
AmazonS3ConfigKey::SecretAccessKey,
creds.secret_access_key().into(),
),
];
if let Some(token) = creds.session_token() {
configs.push((AmazonS3ConfigKey::Token, token.into()))
}
Ok(configs)
}
fn build_aws_cloud_options(plugin: &PolarsPlugin) -> Result<CloudOptions, ShellError> {
let configs = plugin.runtime.block_on(build_aws_cloud_configs())?;
Ok(CloudOptions::default().with_aws(configs.into_iter()))
}