use crate::cloud_sources::{S3Settings, Signing, Source};
use crate::config::CloudConfig;
use std::path::{Path, PathBuf};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProviderKind {
Gcs,
S3,
Azure,
}
impl ProviderKind {
pub fn scheme(self) -> &'static str {
match self {
ProviderKind::Gcs => "gs",
ProviderKind::S3 => "s3",
ProviderKind::Azure => "abfss",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Provider {
pub kind: ProviderKind,
pub label: String,
pub note: String,
pub project: Option<String>,
pub profile: Option<String>,
pub endpoint: Option<String>,
}
impl Provider {
pub fn detail(&self) -> Option<String> {
match (&self.project, &self.profile) {
(Some(project), _) => Some(format!("project: {project}")),
(None, Some(profile)) => Some(format!("profile: {profile}")),
(None, None) => None,
}
}
pub fn can_list_buckets(&self) -> bool {
match self.kind {
ProviderKind::Gcs => self.project.is_some(),
ProviderKind::S3 | ProviderKind::Azure => true,
}
}
}
pub struct Environment<'a> {
pub var: &'a dyn Fn(&str) -> Option<String>,
pub exists: &'a dyn Fn(&Path) -> bool,
pub read: &'a dyn Fn(&Path) -> Option<String>,
pub home: Option<PathBuf>,
pub windows: bool,
pub run: &'a crate::cloud_command::Runner<'a>,
pub all_vars: &'a dyn Fn() -> Vec<(String, String)>,
pub list: &'a dyn Fn(&Path) -> Vec<PathBuf>,
}
impl Environment<'_> {
pub fn current() -> Environment<'static> {
Environment {
var: &crate::cloud_env::var,
exists: &|path| path.exists(),
read: &|path| std::fs::read_to_string(path).ok(),
home: dirs::home_dir(),
windows: cfg!(windows),
run: &|program, args| {
crate::cloud_command::run(program, args, crate::cloud_command::CREDENTIAL_TIMEOUT)
},
all_vars: &crate::cloud_env::vars,
list: &|dir| {
std::fs::read_dir(dir)
.map(|entries| entries.flatten().map(|e| e.path()).collect())
.unwrap_or_default()
},
}
}
}
pub fn detect(config: &CloudConfig, env: &Environment<'_>) -> Vec<Provider> {
let mut providers = Vec::new();
if let Some(gcs) = detect_gcs(config, env) {
providers.push(gcs);
}
if let Some(s3) = detect_s3(config, env) {
providers.push(s3);
}
providers
}
fn detect_gcs(config: &CloudConfig, env: &Environment<'_>) -> Option<Provider> {
let note = if (env.var)("GOOGLE_SERVICE_ACCOUNT").is_some()
|| (env.var)("GOOGLE_SERVICE_ACCOUNT_PATH").is_some()
{
"service account"
} else if (env.var)("GOOGLE_SERVICE_ACCOUNT_KEY").is_some() {
"service account key"
} else if (env.var)("GOOGLE_APPLICATION_CREDENTIALS").is_some() {
"GOOGLE_APPLICATION_CREDENTIALS"
} else if adc_path(env).is_some() {
"gcloud"
} else if instance_identity(config, env).gcp {
"instance identity"
} else {
return None;
};
Some(Provider {
kind: ProviderKind::Gcs,
label: "Google Cloud Storage".to_string(),
note: note.to_string(),
project: gcp_project(env),
profile: None,
endpoint: None,
})
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct InstanceIdentity {
pub aws: bool,
pub gcp: bool,
pub azure: bool,
}
pub fn instance_identity(config: &CloudConfig, env: &Environment<'_>) -> InstanceIdentity {
let opted_in = config.instance_identity;
let set = |key: &str| (env.var)(key).is_some_and(|v| !v.trim().is_empty());
InstanceIdentity {
aws: opted_in,
gcp: opted_in || set("K_SERVICE"),
azure: opted_in || set("IDENTITY_ENDPOINT") || set("MSI_ENDPOINT"),
}
}
pub fn unreadable_google_login(env: &Environment<'_>) -> Option<String> {
let path = (env.var)("GOOGLE_APPLICATION_CREDENTIALS")
.map(PathBuf::from)
.or_else(|| adc_path(env))?;
crate::gcloud::unsupported_credential_type(&(env.read)(&path)?)
}
pub fn adc_path(env: &Environment<'_>) -> Option<PathBuf> {
const FILE: &str = "application_default_credentials.json";
let path = if env.windows {
PathBuf::from((env.var)("APPDATA")?)
.join("gcloud")
.join(FILE)
} else {
env.home.as_ref()?.join(".config").join("gcloud").join(FILE)
};
(env.exists)(&path).then_some(path)
}
pub(crate) fn gcp_project(env: &Environment<'_>) -> Option<String> {
for key in [
"DATUI_GCP_PROJECT",
"GOOGLE_CLOUD_PROJECT",
"GCLOUD_PROJECT",
"CLOUDSDK_CORE_PROJECT",
"GCP_PROJECT",
] {
if let Some(value) = (env.var)(key) {
let value = value.trim().to_string();
if !value.is_empty() {
return Some(value);
}
}
}
adc_quota_project(env)
}
fn adc_quota_project(env: &Environment<'_>) -> Option<String> {
let path = adc_path(env)?;
let contents = (env.read)(&path)?;
let value: serde_json::Value = serde_json::from_str(&contents).ok()?;
let project = value.get("quota_project_id")?.as_str()?.trim();
if project.is_empty() {
return None;
}
Some(project.to_string())
}
fn detect_s3(config: &CloudConfig, env: &Environment<'_>) -> Option<Provider> {
let endpoint = config.s3_endpoint_url.clone();
let configured_keys = config.s3_access_key_id.is_some();
let env_keys = (env.var)("AWS_ACCESS_KEY_ID").is_some();
let profile = (env.var)("AWS_PROFILE");
let shared_credentials = env.home.as_ref().is_some_and(|home| {
(env.exists)(&home.join(".aws/credentials")) || (env.exists)(&home.join(".aws/config"))
});
let container_role = (env.var)("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI").is_some()
|| (env.var)("AWS_CONTAINER_CREDENTIALS_FULL_URI").is_some();
let web_identity = (env.var)("AWS_WEB_IDENTITY_TOKEN_FILE").is_some();
let note = if configured_keys {
"datui config"
} else if env_keys {
"AWS_ACCESS_KEY_ID"
} else if profile.is_some() {
"AWS_PROFILE"
} else if container_role {
"container role"
} else if web_identity {
"web identity"
} else if shared_credentials {
"~/.aws"
} else if instance_identity(config, env).aws {
"instance role"
} else {
return None;
};
let label = match endpoint.as_deref().and_then(endpoint_host) {
Some(host) => format!("S3-compatible ({host})"),
None => "Amazon S3".to_string(),
};
Some(Provider {
kind: ProviderKind::S3,
label,
note: note.to_string(),
project: None,
profile,
endpoint,
})
}
fn endpoint_host(endpoint: &str) -> Option<String> {
let rest = endpoint
.split_once("://")
.map(|(_, rest)| rest)
.unwrap_or(endpoint);
let host = rest.split(['/', '?', '#']).next()?.trim();
if host.is_empty() {
return None;
}
Some(host.to_string())
}
pub fn parse_gcs_buckets(body: &str) -> Result<Vec<String>, String> {
let value: serde_json::Value =
serde_json::from_str(body).map_err(|e| format!("not JSON: {e}"))?;
if let Some(message) = value
.get("error")
.and_then(|e| e.get("message"))
.and_then(|m| m.as_str())
{
return Err(message.to_string());
}
let Some(items) = value.get("items").and_then(|i| i.as_array()) else {
return Ok(Vec::new());
};
Ok(items
.iter()
.filter_map(|item| item.get("name").and_then(|n| n.as_str()))
.filter(|name| !name.is_empty())
.map(str::to_string)
.collect())
}
pub fn gcs_next_page_token(body: &str) -> Option<String> {
serde_json::from_str::<serde_json::Value>(body)
.ok()?
.get("nextPageToken")?
.as_str()
.map(str::to_string)
}
pub fn parse_s3_buckets(body: &str) -> Result<Vec<String>, String> {
use quick_xml::events::Event;
let mut reader = quick_xml::Reader::from_str(body);
reader.config_mut().trim_text(true);
let mut buckets = Vec::new();
let mut path: Vec<Vec<u8>> = Vec::new();
let mut error_message: Option<String> = None;
loop {
match reader.read_event() {
Ok(Event::Start(tag)) => {
if path.len() >= MAX_XML_DEPTH {
return Err("response nested implausibly deeply".to_string());
}
path.push(tag.local_name().as_ref().as_bytes().to_vec());
}
Ok(Event::End(_)) => {
path.pop();
}
Ok(Event::Text(text)) => {
let value = text.xml10_content().into_owned();
match path_tail(&path) {
(Some(b"Name"), Some(b"Bucket")) if !value.is_empty() => {
buckets.push(value);
}
(Some(b"Message"), Some(b"Error")) => error_message = Some(value),
_ => {}
}
}
Ok(Event::Eof) => {
if !path.is_empty() {
return Err("response ended inside an element".to_string());
}
break;
}
Err(e) => return Err(format!("malformed XML: {e}")),
_ => {}
}
}
if let Some(message) = error_message {
return Err(message);
}
Ok(buckets)
}
const MAX_XML_DEPTH: usize = 32;
fn path_tail(path: &[Vec<u8>]) -> (Option<&[u8]>, Option<&[u8]>) {
let len = path.len();
let last = len.checked_sub(1).map(|i| path[i].as_slice());
let parent = len.checked_sub(2).map(|i| path[i].as_slice());
(last, parent)
}
const REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20);
pub(crate) fn http_agent() -> ureq::Agent {
crate::user_agent::ureq_config()
.timeout_global(Some(REQUEST_TIMEOUT))
.build()
.into()
}
pub fn s3_builder(bucket: &str, settings: &S3Settings) -> object_store::aws::AmazonS3Builder {
let builder = if settings.from_env && !settings.skip_signature {
object_store::aws::AmazonS3Builder::from_env()
} else {
object_store::aws::AmazonS3Builder::new()
};
let mut builder = builder.with_bucket_name(bucket).with_config(
object_store::aws::AmazonS3ConfigKey::Client(crate::user_agent::CLIENT_KEY),
crate::user_agent::get(),
);
if settings.skip_signature {
builder = builder.with_skip_signature(true);
}
if let Some(endpoint) = &settings.endpoint {
builder = builder.with_endpoint(endpoint.clone());
if endpoint.starts_with("http://") {
builder = builder.with_allow_http(true);
}
}
if settings.endpoint.is_some() || settings.virtual_hosted.is_some() {
builder = builder.with_virtual_hosted_style_request(settings.virtual_hosted_style());
}
if let Some(region) = &settings.region {
builder = builder.with_region(region.clone());
}
if let Some(key) = &settings.access_key_id {
builder = builder.with_access_key_id(key.clone());
}
if let Some(secret) = &settings.secret_access_key {
builder = builder.with_secret_access_key(secret.clone());
}
if let Some(token) = &settings.session_token {
builder = builder.with_token(token.clone());
}
builder
}
pub fn object_path(key: &str) -> object_store::path::Path {
object_store::path::Path::parse(key).unwrap_or_else(|_| object_store::path::Path::from(key))
}
pub fn store_for_bucket(
kind: ProviderKind,
bucket: &str,
settings: &S3Settings,
unsigned: bool,
google_token: Option<&str>,
google_credentials: Option<&Path>,
) -> Result<std::sync::Arc<dyn object_store::ObjectStore>, String> {
match kind {
ProviderKind::Gcs => Ok(std::sync::Arc::new(gcs_store(
bucket,
unsigned,
google_token,
google_credentials,
)?)),
ProviderKind::S3 => {
let store = s3_builder(bucket, settings)
.build()
.map_err(|e| format!("S3 is not configured: {e}"))?;
Ok(std::sync::Arc::new(store))
}
ProviderKind::Azure => Err("an Azure container needs its account".to_string()),
}
}
pub fn pager_for_bucket(
kind: ProviderKind,
bucket: &str,
settings: &S3Settings,
unsigned: bool,
google_token: Option<&str>,
google_credentials: Option<&Path>,
) -> Result<std::sync::Arc<dyn object_store::list::PaginatedListStore>, String> {
match kind {
ProviderKind::Gcs => Ok(std::sync::Arc::new(gcs_store(
bucket,
unsigned,
google_token,
google_credentials,
)?)),
ProviderKind::S3 => {
let store = s3_builder(bucket, settings)
.build()
.map_err(|e| format!("S3 is not configured: {e}"))?;
Ok(std::sync::Arc::new(store))
}
ProviderKind::Azure => Err("an Azure container needs its account".to_string()),
}
}
fn gcs_store(
bucket: &str,
unsigned: bool,
google_token: Option<&str>,
google_credentials: Option<&Path>,
) -> Result<object_store::gcp::GoogleCloudStorage, String> {
let builder = match (unsigned, google_token) {
(true, _) => object_store::gcp::GoogleCloudStorageBuilder::new().with_skip_signature(true),
(false, Some(token)) => object_store::gcp::GoogleCloudStorageBuilder::new()
.with_credentials(std::sync::Arc::new(
object_store::StaticCredentialProvider::new(object_store::gcp::GcpCredential {
bearer: token.to_string(),
}),
)),
(false, None) => match google_credentials {
Some(file) => object_store::gcp::GoogleCloudStorageBuilder::new()
.with_application_credentials(file.to_string_lossy()),
None => object_store::gcp::GoogleCloudStorageBuilder::from_env(),
},
};
builder
.with_bucket_name(bucket)
.with_config(
object_store::gcp::GoogleConfigKey::Client(crate::user_agent::CLIENT_KEY),
crate::user_agent::get(),
)
.build()
.map_err(|e| format!("Google Cloud Storage is not configured: {e}"))
}
const PEEK_KEYS: usize = 100;
pub async fn peek_kind(
url: &str,
config: &CloudConfig,
) -> Result<(crate::discover::EntryKind, crate::discover::Holds), String> {
let resolved = {
let (url, config) = (url.to_string(), config.clone());
tokio::task::spawn_blocking(move || crate::cloud_sources::resolve(&url, &config))
.await
.map_err(|e| format!("{e}"))??
};
match peek_page(&resolved).await {
Err(refused) if resolved.signing == Signing::Try && is_refusal(&refused) => {
peek_page(&resolved.unsigned()).await.map_err(|_| refused)
}
Err(refused) if is_refusal(&refused) && resolved.login_error.is_some() => {
Err(resolved.login_error.clone().unwrap_or(refused))
}
other => other,
}
}
async fn peek_page(
resolved: &crate::cloud_sources::Resolved,
) -> Result<(crate::discover::EntryKind, crate::discover::Holds), String> {
use object_store::list::{PaginatedListOptions, PaginatedListStore};
let (store, prefix): (std::sync::Arc<dyn PaginatedListStore>, String) =
if let Some((account, container, key)) = crate::source::azure_parts(&resolved.url) {
(
crate::azure::paginated_store(&account, &container, &resolved.azure)?,
key,
)
} else {
let (kind, bucket, key) = split_bucket_url(&resolved.url)
.ok_or_else(|| format!("not an object-store URL: {}", resolved.url))?;
let store = pager_for_bucket(
kind,
&bucket,
&resolved.s3,
resolved.signing == Signing::Unsigned,
resolved.gcloud.as_ref().map(|(_, token)| token.as_str()),
resolved.google_credentials.as_deref(),
)?;
(store, key)
};
let prefix = format!("{}/", prefix.trim_matches('/'));
let page = store
.list_paginated(
Some(&prefix),
PaginatedListOptions {
delimiter: Some("/".into()),
max_keys: Some(PEEK_KEYS),
..Default::default()
},
)
.await
.map_err(|e| format!("{e}"))?;
let directories: Vec<String> = page
.result
.common_prefixes
.iter()
.map(|p| p.as_ref().to_string())
.collect();
let objects: Vec<(String, u64)> = page
.result
.objects
.iter()
.map(|o| (o.location.as_ref().to_string(), o.size))
.collect();
let (kind, holds) = look_at_page(&prefix, &directories, &objects, page.page_token.as_deref());
if kind != crate::discover::EntryKind::MultiFile {
return Ok((kind, holds));
}
Ok((
verified_kind(resolved, &objects).await.unwrap_or(kind),
holds,
))
}
const VERIFY_FOOTERS: usize = 3;
fn footers_to_verify(files: usize) -> Vec<usize> {
if files <= VERIFY_FOOTERS {
(0..files).collect()
} else {
vec![0, files / 2, files - 1]
}
}
async fn verified_kind(
resolved: &crate::cloud_sources::Resolved,
objects: &[(String, u64)],
) -> Option<crate::discover::EntryKind> {
let store = match crate::source::azure_parts(&resolved.url) {
Some((account, container, _)) => {
crate::azure::store(&account, &container, &resolved.azure).ok()?
}
None => {
let (provider, bucket, _) = split_bucket_url(&resolved.url)?;
store_for_bucket(
provider,
&bucket,
&resolved.s3,
resolved.signing == Signing::Unsigned,
resolved.gcloud.as_ref().map(|(_, token)| token.as_str()),
resolved.google_credentials.as_deref(),
)
.ok()?
}
};
kind_from_footers(&store, objects).await
}
async fn kind_from_footers(
store: &std::sync::Arc<dyn object_store::ObjectStore>,
objects: &[(String, u64)],
) -> Option<crate::discover::EntryKind> {
let parquet: Vec<&(String, u64)> = objects
.iter()
.filter(|(key, _)| crate::discover::is_parquet_key(key))
.collect();
if parquet.len() < 2 {
return None;
}
let picks = footers_to_verify(parquet.len());
let meter = std::sync::Arc::new(crate::measurements::Meter::default());
let store = store.clone();
let mut reads = tokio::task::JoinSet::new();
for index in picks {
let (key, size) = parquet[index].clone();
let (store, meter) = (store.clone(), meter.clone());
reads.spawn(async move {
let file = crate::dataset_files::DatasetFile {
key,
size,
stamp: 0,
etag: None,
};
crate::cloud_hive::footer_of_file(&store, &file, &meter)
.await
.ok()
});
}
let mut per_file: Vec<Vec<String>> = Vec::new();
while let Some(joined) = reads.join_next().await {
if let Ok(Some(footer)) = joined {
per_file.push(footer.schema.iter_names().map(|n| n.to_string()).collect());
}
}
if per_file.len() < 2 {
return None;
}
Some(if crate::schema_union::is_nested(&per_file) {
crate::discover::EntryKind::MultiFile
} else {
crate::discover::EntryKind::Directory
})
}
fn look_at_page(
prefix: &str,
directories: &[String],
objects: &[(String, u64)],
next_page: Option<&str>,
) -> (crate::discover::EntryKind, crate::discover::Holds) {
let (kind, mut holds) = look_at_listing(prefix, directories, objects);
holds.truncated = next_page.is_some();
(kind, holds)
}
pub fn look_at_listing(
prefix: &str,
directories: &[String],
objects: &[(String, u64)],
) -> (crate::discover::EntryKind, crate::discover::Holds) {
use crate::discover::EntryKind;
let last = |key: &str| {
key.trim_end_matches('/')
.rsplit('/')
.next()
.unwrap_or("")
.to_string()
};
let here = prefix.trim_matches('/');
let objects: Vec<(String, u64)> = objects
.iter()
.filter(|(key, _)| here.is_empty() || key.trim_matches('/') != here)
.cloned()
.collect();
let objects = objects.as_slice();
let counted: Vec<&String> = directories
.iter()
.filter(|f| !crate::discover::is_bookkeeping(&last(f)))
.collect();
let partitions = counted
.iter()
.filter(|f| crate::discover::is_partition_name(&last(f)))
.count();
let present: Vec<&String> = objects
.iter()
.filter(|(key, size)| {
let name = last(key);
!name.is_empty()
&& !crate::discover::is_bookkeeping(&name)
&& !is_empty_marker(&name, *size)
&& !(*size == 0 && directories.iter().any(|f| last(f) == name))
})
.map(|(key, _)| key)
.collect();
let format_of = |key: &str| crate::discover::data_format(std::path::Path::new(key));
let other_data = present.iter().any(|key| {
format_of(key).is_some_and(|f| !f.is_lines()) || crate::discover::is_parquet_key(key)
});
let files: Vec<&&String> = present
.iter()
.filter(|key| {
format_of(key).is_some_and(|f| !(other_data && f.is_lines()))
|| crate::discover::is_parquet_key(key)
})
.collect();
let mut counts: Vec<(&'static str, usize)> = Vec::new();
for key in &files {
let name = last(key);
let format = crate::discover::data_format(std::path::Path::new(name.as_str()))
.map(|f| f.name())
.unwrap_or("parquet");
match counts.iter_mut().find(|(f, _)| *f == format) {
Some((_, n)) => *n += 1,
None => counts.push((format, 1)),
}
}
let arrow = counts
.iter()
.any(|(f, _)| *f == crate::FileFormat::Arrow.name());
let dataset_dict = !counted.is_empty()
&& files
.iter()
.any(|key| last(key) == crate::hf_splits::DATASET_DICT);
let hugging_face: Vec<String> = files
.iter()
.map(|key| last(key))
.filter(|name| {
(arrow && crate::discover::is_hugging_face_metadata(name))
|| (dataset_dict && name == crate::hf_splits::DATASET_DICT)
})
.collect();
let json = crate::FileFormat::Json.name();
for (format, n) in &mut counts {
if *format == json {
*n -= hugging_face.len();
}
}
counts.retain(|(_, n)| *n > 0);
counts.sort_by(|a, b| crate::discover::rank_formats((a.0, a.1), (b.0, b.1)));
let orphan_markers: Vec<String> = objects
.iter()
.filter(|(key, size)| {
let name = last(key);
!name.is_empty()
&& is_empty_marker(&name, *size)
&& !crate::discover::is_bookkeeping(&name)
&& !directories.iter().any(|f| last(f) == name)
})
.map(|(key, _)| last(key))
.collect();
let mut skipped_names: Vec<String> = objects
.iter()
.filter(|(key, size)| !(*size == 0 && directories.iter().any(|f| last(f) == last(key))))
.map(|(key, _)| last(key))
.chain(directories.iter().map(|f| last(f)))
.filter(|name| !name.is_empty() && crate::discover::is_bookkeeping(name))
.chain(hugging_face)
.collect();
skipped_names.sort();
skipped_names.dedup();
let skipped = skipped_names.len();
skipped_names.truncate(crate::discover::SKIPPED_NAMES_SHOWN);
let unnamed = present
.iter()
.filter(|key| !files.contains(key))
.filter(|key| crate::discover::has_no_extension(std::path::Path::new(&last(key))))
.count();
let holds = crate::discover::Holds {
formats: counts
.into_iter()
.map(|(f, n)| (f.to_string(), n))
.collect(),
directories: counted.len(),
partitions,
not_read: present.len() - files.len() - unnamed + orphan_markers.len(),
unnamed,
skipped,
skipped_names,
truncated: false,
dataset_dict,
};
let directory = |name: &str| directories.iter().any(|f| last(f) == name);
if directory("_delta_log") {
return (EntryKind::Delta, holds);
}
if directory(".hoodie") {
return (EntryKind::Hudi, holds);
}
let parquet = files
.iter()
.filter(|key| crate::discover::is_parquet_key(key))
.count();
if directory("metadata") && directory("data") && parquet == 0 {
return (EntryKind::Iceberg, holds);
}
let seen = counted.len() + present.len() + orphan_markers.len();
if partitions > 0 && partitions >= files.len() {
return (EntryKind::Hive, holds);
}
let kind = if parquet > 1 && parquet == files.len() && parquet * 2 >= seen {
EntryKind::MultiFile
} else {
EntryKind::Directory
};
(kind, holds)
}
pub fn split_bucket_url(url: &str) -> Option<(ProviderKind, String, String)> {
let (_, plain) = crate::source::split_source_id(url);
let (scheme, rest) = plain.split_once("://")?;
let kind = match scheme {
"gs" | "gcs" => ProviderKind::Gcs,
"s3" | "s3a" => ProviderKind::S3,
_ => return None,
};
let rest = rest.trim_end_matches('/');
let (bucket, prefix) = match rest.split_once('/') {
Some((bucket, prefix)) => (bucket, prefix),
None => (rest, ""),
};
if bucket.is_empty() {
return None;
}
Some((
kind,
bucket.to_string(),
prefix.trim_matches('/').to_string(),
))
}
pub const MAX_LEVEL_ROWS: usize = crate::discover::MAX_ENTRIES_PER_DIR;
#[derive(Debug, Clone, Default)]
pub struct Level {
pub rows: Vec<crate::discover::Entry>,
pub truncated: bool,
pub cancelled: bool,
}
pub type Progress = std::sync::Arc<dyn Fn(&[crate::discover::Entry]) + Send + Sync>;
#[derive(Clone, Default)]
pub struct Watch {
pub progress: Option<Progress>,
pub cancelled: std::sync::Arc<std::sync::atomic::AtomicBool>,
pub names_from: Option<String>,
}
impl Watch {
fn cancelled(&self) -> bool {
self.cancelled.load(std::sync::atomic::Ordering::Relaxed)
}
}
pub async fn list_objects(
url: &str,
config: &CloudConfig,
) -> Result<Vec<crate::discover::Entry>, String> {
list_objects_watched(url, config, &Watch::default())
.await
.map(|level| level.rows)
}
pub async fn list_objects_watched(
url: &str,
config: &CloudConfig,
watch: &Watch,
) -> Result<Level, String> {
let resolved = {
let (url, config) = (url.to_string(), config.clone());
tokio::task::spawn_blocking(move || crate::cloud_sources::resolve(&url, &config))
.await
.map_err(|e| format!("{e}"))??
};
let signing = resolved.signing;
let place = resolved.place.clone();
let listed = list_level(url, &resolved, watch).await;
let (listed, resolved) = match listed {
Err(refusal)
if resolved.kind == ProviderKind::Azure
&& crate::azure::is_permission_mismatch(&refusal)
&& resolved.azure.identity.is_some() =>
{
let enabled = config.use_azure_account_keys;
let keyed = {
let (resolved, refusal) = (resolved.clone(), refusal.clone());
tokio::task::spawn_blocking(move || {
let (account, _, _) =
crate::source::azure_parts(&resolved.url).ok_or_else(|| refusal.clone())?;
crate::azure::with_account_key(
&account,
&resolved.azure,
&refusal,
enabled,
&Environment::current(),
)
.map(|azure| crate::cloud_sources::Resolved { azure, ..resolved })
})
.await
.map_err(|e| format!("{e}"))?
};
match keyed {
Ok(keyed) => (list_level(url, &keyed, watch).await, keyed),
Err(why) => (Err(why), resolved),
}
}
other => (other, resolved),
};
if resolved.kind == ProviderKind::Azure
&& listed.is_ok()
&& matches!(resolved.azure.auth, crate::azure::AzureAuth::Bearer(_))
&& let Some((account, _, _)) = crate::source::azure_parts(&resolved.url)
{
crate::azure::remember_token_reads(&account);
}
match listed {
Err(refused) if signing == Signing::Try && is_refusal(&refused) => {
let level = list_level(url, &resolved.unsigned(), watch)
.await
.map_err(|_| refused)?;
crate::cloud_sources::remember_access(&place, true);
Ok(level)
}
Ok(level) => {
if signing == Signing::Try {
crate::cloud_sources::remember_access(&place, false);
}
Ok(level)
}
Err(refused) if is_refusal(&refused) && resolved.login_error.is_some() => {
Err(resolved.login_error.clone().unwrap_or(refused))
}
Err(e) => Err(e),
}
}
pub fn narrowing_prefix(filter: &str, names: &[&str]) -> Option<String> {
let filter = filter.trim();
if filter.is_empty() || filter.contains('/') {
return None;
}
let first = names.first()?;
let mut common = first.len();
for name in &names[1..] {
common = common.min(
first
.bytes()
.zip(name.bytes())
.take_while(|(a, b)| a == b)
.count(),
);
}
while !first.is_char_boundary(common) {
common -= 1;
}
let shared = first[..common]
.rfind(|c: char| !c.is_alphanumeric())
.map_or("", |at| &first[..=at]);
let typed = match filter.get(..shared.len()) {
Some(head) if !shared.is_empty() && head.eq_ignore_ascii_case(shared) => {
&filter[shared.len()..]
}
_ => filter,
};
let rest = names.iter().flat_map(|n| n[shared.len()..].chars());
let (mut upper, mut lower) = (false, false);
for c in rest {
upper |= c.is_uppercase();
lower |= c.is_lowercase();
}
let typed = match (upper, lower) {
(true, false) => typed.to_uppercase(),
(false, true) => typed.to_lowercase(),
_ => typed.to_string(),
};
Some(format!("{shared}{typed}"))
}
pub fn is_refusal(error: &str) -> bool {
let lower = error.to_ascii_lowercase();
[
"403",
"401",
"forbidden",
"unauthorized",
"accessdenied",
"access denied",
"permissiondenied",
"authorizationfailure",
"authenticationfailed",
"invalidauthenticationinfo",
"noauthenticationinformation",
]
.iter()
.any(|word| lower.contains(word))
}
pub fn is_marker(name: &str) -> bool {
name == "_SUCCESS"
|| name.starts_with("_committed_")
|| name.starts_with("_started_")
|| name.ends_with("_$folder$")
}
pub fn is_empty_marker(name: &str, size: u64) -> bool {
size == 0 && !name.contains('.')
}
fn is_listed_object(location: &str, size: u64, prefix: &str, prefixes: &[String]) -> bool {
let name = location.rsplit('/').next().unwrap_or(location);
!(name.is_empty()
|| is_marker(name)
|| crate::azure::is_folder_marker(location, size, prefixes)
|| is_empty_marker(name, size)
|| location.trim_end_matches('/') == prefix)
}
async fn list_level(
url: &str,
resolved: &crate::cloud_sources::Resolved,
watch: &Watch,
) -> Result<Level, String> {
if resolved.kind == ProviderKind::Azure {
return list_azure_objects(resolved, watch).await;
}
let (kind, bucket, prefix) =
split_bucket_url(&resolved.url).ok_or_else(|| format!("not an object-store URL: {url}"))?;
let pager = pager_for_bucket(
kind,
&bucket,
&resolved.s3,
resolved.signing == Signing::Unsigned,
resolved.gcloud.as_ref().map(|(_, token)| token.as_str()),
resolved.google_credentials.as_deref(),
)?;
let base = match crate::source::split_source_id(url).0 {
Some(id) => format!("{}://{id}@{bucket}", kind.scheme()),
None => format!("{}://{bucket}", kind.scheme()),
};
list_pages(pager.as_ref(), &prefix, watch, |result| {
let prefixes: Vec<String> = result
.common_prefixes
.iter()
.map(|p| p.as_ref().to_string())
.collect();
let directories = result
.common_prefixes
.iter()
.map(|common| {
let name = common
.as_ref()
.rsplit('/')
.find(|part| !part.is_empty())
.unwrap_or(common.as_ref())
.to_string();
let mut row = crate::discover::Entry::directory(Path::new(&format!(
"{base}/{}",
common.as_ref()
)));
row.name = name;
row
})
.collect();
let objects = result
.objects
.into_iter()
.filter(|object| {
is_listed_object(object.location.as_ref(), object.size, &prefix, &prefixes)
})
.map(|object| {
let location = object.location.as_ref().to_string();
let name = location.rsplit('/').next().unwrap_or(&location).to_string();
let path = PathBuf::from(format!("{base}/{location}"));
let kind = if crate::discover::unreadable_by_name(&path) {
crate::discover::EntryKind::Other
} else {
crate::discover::EntryKind::File
};
object_row(path, kind, name, &object)
})
.collect();
(directories, objects)
})
.await
}
fn object_row(
path: PathBuf,
kind: crate::discover::EntryKind,
name: String,
object: &object_store::ObjectMeta,
) -> crate::discover::Entry {
crate::discover::Entry {
path,
kind,
name,
size: Some(object.size),
modified: Some(object.last_modified.into()),
rows: None,
cols: None,
cols_sampled: false,
columns: Vec::new(),
cost: Default::default(),
holds: Default::default(),
opens_whole_directory: false,
format_spec: None,
table: None,
}
}
async fn list_pages(
pager: &dyn object_store::list::PaginatedListStore,
prefix: &str,
watch: &Watch,
mut rows_of: impl FnMut(
object_store::ListResult,
) -> (Vec<crate::discover::Entry>, Vec<crate::discover::Entry>),
) -> Result<Level, String> {
let mut key_prefix = if prefix.is_empty() {
String::new()
} else {
format!("{prefix}/")
};
if let Some(names) = &watch.names_from {
key_prefix.push_str(names);
}
let (mut directories, mut objects) = (Vec::new(), Vec::new());
let mut token = None;
let rows = |directories: &[crate::discover::Entry], objects: &[crate::discover::Entry]| {
let mut rows = directories.to_vec();
rows.extend_from_slice(objects);
rows
};
loop {
if watch.cancelled() {
return Ok(Level {
rows: rows(&directories, &objects),
truncated: false,
cancelled: true,
});
}
let page = pager
.list_paginated(
(!key_prefix.is_empty()).then_some(key_prefix.as_str()),
object_store::list::PaginatedListOptions {
delimiter: Some("/".into()),
page_token: token.take(),
..Default::default()
},
)
.await
.map_err(|e| format!("{e}"))?;
let (more_directories, more_objects) = rows_of(page.result);
directories.extend(more_directories);
objects.extend(more_objects);
let mut listed = rows(&directories, &objects);
if listed.len() > MAX_LEVEL_ROWS {
listed.truncate(MAX_LEVEL_ROWS);
return Ok(Level {
rows: listed,
truncated: true,
cancelled: false,
});
}
match page.page_token {
Some(next) => {
if let Some(progress) = &watch.progress {
progress(&listed);
}
token = Some(next);
}
None => {
return Ok(Level {
rows: listed,
truncated: false,
cancelled: false,
});
}
}
}
}
async fn list_azure_objects(
resolved: &crate::cloud_sources::Resolved,
watch: &Watch,
) -> Result<Level, String> {
let (account, container, prefix) = crate::source::azure_parts(&resolved.url)
.ok_or_else(|| format!("not an Azure URL: {}", resolved.url))?;
let pager = crate::azure::paginated_store(&account, &container, &resolved.azure)?;
let prefix = prefix.trim_matches('/').to_string();
list_pages(pager.as_ref(), &prefix, watch, |result| {
let prefixes: Vec<String> = result
.common_prefixes
.iter()
.map(|p| p.as_ref().to_string())
.collect();
let directories = prefixes
.iter()
.map(|common| {
let name = common
.rsplit('/')
.find(|part| !part.is_empty())
.unwrap_or(common)
.to_string();
let mut row = crate::discover::Entry::directory(Path::new(
&crate::source::azure_url(&account, &container, &format!("{common}/")),
));
row.name = name;
row
})
.collect();
let objects = result
.objects
.into_iter()
.filter_map(|object| {
let location = object.location.as_ref().to_string();
let name = location.rsplit('/').next().unwrap_or(&location).to_string();
if name.is_empty()
|| is_marker(&name)
|| crate::azure::is_folder_marker(&location, object.size, &prefixes)
|| is_empty_marker(&name, object.size)
|| (object.size == 0 && location.trim_end_matches('/') == prefix)
{
return None;
}
let path = PathBuf::from(crate::source::azure_url(&account, &container, &location));
Some(object_row(
path,
crate::discover::EntryKind::File,
name,
&object,
))
})
.collect();
(directories, objects)
})
.await
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Listed {
pub name: String,
pub place: PathBuf,
pub details: Vec<(String, String)>,
}
pub async fn list_first_level(source: &Source) -> Result<Vec<Listed>, String> {
if source.kind == ProviderKind::Gcs {
let source = source.clone();
return tokio::task::spawn_blocking(move || {
tokio::runtime::Handle::current().block_on(list_gcs_projects(&source))
})
.await
.map_err(|e| format!("{e}"))?;
}
if source.kind != ProviderKind::Azure {
let blocking = source.clone();
let names = tokio::task::spawn_blocking(move || {
tokio::runtime::Handle::current().block_on(list_buckets(&blocking))
})
.await
.map_err(|e| format!("{e}"))??;
return Ok(names
.into_iter()
.map(|name| Listed {
place: PathBuf::from(source.bucket_url(&name)),
name,
details: Vec::new(),
})
.collect());
}
let source = source.clone();
tokio::task::spawn_blocking(move || {
if let Some(problem) = &source.problem {
return Err(problem.clone());
}
if let Some(account) = &source.azure.account {
return Ok(vec![Listed {
name: account.clone(),
place: PathBuf::from(source.bucket_url(account)),
details: Vec::new(),
}]);
}
let accounts =
crate::azure::discover_accounts(&source.azure.auth, &Environment::current())?;
Ok(accounts
.into_iter()
.map(|account| {
let mut details = Vec::new();
if let Some(subscription) = account.subscription {
details.push(("subscription".to_string(), subscription));
}
if let Some(location) = account.location {
details.push(("region".to_string(), location));
}
let namespace = if account.hierarchical_namespace {
"hierarchical"
} else {
"flat"
};
details.push(("namespace".to_string(), namespace.to_string()));
if account.private_network {
details.push(("network".to_string(), "private".to_string()));
}
if !account.shared_key_access {
details.push(("shared keys".to_string(), "disabled".to_string()));
}
Listed {
place: PathBuf::from(source.bucket_url(&account.name)),
name: account.name,
details,
}
})
.collect())
})
.await
.map_err(|e| format!("{e}"))?
}
pub fn s3_bucket_region(bucket: &str) -> Option<String> {
static REGIONS: std::sync::OnceLock<
std::sync::Mutex<std::collections::HashMap<String, Option<String>>>,
> = std::sync::OnceLock::new();
let regions = REGIONS.get_or_init(Default::default);
if let Some(known) = regions.lock().ok()?.get(bucket) {
return known.clone();
}
let url = if bucket.contains('.') {
format!("https://s3.amazonaws.com/{bucket}")
} else {
format!("https://{bucket}.s3.amazonaws.com/")
};
let response = probe_agent().head(&url).call();
let region = match response {
Ok(response) => response
.headers()
.get("x-amz-bucket-region")
.and_then(|v| v.to_str().ok())
.map(str::to_string),
Err(_) => return None,
};
if let Ok(mut map) = regions.lock() {
map.insert(bucket.to_string(), region.clone());
}
region
}
fn probe_agent() -> ureq::Agent {
crate::user_agent::ureq_config()
.timeout_global(Some(std::time::Duration::from_secs(10)))
.http_status_as_error(false)
.max_redirects(0)
.build()
.into()
}
pub fn probe_unsigned(resolved: &crate::cloud_sources::Resolved) -> Option<bool> {
let url = probe_url(resolved)?;
let agent = probe_agent();
let mut request = if url.contains('?') {
agent.get(&url)
} else {
agent.head(&url)
};
if resolved.kind == ProviderKind::Azure {
request = request.header("x-ms-version", crate::azure::API_VERSION);
}
let response = request.call().ok()?;
match response.status().as_u16() {
200..=299 => Some(true),
401 | 403 => Some(false),
_ => None,
}
}
fn probe_url(resolved: &crate::cloud_sources::Resolved) -> Option<String> {
let encode = |key: &str| key.split('/').map(urlencode).collect::<Vec<_>>().join("/");
let split = |key: &str| -> (String, bool) {
let before_glob = key.split('*').next().unwrap_or("");
if key.contains('*') || key.is_empty() || key.ends_with('/') {
let directory = match before_glob.rsplit_once('/') {
Some((directory, _)) => format!("{directory}/"),
None => String::new(),
};
(directory, true)
} else {
(key.to_string(), false)
}
};
match resolved.kind {
ProviderKind::S3 => {
if resolved.s3.endpoint.is_some() {
return None;
}
let (_, bucket, _) = split_bucket_url(&resolved.url)?;
let key = resolved.url.split_once("://")?.1;
let key = key.split_once('/').map_or("", |(_, key)| key);
let region = resolved.s3.region.as_deref().unwrap_or("us-east-1");
let base = if bucket.contains('.') {
format!("https://s3.{region}.amazonaws.com/{bucket}")
} else {
format!("https://{bucket}.s3.{region}.amazonaws.com")
};
Some(match split(key) {
(directory, true) => format!(
"{base}/?list-type=2&max-keys=1&prefix={}",
urlencode(&directory)
),
(object, false) => format!("{base}/{}", encode(&object)),
})
}
ProviderKind::Gcs => {
let (_, bucket, _) = split_bucket_url(&resolved.url)?;
let key = resolved.url.split_once("://")?.1;
let key = key.split_once('/').map_or("", |(_, key)| key);
Some(match split(key) {
(directory, true) => format!(
"https://storage.googleapis.com/storage/v1/b/{bucket}/o?maxResults=1&prefix={}",
urlencode(&directory)
),
(object, false) => {
format!(
"https://storage.googleapis.com/{bucket}/{}",
encode(&object)
)
}
})
}
ProviderKind::Azure => {
if resolved.azure.blob_endpoint.is_some() || resolved.azure.use_emulator {
return None;
}
let (account, container, key) = crate::source::azure_parts(&resolved.url)?;
let base = format!("https://{account}.blob.core.windows.net/{container}");
Some(match split(&key) {
(directory, true) => format!(
"{base}?restype=container&comp=list&maxresults=1&prefix={}",
urlencode(&directory)
),
(object, false) => format!("{base}/{}", encode(&object)),
})
}
}
}
pub async fn list_account(
source_id: &str,
account: &str,
config: &CloudConfig,
) -> Result<Vec<crate::discover::Entry>, String> {
let (source_id, account, config) = (source_id.to_string(), account.to_string(), config.clone());
let source = {
let (source_id, config) = (source_id.clone(), config.clone());
tokio::task::spawn_blocking(move || {
crate::cloud_sources::discover(&config, &Environment::current())
.into_iter()
.find(|s| s.id == source_id)
.ok_or_else(|| format!("source not found: {source_id}"))
})
.await
.map_err(|e| format!("{e}"))??
};
if source.kind == ProviderKind::Gcs {
let blocking = Source {
project: Some(account.clone()),
..source.clone()
};
let buckets = tokio::task::spawn_blocking(move || {
tokio::runtime::Handle::current().block_on(list_gcs_buckets(&blocking))
})
.await
.map_err(|e| format!("{e}"))??;
return Ok(buckets
.into_iter()
.map(|bucket| {
crate::cloud_sources::remember_bucket(&source, &bucket);
let mut entry =
crate::discover::Entry::directory(Path::new(&format!("gs://{bucket}")));
entry.name = bucket;
entry
})
.collect());
}
tokio::task::spawn_blocking(move || {
let env = Environment::current();
if source.kind != ProviderKind::Azure {
return Err(format!("{source_id} has no accounts"));
}
let settings = source.azure.with_token(&env)?;
let containers = crate::azure::list_containers(&account, &settings)?;
Ok(containers
.into_iter()
.map(|container| {
let mut entry = crate::discover::Entry::directory(Path::new(
&crate::source::azure_url(&account, &container, ""),
));
entry.name = container;
entry
})
.collect())
})
.await
.map_err(|e| format!("{e}"))?
}
pub async fn list_buckets(source: &Source) -> Result<Vec<String>, String> {
if let Some(problem) = &source.problem {
return Err(problem.clone());
}
let source = {
let source = source.clone();
tokio::task::spawn_blocking(move || source.with_credentials(&Environment::current()))
.await
.map_err(|e| format!("{e}"))??
};
let source = &source;
match source.kind {
ProviderKind::Gcs => list_gcs_buckets(source).await,
ProviderKind::S3 => list_s3_buckets(&source.s3).await,
ProviderKind::Azure => Err("Azure lists storage accounts, not buckets".to_string()),
}
}
async fn list_gcs_buckets(source: &Source) -> Result<Vec<String>, String> {
let project = source.project.as_deref().ok_or_else(|| {
"no GCP project is set, so there is nothing to list buckets for. Set \
GOOGLE_CLOUD_PROJECT or DATUI_GCP_PROJECT."
.to_string()
})?;
let bearer = google_bearer(source).await?;
let mut buckets = Vec::new();
let mut page_token: Option<String> = None;
for _ in 0..MAX_BUCKET_PAGES {
let mut url = format!(
"https://storage.googleapis.com/storage/v1/b?project={}&maxResults=1000",
urlencode(project)
);
if let Some(token) = &page_token {
url.push_str(&format!("&pageToken={}", urlencode(token)));
}
let mut response = http_agent()
.get(&url)
.config()
.http_status_as_error(false)
.build()
.header("Authorization", &format!("Bearer {bearer}"))
.call()
.map_err(|e| format!("{e}"))?;
let status = response.status().as_u16();
let body = response
.body_mut()
.read_to_string()
.map_err(|e| format!("could not read the response: {e}"))?;
if status != 200 {
return Err(crate::gcloud::describe_error(status, &body));
}
buckets.extend(parse_gcs_buckets(&body)?);
match gcs_next_page_token(&body) {
Some(token) => page_token = Some(token),
None => break,
}
}
buckets.sort();
Ok(buckets)
}
async fn google_bearer(source: &Source) -> Result<String, String> {
if let Some(problem) = &source.problem {
return Err(problem.clone());
}
if let Some(configuration) = source.gcloud.clone() {
return tokio::task::spawn_blocking(move || {
crate::gcloud::token(&configuration, &Environment::current()).map(|(token, _)| token)
})
.await
.map_err(|e| format!("{e}"))?;
}
let builder = match &source.google_credentials {
Some(file) => object_store::gcp::GoogleCloudStorageBuilder::new()
.with_application_credentials(file.to_string_lossy()),
None => object_store::gcp::GoogleCloudStorageBuilder::from_env(),
};
let store = builder
.with_bucket_name("datui-credential-probe")
.with_config(
object_store::gcp::GoogleConfigKey::Client(crate::user_agent::CLIENT_KEY),
crate::user_agent::get(),
)
.build()
.map_err(|e| format!("Google Cloud Storage is not configured: {e}"))?;
store
.credentials()
.get_credential()
.await
.map(|credential| credential.bearer.clone())
.map_err(|e| format!("could not obtain Google credentials: {e}"))
}
async fn list_gcs_projects(source: &Source) -> Result<Vec<Listed>, String> {
let bearer = google_bearer(source).await?;
let searched = {
let bearer = bearer.clone();
tokio::task::spawn_blocking(move || crate::gcloud::search_projects(&bearer))
.await
.map_err(|e| format!("{e}"))?
};
let mut projects = match (searched, &source.project) {
(Ok(projects), _) => projects,
(Err(_), Some(project)) => vec![crate::gcloud::Project {
id: project.clone(),
name: None,
}],
(Err(e), None) => return Err(e),
};
if let Some(configured) = &source.project {
match projects.iter().position(|p| &p.id == configured) {
Some(i) => {
let first = projects.remove(i);
projects.insert(0, first);
}
None => projects.insert(
0,
crate::gcloud::Project {
id: configured.clone(),
name: None,
},
),
}
}
Ok(projects
.into_iter()
.map(|project| {
let mut details = Vec::new();
if let Some(name) = project.name.filter(|n| n != &project.id) {
details.push(("name".to_string(), name));
}
if source.project.as_deref() == Some(project.id.as_str()) {
details.push(("project".to_string(), "configured".to_string()));
}
Listed {
place: PathBuf::from(source.bucket_url(&project.id)),
name: project.id,
details,
}
})
.collect())
}
async fn list_s3_buckets(settings: &S3Settings) -> Result<Vec<String>, String> {
use object_store::aws::AwsAuthorizer;
let s3 = s3_builder("datui-credential-probe", settings)
.build()
.map_err(|e| format!("S3 is not configured: {e}"))?;
let credential = s3
.credentials()
.get_credential()
.await
.map_err(|e| format!("could not obtain AWS credentials: {e}"))?;
let region = settings
.region
.clone()
.unwrap_or_else(|| "us-east-1".to_string());
let url = s3_list_buckets_url(settings);
let mut signed = http::Request::builder()
.method("GET")
.uri(&url)
.body(object_store::client::HttpRequestBody::empty())
.map_err(|e| format!("could not build the request: {e}"))?;
AwsAuthorizer::new(&credential, "s3", ®ion).authorize(&mut signed, None);
let mut request = http_agent().get(&url);
for (name, value) in signed.headers() {
if let Ok(value) = value.to_str() {
request = request.header(name.as_str(), value);
}
}
let body = request
.call()
.map_err(|e| format!("{e}"))?
.body_mut()
.read_to_string()
.map_err(|e| format!("could not read the response: {e}"))?;
let mut buckets = parse_s3_buckets(&body)?;
buckets.sort();
Ok(buckets)
}
fn s3_list_buckets_url(settings: &S3Settings) -> String {
let endpoint = settings
.endpoint
.as_deref()
.unwrap_or("https://s3.amazonaws.com");
format!("{}/", endpoint.trim_end_matches('/'))
}
const MAX_BUCKET_PAGES: usize = 20;
pub(crate) fn urlencode(value: &str) -> String {
let mut out = String::with_capacity(value.len());
for byte in value.as_bytes() {
match byte {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
out.push(*byte as char)
}
_ => out.push_str(&format!("%{byte:02X}")),
}
}
out
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_filter_narrows_by_the_prefix_the_names_share() {
let stations = [
"STATION=ACW00011604",
"STATION=AE000041196",
"STATION=AFM00040938",
];
assert_eq!(
narrowing_prefix("usw", &stations).as_deref(),
Some("STATION=USW")
);
assert_eq!(
narrowing_prefix("station=usw", &stations).as_deref(),
Some("STATION=USW")
);
let years = ["year=2019", "year=2020", "year=2021"];
assert_eq!(narrowing_prefix("202", &years).as_deref(), Some("year=202"));
let files = ["Sales.csv", "returns.csv"];
assert_eq!(narrowing_prefix("Sal", &files).as_deref(), Some("Sal"));
assert_eq!(narrowing_prefix("", &files), None);
assert_eq!(narrowing_prefix("a/b", &files), None);
assert_eq!(narrowing_prefix("x", &[]), None);
}
fn resolved(url: &str, kind: ProviderKind) -> crate::cloud_sources::Resolved {
crate::cloud_sources::Resolved {
url: url.to_string(),
kind,
source_id: String::new(),
s3: S3Settings::default(),
azure: Default::default(),
signing: Signing::Try,
place: crate::cloud_sources::access_key(url).unwrap(),
gcloud: None,
google_credentials: None,
login_error: None,
}
}
#[test]
fn the_unsigned_look_is_one_small_request() {
let mut s3 = resolved("s3://aws-public-blockchain/v1.0/btc/", ProviderKind::S3);
s3.s3.region = Some("us-east-2".to_string());
assert_eq!(
probe_url(&s3).unwrap(),
"https://aws-public-blockchain.s3.us-east-2.amazonaws.com/?list-type=2&max-keys=1&prefix=v1.0%2Fbtc%2F"
);
let object = resolved("s3://my.dotted.bucket/a b/100%.parquet", ProviderKind::S3);
assert_eq!(
probe_url(&object).unwrap(),
"https://s3.us-east-1.amazonaws.com/my.dotted.bucket/a%20b/100%25.parquet"
);
let glob = resolved("gs://bucket/year=*/part-*.parquet", ProviderKind::Gcs);
assert_eq!(
probe_url(&glob).unwrap(),
"https://storage.googleapis.com/storage/v1/b/bucket/o?maxResults=1&prefix="
);
let azure = resolved(
"abfss://release@overturemapswestus2.dfs.core.windows.net/2026-08-19.0/",
ProviderKind::Azure,
);
assert_eq!(
probe_url(&azure).unwrap(),
"https://overturemapswestus2.blob.core.windows.net/release?restype=container&comp=list&maxresults=1&prefix=2026-08-19.0%2F"
);
let mut minio = resolved("s3://data/x.parquet", ProviderKind::S3);
minio.s3.endpoint = Some("http://127.0.0.1:9000".to_string());
assert_eq!(probe_url(&minio), None, "a custom endpoint is not probed");
}
fn directories(names: &[&str]) -> Vec<String> {
names.iter().map(|n| n.to_string()).collect()
}
fn objects(names: &[(&str, u64)]) -> Vec<(String, u64)> {
names.iter().map(|(n, s)| (n.to_string(), *s)).collect()
}
#[test]
fn the_prefix_being_listed_is_not_something_it_holds() {
let holds = look_at_listing(
"out/sub/",
&directories(&[]),
&objects(&[("out/sub/", 0), ("out/sub/a.parquet", 5)]),
)
.1;
assert_eq!(holds.formats, vec![("parquet".to_string(), 1)]);
assert_eq!(holds.not_read, 0, "the prefix is not a file it cannot read");
assert_eq!(holds.skipped, 0);
assert_eq!(holds.directories, 0);
let holds = look_at_listing(
"bigquery/census/",
&directories(&["bigquery/census/data/"]),
&objects(&[("bigquery/census", 11)]),
)
.1;
assert_eq!(holds.not_read, 0);
assert_eq!(holds.label(), "1 dir");
}
#[test]
fn a_hugging_face_prefix_is_arrow() {
let holds = look_at_listing(
"hf/",
&directories(&[]),
&objects(&[
("hf/data-00000-of-00001.arrow", 5),
("hf/dataset_info.json", 5),
("hf/state.json", 5),
]),
)
.1;
assert_eq!(holds.formats, vec![("arrow".to_string(), 1)]);
assert_eq!(holds.skipped_names, ["dataset_info.json", "state.json"]);
let holds = look_at_listing("j/", &directories(&[]), &objects(&[("j/state.json", 5)])).1;
assert_eq!(holds.formats, vec![("json".to_string(), 1)], "no Arrow");
}
#[test]
fn a_saved_dataset_dict_prefix_is_arrow() {
let (kind, holds) = look_at_listing(
"dd/",
&directories(&["dd/test/", "dd/train/"]),
&objects(&[("dd/dataset_dict.json", 30)]),
);
assert_eq!(kind, crate::discover::EntryKind::Directory);
assert!(holds.dataset_dict);
assert!(holds.formats.is_empty(), "{:?}", holds.formats);
assert_eq!(holds.skipped_names, ["dataset_dict.json"]);
assert_eq!(holds.directories, 2);
let holds = look_at_listing(
"j/",
&directories(&[]),
&objects(&[("j/dataset_dict.json", 30)]),
)
.1;
assert!(!holds.dataset_dict);
assert_eq!(holds.formats, vec![("json".to_string(), 1)]);
}
#[test]
fn the_prefix_being_listed_is_not_a_row_in_it() {
let prefixes = directories(&["bigquery/census/data/"]);
assert!(!is_listed_object(
"bigquery/census",
11,
"bigquery/census",
&prefixes
));
assert!(!is_listed_object(
"bigquery/census",
0,
"bigquery/census",
&prefixes
));
assert!(is_listed_object(
"bigquery/census/test.csv",
240,
"bigquery/census",
&prefixes
));
assert!(is_listed_object("bigquery/census", 11, "bigquery", &[]));
}
#[test]
fn sub_prefixes_count_against_a_prefix_being_one_table() {
let subs: Vec<String> = (0..10).map(|i| format!("out/sub{i}/")).collect();
let (kind, holds) = look_at_listing(
"out/",
&subs,
&objects(&[("out/a.parquet", 5), ("out/b.parquet", 5)]),
);
assert_eq!(kind, crate::discover::EntryKind::Directory);
assert_eq!(holds.label(), "2 parquet");
assert_eq!(holds.directories, 10);
}
#[test]
fn a_consoles_folder_placeholder_is_not_also_a_file_it_cannot_read() {
let holds = look_at_listing(
"out/",
&directories(&["out/sub/"]),
&objects(&[("out/sub", 0)]),
)
.1;
assert_eq!(holds.directories, 1);
assert_eq!(holds.not_read, 0, "the placeholder is the directory itself");
assert_eq!(holds.label(), "1 dir");
}
#[test]
fn an_emr_folder_marker_beside_a_partition_is_a_writers_own_file() {
let holds = look_at_listing(
"out/",
&directories(&["out/year=2024/", "out/year=2025/"]),
&objects(&[("out/year=2024_$folder$", 0), ("out/year=2025_$folder$", 0)]),
)
.1;
assert_eq!(holds.not_read, 0);
assert_eq!(holds.partitions, 2);
assert_eq!(holds.skipped, 2);
}
#[test]
fn partitions_carry_a_prefix_only_while_they_are_the_most_of_it() {
let parts = directories(&["out/year=2024/", "out/year=2025/"]);
let two = objects(&[("out/a.parquet", 5), ("out/b.parquet", 5)]);
assert_eq!(
look_at_listing("out/", &parts, &two).0,
crate::discover::EntryKind::Hive,
"two partitions against two files"
);
let three = objects(&[
("out/a.parquet", 5),
("out/b.parquet", 5),
("out/c.parquet", 5),
]);
assert_ne!(
look_at_listing("out/", &parts, &three).0,
crate::discover::EntryKind::Hive,
"one more file than partitions"
);
}
#[test]
fn a_stray_it_cannot_read_counts_against_a_prefix_the_way_it_does_on_disk() {
let (kind, holds) = look_at_listing(
"yellow/",
&directories(&[]),
&objects(&[
("yellow/a.parquet", 5),
("yellow/b.parquet", 5),
("yellow/year=2028", 0),
("yellow/year=2029", 0),
("yellow/year=2030", 0),
("yellow/year=2031", 0),
("yellow/year=2032", 0),
]),
);
assert_eq!(holds.not_read, 5);
assert_eq!(
kind,
crate::discover::EntryKind::Directory,
"five it cannot read outvote two it can"
);
}
#[test]
fn a_prefix_with_another_page_behind_it_says_so() {
let keys: Vec<(String, u64)> = (0..100)
.map(|i| (format!("out/part-{i:05}.parquet"), 100u64))
.collect();
let holds = look_at_page("out/", &directories(&[]), &keys, Some("next-page-token")).1;
assert!(holds.truncated);
assert_eq!(holds.label(), "100+ parquet");
let holds = look_at_page("out/", &directories(&[]), &keys, None).1;
assert!(!holds.truncated);
assert_eq!(holds.label(), "100 parquet");
}
#[test]
fn a_pane_never_lists_more_skipped_names_than_it_promised() {
let keys: Vec<(String, u64)> = (0..30)
.map(|i| (format!("out/.part-{i:05}.parquet.crc"), 8u64))
.collect();
let holds = look_at_listing("out/", &directories(&[]), &keys).1;
assert_eq!(holds.skipped, 30, "all of them are counted");
assert_eq!(
holds.skipped_names.len(),
crate::discover::SKIPPED_NAMES_SHOWN,
"but only a few are named"
);
let holds = look_at_listing(
"out/",
&directories(&["out/_temporary/"]),
&objects(&[("out/_SUCCESS", 0), ("out/.part.crc", 8)]),
)
.1;
assert_eq!(
holds.skipped_names,
vec![
".part.crc".to_string(),
"_SUCCESS".to_string(),
"_temporary".to_string()
]
);
}
#[test]
fn two_formats_that_tie_are_named_in_the_same_order_every_time() {
let holds = look_at_listing(
"out/",
&directories(&[]),
&objects(&[
("out/a.json", 5),
("out/b.json", 5),
("out/y.csv", 5),
("out/z.csv", 5),
]),
)
.1;
assert_eq!(holds.line(false).unwrap(), "2 csv · 2 json");
assert_eq!(holds.label(), "mixed");
}
#[test]
fn a_directory_is_classified_by_one_page_of_its_listing() {
use crate::discover::EntryKind;
assert_eq!(
look_at_listing(
"v1.0/btc/blocks/",
&directories(&[
"v1.0/btc/blocks/date=2009-01-03/",
"v1.0/btc/blocks/date=2009-01-09/"
]),
&objects(&[("v1.0/btc/blocks/_SUCCESS", 0)]),
)
.0,
EntryKind::Hive
);
assert_eq!(
look_at_listing(
"gbif/occurrence.parquet/",
&directories(&[]),
&objects(&[
("gbif/occurrence.parquet/000001", 10),
("gbif/occurrence.parquet/000002", 10)
]),
)
.0,
EntryKind::MultiFile,
"part files with no extension"
);
assert_eq!(
look_at_listing(
"a/",
&directories(&[]),
&objects(&[("a/x.csv", 5), ("a/y.csv", 5)])
)
.0,
EntryKind::Directory,
"CSV cannot be read in place as one table"
);
assert_eq!(
look_at_listing(
"a/",
&directories(&["a/by_year/", "a/by_station/"]),
&objects(&[])
)
.0,
EntryKind::Directory
);
assert_eq!(
look_at_listing(
"a/",
&directories(&["a/b/"]),
&objects(&[("a/one.parquet", 5)])
)
.0,
EntryKind::Directory,
"one file is a file to open, not a dataset"
);
}
#[test]
fn a_lake_table_is_not_a_directory_of_parquet_files() {
use crate::discover::EntryKind;
let parts = objects(&[
("t/part-00000.parquet", 10),
("t/part-00001.parquet", 10),
("t/part-00002.parquet", 10),
]);
assert_eq!(
look_at_listing("t/", &directories(&["t/_delta_log/"]), &parts).0,
EntryKind::Delta
);
assert_eq!(
look_at_listing("t/", &directories(&["t/.hoodie/"]), &parts).0,
EntryKind::Hudi
);
assert_eq!(
look_at_listing(
"t/",
&directories(&["t/metadata/", "t/data/"]),
&objects(&[])
)
.0,
EntryKind::Iceberg
);
assert_eq!(
look_at_listing("t/", &directories(&["t/metadata/"]), &parts).0,
EntryKind::MultiFile,
"a directory called metadata beside part files is not an Iceberg table"
);
assert_eq!(
look_at_listing("t/", &directories(&["t/metadata/", "t/data/"]), &parts).0,
EntryKind::MultiFile,
"an Iceberg root holds its data under data/, not beside it"
);
assert_eq!(
look_at_listing(
"t/",
&directories(&["t/metadata/", "t/data/"]),
&objects(&[("t/README.md", 20)])
)
.0,
EntryKind::Iceberg,
"but something else beside them does not disqualify it"
);
assert_eq!(
look_at_listing("t/", &directories(&[]), &parts).0,
EntryKind::MultiFile,
"and a directory of part files with no log is still one table"
);
}
#[test]
fn refusals_and_job_files() {
assert!(is_refusal(
"Client error with status 403 Forbidden: <Code>AccessDenied</Code>"
));
assert!(is_refusal(
"Server returned 401 NoAuthenticationInformation"
));
assert!(!is_refusal("error sending request: connection refused"));
for name in [
"_SUCCESS",
"_committed_123",
"_started_123",
"yellow_$folder$",
"_metadata.json",
".crc",
] {
assert!(
crate::discover::is_bookkeeping(name),
"{name} is a writer's own file"
);
}
assert!(!crate::discover::is_bookkeeping("part-0000.parquet"));
let directories = ["out/_temporary/".to_string()];
let objects: Vec<(String, u64)> = [
("out/_temporary", 12u64),
("out/NOTES", 0),
("out/a.parquet", 100),
("out/b.parquet", 100),
]
.iter()
.map(|(k, s)| ((*k).to_string(), *s))
.collect();
let holds = look_at_listing("out", &directories, &objects).1;
assert_eq!(holds.line(true).as_deref(), Some("2 parquet"));
let own_marker: Vec<(String, u64)> = [("out", 0u64), ("out/a.parquet", 100)]
.iter()
.map(|(k, s)| ((*k).to_string(), *s))
.collect();
let holds = look_at_listing("out/", &[], &own_marker).1;
assert_eq!(holds.line(true).as_deref(), Some("1 parquet"));
for (prefix, key) in [("v1.0/", "v1.0"), ("out/_temporary/", "out/_temporary")] {
let objects: Vec<(String, u64)> =
[(key.to_string(), 0u64), (format!("{key}/a.parquet"), 100)]
.into_iter()
.collect();
let holds = look_at_listing(prefix, &[], &objects).1;
assert_eq!(
holds.line(true).as_deref(),
Some("1 parquet"),
"{prefix} counted its own key"
);
}
assert!(crate::discover::is_bookkeeping("_manifest.parquet"));
assert!(!is_marker("_manifest.parquet"));
assert!(!is_marker("_2024_sales.csv"));
for name in ["_SUCCESS", "_committed_1", "_started_1", "yellow_$folder$"] {
assert!(is_marker(name), "{name} stands for no data at all");
}
assert!(is_empty_marker("year=2032", 0));
assert!(!is_empty_marker("year=2032", 10));
assert!(!is_empty_marker("empty.csv", 0));
}
#[test]
fn a_key_is_taken_as_the_service_stores_it() {
assert_eq!(object_path("edge/100%.csv.gz").as_ref(), "edge/100%.csv.gz");
assert_eq!(
object_path("edge/a+b=c&d#e.parquet").as_ref(),
"edge/a+b=c&d#e.parquet"
);
assert_eq!(
object_path("edge/name with spaces").as_ref(),
"edge/name with spaces"
);
assert_eq!(
object_path("a/../b").as_ref(),
object_store::path::Path::from("a/../b").as_ref()
);
}
use std::collections::HashMap;
fn env_of(
vars: &[(&str, &str)],
files: &[&str],
home: Option<&str>,
) -> (HashMap<String, String>, Vec<PathBuf>, Option<PathBuf>) {
(
vars.iter()
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect(),
files.iter().map(PathBuf::from).collect(),
home.map(PathBuf::from),
)
}
macro_rules! environment {
($vars:expr_2021, $files:expr_2021, $home:expr_2021) => {
Environment {
var: &|key| $vars.get(key).cloned(),
exists: &|path| $files.iter().any(|f: &PathBuf| f == path),
read: &|_| None,
home: $home.clone(),
windows: false,
run: &|_, _| {
Err(crate::cloud_command::CommandError::Missing(
"test".to_string(),
))
},
all_vars: &|| Vec::new(),
list: &|_| Vec::new(),
}
};
($vars:expr_2021, $files:expr_2021, $home:expr_2021, $contents:expr_2021) => {
Environment {
var: &|key| $vars.get(key).cloned(),
exists: &|path| $files.iter().any(|f: &PathBuf| f == path),
read: &|_| Some($contents.to_string()),
home: $home.clone(),
windows: false,
run: &|_, _| {
Err(crate::cloud_command::CommandError::Missing(
"test".to_string(),
))
},
all_vars: &|| Vec::new(),
list: &|_| Vec::new(),
}
};
}
#[test]
fn nothing_configured_finds_nothing() {
let (vars, files, home) = env_of(&[], &[], Some("/home/u"));
let env = environment!(vars, files, home);
assert!(detect(&CloudConfig::default(), &env).is_empty());
}
#[test]
fn gcloud_default_credentials_are_enough_to_list_gcs() {
let (vars, files, home) = env_of(
&[("GOOGLE_CLOUD_PROJECT", "example-project")],
&["/home/u/.config/gcloud/application_default_credentials.json"],
Some("/home/u"),
);
let env = environment!(vars, files, home);
let found = detect(&CloudConfig::default(), &env);
assert_eq!(found.len(), 1);
assert_eq!(found[0].kind, ProviderKind::Gcs);
assert_eq!(found[0].project.as_deref(), Some("example-project"));
assert!(found[0].can_list_buckets());
assert_eq!(found[0].note, "gcloud");
assert_eq!(
found[0].detail().as_deref(),
Some("project: example-project")
);
}
#[test]
fn an_aws_profile_is_the_detail_for_s3() {
let (vars, files, home) = env_of(&[("AWS_PROFILE", "research")], &[], Some("/home/u"));
let env = environment!(vars, files, home);
let found = detect(&CloudConfig::default(), &env);
assert_eq!(found[0].detail().as_deref(), Some("profile: research"));
}
#[test]
fn gcs_without_a_project_is_shown_but_cannot_enumerate() {
let (vars, files, home) = env_of(
&[],
&["/home/u/.config/gcloud/application_default_credentials.json"],
Some("/home/u"),
);
let env = environment!(vars, files, home);
let found = detect(&CloudConfig::default(), &env);
assert_eq!(found.len(), 1);
assert!(found[0].project.is_none());
assert!(!found[0].can_list_buckets());
}
#[test]
fn a_service_account_outranks_the_developer_login_in_the_note() {
let (vars, files, home) = env_of(
&[("GOOGLE_SERVICE_ACCOUNT", "/keys/sa.json")],
&["/home/u/.config/gcloud/application_default_credentials.json"],
Some("/home/u"),
);
let env = environment!(vars, files, home);
let found = detect(&CloudConfig::default(), &env);
assert_eq!(found[0].note, "service account");
}
#[test]
fn aws_keys_in_the_environment_are_amazon_until_an_endpoint_says_otherwise() {
let (vars, files, home) = env_of(&[("AWS_ACCESS_KEY_ID", "AKIA")], &[], Some("/home/u"));
let env = environment!(vars, files, home);
let found = detect(&CloudConfig::default(), &env);
assert_eq!(found.len(), 1);
assert_eq!(found[0].label, "Amazon S3");
assert_eq!(found[0].note, "AWS_ACCESS_KEY_ID");
}
#[test]
fn a_custom_endpoint_is_named_by_its_host_and_not_guessed_at() {
let config = CloudConfig {
s3_endpoint_url: Some("http://localhost:9000".to_string()),
s3_access_key_id: Some("minioadmin".to_string()),
..CloudConfig::default()
};
let (vars, files, home) = env_of(&[], &[], Some("/home/u"));
let env = environment!(vars, files, home);
let found = detect(&config, &env);
assert_eq!(found.len(), 1);
assert_eq!(found[0].label, "S3-compatible (localhost:9000)");
assert_eq!(found[0].note, "datui config");
assert_eq!(found[0].endpoint.as_deref(), Some("http://localhost:9000"));
}
fn effective(config: &CloudConfig, env: &Environment<'_>) -> CloudConfig {
let mut merged = config.clone();
merged.overlay(CloudConfig::from_env(env.var));
merged
}
#[test]
fn an_endpoint_from_the_environment_counts_too() {
let (vars, files, home) = env_of(
&[
("AWS_ACCESS_KEY_ID", "minioadmin"),
("AWS_ENDPOINT_URL", "https://minio.internal:9000/"),
],
&[],
Some("/home/u"),
);
let env = environment!(vars, files, home);
let config = effective(&CloudConfig::default(), &env);
let found = detect(&config, &env);
assert_eq!(found[0].label, "S3-compatible (minio.internal:9000)");
assert_eq!(
s3_list_buckets_url(&S3Settings::from_config(&config)),
"https://minio.internal:9000/"
);
}
#[test]
fn the_service_specific_endpoint_variable_outranks_the_general_one() {
let (vars, files, home) = env_of(
&[
("AWS_ACCESS_KEY_ID", "k"),
("AWS_ENDPOINT", "http://third:1"),
("AWS_ENDPOINT_URL", "http://second:2"),
("AWS_ENDPOINT_URL_S3", "http://first:3"),
],
&[],
None,
);
let env = environment!(vars, files, home);
let config = effective(&CloudConfig::default(), &env);
assert_eq!(
s3_list_buckets_url(&S3Settings::from_config(&config)),
"http://first:3/"
);
assert_eq!(detect(&config, &env)[0].label, "S3-compatible (first:3)");
}
#[test]
fn a_blank_endpoint_variable_does_not_erase_the_configured_one() {
let (vars, files, home) = env_of(
&[("AWS_ACCESS_KEY_ID", "k"), ("AWS_ENDPOINT_URL", " ")],
&[],
None,
);
let env = environment!(vars, files, home);
let file = CloudConfig {
s3_endpoint_url: Some("http://localhost:9000".to_string()),
..CloudConfig::default()
};
let config = effective(&file, &env);
assert_eq!(
s3_list_buckets_url(&S3Settings::from_config(&config)),
"http://localhost:9000/"
);
assert_eq!(
detect(&config, &env)[0].label,
"S3-compatible (localhost:9000)"
);
}
#[test]
fn a_key_without_a_secret_still_reaches_the_builder() {
use object_store::aws::AmazonS3ConfigKey;
let config = CloudConfig {
s3_access_key_id: Some("from-config".to_string()),
..CloudConfig::default()
};
let builder = s3_builder("bucket", &S3Settings::from_config(&config));
assert_eq!(
builder.get_config_value(&AmazonS3ConfigKey::AccessKeyId),
Some("from-config".to_string())
);
let config = CloudConfig {
s3_secret_access_key: Some("from-env".to_string()),
..CloudConfig::default()
};
let builder = s3_builder("bucket", &S3Settings::from_config(&config));
assert_eq!(
builder.get_config_value(&AmazonS3ConfigKey::SecretAccessKey),
Some("from-env".to_string())
);
}
#[test]
fn a_shared_credentials_file_is_enough() {
let (vars, files, home) = env_of(&[], &["/home/u/.aws/credentials"], Some("/home/u"));
let env = environment!(vars, files, home);
let found = detect(&CloudConfig::default(), &env);
assert_eq!(found.len(), 1);
assert_eq!(found[0].note, "~/.aws");
}
#[test]
fn both_providers_appear_when_both_are_usable() {
let (vars, files, home) = env_of(
&[("AWS_ACCESS_KEY_ID", "AKIA"), ("GOOGLE_CLOUD_PROJECT", "p")],
&["/home/u/.config/gcloud/application_default_credentials.json"],
Some("/home/u"),
);
let env = environment!(vars, files, home);
let found = detect(&CloudConfig::default(), &env);
assert_eq!(found.len(), 2);
assert_eq!(found[0].kind, ProviderKind::Gcs);
assert_eq!(found[1].kind, ProviderKind::S3);
}
#[test]
fn a_blank_project_variable_is_not_a_project() {
let (vars, files, home) = env_of(
&[("GOOGLE_CLOUD_PROJECT", " ")],
&["/home/u/.config/gcloud/application_default_credentials.json"],
Some("/home/u"),
);
let env = environment!(vars, files, home);
let found = detect(&CloudConfig::default(), &env);
assert!(found[0].project.is_none());
}
#[test]
fn the_project_comes_from_the_credentials_file_when_nothing_else_says() {
let adc = r#"{
"type": "authorized_user",
"client_id": "x.apps.googleusercontent.com",
"refresh_token": "secret-and-not-read-here",
"quota_project_id": "example-project"
}"#;
let (vars, files, home) = env_of(
&[],
&["/home/u/.config/gcloud/application_default_credentials.json"],
Some("/home/u"),
);
let env = environment!(vars, files, home, adc);
let found = detect(&CloudConfig::default(), &env);
assert_eq!(found[0].project.as_deref(), Some("example-project"));
assert!(found[0].can_list_buckets());
}
#[test]
fn an_environment_variable_outranks_the_credentials_file() {
let adc = r#"{"quota_project_id": "from-the-file"}"#;
let (vars, files, home) = env_of(
&[("GOOGLE_CLOUD_PROJECT", "from-the-shell")],
&["/home/u/.config/gcloud/application_default_credentials.json"],
Some("/home/u"),
);
let env = environment!(vars, files, home, adc);
let found = detect(&CloudConfig::default(), &env);
assert_eq!(found[0].project.as_deref(), Some("from-the-shell"));
}
#[test]
fn an_unparseable_credentials_file_costs_the_project_and_nothing_else() {
let (vars, files, home) = env_of(
&[],
&["/home/u/.config/gcloud/application_default_credentials.json"],
Some("/home/u"),
);
let env = environment!(vars, files, home, "{ not json");
let found = detect(&CloudConfig::default(), &env);
assert_eq!(found.len(), 1, "the provider is still usable");
assert!(found[0].project.is_none());
}
#[test]
fn gcs_buckets_come_out_of_a_real_shaped_response() {
let body = r#"{
"kind": "storage#buckets",
"items": [
{"kind": "storage#bucket", "name": "example-data", "location": "US-CENTRAL1"},
{"kind": "storage#bucket", "name": "example-backups"}
]
}"#;
assert_eq!(
parse_gcs_buckets(body).unwrap(),
vec!["example-data", "example-backups"]
);
}
#[test]
fn a_project_with_no_buckets_is_an_answer_not_an_error() {
assert_eq!(
parse_gcs_buckets(r#"{"kind": "storage#buckets"}"#).unwrap(),
Vec::<String>::new()
);
}
#[test]
fn a_gcs_error_body_is_reported_rather_than_read_as_emptiness() {
let body =
r#"{"error": {"code": 403, "message": "does not have storage.buckets.list access"}}"#;
let err = parse_gcs_buckets(body).unwrap_err();
assert!(err.contains("storage.buckets.list"), "{err}");
}
#[test]
fn a_malformed_gcs_entry_does_not_discard_the_page() {
let body = r#"{"items": [{"name": ""}, {"nome": "typo"}, {"name": "good"}]}"#;
assert_eq!(parse_gcs_buckets(body).unwrap(), vec!["good"]);
}
#[test]
fn gcs_pagination_token_is_found_when_present() {
assert_eq!(
gcs_next_page_token(r#"{"nextPageToken": "abc", "items": []}"#).as_deref(),
Some("abc")
);
assert!(gcs_next_page_token(r#"{"items": []}"#).is_none());
}
#[test]
fn the_listing_goes_to_amazon_when_no_endpoint_is_set() {
assert_eq!(
s3_list_buckets_url(&S3Settings::from_config(&CloudConfig::default())),
"https://s3.amazonaws.com/"
);
}
#[test]
fn the_listing_honours_the_endpoint_override() {
let config = CloudConfig {
s3_endpoint_url: Some("http://localhost:9000/".to_string()),
..CloudConfig::default()
};
let mut effective = config.clone();
effective.overlay(CloudConfig::from_env(&|name| {
(name == "AWS_ENDPOINT_URL").then(|| "http://127.0.0.1:9101".to_string())
}));
assert_eq!(
s3_list_buckets_url(&S3Settings::from_config(&effective)),
"http://127.0.0.1:9101/"
);
let (vars, files, home) = env_of(&[("AWS_ACCESS_KEY_ID", "testing")], &[], None);
let env = environment!(vars, files, home);
let found = detect(&effective, &env);
assert_eq!(found[0].label, "S3-compatible (127.0.0.1:9101)");
let mut effective = config.clone();
effective.overlay(CloudConfig::from_env(&|_| None));
assert_eq!(
s3_list_buckets_url(&S3Settings::from_config(&effective)),
"http://localhost:9000/"
);
}
#[test]
fn the_override_carries_keys_and_region_too() {
let config = CloudConfig {
s3_access_key_id: Some("from-config".to_string()),
s3_region: Some("eu-west-1".to_string()),
..CloudConfig::default()
};
let mut effective = config.clone();
effective.overlay(CloudConfig::from_env(&|name| match name {
"AWS_ACCESS_KEY_ID" => Some("from-env".to_string()),
"AWS_SECRET_ACCESS_KEY" => Some("secret".to_string()),
_ => None,
}));
assert_eq!(effective.s3_access_key_id.as_deref(), Some("from-env"));
assert_eq!(effective.s3_secret_access_key.as_deref(), Some("secret"));
assert_eq!(effective.s3_region.as_deref(), Some("eu-west-1"));
}
#[test]
fn s3_buckets_come_out_of_a_real_shaped_response() {
let body = r#"<?xml version="1.0" encoding="UTF-8"?>
<ListAllMyBucketsResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">
<Owner><ID>abc</ID><DisplayName>owner</DisplayName></Owner>
<Buckets>
<Bucket><Name>first-bucket</Name><CreationDate>2024-01-01T00:00:00.000Z</CreationDate></Bucket>
<Bucket><Name>second-bucket</Name><CreationDate>2024-02-01T00:00:00.000Z</CreationDate></Bucket>
</Buckets>
</ListAllMyBucketsResult>"#;
assert_eq!(
parse_s3_buckets(body).unwrap(),
vec!["first-bucket", "second-bucket"]
);
}
#[test]
fn the_owner_display_name_is_not_mistaken_for_a_bucket() {
let body = r#"<ListAllMyBucketsResult>
<Owner><ID>x</ID><DisplayName>Name</DisplayName></Owner>
<Buckets><Bucket><Name>only-bucket</Name></Bucket></Buckets>
</ListAllMyBucketsResult>"#;
assert_eq!(parse_s3_buckets(body).unwrap(), vec!["only-bucket"]);
}
#[test]
fn an_s3_error_document_is_reported() {
let body = r#"<Error><Code>InvalidAccessKeyId</Code><Message>The key is not valid</Message></Error>"#;
let err = parse_s3_buckets(body).unwrap_err();
assert!(err.contains("not valid"), "{err}");
}
#[test]
fn no_buckets_is_not_an_error() {
let body = r#"<ListAllMyBucketsResult><Buckets></Buckets></ListAllMyBucketsResult>"#;
assert_eq!(parse_s3_buckets(body).unwrap(), Vec::<String>::new());
}
#[test]
fn implausibly_deep_xml_is_refused_rather_than_followed() {
let body = "<a>".repeat(MAX_XML_DEPTH + 2);
assert!(parse_s3_buckets(&body).is_err());
}
#[test]
fn malformed_xml_is_an_error_and_not_a_panic() {
assert!(parse_s3_buckets("<Buckets><Bucket><Name>x").is_err());
}
#[test]
fn an_endpoint_host_is_extracted_or_declined() {
assert_eq!(
endpoint_host("http://localhost:9000"),
Some("localhost:9000".into())
);
assert_eq!(endpoint_host("https://a.b/c/d"), Some("a.b".into()));
assert_eq!(endpoint_host("minio:9000"), Some("minio:9000".into()));
assert_eq!(endpoint_host(""), None);
assert_eq!(endpoint_host("http://"), None);
}
}
#[cfg(test)]
mod aws_role_tests {
use super::*;
use std::collections::HashMap;
fn detect_with(vars: &[(&str, &str)]) -> Vec<Provider> {
let vars: HashMap<String, String> = vars
.iter()
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect();
let env = Environment {
var: &|key| vars.get(key).cloned(),
exists: &|_| false,
read: &|_| None,
home: Some(PathBuf::from("/home/u")),
windows: false,
run: &|_, _| {
Err(crate::cloud_command::CommandError::Missing(
"test".to_string(),
))
},
all_vars: &|| Vec::new(),
list: &|_| Vec::new(),
};
detect(&CloudConfig::default(), &env)
}
#[test]
fn a_gcloud_login_on_windows_is_found_under_appdata() {
let vars: HashMap<String, String> = [("APPDATA", r"C:\Users\u\AppData\Roaming")]
.iter()
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect();
let adc = PathBuf::from(r"C:\Users\u\AppData\Roaming")
.join("gcloud")
.join("application_default_credentials.json");
let under_home = PathBuf::from(r"C:\Users\u")
.join(".config")
.join("gcloud")
.join("application_default_credentials.json");
let windows = |exists: PathBuf, windows: bool| {
let env = Environment {
var: &|key| vars.get(key).cloned(),
exists: &|path| path == exists,
read: &|_| None,
home: Some(PathBuf::from(r"C:\Users\u")),
windows,
run: &|_, _| {
Err(crate::cloud_command::CommandError::Missing(
"test".to_string(),
))
},
all_vars: &|| Vec::new(),
list: &|_| Vec::new(),
};
detect(&CloudConfig::default(), &env)
.iter()
.any(|p| p.kind == ProviderKind::Gcs)
};
assert!(windows(adc.clone(), true));
assert!(!windows(under_home, true));
assert!(!windows(adc, false));
}
#[test]
fn an_ecs_task_role_is_enough() {
for key in [
"AWS_CONTAINER_CREDENTIALS_RELATIVE_URI",
"AWS_CONTAINER_CREDENTIALS_FULL_URI",
] {
let found = detect_with(&[(key, "/v2/credentials/abc")]);
assert_eq!(found.len(), 1, "{key} should be enough");
assert_eq!(found[0].kind, ProviderKind::S3);
assert_eq!(found[0].label, "Amazon S3");
assert_eq!(found[0].note, "container role");
}
}
#[test]
fn an_eks_web_identity_is_enough() {
let found = detect_with(&[
("AWS_WEB_IDENTITY_TOKEN_FILE", "/var/run/secrets/token"),
("AWS_ROLE_ARN", "arn:aws:iam::1:role/r"),
]);
assert_eq!(found.len(), 1);
assert_eq!(found[0].note, "web identity");
}
#[test]
fn a_region_alone_is_not_a_credential() {
assert!(detect_with(&[("AWS_REGION", "us-east-1")]).is_empty());
assert!(detect_with(&[("AWS_DEFAULT_REGION", "us-east-1")]).is_empty());
}
#[test]
fn a_key_still_outranks_a_role_in_the_note() {
let found = detect_with(&[
("AWS_ACCESS_KEY_ID", "AKIA"),
("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", "/v2/creds"),
]);
assert_eq!(found[0].note, "AWS_ACCESS_KEY_ID");
}
}
#[cfg(test)]
mod one_table_tests {
use super::*;
use object_store::{ObjectStore, ObjectStoreExt, PutPayload, path::Path as OsPath};
use polars::prelude::*;
use std::sync::Arc;
fn parquet(columns: &[&str]) -> Vec<u8> {
let mut frame = DataFrame::new(
1,
columns
.iter()
.map(|c| Column::new((*c).into(), &[1i32]))
.collect::<Vec<_>>(),
)
.unwrap();
let mut bytes = Vec::new();
ParquetWriter::new(&mut bytes).finish(&mut frame).unwrap();
bytes
}
fn kind_of(files: &[(&str, &[&str])]) -> Option<crate::discover::EntryKind> {
let rt = tokio::runtime::Runtime::new().unwrap();
let store: Arc<dyn ObjectStore> = Arc::new(object_store::memory::InMemory::new());
rt.block_on(async {
let mut objects = Vec::new();
for (key, columns) in files {
let bytes = parquet(columns);
objects.push((format!("data/{key}"), bytes.len() as u64));
store
.put(
&OsPath::from(format!("data/{key}")),
PutPayload::from(bytes),
)
.await
.unwrap();
}
kind_from_footers(&store, &objects).await
})
}
#[test]
fn separate_tables_in_a_bucket_are_a_directory() {
let kind = kind_of(&[
("circuits.parquet", &["circuit_id", "lat", "lng"]),
("drivers.parquet", &["driver_id", "code", "nationality"]),
("laps.parquet", &["lap", "position", "time_millis"]),
]);
assert_eq!(kind, Some(crate::discover::EntryKind::Directory));
}
#[test]
fn parts_of_one_table_stay_one_dataset() {
let kind = kind_of(&[
("part-00000.parquet", &["id", "ts", "amount"]),
("part-00001.parquet", &["id", "ts", "amount"]),
("part-00002.parquet", &["id", "ts", "amount"]),
]);
assert_eq!(kind, Some(crate::discover::EntryKind::MultiFile));
}
#[test]
fn a_dataset_that_gained_columns_stays_one_dataset() {
let kind = kind_of(&[
("date=2009-01-03.parquet", &["id", "ts"]),
("date=2015-06-01.parquet", &["id", "ts", "fee"]),
(
"date=2025-06-01.parquet",
&["id", "ts", "fee", "witness", "address"],
),
]);
assert_eq!(kind, Some(crate::discover::EntryKind::MultiFile));
}
#[test]
fn a_directory_that_cannot_be_read_is_left_as_it_was() {
let rt = tokio::runtime::Runtime::new().unwrap();
let store: Arc<dyn ObjectStore> = Arc::new(object_store::memory::InMemory::new());
let objects = vec![
("data/a.parquet".to_string(), 10),
("data/b.parquet".to_string(), 10),
];
assert_eq!(rt.block_on(kind_from_footers(&store, &objects)), None);
}
#[test]
fn a_single_file_decides_nothing() {
assert_eq!(kind_of(&[("only.parquet", &["a", "b"])]), None);
}
#[test]
fn the_files_read_span_the_listing() {
assert_eq!(footers_to_verify(2), vec![0, 1]);
assert_eq!(footers_to_verify(3), vec![0, 1, 2]);
assert_eq!(footers_to_verify(15), vec![0, 7, 14]);
}
}