From 6d51d2368164b692ea2f88a4da28f413c6da9eed Mon Sep 17 00:00:00 2001
From: Jack Wright <ayax79@gmail.com>
Date: Thu, 5 Dec 2024 16:58:21 -0800
Subject: [PATCH] working s3 cloud support

---
 Cargo.lock                               | 60 ++++++++++++++++++++++--
 crates/nu_plugin_polars/Cargo.toml       |  3 +-
 crates/nu_plugin_polars/src/cloud/aws.rs | 37 +++++++++++++--
 3 files changed, 92 insertions(+), 8 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 1a1bb78e85..3d206f3966 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -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"
diff --git a/crates/nu_plugin_polars/Cargo.toml b/crates/nu_plugin_polars/Cargo.toml
index 36dce3e975..aac4f458f4 100644
--- a/crates/nu_plugin_polars/Cargo.toml
+++ b/crates/nu_plugin_polars/Cargo.toml
@@ -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 = [
diff --git a/crates/nu_plugin_polars/src/cloud/aws.rs b/crates/nu_plugin_polars/src/cloud/aws.rs
index c9872d1173..55a3c82e38 100644
--- a/crates/nu_plugin_polars/src/cloud/aws.rs
+++ b/crates/nu_plugin_polars/src/cloud/aws.rs
@@ -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()))
 }