mirror of
https://github.com/windmill-labs/windmill.git
synced 2026-09-06 16:02:23 +00:00
8b80b09f33
* fix: restrict filesystem workspace storage to debug builds Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Q7p2VbtYqaXHGaAskgwVk5 * chore: update ee-repo-ref to b58ad414b098d3d7787001a352bfbb13e43a335f This commit updates the EE repository reference after PR #747 was merged in windmill-ee-private. Previous ee-repo-ref: 1b4dada77a8fe2224579c643550c63b1ac2616de New ee-repo-ref: b58ad414b098d3d7787001a352bfbb13e43a335f Automated by sync-ee-ref workflow. --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: windmill-internal-app[bot] <windmill-internal-app[bot]@users.noreply.github.com>
2472 lines
88 KiB
Rust
2472 lines
88 KiB
Rust
pub use windmill_types::s3::*;
|
|
|
|
use std::collections::HashMap;
|
|
|
|
use quick_cache::sync::Cache;
|
|
use windmill_common::error::{self};
|
|
|
|
#[cfg(feature = "parquet")]
|
|
use async_trait::async_trait;
|
|
#[cfg(feature = "parquet")]
|
|
use aws_config::{default_provider::credentials::DefaultCredentialsChain, Region};
|
|
#[cfg(feature = "parquet")]
|
|
use aws_sdk_sts::config::ProvideCredentials;
|
|
#[cfg(feature = "parquet")]
|
|
use bytes::Bytes;
|
|
#[cfg(feature = "parquet")]
|
|
use chrono::{DateTime, Utc};
|
|
#[cfg(feature = "parquet")]
|
|
use datafusion::arrow::array::{RecordBatch, RecordBatchWriter};
|
|
#[cfg(feature = "parquet")]
|
|
use datafusion::arrow::error::ArrowError;
|
|
#[cfg(feature = "parquet")]
|
|
use datafusion::arrow::json::writer::JsonArray;
|
|
#[cfg(feature = "parquet")]
|
|
use datafusion::arrow::{csv, json};
|
|
#[cfg(feature = "parquet")]
|
|
use datafusion::parquet::arrow::ArrowWriter;
|
|
#[cfg(feature = "parquet")]
|
|
use futures::TryStreamExt;
|
|
#[cfg(feature = "parquet")]
|
|
use object_store::aws::AwsCredential;
|
|
#[cfg(feature = "parquet")]
|
|
use object_store::azure::MicrosoftAzureBuilder;
|
|
#[cfg(feature = "parquet")]
|
|
use object_store::gcp::GoogleCloudStorageBuilder;
|
|
#[cfg(feature = "parquet")]
|
|
use object_store::CredentialProvider;
|
|
#[cfg(feature = "parquet")]
|
|
use object_store::ObjectStore;
|
|
#[cfg(feature = "parquet")]
|
|
use object_store::{aws::AmazonS3Builder, ClientOptions};
|
|
#[cfg(feature = "parquet")]
|
|
use reqwest::header::HeaderMap;
|
|
#[cfg(feature = "parquet")]
|
|
use std::io::Write;
|
|
#[cfg(feature = "parquet")]
|
|
use std::sync::{Arc, Mutex};
|
|
#[cfg(feature = "parquet")]
|
|
use tokio::sync::RwLock;
|
|
#[cfg(feature = "parquet")]
|
|
use tokio::task;
|
|
#[cfg(feature = "parquet")]
|
|
use windmill_common::error::to_anyhow;
|
|
#[cfg(feature = "parquet")]
|
|
use windmill_common::jobs::is_safe_log_file_path;
|
|
#[cfg(feature = "parquet")]
|
|
use windmill_common::utils::rd_string;
|
|
#[cfg(all(feature = "parquet", feature = "private"))]
|
|
pub mod job_s3_helpers_ee;
|
|
#[cfg(feature = "parquet")]
|
|
pub mod job_s3_helpers_oss;
|
|
|
|
// Re-export object_store types so consumers don't need a direct object_store dep
|
|
#[cfg(feature = "parquet")]
|
|
pub mod object_store_reexports {
|
|
pub use object_store::local::LocalFileSystem;
|
|
pub use object_store::memory::InMemory;
|
|
pub use object_store::path::Path;
|
|
pub use object_store::{
|
|
Attribute, Attributes, Error as ObjectStoreError, GetOptions, GetRange, GetResult,
|
|
ObjectStore, PutMultipartOpts, PutPayload, PutResult, Result as ObjectStoreResult,
|
|
WriteMultipart,
|
|
};
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
pub fn object_store_error_to_error(err: object_store::Error) -> error::Error {
|
|
use object_store::Error::*;
|
|
match err {
|
|
Generic { store, source } => error::Error::Generic(
|
|
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
|
|
format!("Generic {} error: {}", store, source),
|
|
),
|
|
NotFound { path, source } => error::Error::NotFound(format!("{}: {}", path, source)),
|
|
InvalidPath { source } => error::Error::BadRequest(format!("Invalid path: {}", source)),
|
|
JoinError { source } => error::Error::InternalErr(format!("Join error: {}", source)),
|
|
NotSupported { source } => {
|
|
error::Error::BadRequest(format!("Operation not supported: {}", source))
|
|
}
|
|
AlreadyExists { path, source } => {
|
|
error::Error::BadRequest(format!("Object at {} already exists: {}", path, source))
|
|
}
|
|
Precondition { path, source } => {
|
|
error::Error::BadRequest(format!("Precondition failed at {}: {}", path, source))
|
|
}
|
|
NotModified { path, source } => {
|
|
error::Error::ExecutionErr(format!("Not modified at {}: {}", path, source))
|
|
}
|
|
NotImplemented => error::Error::BadRequest("Operation not yet implemented.".to_string()),
|
|
PermissionDenied { path, source } => {
|
|
error::Error::PermissionDenied(format!("Permission denied at {}: {}", path, source))
|
|
}
|
|
Unauthenticated { path, source } => {
|
|
error::Error::NotAuthorized(format!("Unauthenticated for {}: {}", path, source))
|
|
}
|
|
UnknownConfigurationKey { store, key } => error::Error::BadConfig(format!(
|
|
"Invalid config key '{}' for store '{}'",
|
|
key, store
|
|
)),
|
|
_ => error::Error::InternalErr(format!("Object store error: {}", err)),
|
|
}
|
|
}
|
|
|
|
// --- Object store builder infrastructure (moved from windmill-common/src/s3_helpers.rs) ---
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[derive(Clone)]
|
|
pub struct ExpirableObjectStore {
|
|
pub store: Arc<dyn ObjectStore>,
|
|
pub refresh: Option<ObjectStoreRefresh>,
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[derive(Clone)]
|
|
pub struct ObjectStoreRefresh {
|
|
refresh: Option<DateTime<Utc>>,
|
|
settings: ObjectSettings,
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
impl ObjectStoreRefresh {
|
|
pub fn new(settings: ObjectSettings, refresh: Option<DateTime<Utc>>) -> Self {
|
|
Self { settings, refresh }
|
|
}
|
|
fn refresh_needed(&self) -> bool {
|
|
if let Some(refresh) = self.refresh {
|
|
if refresh < Utc::now() - chrono::Duration::minutes(1) {
|
|
return true;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
async fn refresh(&self) -> Option<ExpirableObjectStore> {
|
|
return build_object_store_from_settings(self.settings.clone(), None)
|
|
.await
|
|
.map_err(|e| {
|
|
tracing::error!("Error building s3 client from settings: {:?}", e);
|
|
e
|
|
})
|
|
.ok();
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
impl From<Arc<dyn ObjectStore>> for ExpirableObjectStore {
|
|
fn from(store: Arc<dyn ObjectStore>) -> Self {
|
|
Self { store, refresh: None }
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
lazy_static::lazy_static! {
|
|
pub static ref OBJECT_STORE_SETTINGS: Arc<RwLock<Option<ExpirableObjectStore>>> = Arc::new(RwLock::new(None));
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
pub async fn get_object_store() -> Option<Arc<dyn ObjectStore>> {
|
|
let settings = OBJECT_STORE_SETTINGS.read().await;
|
|
if let Some(s) = settings.as_ref() {
|
|
match &s.refresh {
|
|
Some(refresh) => {
|
|
if refresh.refresh_needed() {
|
|
let refresh = refresh.clone();
|
|
drop(settings);
|
|
let new_store = refresh.refresh().await;
|
|
if let Some(new_store) = new_store {
|
|
let mut s3_cache_settings = OBJECT_STORE_SETTINGS.write().await;
|
|
let arc = new_store.store.clone();
|
|
*s3_cache_settings = Some(new_store);
|
|
return Some(arc);
|
|
} else {
|
|
return None;
|
|
}
|
|
} else {
|
|
return Some(s.store.clone());
|
|
}
|
|
}
|
|
None => {
|
|
return Some(s.store.clone());
|
|
}
|
|
}
|
|
} else {
|
|
return None;
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
pub enum ObjectStoreReload {
|
|
Later,
|
|
Never,
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
pub async fn reload_object_store_setting(db: &windmill_common::DB) -> ObjectStoreReload {
|
|
use windmill_common::{
|
|
ee_oss::{get_license_plan, LicensePlan},
|
|
global_settings::{load_value_from_global_settings, OBJECT_STORE_CONFIG_SETTING},
|
|
};
|
|
|
|
let s3_config = load_value_from_global_settings(db, OBJECT_STORE_CONFIG_SETTING).await;
|
|
if let Err(e) = s3_config {
|
|
tracing::error!("Error reloading s3 cache config: {:?}", e)
|
|
} else {
|
|
if let Some(v) = s3_config.unwrap() {
|
|
if matches!(get_license_plan().await, LicensePlan::Pro) {
|
|
tracing::error!("S3 cache is not available for pro plan");
|
|
return ObjectStoreReload::Never;
|
|
}
|
|
let setting = serde_json::from_value::<ObjectSettings>(v);
|
|
match setting {
|
|
Ok(setting) => {
|
|
let is_oidc = matches!(setting, ObjectSettings::AwsOidc(_));
|
|
let s3_client = build_object_store_from_settings(setting, Some(db)).await;
|
|
match s3_client {
|
|
Ok(s3_client) => {
|
|
let mut s3_cache_settings = OBJECT_STORE_SETTINGS.write().await;
|
|
*s3_cache_settings = Some(s3_client);
|
|
}
|
|
Err(e) => {
|
|
if is_oidc {
|
|
tracing::error!("Error building s3 client from oidc settings. It may be due to the jwks endpoints not being up yet, it will be attempted again in 10s to leave time for the server to be ready: {:?}", e);
|
|
return ObjectStoreReload::Later;
|
|
} else {
|
|
tracing::error!("Error building s3 client from settings: {:?}", e);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Err(e) => {
|
|
tracing::error!("Error parsing s3 cache config: {:?}", e)
|
|
}
|
|
}
|
|
} else {
|
|
let mut s3_cache_settings = OBJECT_STORE_SETTINGS.write().await;
|
|
if std::env::var("S3_CACHE_BUCKET").is_ok() {
|
|
if matches!(get_license_plan().await, LicensePlan::Pro) {
|
|
tracing::error!("S3 cache is not available for pro plan");
|
|
return ObjectStoreReload::Never;
|
|
}
|
|
*s3_cache_settings = build_s3_client_from_settings(S3Settings {
|
|
bucket: None,
|
|
region: None,
|
|
access_key: None,
|
|
secret_key: None,
|
|
endpoint: None,
|
|
store_logs: None,
|
|
path_style: None,
|
|
allow_http: None,
|
|
port: None,
|
|
})
|
|
.await
|
|
.ok()
|
|
.map(|x| ExpirableObjectStore::from(x))
|
|
} else {
|
|
*s3_cache_settings = None;
|
|
}
|
|
}
|
|
}
|
|
return ObjectStoreReload::Never;
|
|
}
|
|
|
|
pub fn render_endpoint(
|
|
raw_endpoint: String,
|
|
use_ssl: bool,
|
|
port: Option<u16>,
|
|
path_style: Option<bool>,
|
|
bucket: String,
|
|
) -> String {
|
|
let url_with_prefix =
|
|
if raw_endpoint.starts_with("http://") || raw_endpoint.starts_with("https://") {
|
|
raw_endpoint.clone()
|
|
} else {
|
|
let scheme = if use_ssl { "https" } else { "http" };
|
|
format!(
|
|
"{}://{}",
|
|
scheme,
|
|
if path_style.unwrap_or(true) {
|
|
raw_endpoint
|
|
} else {
|
|
format!("{}.{}", bucket, raw_endpoint)
|
|
}
|
|
)
|
|
};
|
|
if port.is_some() {
|
|
format!("{}:{}", url_with_prefix, port.unwrap())
|
|
} else {
|
|
url_with_prefix
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
pub async fn build_object_store_client(
|
|
resource_ref: &ObjectStoreResource,
|
|
) -> error::Result<Arc<dyn ObjectStore>> {
|
|
match resource_ref {
|
|
ObjectStoreResource::S3(s3_resource_ref) => build_s3_client(&s3_resource_ref).await,
|
|
ObjectStoreResource::Azure(azure_blob_resource_ref) => {
|
|
build_azure_blob_client(&azure_blob_resource_ref)
|
|
}
|
|
ObjectStoreResource::Gcs(gcs_resource_ref) => build_gcs_client(&gcs_resource_ref).await,
|
|
ObjectStoreResource::Filesystem(fs) => build_filesystem_client(&fs.root_path),
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
pub async fn attempt_fetch_bytes(
|
|
client: Arc<dyn ObjectStore>,
|
|
path: &str,
|
|
) -> error::Result<bytes::Bytes> {
|
|
use object_store::path::Path;
|
|
|
|
let object = client.get(&Path::from(path)).await;
|
|
if let Err(e) = object {
|
|
tracing::info!(
|
|
"Failed to pull bytes from object store at path {path}. Error: {:?}",
|
|
e
|
|
);
|
|
return Err(error::Error::ExecutionErr(format!(
|
|
"Failed to pull bytes from object store: {path}"
|
|
)));
|
|
}
|
|
|
|
let bytes = object.unwrap().bytes().await;
|
|
if bytes.is_err() {
|
|
tracing::info!(
|
|
"Failed to read bytes from object store: {path}. Error: {:?}",
|
|
bytes.err()
|
|
);
|
|
return Err(error::Error::ExecutionErr(format!(
|
|
"Failed to read bytes from object store: {path}"
|
|
)));
|
|
}
|
|
let bytes = bytes.unwrap();
|
|
|
|
tracing::info!("{path} len: {}", bytes.len());
|
|
|
|
if bytes.len() == 0 {
|
|
tracing::info!("object {path} not found in bucket, bytes empty",);
|
|
return Err(error::Error::ExecutionErr(format!(
|
|
"object {path} does not exist in bucket"
|
|
)));
|
|
}
|
|
|
|
return Ok(bytes);
|
|
}
|
|
|
|
/// Whether an S3 resource carries static credentials. When it does not, the
|
|
/// ambient AWS chain (env, profile, ECS/EC2 instance role) is used instead.
|
|
/// Shared so callers that sign requests by hand resolve credentials on exactly the
|
|
/// same condition as `build_s3_client`.
|
|
#[cfg(feature = "parquet")]
|
|
pub fn s3_resource_has_static_credentials(s3_resource: &S3Resource) -> bool {
|
|
s3_resource.access_key.as_ref().is_some_and(|x| x != "")
|
|
|| s3_resource.secret_key.as_ref().is_some_and(|x| x != "")
|
|
}
|
|
|
|
/// Ambient AWS credentials from the shared, cached provider backing
|
|
/// `build_s3_client`. Callers that sign their own requests must go through this
|
|
/// rather than resolving the default chain themselves: the cache is what keeps a
|
|
/// burst of requests from hitting the instance metadata service once each.
|
|
///
|
|
/// These are the **instance's own** credentials, not any caller's, and they are
|
|
/// returned in the clear. A caller therefore MUST:
|
|
/// - authorize the request target itself — reaching this function implies no
|
|
/// permission check, and the credentials typically outrank the requesting user;
|
|
/// - use them only to sign a request it has already authorized, never surface them
|
|
/// in a response, log, or error message, and never hand them to a caller-supplied
|
|
/// endpoint.
|
|
///
|
|
/// Prefer `build_s3_client`, which confines them to the object-store client; reach
|
|
/// for this only where a request must be signed by hand.
|
|
#[cfg(feature = "parquet")]
|
|
pub async fn ambient_aws_credentials(
|
|
region: &str,
|
|
) -> anyhow::Result<aws_sdk_sts::config::Credentials> {
|
|
ambient_aws_credentials_provider(region).await.get().await
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
pub async fn build_s3_client(s3_resource_ref: &S3Resource) -> error::Result<Arc<dyn ObjectStore>> {
|
|
let static_creds = s3_resource_has_static_credentials(s3_resource_ref);
|
|
|
|
let credentials_provider = if !static_creds {
|
|
Some(ambient_aws_credentials_provider(&s3_resource_ref.region).await)
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let s3_resource = s3_resource_ref.clone();
|
|
let endpoint = render_endpoint(
|
|
s3_resource.endpoint_with_region_fallback(None),
|
|
s3_resource.use_ssl,
|
|
s3_resource.port,
|
|
s3_resource.path_style,
|
|
s3_resource.bucket.clone(),
|
|
);
|
|
let mut store_builder = AmazonS3Builder::new()
|
|
.with_client_options(
|
|
ClientOptions::new()
|
|
.with_timeout_disabled()
|
|
.with_default_headers(HeaderMap::from_iter(vec![(
|
|
"Accept-Encoding".parse().unwrap(),
|
|
"".parse().unwrap(),
|
|
)])),
|
|
)
|
|
.with_region(s3_resource.region)
|
|
.with_bucket_name(s3_resource.bucket)
|
|
.with_endpoint(endpoint);
|
|
|
|
if let Some(credentials_provider) = credentials_provider {
|
|
store_builder = store_builder.with_credentials(Arc::new(AwsCredentialAdapter {
|
|
inner: credentials_provider,
|
|
}));
|
|
}
|
|
|
|
if !s3_resource.use_ssl {
|
|
store_builder = store_builder.with_allow_http(true)
|
|
}
|
|
|
|
if let Some(key) = s3_resource.access_key {
|
|
if key != "" {
|
|
store_builder = store_builder.with_access_key_id(key);
|
|
}
|
|
}
|
|
|
|
if let Some(token) = s3_resource.token {
|
|
if token != "" {
|
|
store_builder = store_builder.with_token(token);
|
|
}
|
|
}
|
|
if let Some(secret_key) = s3_resource.secret_key {
|
|
if secret_key != "" {
|
|
store_builder = store_builder.with_secret_access_key(secret_key);
|
|
}
|
|
}
|
|
if !s3_resource.path_style.unwrap_or(true) {
|
|
store_builder = store_builder.with_virtual_hosted_style_request(true);
|
|
}
|
|
|
|
let store = store_builder.build().map_err(|err| {
|
|
tracing::error!("Error building object store client: {:?}", err);
|
|
error::Error::internal_err(format!("Error building object store client: {:?}", err))
|
|
})?;
|
|
|
|
return Ok(Arc::new(store));
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
fn build_azure_blob_client(
|
|
azure_blob_resource_ref: &AzureBlobResource,
|
|
) -> error::Result<Arc<dyn ObjectStore>> {
|
|
let blob_resource = azure_blob_resource_ref.clone();
|
|
|
|
let mut store_builder = MicrosoftAzureBuilder::new()
|
|
.with_client_options(
|
|
ClientOptions::new()
|
|
.with_timeout_disabled()
|
|
.with_default_headers(HeaderMap::from_iter(vec![(
|
|
"Accept-Encoding".parse().unwrap(),
|
|
"".parse().unwrap(),
|
|
)])),
|
|
)
|
|
.with_account(blob_resource.account_name)
|
|
.with_container_name(blob_resource.container_name);
|
|
|
|
if let Some(federated_token_file) = blob_resource.federated_token_file {
|
|
if federated_token_file != "" {
|
|
store_builder = store_builder.with_federated_token_file(federated_token_file);
|
|
}
|
|
}
|
|
if let Some(tenant_id) = blob_resource.tenant_id {
|
|
if tenant_id != "" {
|
|
store_builder = store_builder.with_tenant_id(tenant_id);
|
|
}
|
|
}
|
|
if let Some(client_id) = blob_resource.client_id {
|
|
if client_id != "" {
|
|
store_builder = store_builder.with_client_id(client_id);
|
|
}
|
|
}
|
|
if let Some(endpoint) = blob_resource.endpoint {
|
|
if endpoint != "" {
|
|
let endpoint = render_endpoint(
|
|
endpoint,
|
|
blob_resource.use_ssl.unwrap_or(false),
|
|
None,
|
|
None,
|
|
"".to_string(),
|
|
);
|
|
store_builder = store_builder.with_endpoint(endpoint)
|
|
}
|
|
}
|
|
|
|
if !blob_resource.use_ssl.unwrap_or(false) {
|
|
store_builder = store_builder.with_allow_http(true)
|
|
}
|
|
|
|
if let Some(key) = blob_resource.access_key {
|
|
if key != "" {
|
|
store_builder = store_builder.with_access_key(key);
|
|
}
|
|
}
|
|
|
|
let store = store_builder.build().map_err(|err| {
|
|
tracing::error!("Error building object store client: {:?}", err);
|
|
error::Error::internal_err(format!("Error building object store client: {:?}", err))
|
|
})?;
|
|
|
|
return Ok(Arc::new(store));
|
|
}
|
|
|
|
/// Whether a GCS `service_account_key` carries no static credentials, in which case the client
|
|
/// should fall back to the instance's ambient credentials (GKE Workload Identity / metadata server)
|
|
/// instead of being handed an unparseable key. Besides an empty/whitespace string, the settings UI
|
|
/// stores "no key" as an empty JSON object `{}` (and `serde_json` may yield `null`), so treat those
|
|
/// as absent too. Shared with the connectivity-test SSRF guard so both agree on what "no key" means.
|
|
pub fn gcs_service_account_key_is_blank(service_account_key: &str) -> bool {
|
|
let trimmed = service_account_key.trim();
|
|
if trimmed.is_empty() {
|
|
return true;
|
|
}
|
|
match serde_json::from_str::<serde_json::Value>(trimmed) {
|
|
Ok(serde_json::Value::Null) => true,
|
|
Ok(serde_json::Value::Object(map)) => map.is_empty(),
|
|
_ => false,
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
async fn build_gcs_client(gcs_resource_ref: &GcsResource) -> error::Result<Arc<dyn ObjectStore>> {
|
|
let gcs_resource = gcs_resource_ref.clone();
|
|
|
|
let mut store_builder = GoogleCloudStorageBuilder::new()
|
|
.with_client_options(
|
|
ClientOptions::new()
|
|
.with_timeout_disabled()
|
|
.with_default_headers(HeaderMap::from_iter(vec![(
|
|
"Accept-Encoding".parse().unwrap(),
|
|
"".parse().unwrap(),
|
|
)])),
|
|
)
|
|
.with_bucket_name(gcs_resource.bucket);
|
|
|
|
// A blank key means no static credentials: let the builder fall back to the metadata server
|
|
// (InstanceCredentialProvider) so GKE Workload Identity / ambient credentials work. Passing a
|
|
// blank/`{}` key to `with_service_account_key` would instead fail to parse.
|
|
if !gcs_service_account_key_is_blank(&gcs_resource.service_account_key) {
|
|
store_builder = store_builder.with_service_account_key(gcs_resource.service_account_key);
|
|
}
|
|
|
|
let store = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| store_builder.build()))
|
|
.map_err(|panic_info| {
|
|
tracing::error!(
|
|
"Panic while building GCS object store client: {:?}",
|
|
panic_info
|
|
);
|
|
error::Error::internal_err(format!(
|
|
"Panic while building GCS object store client: {:?}",
|
|
panic_info
|
|
))
|
|
})?
|
|
.map_err(|err| {
|
|
tracing::error!("Error building GCS object store client: {:?}", err);
|
|
error::Error::internal_err(format!("Error building GCS object store client: {:?}", err))
|
|
})?;
|
|
|
|
return Ok(Arc::new(store));
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
pub fn build_filesystem_client(root_path: &str) -> error::Result<Arc<dyn ObjectStore>> {
|
|
let store = object_store::local::LocalFileSystem::new_with_prefix(root_path).map_err(|e| {
|
|
error::Error::internal_err(format!("Error building filesystem object store: {:?}", e))
|
|
})?;
|
|
Ok(Arc::new(FilesystemStoreIgnoringAttributes(store)))
|
|
}
|
|
|
|
/// `LocalFileSystem` rejects put/multipart uploads whose options carry
|
|
/// attributes (content-type, content-disposition, ...) with `NotImplemented`.
|
|
/// Attributes are advisory metadata a plain filesystem cannot persist, so
|
|
/// drop them instead of failing the upload.
|
|
#[cfg(feature = "parquet")]
|
|
#[derive(Debug)]
|
|
struct FilesystemStoreIgnoringAttributes(object_store::local::LocalFileSystem);
|
|
|
|
#[cfg(feature = "parquet")]
|
|
impl std::fmt::Display for FilesystemStoreIgnoringAttributes {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
self.0.fmt(f)
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[async_trait]
|
|
impl ObjectStore for FilesystemStoreIgnoringAttributes {
|
|
async fn put_opts(
|
|
&self,
|
|
location: &object_store::path::Path,
|
|
payload: object_store::PutPayload,
|
|
mut opts: object_store::PutOptions,
|
|
) -> object_store::Result<object_store::PutResult> {
|
|
opts.attributes = Default::default();
|
|
self.0.put_opts(location, payload, opts).await
|
|
}
|
|
|
|
async fn put_multipart_opts(
|
|
&self,
|
|
location: &object_store::path::Path,
|
|
mut opts: object_store::PutMultipartOpts,
|
|
) -> object_store::Result<Box<dyn object_store::MultipartUpload>> {
|
|
opts.attributes = Default::default();
|
|
self.0.put_multipart_opts(location, opts).await
|
|
}
|
|
|
|
async fn get_opts(
|
|
&self,
|
|
location: &object_store::path::Path,
|
|
options: object_store::GetOptions,
|
|
) -> object_store::Result<object_store::GetResult> {
|
|
self.0.get_opts(location, options).await
|
|
}
|
|
|
|
async fn get_range(
|
|
&self,
|
|
location: &object_store::path::Path,
|
|
range: std::ops::Range<u64>,
|
|
) -> object_store::Result<Bytes> {
|
|
self.0.get_range(location, range).await
|
|
}
|
|
|
|
async fn get_ranges(
|
|
&self,
|
|
location: &object_store::path::Path,
|
|
ranges: &[std::ops::Range<u64>],
|
|
) -> object_store::Result<Vec<Bytes>> {
|
|
self.0.get_ranges(location, ranges).await
|
|
}
|
|
|
|
async fn head(
|
|
&self,
|
|
location: &object_store::path::Path,
|
|
) -> object_store::Result<object_store::ObjectMeta> {
|
|
self.0.head(location).await
|
|
}
|
|
|
|
async fn delete(&self, location: &object_store::path::Path) -> object_store::Result<()> {
|
|
self.0.delete(location).await
|
|
}
|
|
|
|
fn list(
|
|
&self,
|
|
prefix: Option<&object_store::path::Path>,
|
|
) -> futures::stream::BoxStream<'static, object_store::Result<object_store::ObjectMeta>> {
|
|
self.0.list(prefix)
|
|
}
|
|
|
|
fn list_with_offset(
|
|
&self,
|
|
prefix: Option<&object_store::path::Path>,
|
|
offset: &object_store::path::Path,
|
|
) -> futures::stream::BoxStream<'static, object_store::Result<object_store::ObjectMeta>> {
|
|
self.0.list_with_offset(prefix, offset)
|
|
}
|
|
|
|
async fn list_with_delimiter(
|
|
&self,
|
|
prefix: Option<&object_store::path::Path>,
|
|
) -> object_store::Result<object_store::ListResult> {
|
|
self.0.list_with_delimiter(prefix).await
|
|
}
|
|
|
|
async fn copy(
|
|
&self,
|
|
from: &object_store::path::Path,
|
|
to: &object_store::path::Path,
|
|
) -> object_store::Result<()> {
|
|
self.0.copy(from, to).await
|
|
}
|
|
|
|
async fn rename(
|
|
&self,
|
|
from: &object_store::path::Path,
|
|
to: &object_store::path::Path,
|
|
) -> object_store::Result<()> {
|
|
self.0.rename(from, to).await
|
|
}
|
|
|
|
async fn copy_if_not_exists(
|
|
&self,
|
|
from: &object_store::path::Path,
|
|
to: &object_store::path::Path,
|
|
) -> object_store::Result<()> {
|
|
self.0.copy_if_not_exists(from, to).await
|
|
}
|
|
|
|
async fn rename_if_not_exists(
|
|
&self,
|
|
from: &object_store::path::Path,
|
|
to: &object_store::path::Path,
|
|
) -> object_store::Result<()> {
|
|
self.0.rename_if_not_exists(from, to).await
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
pub async fn build_object_store_from_settings(
|
|
settings: ObjectSettings,
|
|
init_private_key: Option<&windmill_common::DB>,
|
|
) -> error::Result<ExpirableObjectStore> {
|
|
match settings {
|
|
ObjectSettings::S3(s3_settings) => build_s3_client_from_settings(s3_settings)
|
|
.await
|
|
.map(|x| ExpirableObjectStore::from(x)),
|
|
ObjectSettings::Azure(azure_settings) => {
|
|
let azure_blob_resource = azure_settings;
|
|
build_azure_blob_client(&azure_blob_resource).map(|x| ExpirableObjectStore::from(x))
|
|
}
|
|
ObjectSettings::AwsOidc(ref s3_aws_oidc_settings) => {
|
|
let token_generator = crate::job_s3_helpers_oss::TokenGenerator::AsServerInstance();
|
|
let res = crate::job_s3_helpers_oss::generate_s3_aws_oidc_resource(
|
|
s3_aws_oidc_settings.clone(),
|
|
token_generator,
|
|
init_private_key,
|
|
)
|
|
.await?;
|
|
|
|
build_object_store_client(&res)
|
|
.await
|
|
.map(|x| ExpirableObjectStore {
|
|
store: x,
|
|
refresh: Some(ObjectStoreRefresh::new(settings.clone(), res.expiration())),
|
|
})
|
|
}
|
|
ObjectSettings::Gcs(gcs_settings) => {
|
|
let gcs_resource = gcs_settings;
|
|
build_gcs_client(&gcs_resource)
|
|
.await
|
|
.map(|x| ExpirableObjectStore::from(x))
|
|
}
|
|
ObjectSettings::Filesystem(fs) => {
|
|
build_filesystem_client(&fs.root_path).map(|x| ExpirableObjectStore::from(x))
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
fn none_if_empty(s: Option<String>) -> Option<String> {
|
|
if s.is_none() || s.as_ref().unwrap().is_empty() {
|
|
None
|
|
} else {
|
|
s
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
pub async fn build_s3_client_from_settings(
|
|
settings: S3Settings,
|
|
) -> error::Result<Arc<dyn ObjectStore>> {
|
|
let region = none_if_empty(settings.region)
|
|
.unwrap_or_else(|| std::env::var("AWS_REGION").unwrap_or_else(|_| "us-east-1".to_string()));
|
|
|
|
let s3_resource = S3Resource {
|
|
endpoint: none_if_empty(settings.endpoint).unwrap_or_else(|| {
|
|
std::env::var("S3_ENDPOINT").unwrap_or_else(|_| format!("s3.{region}.amazonaws.com"))
|
|
}),
|
|
bucket: settings.bucket.clone().unwrap_or_else(|| {
|
|
std::env::var("S3_CACHE_BUCKET").unwrap_or_else(|_| "missingbucket".to_string())
|
|
}),
|
|
region,
|
|
access_key: settings.access_key,
|
|
secret_key: settings.secret_key,
|
|
use_ssl: !settings.allow_http.unwrap_or(true),
|
|
path_style: settings.path_style,
|
|
port: settings.port,
|
|
token: None,
|
|
expiration: None,
|
|
};
|
|
|
|
build_s3_client(&s3_resource).await
|
|
}
|
|
|
|
// Resolving the default chain goes over the network (ECS/IMDS) on instances relying on an
|
|
// instance role, and object_store asks its CredentialProvider on every request — so resolved
|
|
// credentials must be cached and only re-fetched when close to expiring.
|
|
#[cfg(feature = "parquet")]
|
|
#[derive(Debug)]
|
|
struct AmbientAwsCredentials {
|
|
chain: DefaultCredentialsChain,
|
|
cached: RwLock<Option<(aws_sdk_sts::config::Credentials, std::time::Instant)>>,
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
impl AmbientAwsCredentials {
|
|
// Credentials without an expiry (env vars, static profile) are still re-resolved
|
|
// periodically so runtime changes to the environment are eventually picked up.
|
|
const NO_EXPIRY_TTL: std::time::Duration = std::time::Duration::from_secs(300);
|
|
const EXPIRY_MARGIN: std::time::Duration = std::time::Duration::from_secs(120);
|
|
|
|
fn still_valid(creds: &aws_sdk_sts::config::Credentials, age: std::time::Duration) -> bool {
|
|
match creds.expiry() {
|
|
Some(expiry) => std::time::SystemTime::now() + Self::EXPIRY_MARGIN < expiry,
|
|
None => age < Self::NO_EXPIRY_TTL,
|
|
}
|
|
}
|
|
|
|
async fn get(&self) -> anyhow::Result<aws_sdk_sts::config::Credentials> {
|
|
if let Some((creds, fetched_at)) = self.cached.read().await.as_ref() {
|
|
if Self::still_valid(creds, fetched_at.elapsed()) {
|
|
return Ok(creds.clone());
|
|
}
|
|
}
|
|
// The write lock is held across the chain resolution so concurrent requests don't all
|
|
// hit the metadata service at once.
|
|
let mut guard = self.cached.write().await;
|
|
if let Some((creds, fetched_at)) = guard.as_ref() {
|
|
if Self::still_valid(creds, fetched_at.elapsed()) {
|
|
return Ok(creds.clone());
|
|
}
|
|
}
|
|
let creds = self.chain.provide_credentials().await.map_err(|e| {
|
|
anyhow::anyhow!(
|
|
"no S3 access key/secret key is configured and no ambient AWS credentials could \
|
|
be loaded through the AWS SDK default chain (env vars, profile, ECS/EC2 instance \
|
|
role): {cause}. If an EC2/ECS instance role is expected to be used, the instance \
|
|
metadata service must be reachable from the process running Windmill — on EC2 the \
|
|
AWS Rust SDK only supports IMDSv2, so when Windmill runs in a Docker container \
|
|
the instance metadata hop limit (HttpPutResponseHopLimit) must be at least 2",
|
|
cause = format!("{:#}", anyhow::Error::new(e))
|
|
)
|
|
})?;
|
|
*guard = Some((creds.clone(), std::time::Instant::now()));
|
|
Ok(creds)
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
lazy_static::lazy_static! {
|
|
static ref AMBIENT_AWS_CREDS_PROVIDERS: Cache<String, Arc<AmbientAwsCredentials>> =
|
|
Cache::new(20);
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
async fn ambient_aws_credentials_provider(region: &str) -> Arc<AmbientAwsCredentials> {
|
|
// Single-flight: concurrent cold misses for the same region must share one provider,
|
|
// otherwise each gets its own instance and their per-instance refresh locks can't serialize
|
|
// the initial credential resolution — every caller would hit the metadata service.
|
|
match AMBIENT_AWS_CREDS_PROVIDERS
|
|
.get_value_or_guard_async(region)
|
|
.await
|
|
{
|
|
Ok(provider) => provider,
|
|
Err(guard) => {
|
|
let chain = DefaultCredentialsChain::builder()
|
|
.region(Region::new(region.to_string()))
|
|
.build()
|
|
.await;
|
|
let provider = Arc::new(AmbientAwsCredentials { chain, cached: RwLock::new(None) });
|
|
let _ = guard.insert(provider.clone());
|
|
provider
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[derive(Debug)]
|
|
struct AwsCredentialAdapter {
|
|
pub inner: Arc<AmbientAwsCredentials>,
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[async_trait]
|
|
impl CredentialProvider for AwsCredentialAdapter {
|
|
type Credential = AwsCredential;
|
|
async fn get_credential(&self) -> object_store::Result<Arc<Self::Credential>> {
|
|
let creds = self.inner.get().await.map_err(|e| {
|
|
tracing::error!("Error getting AWS credentials: {e:#}");
|
|
object_store::Error::Generic { store: "AWS", source: e.into() }
|
|
})?;
|
|
Ok(Arc::new(Self::Credential {
|
|
key_id: creds.access_key_id().to_string(),
|
|
secret_key: creds.secret_access_key().to_string(),
|
|
token: creds.session_token().map(|s| s.to_string()),
|
|
}))
|
|
}
|
|
}
|
|
|
|
// --- End moved code ---
|
|
|
|
lazy_static::lazy_static! {
|
|
static ref S3_BUCKET_RESTRICTIONS: Option<HashMap<String, Vec<String>>> = {
|
|
parse_bucket_restrictions()
|
|
};
|
|
static ref AZ_ACCOUNT_NAME_RESTRICTIONS: Option<HashMap<String, Vec<String>>> = {
|
|
parse_az_account_name_restrictions()
|
|
};
|
|
}
|
|
|
|
fn parse_bucket_restrictions() -> Option<HashMap<String, Vec<String>>> {
|
|
let env_var = std::env::var("S3_BUCKETS_WORKSPACE_RESTRICTIONS").ok()?;
|
|
parse_restrictions_from_str(&env_var, "S3 bucket")
|
|
}
|
|
|
|
fn parse_az_account_name_restrictions() -> Option<HashMap<String, Vec<String>>> {
|
|
let env_var = std::env::var("AZ_ACCOUNT_NAME_WORKSPACE_RESTRICTIONS").ok()?;
|
|
parse_restrictions_from_str(&env_var, "Azure account name")
|
|
}
|
|
|
|
fn parse_restrictions_from_str(input: &str, label: &str) -> Option<HashMap<String, Vec<String>>> {
|
|
if input.trim().is_empty() {
|
|
return None;
|
|
}
|
|
|
|
let mut restrictions = HashMap::new();
|
|
|
|
for rule in input.split(';') {
|
|
let rule = rule.trim();
|
|
if rule.is_empty() {
|
|
continue;
|
|
}
|
|
|
|
let parts: Vec<&str> = rule.splitn(2, ':').collect();
|
|
if parts.len() != 2 {
|
|
tracing::warn!(
|
|
"Invalid {} restriction format: '{}'. Expected 'name:workspace1,workspace2'",
|
|
label,
|
|
rule
|
|
);
|
|
continue;
|
|
}
|
|
|
|
let name = parts[0].trim().to_string();
|
|
let workspaces: Vec<String> = parts[1]
|
|
.split(',')
|
|
.map(|w| w.trim().to_string())
|
|
.filter(|w| !w.is_empty())
|
|
.collect();
|
|
|
|
if workspaces.is_empty() {
|
|
tracing::warn!(
|
|
"No workspaces specified for {} '{}', skipping restriction",
|
|
label,
|
|
name
|
|
);
|
|
continue;
|
|
}
|
|
|
|
restrictions.insert(name, workspaces);
|
|
}
|
|
|
|
if restrictions.is_empty() {
|
|
None
|
|
} else {
|
|
tracing::info!(
|
|
"{} restrictions loaded for {} entries",
|
|
label,
|
|
restrictions.len()
|
|
);
|
|
Some(restrictions)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn parse_bucket_restrictions_from_str(input: &str) -> Option<HashMap<String, Vec<String>>> {
|
|
parse_restrictions_from_str(input, "S3 bucket")
|
|
}
|
|
|
|
pub fn check_bucket_workspace_restriction(
|
|
bucket_name: &str,
|
|
workspace_id: &str,
|
|
) -> error::Result<()> {
|
|
if let Some(ref restrictions) = *S3_BUCKET_RESTRICTIONS {
|
|
if let Some(allowed_workspaces) = restrictions.get(bucket_name) {
|
|
if !allowed_workspaces.contains(&workspace_id.to_string()) {
|
|
return Err(error::Error::NotAuthorized(format!(
|
|
"Workspace '{}' is not authorized to access bucket '{}'",
|
|
workspace_id, bucket_name
|
|
)));
|
|
}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub fn check_az_account_name_workspace_restriction(
|
|
account_name: &str,
|
|
workspace_id: &str,
|
|
) -> error::Result<()> {
|
|
if let Some(ref restrictions) = *AZ_ACCOUNT_NAME_RESTRICTIONS {
|
|
if let Some(allowed_workspaces) = restrictions.get(account_name) {
|
|
if !allowed_workspaces.contains(&workspace_id.to_string()) {
|
|
return Err(error::Error::NotAuthorized(format!(
|
|
"Workspace '{}' is not authorized to access Azure account '{}'",
|
|
workspace_id, account_name
|
|
)));
|
|
}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// The `&storage=` fragment of a presigned s3 signature's HMAC message. `_default_` and an unset
|
|
/// storage name the same storage and must fold to the same fragment: a signature minted for one
|
|
/// is redeemed through a URL carrying the other, and a disagreement between signer and validator
|
|
/// surfaces only as `Invalid signature`. Both must build the message through this.
|
|
pub fn s3_signature_storage_fragment(storage: Option<&str>) -> String {
|
|
match storage.filter(|s| *s != DEFAULT_STORAGE) {
|
|
Some(name) => format!("&storage={name}"),
|
|
None => String::new(),
|
|
}
|
|
}
|
|
|
|
/// Error for a workspace file-storage lookup that resolved to nothing. Naming the requested
|
|
/// storage separates the two causes — a name this workspace has no storage for, versus a
|
|
/// workspace with no storage configured at all — which the reader cannot otherwise tell apart.
|
|
/// The wording stays neutral about where the name came from: callers pass an s3 object's
|
|
/// `storage`, a request field, or a trigger's stored config.
|
|
///
|
|
/// The asset previewer renders "this object has not been written yet" for a 404, and for any
|
|
/// other non-400 whose body contains "not found" (`S3FilePreview.svelte`, `isNotFoundError`).
|
|
/// So the named variant must stay a **400** — its message echoes a caller-supplied name, which
|
|
/// may itself contain "not found" — and the unnamed one must keep a message that does not.
|
|
pub fn workspace_storage_not_found(storage: Option<&str>) -> error::Error {
|
|
match workspace_storage_not_found_message(storage) {
|
|
Some(msg) => error::Error::BadRequest(msg),
|
|
None => error::Error::InternalErr(
|
|
"No files storage resource defined at the workspace level".to_string(),
|
|
),
|
|
}
|
|
}
|
|
|
|
/// [`workspace_storage_not_found`] for a caller whose storage name comes from stored
|
|
/// configuration rather than from the request — a trigger's static-asset config, say. Same
|
|
/// message, but a server-side class: the requester cannot correct a name they never supplied.
|
|
///
|
|
/// Not for any route the asset previewer reads: this is the 500-with-an-interpolated-name shape
|
|
/// that `isNotFoundError` falls through to its "not found" substring test for, so a storage
|
|
/// named `archive not found` would render there as "asset not yet materialized".
|
|
#[track_caller]
|
|
pub fn workspace_storage_misconfigured(storage: Option<&str>) -> error::Error {
|
|
// `internal_err` on both arms: it is `#[track_caller]`, so the `@file:line` stamp lands on
|
|
// the handler that misconfigured the storage rather than on this helper.
|
|
error::Error::internal_err(
|
|
workspace_storage_not_found_message(storage).unwrap_or_else(|| {
|
|
"No files storage resource defined at the workspace level".to_string()
|
|
}),
|
|
)
|
|
}
|
|
|
|
/// `None` when the request named no storage (or named the primary one), so the caller reports
|
|
/// the workspace as having no storage configured at all.
|
|
fn workspace_storage_not_found_message(storage: Option<&str>) -> Option<String> {
|
|
storage.filter(|s| *s != DEFAULT_STORAGE).map(|name| {
|
|
format!(
|
|
"No files storage named '{name}' is defined at the workspace level. A storage name \
|
|
must be one of this workspace's secondary storages, or `{DEFAULT_STORAGE}` \
|
|
(equivalently, nothing at all) for the primary one."
|
|
)
|
|
})
|
|
}
|
|
|
|
pub fn bundle(w_id: &str, hash: &str) -> String {
|
|
format!("script_bundle/{}/{}", w_id, hash)
|
|
}
|
|
|
|
pub fn raw_app(w_id: &str, version: &i64) -> String {
|
|
format!("raw_app/{}/{}", w_id, version)
|
|
}
|
|
|
|
pub async fn upload_artifact_to_store(
|
|
path: &str,
|
|
data: bytes::Bytes,
|
|
standalone_dir: &str,
|
|
) -> error::Result<()> {
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
let object_store = get_object_store().await;
|
|
#[cfg(not(all(feature = "enterprise", feature = "parquet")))]
|
|
let object_store: Option<()> = None;
|
|
Ok(
|
|
if &windmill_common::utils::MODE_AND_ADDONS.mode
|
|
== &windmill_common::utils::Mode::Standalone
|
|
&& object_store.is_none()
|
|
{
|
|
let path = format!("{}/{}", standalone_dir, path);
|
|
tracing::info!("Writing file to path {path}");
|
|
|
|
let split_path = path.split("/").collect::<Vec<&str>>();
|
|
std::fs::create_dir_all(split_path[..split_path.len() - 1].join("/"))?;
|
|
|
|
windmill_common::worker::write_file_bytes(&path, &data)?;
|
|
} else {
|
|
#[cfg(not(all(feature = "enterprise", feature = "parquet")))]
|
|
{
|
|
return Err(error::Error::ExecutionErr(
|
|
"codebase is an EE feature".to_string(),
|
|
));
|
|
}
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
if let Some(os) = object_store {
|
|
if let Err(e) = os
|
|
.put(&object_store::path::Path::from(path), data.into())
|
|
.await
|
|
{
|
|
tracing::info!("Failed to put snapshot to s3 at {path}: {:?}", e);
|
|
return Err(error::Error::ExecutionErr(format!(
|
|
"Failed to put {path} to s3"
|
|
)));
|
|
}
|
|
} else {
|
|
return Err(error::Error::BadConfig("Object store is required for snapshot script and is not configured for servers".to_string()));
|
|
}
|
|
},
|
|
)
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
pub async fn get_etag_or_empty(
|
|
object_store_resource: &ObjectStoreResource,
|
|
s3_object: S3Object,
|
|
) -> Option<String> {
|
|
let object_store_client = build_object_store_client(object_store_resource).await;
|
|
if object_store_client.is_err() {
|
|
return None;
|
|
}
|
|
|
|
let object_key = object_store::path::Path::from(s3_object.s3);
|
|
|
|
return object_store_client
|
|
.unwrap()
|
|
.head(&object_key)
|
|
.await
|
|
.ok()
|
|
.map(|meta| meta.e_tag)
|
|
.flatten();
|
|
}
|
|
|
|
pub fn lfs_to_object_store_resource(
|
|
lfs: &LargeFileStorage,
|
|
resource_value: serde_json::Value,
|
|
) -> error::Result<ObjectStoreResource> {
|
|
match lfs {
|
|
LargeFileStorage::S3Storage(_) | LargeFileStorage::S3AwsOidc(_) => {
|
|
let s3_resource: S3Resource = serde_json::from_value(resource_value).map_err(|e| {
|
|
error::Error::internal_err(format!("Error parsing S3 resource: {e:?}"))
|
|
})?;
|
|
Ok(ObjectStoreResource::S3(s3_resource))
|
|
}
|
|
LargeFileStorage::AzureBlobStorage(_) | LargeFileStorage::AzureWorkloadIdentity(_) => {
|
|
let azure_blob_resource: AzureBlobResource = serde_json::from_value(resource_value)
|
|
.map_err(|e| {
|
|
error::Error::internal_err(format!("Error parsing Azure Blob resource: {e:?}"))
|
|
})?;
|
|
Ok(ObjectStoreResource::Azure(azure_blob_resource))
|
|
}
|
|
LargeFileStorage::GoogleCloudStorage(_) => {
|
|
let gcs_resource: GcsResource =
|
|
serde_json::from_value(resource_value).map_err(|e| {
|
|
error::Error::internal_err(format!("Error parsing GCS resource: {e:?}"))
|
|
})?;
|
|
Ok(ObjectStoreResource::Gcs(gcs_resource))
|
|
}
|
|
LargeFileStorage::FilesystemStorage(fs) => {
|
|
windmill_common::workspaces::ensure_filesystem_storage_allowed()?;
|
|
Ok(ObjectStoreResource::Filesystem(FilesystemSettings {
|
|
root_path: fs.root_path.clone(),
|
|
}))
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn format_duckdb_connection_settings(
|
|
object_store_resource: ObjectStoreResource,
|
|
) -> error::Result<DuckdbConnectionSettingsResponse> {
|
|
match object_store_resource {
|
|
ObjectStoreResource::S3(s3_resource) => duckdb_connection_settings_internal(s3_resource),
|
|
ObjectStoreResource::Azure(azure_resource) => {
|
|
let connection_string = format!(
|
|
"CREATE SECRET az_secret (TYPE AZURE, CONNECTION_STRING 'DefaultEndpointsProtocol=https;AccountName={};AccountKey={};EndpointSuffix=core.windows.net');",
|
|
azure_resource.account_name,
|
|
azure_resource.access_key.unwrap_or_default()
|
|
);
|
|
let response = DuckdbConnectionSettingsResponse {
|
|
connection_settings_str: connection_string,
|
|
azure_container_path: Some(format!("az://{}", azure_resource.container_name)),
|
|
s3_bucket: None,
|
|
};
|
|
Ok(response)
|
|
}
|
|
ObjectStoreResource::Gcs(_) => {
|
|
return Err(error::Error::BadRequest(
|
|
"GCS is not supported in DuckDB".to_string(),
|
|
));
|
|
}
|
|
ObjectStoreResource::Filesystem(_) => {
|
|
return Err(error::Error::BadRequest(
|
|
"Filesystem is not supported in DuckDB".to_string(),
|
|
));
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn duckdb_connection_settings_internal(
|
|
s3_resource: S3Resource,
|
|
) -> error::Result<DuckdbConnectionSettingsResponse> {
|
|
let mut duckdb_settings: String = String::new();
|
|
|
|
duckdb_settings.push_str("SET home_directory='./';\n");
|
|
duckdb_settings.push_str("INSTALL 'httpfs';\n");
|
|
if s3_resource.path_style.unwrap_or(true) {
|
|
duckdb_settings.push_str("SET s3_url_style='path';\n");
|
|
}
|
|
duckdb_settings.push_str(format!("SET s3_region='{}';\n", s3_resource.region).as_str());
|
|
duckdb_settings.push_str(
|
|
format!(
|
|
"SET s3_endpoint='{}';\n",
|
|
s3_resource.endpoint_with_region_fallback(None)
|
|
)
|
|
.as_str(),
|
|
);
|
|
if !s3_resource.use_ssl {
|
|
duckdb_settings.push_str("SET s3_use_ssl=0;\n");
|
|
}
|
|
if let Some(access_key_id) = s3_resource.access_key {
|
|
duckdb_settings.push_str(format!("SET s3_access_key_id='{}';\n", access_key_id).as_str());
|
|
}
|
|
if let Some(secret_access_key) = s3_resource.secret_key {
|
|
duckdb_settings
|
|
.push_str(format!("SET s3_secret_access_key='{}';\n", secret_access_key).as_str());
|
|
}
|
|
|
|
let response = DuckdbConnectionSettingsResponse {
|
|
connection_settings_str: duckdb_settings,
|
|
azure_container_path: None,
|
|
s3_bucket: Some(s3_resource.bucket),
|
|
};
|
|
return Ok(response);
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
enum RecordBatchWriterEnum {
|
|
Parquet(ArrowWriter<ChannelWriter>),
|
|
Csv(csv::Writer<ChannelWriter>),
|
|
Json(json::Writer<ChannelWriter, JsonArray>),
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
impl RecordBatchWriter for RecordBatchWriterEnum {
|
|
fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
|
|
match self {
|
|
RecordBatchWriterEnum::Parquet(w) => w.write(batch).map_err(|e| e.into()),
|
|
RecordBatchWriterEnum::Csv(w) => w.write(batch),
|
|
RecordBatchWriterEnum::Json(w) => w.write(batch),
|
|
}
|
|
}
|
|
|
|
fn close(self) -> Result<(), ArrowError> {
|
|
match self {
|
|
RecordBatchWriterEnum::Parquet(w) => w.close().map_err(|e| e.into()).map(drop),
|
|
RecordBatchWriterEnum::Csv(w) => w.close(),
|
|
RecordBatchWriterEnum::Json(w) => w.close(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Infer the Arrow schema of a newline-delimited JSON file.
|
|
///
|
|
/// Inference only looks at the first `DEFAULT_SCHEMA_INFER_MAX_RECORD` rows, and a column
|
|
/// that holds nothing but JSON `null` across that sample is typed `DataType::Null`, which
|
|
/// makes the reader reject the first real value further down the file. When a longer file
|
|
/// leaves such a column behind, re-infer over all of it so the column's type comes from
|
|
/// wherever its first non-null value is.
|
|
#[cfg(feature = "parquet")]
|
|
fn infer_ndjson_schema(
|
|
path: &std::path::Path,
|
|
row_count: u64,
|
|
) -> anyhow::Result<datafusion::arrow::datatypes::Schema> {
|
|
use datafusion::arrow::datatypes::{DataType, Schema};
|
|
use datafusion::arrow::json::reader::infer_json_schema;
|
|
use datafusion::datasource::file_format::DEFAULT_SCHEMA_INFER_MAX_RECORD;
|
|
|
|
fn is_untyped(data_type: &DataType) -> bool {
|
|
match data_type {
|
|
DataType::Null => true,
|
|
DataType::Struct(fields) => fields.iter().any(|f| is_untyped(f.data_type())),
|
|
DataType::List(field) | DataType::LargeList(field) => is_untyped(field.data_type()),
|
|
_ => false,
|
|
}
|
|
}
|
|
|
|
let infer = |max_records: Option<usize>| -> anyhow::Result<Schema> {
|
|
let reader = std::io::BufReader::new(std::fs::File::open(path)?);
|
|
Ok(infer_json_schema(reader, max_records).map_err(to_anyhow)?.0)
|
|
};
|
|
|
|
let schema = infer(Some(DEFAULT_SCHEMA_INFER_MAX_RECORD))?;
|
|
if row_count > DEFAULT_SCHEMA_INFER_MAX_RECORD as u64
|
|
&& schema.fields().iter().any(|f| is_untyped(f.data_type()))
|
|
{
|
|
return infer(None);
|
|
}
|
|
Ok(schema)
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
struct ChannelWriter {
|
|
sender: tokio::sync::mpsc::Sender<anyhow::Result<Bytes>>,
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
impl Write for ChannelWriter {
|
|
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
|
|
let data: Bytes = buf.to_vec().into();
|
|
self.sender.blocking_send(Ok(data)).map_err(|e| {
|
|
std::io::Error::new(
|
|
std::io::ErrorKind::BrokenPipe,
|
|
format!("Channel send error: {}", e),
|
|
)
|
|
})?;
|
|
Ok(buf.len())
|
|
}
|
|
|
|
fn flush(&mut self) -> std::io::Result<()> {
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[cfg(not(feature = "parquet"))]
|
|
#[derive(Debug, Clone, Copy, Default)]
|
|
pub struct IngestStats {
|
|
pub rows: u64,
|
|
pub bytes: u64,
|
|
pub elapsed: std::time::Duration,
|
|
pub fetch_wait: std::time::Duration,
|
|
pub write_time: std::time::Duration,
|
|
pub first_row_latency: Option<std::time::Duration>,
|
|
}
|
|
|
|
#[cfg(not(feature = "parquet"))]
|
|
pub async fn convert_json_line_stream<V, E>(
|
|
mut _stream: impl futures::TryStreamExt<Item = Result<V, E>> + Unpin,
|
|
_output_format: S3ModeFormat,
|
|
_progress: Option<tokio::sync::mpsc::Sender<IngestStats>>,
|
|
) -> anyhow::Result<(
|
|
futures::stream::BoxStream<'static, anyhow::Result<bytes::Bytes>>,
|
|
IngestStats,
|
|
)>
|
|
where
|
|
V: serde::Serialize,
|
|
E: Into<anyhow::Error>,
|
|
{
|
|
use futures::StreamExt;
|
|
let stream = async_stream::stream! {
|
|
yield Err(anyhow::anyhow!("Parquet feature is not enabled. Cannot convert JSON line stream."));
|
|
};
|
|
Ok((stream.boxed(), IngestStats::default()))
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub struct IngestStats {
|
|
pub rows: u64,
|
|
pub bytes: u64,
|
|
pub elapsed: std::time::Duration,
|
|
pub fetch_wait: std::time::Duration,
|
|
pub write_time: std::time::Duration,
|
|
pub first_row_latency: Option<std::time::Duration>,
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
pub async fn convert_json_line_stream<V, E>(
|
|
mut stream: impl TryStreamExt<Item = Result<V, E>> + Unpin,
|
|
output_format: S3ModeFormat,
|
|
progress: Option<tokio::sync::mpsc::Sender<IngestStats>>,
|
|
) -> anyhow::Result<(
|
|
futures::stream::BoxStream<'static, anyhow::Result<bytes::Bytes>>,
|
|
IngestStats,
|
|
)>
|
|
where
|
|
V: serde::Serialize,
|
|
E: Into<anyhow::Error>,
|
|
{
|
|
const MAX_MPSC_SIZE: usize = 1000;
|
|
const WRITE_BUF_CAPACITY: usize = 256 * 1024;
|
|
const PROGRESS_INTERVAL_SECS: u64 = 10;
|
|
|
|
use datafusion::{execution::context::SessionContext, prelude::NdJsonReadOptions};
|
|
use futures::StreamExt;
|
|
use std::path::PathBuf;
|
|
use std::time::{Duration, Instant};
|
|
use tokio::io::AsyncWriteExt;
|
|
|
|
let mut path = PathBuf::from(std::env::temp_dir());
|
|
path.push(format!("{}.json", rd_string(8)));
|
|
let path_str = path
|
|
.to_str()
|
|
.ok_or_else(|| anyhow::anyhow!("Invalid path"))?;
|
|
|
|
let file: tokio::fs::File = tokio::fs::File::create(&path).await.map_err(to_anyhow)?;
|
|
let mut file = tokio::io::BufWriter::with_capacity(WRITE_BUF_CAPACITY, file);
|
|
|
|
let ingest_start = Instant::now();
|
|
let mut first_row_latency: Option<Duration> = None;
|
|
let mut row_count: u64 = 0;
|
|
let mut bytes_written: u64 = 0;
|
|
let mut write_time = Duration::ZERO;
|
|
let mut progress_timer = tokio::time::interval(Duration::from_secs(PROGRESS_INTERVAL_SECS));
|
|
progress_timer.tick().await; // drop the immediate tick
|
|
let build_stats = |row_count: u64,
|
|
bytes_written: u64,
|
|
write_time: Duration,
|
|
first_row_latency: Option<Duration>,
|
|
elapsed: Duration|
|
|
-> IngestStats {
|
|
IngestStats {
|
|
rows: row_count,
|
|
bytes: bytes_written,
|
|
elapsed,
|
|
fetch_wait: elapsed.saturating_sub(write_time),
|
|
write_time,
|
|
first_row_latency,
|
|
}
|
|
};
|
|
|
|
let mut done = false;
|
|
while !done {
|
|
tokio::select! {
|
|
chunk = stream.next() => {
|
|
match chunk {
|
|
Some(Ok(chunk)) => {
|
|
if first_row_latency.is_none() {
|
|
first_row_latency = Some(ingest_start.elapsed());
|
|
}
|
|
let t_write = Instant::now();
|
|
let mut s = serde_json::to_string(&chunk)?;
|
|
s.push('\n');
|
|
bytes_written += s.len() as u64;
|
|
file.write_all(s.as_bytes()).await?;
|
|
write_time += t_write.elapsed();
|
|
row_count += 1;
|
|
}
|
|
Some(Err(e)) => {
|
|
tokio::fs::remove_file(&path).await?;
|
|
return Err(e.into());
|
|
}
|
|
None => done = true,
|
|
}
|
|
}
|
|
_ = progress_timer.tick(), if progress.is_some() => {
|
|
let stats = build_stats(row_count, bytes_written, write_time, first_row_latency, ingest_start.elapsed());
|
|
if let Some(tx) = &progress {
|
|
let _ = tx.try_send(stats);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
file.flush().await?;
|
|
let file = file.into_inner();
|
|
file.sync_all().await?;
|
|
drop(file);
|
|
|
|
let ingest_stats = build_stats(
|
|
row_count,
|
|
bytes_written,
|
|
write_time,
|
|
first_row_latency,
|
|
ingest_start.elapsed(),
|
|
);
|
|
|
|
let inferred_schema = {
|
|
let path = path.clone();
|
|
task::spawn_blocking(move || infer_ndjson_schema(&path, row_count))
|
|
.await
|
|
.map_err(to_anyhow)??
|
|
};
|
|
|
|
let ctx = SessionContext::new();
|
|
ctx.register_json(
|
|
"my_table",
|
|
path_str,
|
|
NdJsonReadOptions::default().schema(&inferred_schema),
|
|
)
|
|
.await
|
|
.map_err(to_anyhow)?;
|
|
|
|
let df = ctx.sql("SELECT * FROM my_table").await.map_err(to_anyhow)?;
|
|
let schema = df.schema().clone().into();
|
|
let mut datafusion_stream = df.execute_stream().await.map_err(to_anyhow)?;
|
|
|
|
let (tx, rx) = tokio::sync::mpsc::channel(MAX_MPSC_SIZE);
|
|
let writer: Arc<Mutex<Option<RecordBatchWriterEnum>>> =
|
|
Arc::new(Mutex::new(Some(match output_format {
|
|
S3ModeFormat::Parquet => RecordBatchWriterEnum::Parquet(
|
|
ArrowWriter::try_new(ChannelWriter { sender: tx.clone() }, Arc::new(schema), None)
|
|
.map_err(to_anyhow)?,
|
|
),
|
|
|
|
S3ModeFormat::Csv => {
|
|
RecordBatchWriterEnum::Csv(csv::Writer::new(ChannelWriter { sender: tx.clone() }))
|
|
}
|
|
S3ModeFormat::Json => {
|
|
RecordBatchWriterEnum::Json(json::Writer::<_, JsonArray>::new(ChannelWriter {
|
|
sender: tx.clone(),
|
|
}))
|
|
}
|
|
})));
|
|
|
|
task::spawn(async move {
|
|
while let Some(batch_result) = datafusion_stream.next().await {
|
|
let batch: RecordBatch = match batch_result {
|
|
Ok(batch) => batch,
|
|
Err(e) => {
|
|
tracing::error!("Error in datafusion stream: {:?}", &e);
|
|
match tx.send(Err(e.into())).await {
|
|
Ok(_) => {}
|
|
Err(e) => tracing::error!("Failed to write error to channel: {:?}", &e),
|
|
}
|
|
break;
|
|
}
|
|
};
|
|
let writer = writer.clone();
|
|
let write_result = task::spawn_blocking(move || {
|
|
writer.lock().unwrap().as_mut().unwrap().write(&batch)
|
|
})
|
|
.await;
|
|
match write_result {
|
|
Ok(Ok(_)) => {}
|
|
Ok(Err(e)) => {
|
|
tracing::error!("Error writing batch: {:?}", &e);
|
|
match tx.send(Err(e.into())).await {
|
|
Ok(_) => {}
|
|
Err(e) => tracing::error!("Failed to write error to channel: {:?}", &e),
|
|
}
|
|
}
|
|
Err(e) => tracing::error!("Error in blocking task: {:?}", &e),
|
|
};
|
|
}
|
|
let close_result = task::spawn_blocking(move || {
|
|
writer.lock().unwrap().take().unwrap().close()?;
|
|
drop(writer);
|
|
Ok::<_, anyhow::Error>(())
|
|
})
|
|
.await;
|
|
match close_result {
|
|
Ok(Ok(())) => {}
|
|
Ok(Err(e)) => {
|
|
tracing::error!("Error closing S3 stream writer: {:?}", e);
|
|
let _ = tx.send(Err(e)).await;
|
|
}
|
|
Err(e) => {
|
|
tracing::error!("S3 stream writer close task panicked: {:?}", e);
|
|
let _ = tx
|
|
.send(Err(anyhow::anyhow!("writer close task panicked: {e}")))
|
|
.await;
|
|
}
|
|
}
|
|
drop(ctx);
|
|
if let Err(e) = tokio::fs::remove_file(&path).await {
|
|
tracing::error!("Error removing temp file {}: {:?}", path.display(), e);
|
|
}
|
|
Ok::<_, anyhow::Error>(())
|
|
});
|
|
|
|
Ok((
|
|
tokio_stream::wrappers::ReceiverStream::new(rx).boxed(),
|
|
ingest_stats,
|
|
))
|
|
}
|
|
|
|
/// Decode the bytes of a Parquet file into a JSON array text (`[ {...}, {...} ]`)
|
|
/// suitable for binding as a single SQL parameter and consuming with `OPENJSON`,
|
|
/// `jsonb_to_recordset`, `JSON_TABLE`, etc.
|
|
///
|
|
/// Runs the synchronous Arrow parquet reader on a `spawn_blocking` thread, which is
|
|
/// fine for the ~500 MB ceiling we target. Larger files should use a streaming
|
|
/// path (out of scope for the s3-input feature).
|
|
#[cfg(feature = "parquet")]
|
|
pub async fn decode_parquet_bytes_to_json_array(bytes: bytes::Bytes) -> anyhow::Result<String> {
|
|
use datafusion::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
|
|
|
|
task::spawn_blocking(move || {
|
|
let builder = ParquetRecordBatchReaderBuilder::try_new(bytes).map_err(to_anyhow)?;
|
|
let reader = builder.build().map_err(to_anyhow)?;
|
|
|
|
let mut out: Vec<u8> = Vec::new();
|
|
let mut writer = json::Writer::<_, JsonArray>::new(&mut out);
|
|
for batch in reader {
|
|
let batch = batch.map_err(to_anyhow)?;
|
|
writer.write(&batch).map_err(to_anyhow)?;
|
|
}
|
|
writer.finish().map_err(to_anyhow)?;
|
|
drop(writer);
|
|
String::from_utf8(out).map_err(to_anyhow)
|
|
})
|
|
.await
|
|
.map_err(to_anyhow)?
|
|
}
|
|
|
|
#[cfg(not(feature = "parquet"))]
|
|
pub async fn decode_parquet_bytes_to_json_array(_bytes: bytes::Bytes) -> anyhow::Result<String> {
|
|
anyhow::bail!("Parquet S3 input requires the `parquet` feature to be enabled on this build")
|
|
}
|
|
|
|
/// Decode the bytes of a CSV file into a JSON array text using the first row as headers.
|
|
/// Same blocking-thread pattern as the parquet decoder.
|
|
#[cfg(feature = "parquet")]
|
|
pub async fn decode_csv_bytes_to_json_array(bytes: bytes::Bytes) -> anyhow::Result<String> {
|
|
use datafusion::arrow::csv::ReaderBuilder;
|
|
use std::io::Cursor;
|
|
|
|
task::spawn_blocking(move || {
|
|
let cursor = Cursor::new(bytes);
|
|
// Two-pass: infer schema from the bytes, then build the reader. The infer step
|
|
// rewinds the underlying reader for us.
|
|
let (schema, _) = datafusion::arrow::csv::reader::Format::default()
|
|
.with_header(true)
|
|
.infer_schema(Cursor::new(&cursor.get_ref()[..]), Some(1024))
|
|
.map_err(to_anyhow)?;
|
|
|
|
let reader = ReaderBuilder::new(Arc::new(schema))
|
|
.with_header(true)
|
|
.build(cursor)
|
|
.map_err(to_anyhow)?;
|
|
|
|
let mut out: Vec<u8> = Vec::new();
|
|
let mut writer = json::Writer::<_, JsonArray>::new(&mut out);
|
|
for batch in reader {
|
|
let batch = batch.map_err(to_anyhow)?;
|
|
writer.write(&batch).map_err(to_anyhow)?;
|
|
}
|
|
writer.finish().map_err(to_anyhow)?;
|
|
drop(writer);
|
|
String::from_utf8(out).map_err(to_anyhow)
|
|
})
|
|
.await
|
|
.map_err(to_anyhow)?
|
|
}
|
|
|
|
#[cfg(not(feature = "parquet"))]
|
|
pub async fn decode_csv_bytes_to_json_array(_bytes: bytes::Bytes) -> anyhow::Result<String> {
|
|
anyhow::bail!("CSV S3 input requires the `parquet` feature to be enabled on this build")
|
|
}
|
|
|
|
lazy_static::lazy_static! {
|
|
pub static ref S3_PROXY_LAST_ERRORS_CACHE: Cache<String, String> = Cache::new(4);
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
pub async fn get_logs_from_store(
|
|
log_offset: i32,
|
|
logs: &str,
|
|
log_file_index: &Option<Vec<String>>,
|
|
) -> Option<impl futures::Stream<Item = Result<bytes::Bytes, object_store::Error>>> {
|
|
if log_offset > 0 {
|
|
if let Some(file_index) = log_file_index.clone() {
|
|
if file_index.iter().any(|p| !is_safe_log_file_path(p)) {
|
|
return None;
|
|
}
|
|
if let Some(os) = get_object_store().await {
|
|
let logs = logs.to_string();
|
|
let stream = async_stream::stream! {
|
|
for file_p in file_index {
|
|
let file = os.get(&object_store::path::Path::from(file_p.clone())).await;
|
|
match file {
|
|
Ok(file) => {
|
|
if let Ok(bytes) = file.bytes().await {
|
|
yield Ok(bytes::Bytes::from(bytes));
|
|
}
|
|
}
|
|
Err(e) => {
|
|
tracing::debug!("error getting file from store: {file_p}: {e}");
|
|
}
|
|
}
|
|
}
|
|
yield Ok(bytes::Bytes::from(logs))
|
|
};
|
|
return Some(stream);
|
|
} else {
|
|
tracing::debug!("object store client not present");
|
|
}
|
|
}
|
|
}
|
|
None
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
// --- ambient credentials cache tests ---
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[test]
|
|
fn test_ambient_credentials_still_valid() {
|
|
use std::time::{Duration, SystemTime};
|
|
|
|
fn creds(expiry: Option<SystemTime>) -> aws_sdk_sts::config::Credentials {
|
|
let mut builder = aws_sdk_sts::config::Credentials::builder()
|
|
.access_key_id("AK")
|
|
.secret_access_key("SK")
|
|
.provider_name("test");
|
|
if let Some(expiry) = expiry {
|
|
builder = builder.expiry(expiry);
|
|
}
|
|
builder.build()
|
|
}
|
|
|
|
// Expiry far in the future: valid regardless of fetch time
|
|
assert!(AmbientAwsCredentials::still_valid(
|
|
&creds(Some(SystemTime::now() + Duration::from_secs(3600))),
|
|
Duration::ZERO
|
|
));
|
|
// Expiry within the refresh margin: must be re-fetched
|
|
assert!(!AmbientAwsCredentials::still_valid(
|
|
&creds(Some(SystemTime::now() + Duration::from_secs(30))),
|
|
Duration::ZERO
|
|
));
|
|
// No expiry: valid while fresh, re-fetched after the TTL
|
|
assert!(AmbientAwsCredentials::still_valid(
|
|
&creds(None),
|
|
Duration::ZERO
|
|
));
|
|
assert!(!AmbientAwsCredentials::still_valid(
|
|
&creds(None),
|
|
AmbientAwsCredentials::NO_EXPIRY_TTL + Duration::from_secs(1)
|
|
));
|
|
}
|
|
|
|
// --- render_endpoint tests ---
|
|
|
|
#[test]
|
|
fn test_render_endpoint_ssl_prefix() {
|
|
let result = render_endpoint(
|
|
"s3.amazonaws.com".to_string(),
|
|
true,
|
|
None,
|
|
Some(true),
|
|
"mybucket".to_string(),
|
|
);
|
|
assert_eq!(result, "https://s3.amazonaws.com");
|
|
}
|
|
|
|
#[test]
|
|
fn test_render_endpoint_non_ssl_prefix() {
|
|
let result = render_endpoint(
|
|
"minio.local".to_string(),
|
|
false,
|
|
None,
|
|
Some(true),
|
|
"mybucket".to_string(),
|
|
);
|
|
assert_eq!(result, "http://minio.local");
|
|
}
|
|
|
|
#[test]
|
|
fn test_render_endpoint_with_port() {
|
|
let result = render_endpoint(
|
|
"minio.local".to_string(),
|
|
false,
|
|
Some(9000),
|
|
Some(true),
|
|
"mybucket".to_string(),
|
|
);
|
|
assert_eq!(result, "http://minio.local:9000");
|
|
}
|
|
|
|
#[test]
|
|
fn test_render_endpoint_virtual_hosted_style() {
|
|
let result = render_endpoint(
|
|
"s3.amazonaws.com".to_string(),
|
|
true,
|
|
None,
|
|
Some(false),
|
|
"mybucket".to_string(),
|
|
);
|
|
assert_eq!(result, "https://mybucket.s3.amazonaws.com");
|
|
}
|
|
|
|
#[test]
|
|
fn test_render_endpoint_passthrough_with_scheme() {
|
|
let result = render_endpoint(
|
|
"https://custom.endpoint.com".to_string(),
|
|
false, // use_ssl is ignored when scheme already present
|
|
None,
|
|
Some(false), // path_style is also ignored
|
|
"mybucket".to_string(),
|
|
);
|
|
assert_eq!(result, "https://custom.endpoint.com");
|
|
}
|
|
|
|
// --- lfs_to_object_store_resource tests ---
|
|
|
|
#[test]
|
|
fn test_lfs_to_object_store_resource_s3() {
|
|
let lfs = LargeFileStorage::S3Storage(S3Storage {
|
|
s3_resource_path: "u/admin/s3_resource".to_string(),
|
|
public_resource: None,
|
|
advanced_permissions: None,
|
|
});
|
|
let resource_json = serde_json::json!({
|
|
"bucket": "my-bucket",
|
|
"region": "us-east-1",
|
|
"endPoint": "s3.amazonaws.com",
|
|
"useSSL": true,
|
|
"accessKey": "AKIA...",
|
|
"secretKey": "secret"
|
|
});
|
|
let result = lfs_to_object_store_resource(&lfs, resource_json).unwrap();
|
|
match result {
|
|
ObjectStoreResource::S3(s3) => {
|
|
assert_eq!(s3.bucket, "my-bucket");
|
|
assert_eq!(s3.region, "us-east-1");
|
|
assert_eq!(s3.endpoint, "s3.amazonaws.com");
|
|
assert!(s3.use_ssl);
|
|
}
|
|
_ => panic!("Expected S3 resource"),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_lfs_to_object_store_resource_azure() {
|
|
let lfs = LargeFileStorage::AzureBlobStorage(AzureBlobStorage {
|
|
azure_blob_resource_path: "u/admin/azure_resource".to_string(),
|
|
public_resource: None,
|
|
advanced_permissions: None,
|
|
});
|
|
let resource_json = serde_json::json!({
|
|
"accountName": "myaccount",
|
|
"containerName": "mycontainer"
|
|
});
|
|
let result = lfs_to_object_store_resource(&lfs, resource_json).unwrap();
|
|
match result {
|
|
ObjectStoreResource::Azure(az) => {
|
|
assert_eq!(az.account_name, "myaccount");
|
|
assert_eq!(az.container_name, "mycontainer");
|
|
}
|
|
_ => panic!("Expected Azure resource"),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_lfs_to_object_store_resource_gcs() {
|
|
let lfs = LargeFileStorage::GoogleCloudStorage(GoogleCloudStorage {
|
|
gcs_resource_path: "u/admin/gcs_resource".to_string(),
|
|
public_resource: None,
|
|
advanced_permissions: None,
|
|
});
|
|
let resource_json = serde_json::json!({
|
|
"bucket": "gcs-bucket",
|
|
"serviceAccountKey": {"type": "service_account", "project_id": "test"}
|
|
});
|
|
let result = lfs_to_object_store_resource(&lfs, resource_json).unwrap();
|
|
match result {
|
|
ObjectStoreResource::Gcs(gcs) => {
|
|
assert_eq!(gcs.bucket, "gcs-bucket");
|
|
}
|
|
_ => panic!("Expected GCS resource"),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_lfs_to_object_store_resource_filesystem() {
|
|
let lfs = LargeFileStorage::FilesystemStorage(FilesystemStorage {
|
|
root_path: "/tmp/mydata".to_string(),
|
|
public_resource: None,
|
|
advanced_permissions: None,
|
|
});
|
|
// resource_value is ignored for filesystem
|
|
let result = lfs_to_object_store_resource(&lfs, serde_json::Value::Null).unwrap();
|
|
match result {
|
|
ObjectStoreResource::Filesystem(fs) => {
|
|
assert_eq!(fs.root_path, "/tmp/mydata");
|
|
}
|
|
_ => panic!("Expected Filesystem resource"),
|
|
}
|
|
}
|
|
|
|
// --- duckdb_connection_settings tests ---
|
|
|
|
#[test]
|
|
fn test_duckdb_connection_settings_s3_basic() {
|
|
let s3 = S3Resource {
|
|
bucket: "test-bucket".to_string(),
|
|
region: "eu-west-1".to_string(),
|
|
endpoint: "s3.eu-west-1.amazonaws.com".to_string(),
|
|
use_ssl: true,
|
|
access_key: Some("AKIA123".to_string()),
|
|
secret_key: Some("secret456".to_string()),
|
|
path_style: Some(true),
|
|
token: None,
|
|
expiration: None,
|
|
port: None,
|
|
};
|
|
let result = duckdb_connection_settings_internal(s3).unwrap();
|
|
assert!(result
|
|
.connection_settings_str
|
|
.contains("SET s3_region='eu-west-1'"));
|
|
assert!(result
|
|
.connection_settings_str
|
|
.contains("SET s3_access_key_id='AKIA123'"));
|
|
assert!(result
|
|
.connection_settings_str
|
|
.contains("SET s3_secret_access_key='secret456'"));
|
|
assert!(result
|
|
.connection_settings_str
|
|
.contains("SET s3_url_style='path'"));
|
|
assert!(!result.connection_settings_str.contains("SET s3_use_ssl=0"));
|
|
assert_eq!(result.s3_bucket, Some("test-bucket".to_string()));
|
|
assert!(result.azure_container_path.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_duckdb_connection_settings_s3_no_ssl() {
|
|
let s3 = S3Resource {
|
|
bucket: "bucket".to_string(),
|
|
region: "us-east-1".to_string(),
|
|
endpoint: "minio:9000".to_string(),
|
|
use_ssl: false,
|
|
access_key: None,
|
|
secret_key: None,
|
|
path_style: Some(false),
|
|
token: None,
|
|
expiration: None,
|
|
port: None,
|
|
};
|
|
let result = duckdb_connection_settings_internal(s3).unwrap();
|
|
assert!(result.connection_settings_str.contains("SET s3_use_ssl=0"));
|
|
assert!(!result
|
|
.connection_settings_str
|
|
.contains("SET s3_url_style='path'"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_duckdb_connection_settings_azure() {
|
|
let resource = ObjectStoreResource::Azure(AzureBlobResource {
|
|
endpoint: None,
|
|
use_ssl: None,
|
|
account_name: "myaccount".to_string(),
|
|
tenant_id: None,
|
|
client_id: None,
|
|
container_name: "mycontainer".to_string(),
|
|
access_key: Some("base64key==".to_string()),
|
|
federated_token_file: None,
|
|
});
|
|
let result = format_duckdb_connection_settings(resource).unwrap();
|
|
assert!(result
|
|
.connection_settings_str
|
|
.contains("AccountName=myaccount"));
|
|
assert!(result
|
|
.connection_settings_str
|
|
.contains("AccountKey=base64key=="));
|
|
assert_eq!(
|
|
result.azure_container_path,
|
|
Some("az://mycontainer".to_string())
|
|
);
|
|
assert!(result.s3_bucket.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_duckdb_connection_settings_gcs_unsupported() {
|
|
let resource = ObjectStoreResource::Gcs(GcsResource {
|
|
bucket: "bucket".to_string(),
|
|
service_account_key: "{}".to_string(),
|
|
});
|
|
let result = format_duckdb_connection_settings(resource);
|
|
assert!(result.is_err());
|
|
assert!(result
|
|
.unwrap_err()
|
|
.to_string()
|
|
.contains("GCS is not supported"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_gcs_service_account_key_is_blank() {
|
|
for blank in ["", " ", "\n\t", "{}", " {} ", "null"] {
|
|
assert!(
|
|
gcs_service_account_key_is_blank(blank),
|
|
"{blank:?} should be treated as no key"
|
|
);
|
|
}
|
|
for present in ["{\"client_email\":\"x@y.z\"}", "not json"] {
|
|
assert!(
|
|
!gcs_service_account_key_is_blank(present),
|
|
"{present:?} should be treated as a key"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[tokio::test]
|
|
async fn test_build_gcs_client_blank_key_uses_instance_credentials() {
|
|
// A blank service account key must not be passed to `with_service_account_key`
|
|
// (which would fail to parse): the builder should fall back to instance credentials
|
|
// (GKE Workload Identity / metadata server) and construct successfully. `{}` is the
|
|
// settings UI's representation of "no key".
|
|
for key in ["", " ", "{}"] {
|
|
let resource =
|
|
GcsResource { bucket: "bucket".to_string(), service_account_key: key.to_string() };
|
|
assert!(
|
|
build_gcs_client(&resource).await.is_ok(),
|
|
"blank key {:?} should build via instance credentials",
|
|
key
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_duckdb_connection_settings_filesystem_unsupported() {
|
|
let resource = ObjectStoreResource::Filesystem(FilesystemSettings {
|
|
root_path: "/tmp/data".to_string(),
|
|
});
|
|
let result = format_duckdb_connection_settings(resource);
|
|
assert!(result.is_err());
|
|
assert!(result
|
|
.unwrap_err()
|
|
.to_string()
|
|
.contains("Filesystem is not supported"));
|
|
}
|
|
|
|
// --- object_store error mapping ---
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[test]
|
|
fn test_object_store_error_mapping() {
|
|
use windmill_common::error::Error;
|
|
|
|
// NotFound
|
|
let err =
|
|
object_store::Error::NotFound { path: "test/path".into(), source: "missing".into() };
|
|
let mapped = object_store_error_to_error(err);
|
|
assert!(matches!(mapped, Error::NotFound(_)));
|
|
|
|
// PermissionDenied
|
|
let err = object_store::Error::PermissionDenied {
|
|
path: "secret".into(),
|
|
source: "forbidden".into(),
|
|
};
|
|
let mapped = object_store_error_to_error(err);
|
|
assert!(matches!(mapped, Error::PermissionDenied(_)));
|
|
|
|
// InvalidPath
|
|
let err = object_store::Error::InvalidPath {
|
|
source: object_store::path::Error::EmptySegment { path: "".into() },
|
|
};
|
|
let mapped = object_store_error_to_error(err);
|
|
assert!(matches!(mapped, Error::BadRequest(_)));
|
|
|
|
// NotImplemented
|
|
let mapped = object_store_error_to_error(object_store::Error::NotImplemented);
|
|
assert!(matches!(mapped, Error::BadRequest(_)));
|
|
|
|
// Unauthenticated
|
|
let err =
|
|
object_store::Error::Unauthenticated { path: "obj".into(), source: "no creds".into() };
|
|
let mapped = object_store_error_to_error(err);
|
|
assert!(matches!(mapped, Error::NotAuthorized(_)));
|
|
}
|
|
|
|
// --- Filesystem-backed integration tests ---
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[test]
|
|
fn test_build_filesystem_client() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let result = build_filesystem_client(dir.path().to_str().unwrap());
|
|
assert!(result.is_ok());
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[tokio::test]
|
|
async fn test_filesystem_put_get_roundtrip() {
|
|
use object_store::{path::Path, ObjectStore, PutPayload};
|
|
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let client = build_filesystem_client(dir.path().to_str().unwrap()).unwrap();
|
|
let path = Path::from("test.txt");
|
|
let data = bytes::Bytes::from("hello world");
|
|
|
|
client
|
|
.put(&path, PutPayload::from(data.clone()))
|
|
.await
|
|
.unwrap();
|
|
let result = client.get(&path).await.unwrap().bytes().await.unwrap();
|
|
assert_eq!(result, data);
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[tokio::test]
|
|
async fn test_filesystem_nested_paths() {
|
|
use object_store::{path::Path, ObjectStore, PutPayload};
|
|
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let client = build_filesystem_client(dir.path().to_str().unwrap()).unwrap();
|
|
let path = Path::from("a/b/c.txt");
|
|
let data = bytes::Bytes::from("nested content");
|
|
|
|
client
|
|
.put(&path, PutPayload::from(data.clone()))
|
|
.await
|
|
.unwrap();
|
|
let result = client.get(&path).await.unwrap().bytes().await.unwrap();
|
|
assert_eq!(result, data);
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[tokio::test]
|
|
async fn test_filesystem_overwrite() {
|
|
use object_store::{path::Path, ObjectStore, PutPayload};
|
|
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let client = build_filesystem_client(dir.path().to_str().unwrap()).unwrap();
|
|
let path = Path::from("overwrite.txt");
|
|
|
|
client
|
|
.put(&path, PutPayload::from(bytes::Bytes::from("v1")))
|
|
.await
|
|
.unwrap();
|
|
client
|
|
.put(&path, PutPayload::from(bytes::Bytes::from("v2")))
|
|
.await
|
|
.unwrap();
|
|
let result = client.get(&path).await.unwrap().bytes().await.unwrap();
|
|
assert_eq!(result, bytes::Bytes::from("v2"));
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[tokio::test]
|
|
async fn test_attempt_fetch_bytes_success() {
|
|
use object_store::{path::Path, ObjectStore, PutPayload};
|
|
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let client = build_filesystem_client(dir.path().to_str().unwrap()).unwrap();
|
|
let data = bytes::Bytes::from("fetch me");
|
|
client
|
|
.put(&Path::from("fetch.txt"), PutPayload::from(data.clone()))
|
|
.await
|
|
.unwrap();
|
|
|
|
let result = attempt_fetch_bytes(client, "fetch.txt").await.unwrap();
|
|
assert_eq!(result, data);
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[tokio::test]
|
|
async fn test_attempt_fetch_bytes_missing_key() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let client = build_filesystem_client(dir.path().to_str().unwrap()).unwrap();
|
|
|
|
let result = attempt_fetch_bytes(client, "nonexistent.txt").await;
|
|
assert!(result.is_err());
|
|
assert!(matches!(
|
|
result.unwrap_err(),
|
|
windmill_common::error::Error::ExecutionErr(_)
|
|
));
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[tokio::test]
|
|
async fn test_attempt_fetch_bytes_empty_object() {
|
|
use object_store::{path::Path, ObjectStore, PutPayload};
|
|
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let client = build_filesystem_client(dir.path().to_str().unwrap()).unwrap();
|
|
client
|
|
.put(
|
|
&Path::from("empty.txt"),
|
|
PutPayload::from(bytes::Bytes::new()),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
let result = attempt_fetch_bytes(client, "empty.txt").await;
|
|
assert!(result.is_err());
|
|
assert!(result
|
|
.unwrap_err()
|
|
.to_string()
|
|
.contains("does not exist in bucket"));
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[tokio::test]
|
|
async fn test_get_etag_or_empty() {
|
|
use object_store::{path::Path, ObjectStore, PutPayload};
|
|
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let root = dir.path().to_str().unwrap();
|
|
let client = build_filesystem_client(root).unwrap();
|
|
client
|
|
.put(
|
|
&Path::from("etag.txt"),
|
|
PutPayload::from(bytes::Bytes::from("content")),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
let resource =
|
|
ObjectStoreResource::Filesystem(FilesystemSettings { root_path: root.to_string() });
|
|
let s3_obj =
|
|
S3Object { s3: "etag.txt".to_string(), storage: None, filename: None, presigned: None };
|
|
|
|
let etag = get_etag_or_empty(&resource, s3_obj).await;
|
|
// LocalFileSystem should return an etag based on file metadata
|
|
assert!(etag.is_some());
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[tokio::test]
|
|
async fn test_get_etag_or_empty_missing() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let root = dir.path().to_str().unwrap();
|
|
|
|
let resource =
|
|
ObjectStoreResource::Filesystem(FilesystemSettings { root_path: root.to_string() });
|
|
let s3_obj = S3Object {
|
|
s3: "nonexistent.txt".to_string(),
|
|
storage: None,
|
|
filename: None,
|
|
presigned: None,
|
|
};
|
|
|
|
let etag = get_etag_or_empty(&resource, s3_obj).await;
|
|
assert!(etag.is_none());
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[tokio::test]
|
|
async fn test_settings_to_client_end_to_end() {
|
|
use object_store::{path::Path, ObjectStore, PutPayload};
|
|
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let root = dir.path().to_str().unwrap().to_string();
|
|
let settings = ObjectSettings::Filesystem(FilesystemSettings { root_path: root });
|
|
|
|
let expirable = build_object_store_from_settings(settings, None)
|
|
.await
|
|
.unwrap();
|
|
let data = bytes::Bytes::from("end to end via settings");
|
|
expirable
|
|
.store
|
|
.put(&Path::from("e2e.txt"), PutPayload::from(data.clone()))
|
|
.await
|
|
.unwrap();
|
|
let result = expirable
|
|
.store
|
|
.get(&Path::from("e2e.txt"))
|
|
.await
|
|
.unwrap()
|
|
.bytes()
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(result, data);
|
|
assert!(expirable.refresh.is_none());
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[tokio::test]
|
|
async fn test_resource_to_client_end_to_end() {
|
|
use object_store::{path::Path, ObjectStore, PutPayload};
|
|
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let root = dir.path().to_str().unwrap().to_string();
|
|
let resource = ObjectStoreResource::Filesystem(FilesystemSettings { root_path: root });
|
|
|
|
let client = build_object_store_client(&resource).await.unwrap();
|
|
let data = bytes::Bytes::from("end to end via resource");
|
|
client
|
|
.put(&Path::from("e2e.txt"), PutPayload::from(data.clone()))
|
|
.await
|
|
.unwrap();
|
|
let result = client
|
|
.get(&Path::from("e2e.txt"))
|
|
.await
|
|
.unwrap()
|
|
.bytes()
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(result, data);
|
|
}
|
|
|
|
// --- bundle / raw_app path tests ---
|
|
|
|
#[test]
|
|
fn test_bundle_path_format() {
|
|
assert_eq!(
|
|
bundle("my_workspace", "abc123"),
|
|
"script_bundle/my_workspace/abc123"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_raw_app_path_format() {
|
|
assert_eq!(raw_app("my_workspace", &42), "raw_app/my_workspace/42");
|
|
}
|
|
|
|
// --- parse_bucket_restrictions tests ---
|
|
|
|
#[test]
|
|
fn test_parse_bucket_restrictions_single_bucket() {
|
|
let result = parse_bucket_restrictions_from_str("my-bucket:workspace1,workspace2").unwrap();
|
|
assert_eq!(result.len(), 1);
|
|
assert_eq!(
|
|
result.get("my-bucket").unwrap(),
|
|
&vec!["workspace1".to_string(), "workspace2".to_string()]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_parse_bucket_restrictions_multiple_buckets() {
|
|
let result = parse_bucket_restrictions_from_str("bucket-a:ws1,ws2;bucket-b:ws3").unwrap();
|
|
assert_eq!(result.len(), 2);
|
|
assert_eq!(
|
|
result.get("bucket-a").unwrap(),
|
|
&vec!["ws1".to_string(), "ws2".to_string()]
|
|
);
|
|
assert_eq!(result.get("bucket-b").unwrap(), &vec!["ws3".to_string()]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_parse_bucket_restrictions_empty_string() {
|
|
assert!(parse_bucket_restrictions_from_str("").is_none());
|
|
assert!(parse_bucket_restrictions_from_str(" ").is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_parse_bucket_restrictions_invalid_format_skipped() {
|
|
// "no-colon" is invalid, only "valid:ws1" should be parsed
|
|
let result = parse_bucket_restrictions_from_str("no-colon;valid:ws1").unwrap();
|
|
assert_eq!(result.len(), 1);
|
|
assert!(result.contains_key("valid"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_parse_bucket_restrictions_trailing_semicolons() {
|
|
let result = parse_bucket_restrictions_from_str(";bucket:ws1;;").unwrap();
|
|
assert_eq!(result.len(), 1);
|
|
assert!(result.contains_key("bucket"));
|
|
}
|
|
|
|
// --- build_filesystem_client error path ---
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[test]
|
|
fn test_build_filesystem_client_nonexistent_path() {
|
|
// LocalFileSystem::new_with_prefix fails when the directory doesn't exist
|
|
let result = build_filesystem_client("/tmp/windmill_test_nonexistent_dir_12345_xyz");
|
|
assert!(result.is_err());
|
|
assert!(result
|
|
.unwrap_err()
|
|
.to_string()
|
|
.contains("Error building filesystem object store"));
|
|
}
|
|
|
|
// --- get_logs_from_store test ---
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[tokio::test]
|
|
async fn test_get_logs_from_store_with_filesystem() {
|
|
use futures::StreamExt;
|
|
use object_store::{path::Path, ObjectStore, PutPayload};
|
|
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let root = dir.path().to_str().unwrap().to_string();
|
|
let client = build_filesystem_client(&root).unwrap();
|
|
|
|
// Write two log chunk files
|
|
client
|
|
.put(
|
|
&Path::from("logs/chunk1.log"),
|
|
PutPayload::from(bytes::Bytes::from("chunk1 content\n")),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
client
|
|
.put(
|
|
&Path::from("logs/chunk2.log"),
|
|
PutPayload::from(bytes::Bytes::from("chunk2 content\n")),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Set the global OBJECT_STORE_SETTINGS to our filesystem store
|
|
{
|
|
let mut settings = OBJECT_STORE_SETTINGS.write().await;
|
|
*settings = Some(ExpirableObjectStore::from(client));
|
|
}
|
|
|
|
let file_index = Some(vec![
|
|
"logs/chunk1.log".to_string(),
|
|
"logs/chunk2.log".to_string(),
|
|
]);
|
|
let tail_logs = "tail logs here";
|
|
|
|
// log_offset > 0, file_index Some, object store set → should return a stream
|
|
let stream = get_logs_from_store(1, tail_logs, &file_index).await;
|
|
assert!(stream.is_some());
|
|
|
|
let chunks: Vec<bytes::Bytes> = stream
|
|
.unwrap()
|
|
.filter_map(|r| async { r.ok() })
|
|
.collect()
|
|
.await;
|
|
assert_eq!(chunks.len(), 3); // chunk1 + chunk2 + tail_logs
|
|
assert_eq!(chunks[0], bytes::Bytes::from("chunk1 content\n"));
|
|
assert_eq!(chunks[1], bytes::Bytes::from("chunk2 content\n"));
|
|
assert_eq!(chunks[2], bytes::Bytes::from("tail logs here"));
|
|
|
|
// Clean up global state
|
|
{
|
|
let mut settings = OBJECT_STORE_SETTINGS.write().await;
|
|
*settings = None;
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[tokio::test]
|
|
async fn test_get_logs_from_store_zero_offset() {
|
|
// log_offset == 0 → always returns None regardless of other params
|
|
let file_index = Some(vec!["some/file.log".to_string()]);
|
|
let result = get_logs_from_store(0, "logs", &file_index).await;
|
|
assert!(result.is_none());
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[tokio::test]
|
|
async fn test_get_logs_from_store_no_file_index() {
|
|
// file_index is None → returns None
|
|
let result = get_logs_from_store(1, "logs", &None).await;
|
|
assert!(result.is_none());
|
|
}
|
|
|
|
#[cfg(feature = "parquet")]
|
|
#[tokio::test]
|
|
async fn test_convert_json_line_stream_value_after_null_only_sample() {
|
|
use datafusion::datasource::file_format::DEFAULT_SCHEMA_INFER_MAX_RECORD as SAMPLE;
|
|
use futures::StreamExt;
|
|
|
|
// One more all-null row than the inference sample can see, so the column's type has
|
|
// to come from the row that follows it.
|
|
let total = SAMPLE + 1;
|
|
let rows = (0..total).map(|i| {
|
|
Ok::<_, anyhow::Error>(serde_json::json!({
|
|
"col": if i < SAMPLE { serde_json::Value::Null } else { serde_json::json!("2023-11-30") }
|
|
}))
|
|
});
|
|
|
|
let (mut out, stats) =
|
|
convert_json_line_stream(futures::stream::iter(rows), S3ModeFormat::Json, None)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(stats.rows, total as u64);
|
|
|
|
let mut bytes = Vec::new();
|
|
while let Some(chunk) = out.next().await {
|
|
bytes.extend_from_slice(&chunk.expect("row after the null-only sample must decode"));
|
|
}
|
|
let parsed: Vec<serde_json::Value> = serde_json::from_slice(&bytes).unwrap();
|
|
assert_eq!(parsed.len(), total);
|
|
assert_eq!(parsed[total - 1]["col"], serde_json::json!("2023-11-30"));
|
|
}
|
|
}
|