1use crate::cloud_sources::{S3Settings, Signing, Source};
25use crate::config::CloudConfig;
26use std::path::{Path, PathBuf};
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub enum ProviderKind {
32 Gcs,
33 S3,
34 Azure,
35}
36
37impl ProviderKind {
38 pub fn scheme(self) -> &'static str {
40 match self {
41 ProviderKind::Gcs => "gs",
42 ProviderKind::S3 => "s3",
43 ProviderKind::Azure => "abfss",
44 }
45 }
46}
47
48#[derive(Debug, Clone, PartialEq, Eq)]
50pub struct Provider {
51 pub kind: ProviderKind,
52 pub label: String,
54 pub note: String,
56 pub project: Option<String>,
60 pub profile: Option<String>,
62 pub endpoint: Option<String>,
65}
66
67impl Provider {
68 pub fn detail(&self) -> Option<String> {
71 match (&self.project, &self.profile) {
72 (Some(project), _) => Some(format!("project: {project}")),
73 (None, Some(profile)) => Some(format!("profile: {profile}")),
74 (None, None) => None,
75 }
76 }
77
78 pub fn can_list_buckets(&self) -> bool {
84 match self.kind {
85 ProviderKind::Gcs => self.project.is_some(),
86 ProviderKind::S3 | ProviderKind::Azure => true,
87 }
88 }
89}
90
91pub struct Environment<'a> {
98 pub var: &'a dyn Fn(&str) -> Option<String>,
100 pub exists: &'a dyn Fn(&Path) -> bool,
103 pub read: &'a dyn Fn(&Path) -> Option<String>,
107 pub home: Option<PathBuf>,
109 pub windows: bool,
111 pub run: &'a crate::cloud_command::Runner<'a>,
113 pub all_vars: &'a dyn Fn() -> Vec<(String, String)>,
115 pub list: &'a dyn Fn(&Path) -> Vec<PathBuf>,
118}
119
120impl Environment<'_> {
121 pub fn current() -> Environment<'static> {
123 Environment {
124 var: &crate::cloud_env::var,
125 exists: &|path| path.exists(),
126 read: &|path| std::fs::read_to_string(path).ok(),
127 home: dirs::home_dir(),
128 windows: cfg!(windows),
129 run: &|program, args| {
130 crate::cloud_command::run(program, args, crate::cloud_command::CREDENTIAL_TIMEOUT)
131 },
132 all_vars: &crate::cloud_env::vars,
133 list: &|dir| {
134 std::fs::read_dir(dir)
135 .map(|entries| entries.flatten().map(|e| e.path()).collect())
136 .unwrap_or_default()
137 },
138 }
139 }
140}
141
142pub fn detect(config: &CloudConfig, env: &Environment<'_>) -> Vec<Provider> {
147 let mut providers = Vec::new();
148 if let Some(gcs) = detect_gcs(config, env) {
149 providers.push(gcs);
150 }
151 if let Some(s3) = detect_s3(config, env) {
152 providers.push(s3);
153 }
154 providers
155}
156
157fn detect_gcs(config: &CloudConfig, env: &Environment<'_>) -> Option<Provider> {
159 let note = if (env.var)("GOOGLE_SERVICE_ACCOUNT").is_some()
162 || (env.var)("GOOGLE_SERVICE_ACCOUNT_PATH").is_some()
163 {
164 "service account"
165 } else if (env.var)("GOOGLE_SERVICE_ACCOUNT_KEY").is_some() {
166 "service account key"
167 } else if (env.var)("GOOGLE_APPLICATION_CREDENTIALS").is_some() {
168 "GOOGLE_APPLICATION_CREDENTIALS"
169 } else if adc_path(env).is_some() {
170 "gcloud"
174 } else if instance_identity(config, env).gcp {
175 "instance identity"
176 } else {
177 return None;
178 };
179
180 Some(Provider {
181 kind: ProviderKind::Gcs,
182 label: "Google Cloud Storage".to_string(),
183 note: note.to_string(),
184 project: gcp_project(env),
185 profile: None,
186 endpoint: None,
187 })
188}
189
190#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
197pub struct InstanceIdentity {
198 pub aws: bool,
199 pub gcp: bool,
200 pub azure: bool,
201}
202
203pub fn instance_identity(config: &CloudConfig, env: &Environment<'_>) -> InstanceIdentity {
204 let opted_in = config.instance_identity;
205 let set = |key: &str| (env.var)(key).is_some_and(|v| !v.trim().is_empty());
206 InstanceIdentity {
207 aws: opted_in,
208 gcp: opted_in || set("K_SERVICE"),
209 azure: opted_in || set("IDENTITY_ENDPOINT") || set("MSI_ENDPOINT"),
210 }
211}
212
213pub fn unreadable_google_login(env: &Environment<'_>) -> Option<String> {
217 let path = (env.var)("GOOGLE_APPLICATION_CREDENTIALS")
218 .map(PathBuf::from)
219 .or_else(|| adc_path(env))?;
220 crate::gcloud::unsupported_credential_type(&(env.read)(&path)?)
221}
222
223pub fn adc_path(env: &Environment<'_>) -> Option<PathBuf> {
228 const FILE: &str = "application_default_credentials.json";
229 let path = if env.windows {
230 PathBuf::from((env.var)("APPDATA")?)
231 .join("gcloud")
232 .join(FILE)
233 } else {
234 env.home.as_ref()?.join(".config").join("gcloud").join(FILE)
235 };
236 (env.exists)(&path).then_some(path)
237}
238
239pub(crate) fn gcp_project(env: &Environment<'_>) -> Option<String> {
254 for key in [
255 "DATUI_GCP_PROJECT",
256 "GOOGLE_CLOUD_PROJECT",
257 "GCLOUD_PROJECT",
258 "CLOUDSDK_CORE_PROJECT",
259 "GCP_PROJECT",
260 ] {
261 if let Some(value) = (env.var)(key) {
262 let value = value.trim().to_string();
263 if !value.is_empty() {
264 return Some(value);
265 }
266 }
267 }
268 adc_quota_project(env)
269}
270
271fn adc_quota_project(env: &Environment<'_>) -> Option<String> {
277 let path = adc_path(env)?;
278 let contents = (env.read)(&path)?;
279 let value: serde_json::Value = serde_json::from_str(&contents).ok()?;
280 let project = value.get("quota_project_id")?.as_str()?.trim();
281 if project.is_empty() {
282 return None;
283 }
284 Some(project.to_string())
285}
286
287fn detect_s3(config: &CloudConfig, env: &Environment<'_>) -> Option<Provider> {
295 let endpoint = config.s3_endpoint_url.clone();
296
297 let configured_keys = config.s3_access_key_id.is_some();
298 let env_keys = (env.var)("AWS_ACCESS_KEY_ID").is_some();
299 let profile = (env.var)("AWS_PROFILE");
300 let shared_credentials = env.home.as_ref().is_some_and(|home| {
301 (env.exists)(&home.join(".aws/credentials")) || (env.exists)(&home.join(".aws/config"))
302 });
303 let container_role = (env.var)("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI").is_some()
309 || (env.var)("AWS_CONTAINER_CREDENTIALS_FULL_URI").is_some();
310 let web_identity = (env.var)("AWS_WEB_IDENTITY_TOKEN_FILE").is_some();
311
312 let note = if configured_keys {
313 "datui config"
314 } else if env_keys {
315 "AWS_ACCESS_KEY_ID"
316 } else if profile.is_some() {
317 "AWS_PROFILE"
318 } else if container_role {
319 "container role"
320 } else if web_identity {
321 "web identity"
322 } else if shared_credentials {
323 "~/.aws"
324 } else if instance_identity(config, env).aws {
325 "instance role"
326 } else {
327 return None;
335 };
336
337 let label = match endpoint.as_deref().and_then(endpoint_host) {
341 Some(host) => format!("S3-compatible ({host})"),
342 None => "Amazon S3".to_string(),
343 };
344
345 Some(Provider {
346 kind: ProviderKind::S3,
347 label,
348 note: note.to_string(),
349 project: None,
350 profile,
351 endpoint,
352 })
353}
354
355fn endpoint_host(endpoint: &str) -> Option<String> {
359 let rest = endpoint
360 .split_once("://")
361 .map(|(_, rest)| rest)
362 .unwrap_or(endpoint);
363 let host = rest.split(['/', '?', '#']).next()?.trim();
364 if host.is_empty() {
365 return None;
366 }
367 Some(host.to_string())
368}
369
370pub fn parse_gcs_buckets(body: &str) -> Result<Vec<String>, String> {
376 let value: serde_json::Value =
377 serde_json::from_str(body).map_err(|e| format!("not JSON: {e}"))?;
378
379 if let Some(message) = value
382 .get("error")
383 .and_then(|e| e.get("message"))
384 .and_then(|m| m.as_str())
385 {
386 return Err(message.to_string());
387 }
388
389 let Some(items) = value.get("items").and_then(|i| i.as_array()) else {
390 return Ok(Vec::new());
391 };
392 Ok(items
393 .iter()
394 .filter_map(|item| item.get("name").and_then(|n| n.as_str()))
395 .filter(|name| !name.is_empty())
396 .map(str::to_string)
397 .collect())
398}
399
400pub fn gcs_next_page_token(body: &str) -> Option<String> {
402 serde_json::from_str::<serde_json::Value>(body)
403 .ok()?
404 .get("nextPageToken")?
405 .as_str()
406 .map(str::to_string)
407}
408
409pub fn parse_s3_buckets(body: &str) -> Result<Vec<String>, String> {
417 use quick_xml::events::Event;
418
419 let mut reader = quick_xml::Reader::from_str(body);
420 reader.config_mut().trim_text(true);
421
422 let mut buckets = Vec::new();
423 let mut path: Vec<Vec<u8>> = Vec::new();
424 let mut error_message: Option<String> = None;
425
426 loop {
427 match reader.read_event() {
428 Ok(Event::Start(tag)) => {
429 if path.len() >= MAX_XML_DEPTH {
430 return Err("response nested implausibly deeply".to_string());
431 }
432 path.push(tag.local_name().as_ref().as_bytes().to_vec());
433 }
434 Ok(Event::End(_)) => {
435 path.pop();
436 }
437 Ok(Event::Text(text)) => {
438 let value = text.xml10_content().into_owned();
439 match path_tail(&path) {
440 (Some(b"Name"), Some(b"Bucket")) if !value.is_empty() => {
442 buckets.push(value);
443 }
444 (Some(b"Message"), Some(b"Error")) => error_message = Some(value),
447 _ => {}
448 }
449 }
450 Ok(Event::Eof) => {
451 if !path.is_empty() {
456 return Err("response ended inside an element".to_string());
457 }
458 break;
459 }
460 Err(e) => return Err(format!("malformed XML: {e}")),
461 _ => {}
462 }
463 }
464
465 if let Some(message) = error_message {
466 return Err(message);
467 }
468 Ok(buckets)
469}
470
471const MAX_XML_DEPTH: usize = 32;
474
475fn path_tail(path: &[Vec<u8>]) -> (Option<&[u8]>, Option<&[u8]>) {
477 let len = path.len();
478 let last = len.checked_sub(1).map(|i| path[i].as_slice());
479 let parent = len.checked_sub(2).map(|i| path[i].as_slice());
480 (last, parent)
481}
482
483const REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20);
490
491pub(crate) fn http_agent() -> ureq::Agent {
492 crate::user_agent::ureq_config()
493 .timeout_global(Some(REQUEST_TIMEOUT))
494 .build()
495 .into()
496}
497
498pub fn s3_builder(bucket: &str, settings: &S3Settings) -> object_store::aws::AmazonS3Builder {
512 let builder = if settings.from_env && !settings.skip_signature {
516 object_store::aws::AmazonS3Builder::from_env()
517 } else {
518 object_store::aws::AmazonS3Builder::new()
519 };
520 let mut builder = builder.with_bucket_name(bucket).with_config(
521 object_store::aws::AmazonS3ConfigKey::Client(crate::user_agent::CLIENT_KEY),
522 crate::user_agent::get(),
523 );
524 if settings.skip_signature {
525 builder = builder.with_skip_signature(true);
526 }
527 if let Some(endpoint) = &settings.endpoint {
528 builder = builder.with_endpoint(endpoint.clone());
531 if endpoint.starts_with("http://") {
532 builder = builder.with_allow_http(true);
533 }
534 }
535 if settings.endpoint.is_some() || settings.virtual_hosted.is_some() {
536 builder = builder.with_virtual_hosted_style_request(settings.virtual_hosted_style());
537 }
538 if let Some(region) = &settings.region {
539 builder = builder.with_region(region.clone());
540 }
541 if let Some(key) = &settings.access_key_id {
545 builder = builder.with_access_key_id(key.clone());
546 }
547 if let Some(secret) = &settings.secret_access_key {
548 builder = builder.with_secret_access_key(secret.clone());
549 }
550 if let Some(token) = &settings.session_token {
551 builder = builder.with_token(token.clone());
552 }
553 builder
554}
555
556pub fn object_path(key: &str) -> object_store::path::Path {
564 object_store::path::Path::parse(key).unwrap_or_else(|_| object_store::path::Path::from(key))
565}
566
567pub fn store_for_bucket(
573 kind: ProviderKind,
574 bucket: &str,
575 settings: &S3Settings,
576 unsigned: bool,
577 google_token: Option<&str>,
578 google_credentials: Option<&Path>,
579) -> Result<std::sync::Arc<dyn object_store::ObjectStore>, String> {
580 match kind {
581 ProviderKind::Gcs => Ok(std::sync::Arc::new(gcs_store(
582 bucket,
583 unsigned,
584 google_token,
585 google_credentials,
586 )?)),
587 ProviderKind::S3 => {
588 let store = s3_builder(bucket, settings)
589 .build()
590 .map_err(|e| format!("S3 is not configured: {e}"))?;
591 Ok(std::sync::Arc::new(store))
592 }
593 ProviderKind::Azure => Err("an Azure container needs its account".to_string()),
594 }
595}
596
597pub fn pager_for_bucket(
599 kind: ProviderKind,
600 bucket: &str,
601 settings: &S3Settings,
602 unsigned: bool,
603 google_token: Option<&str>,
604 google_credentials: Option<&Path>,
605) -> Result<std::sync::Arc<dyn object_store::list::PaginatedListStore>, String> {
606 match kind {
607 ProviderKind::Gcs => Ok(std::sync::Arc::new(gcs_store(
608 bucket,
609 unsigned,
610 google_token,
611 google_credentials,
612 )?)),
613 ProviderKind::S3 => {
614 let store = s3_builder(bucket, settings)
615 .build()
616 .map_err(|e| format!("S3 is not configured: {e}"))?;
617 Ok(std::sync::Arc::new(store))
618 }
619 ProviderKind::Azure => Err("an Azure container needs its account".to_string()),
620 }
621}
622
623fn gcs_store(
625 bucket: &str,
626 unsigned: bool,
627 google_token: Option<&str>,
628 google_credentials: Option<&Path>,
629) -> Result<object_store::gcp::GoogleCloudStorage, String> {
630 let builder = match (unsigned, google_token) {
633 (true, _) => object_store::gcp::GoogleCloudStorageBuilder::new().with_skip_signature(true),
634 (false, Some(token)) => object_store::gcp::GoogleCloudStorageBuilder::new()
635 .with_credentials(std::sync::Arc::new(
636 object_store::StaticCredentialProvider::new(object_store::gcp::GcpCredential {
637 bearer: token.to_string(),
638 }),
639 )),
640 (false, None) => match google_credentials {
641 Some(file) => object_store::gcp::GoogleCloudStorageBuilder::new()
642 .with_application_credentials(file.to_string_lossy()),
643 None => object_store::gcp::GoogleCloudStorageBuilder::from_env(),
644 },
645 };
646 builder
647 .with_bucket_name(bucket)
648 .with_config(
649 object_store::gcp::GoogleConfigKey::Client(crate::user_agent::CLIENT_KEY),
650 crate::user_agent::get(),
651 )
652 .build()
653 .map_err(|e| format!("Google Cloud Storage is not configured: {e}"))
654}
655
656const PEEK_KEYS: usize = 100;
659
660pub async fn peek_kind(
664 url: &str,
665 config: &CloudConfig,
666) -> Result<(crate::discover::EntryKind, crate::discover::Holds), String> {
667 let resolved = {
668 let (url, config) = (url.to_string(), config.clone());
669 tokio::task::spawn_blocking(move || crate::cloud_sources::resolve(&url, &config))
670 .await
671 .map_err(|e| format!("{e}"))??
672 };
673 match peek_page(&resolved).await {
674 Err(refused) if resolved.signing == Signing::Try && is_refusal(&refused) => {
675 peek_page(&resolved.unsigned()).await.map_err(|_| refused)
676 }
677 Err(refused) if is_refusal(&refused) && resolved.login_error.is_some() => {
678 Err(resolved.login_error.clone().unwrap_or(refused))
679 }
680 other => other,
681 }
682}
683
684async fn peek_page(
685 resolved: &crate::cloud_sources::Resolved,
686) -> Result<(crate::discover::EntryKind, crate::discover::Holds), String> {
687 use object_store::list::{PaginatedListOptions, PaginatedListStore};
688 let (store, prefix): (std::sync::Arc<dyn PaginatedListStore>, String) =
689 if let Some((account, container, key)) = crate::source::azure_parts(&resolved.url) {
690 (
691 crate::azure::paginated_store(&account, &container, &resolved.azure)?,
692 key,
693 )
694 } else {
695 let (kind, bucket, key) = split_bucket_url(&resolved.url)
696 .ok_or_else(|| format!("not an object-store URL: {}", resolved.url))?;
697 let store = pager_for_bucket(
698 kind,
699 &bucket,
700 &resolved.s3,
701 resolved.signing == Signing::Unsigned,
702 resolved.gcloud.as_ref().map(|(_, token)| token.as_str()),
703 resolved.google_credentials.as_deref(),
704 )?;
705 (store, key)
706 };
707 let prefix = format!("{}/", prefix.trim_matches('/'));
708 let page = store
709 .list_paginated(
710 Some(&prefix),
711 PaginatedListOptions {
712 delimiter: Some("/".into()),
713 max_keys: Some(PEEK_KEYS),
714 ..Default::default()
715 },
716 )
717 .await
718 .map_err(|e| format!("{e}"))?;
719 let directories: Vec<String> = page
720 .result
721 .common_prefixes
722 .iter()
723 .map(|p| p.as_ref().to_string())
724 .collect();
725 let objects: Vec<(String, u64)> = page
726 .result
727 .objects
728 .iter()
729 .map(|o| (o.location.as_ref().to_string(), o.size))
730 .collect();
731 let (kind, holds) = look_at_page(&prefix, &directories, &objects, page.page_token.as_deref());
732 if kind != crate::discover::EntryKind::MultiFile {
733 return Ok((kind, holds));
734 }
735 Ok((
740 verified_kind(resolved, &objects).await.unwrap_or(kind),
741 holds,
742 ))
743}
744
745const VERIFY_FOOTERS: usize = 3;
752
753fn footers_to_verify(files: usize) -> Vec<usize> {
760 if files <= VERIFY_FOOTERS {
761 (0..files).collect()
762 } else {
763 vec![0, files / 2, files - 1]
764 }
765}
766
767async fn verified_kind(
771 resolved: &crate::cloud_sources::Resolved,
772 objects: &[(String, u64)],
773) -> Option<crate::discover::EntryKind> {
774 let store = match crate::source::azure_parts(&resolved.url) {
780 Some((account, container, _)) => {
781 crate::azure::store(&account, &container, &resolved.azure).ok()?
782 }
783 None => {
784 let (provider, bucket, _) = split_bucket_url(&resolved.url)?;
785 store_for_bucket(
786 provider,
787 &bucket,
788 &resolved.s3,
789 resolved.signing == Signing::Unsigned,
790 resolved.gcloud.as_ref().map(|(_, token)| token.as_str()),
791 resolved.google_credentials.as_deref(),
792 )
793 .ok()?
794 }
795 };
796 kind_from_footers(&store, objects).await
797}
798
799async fn kind_from_footers(
801 store: &std::sync::Arc<dyn object_store::ObjectStore>,
802 objects: &[(String, u64)],
803) -> Option<crate::discover::EntryKind> {
804 let parquet: Vec<&(String, u64)> = objects
805 .iter()
806 .filter(|(key, _)| crate::discover::is_parquet_key(key))
807 .collect();
808 if parquet.len() < 2 {
809 return None;
810 }
811 let picks = footers_to_verify(parquet.len());
812
813 let meter = std::sync::Arc::new(crate::measurements::Meter::default());
814 let store = store.clone();
815 let mut reads = tokio::task::JoinSet::new();
816 for index in picks {
817 let (key, size) = parquet[index].clone();
818 let (store, meter) = (store.clone(), meter.clone());
819 reads.spawn(async move {
820 let file = crate::dataset_files::DatasetFile {
821 key,
822 size,
823 stamp: 0,
824 etag: None,
825 };
826 crate::cloud_hive::footer_of_file(&store, &file, &meter)
827 .await
828 .ok()
829 });
830 }
831
832 let mut per_file: Vec<Vec<String>> = Vec::new();
833 while let Some(joined) = reads.join_next().await {
834 if let Ok(Some(footer)) = joined {
835 per_file.push(footer.schema.iter_names().map(|n| n.to_string()).collect());
836 }
837 }
838 if per_file.len() < 2 {
840 return None;
841 }
842 Some(if crate::schema_union::is_nested(&per_file) {
843 crate::discover::EntryKind::MultiFile
844 } else {
845 crate::discover::EntryKind::Directory
846 })
847}
848
849fn look_at_page(
860 prefix: &str,
861 directories: &[String],
862 objects: &[(String, u64)],
863 next_page: Option<&str>,
864) -> (crate::discover::EntryKind, crate::discover::Holds) {
865 let (kind, mut holds) = look_at_listing(prefix, directories, objects);
866 holds.truncated = next_page.is_some();
867 (kind, holds)
868}
869
870pub fn look_at_listing(
873 prefix: &str,
874 directories: &[String],
875 objects: &[(String, u64)],
876) -> (crate::discover::EntryKind, crate::discover::Holds) {
877 use crate::discover::EntryKind;
878 let last = |key: &str| {
879 key.trim_end_matches('/')
880 .rsplit('/')
881 .next()
882 .unwrap_or("")
883 .to_string()
884 };
885 let here = prefix.trim_matches('/');
898 let objects: Vec<(String, u64)> = objects
899 .iter()
900 .filter(|(key, _)| here.is_empty() || key.trim_matches('/') != here)
901 .cloned()
902 .collect();
903 let objects = objects.as_slice();
904
905 let counted: Vec<&String> = directories
906 .iter()
907 .filter(|f| !crate::discover::is_bookkeeping(&last(f)))
908 .collect();
909 let partitions = counted
910 .iter()
911 .filter(|f| crate::discover::is_partition_name(&last(f)))
912 .count();
913 let present: Vec<&String> = objects
916 .iter()
917 .filter(|(key, size)| {
918 let name = last(key);
919 !name.is_empty()
920 && !crate::discover::is_bookkeeping(&name)
921 && !is_empty_marker(&name, *size)
922 && !(*size == 0 && directories.iter().any(|f| last(f) == name))
923 })
924 .map(|(key, _)| key)
925 .collect();
926 let format_of = |key: &str| crate::discover::data_format(std::path::Path::new(key));
931 let other_data = present.iter().any(|key| {
933 format_of(key).is_some_and(|f| !f.is_lines()) || crate::discover::is_parquet_key(key)
934 });
935 let files: Vec<&&String> = present
936 .iter()
937 .filter(|key| {
938 format_of(key).is_some_and(|f| !(other_data && f.is_lines()))
941 || crate::discover::is_parquet_key(key)
942 })
943 .collect();
944 let mut counts: Vec<(&'static str, usize)> = Vec::new();
948 for key in &files {
949 let name = last(key);
950 let format = crate::discover::data_format(std::path::Path::new(name.as_str()))
951 .map(|f| f.name())
952 .unwrap_or("parquet");
953 match counts.iter_mut().find(|(f, _)| *f == format) {
954 Some((_, n)) => *n += 1,
955 None => counts.push((format, 1)),
956 }
957 }
958 let arrow = counts
962 .iter()
963 .any(|(f, _)| *f == crate::FileFormat::Arrow.name());
964 let dataset_dict = !counted.is_empty()
965 && files
966 .iter()
967 .any(|key| last(key) == crate::hf_splits::DATASET_DICT);
968 let hugging_face: Vec<String> = files
969 .iter()
970 .map(|key| last(key))
971 .filter(|name| {
972 (arrow && crate::discover::is_hugging_face_metadata(name))
973 || (dataset_dict && name == crate::hf_splits::DATASET_DICT)
974 })
975 .collect();
976 let json = crate::FileFormat::Json.name();
977 for (format, n) in &mut counts {
978 if *format == json {
979 *n -= hugging_face.len();
980 }
981 }
982 counts.retain(|(_, n)| *n > 0);
983 counts.sort_by(|a, b| crate::discover::rank_formats((a.0, a.1), (b.0, b.1)));
986 let orphan_markers: Vec<String> = objects
994 .iter()
995 .filter(|(key, size)| {
996 let name = last(key);
997 !name.is_empty()
998 && is_empty_marker(&name, *size)
999 && !crate::discover::is_bookkeeping(&name)
1000 && !directories.iter().any(|f| last(f) == name)
1001 })
1002 .map(|(key, _)| last(key))
1003 .collect();
1004 let mut skipped_names: Vec<String> = objects
1005 .iter()
1006 .filter(|(key, size)| !(*size == 0 && directories.iter().any(|f| last(f) == last(key))))
1009 .map(|(key, _)| last(key))
1010 .chain(directories.iter().map(|f| last(f)))
1011 .filter(|name| !name.is_empty() && crate::discover::is_bookkeeping(name))
1012 .chain(hugging_face)
1013 .collect();
1014 skipped_names.sort();
1015 skipped_names.dedup();
1018 let skipped = skipped_names.len();
1019 skipped_names.truncate(crate::discover::SKIPPED_NAMES_SHOWN);
1020 let unnamed = present
1023 .iter()
1024 .filter(|key| !files.contains(key))
1025 .filter(|key| crate::discover::has_no_extension(std::path::Path::new(&last(key))))
1026 .count();
1027 let holds = crate::discover::Holds {
1028 formats: counts
1029 .into_iter()
1030 .map(|(f, n)| (f.to_string(), n))
1031 .collect(),
1032 directories: counted.len(),
1033 partitions,
1034 not_read: present.len() - files.len() - unnamed + orphan_markers.len(),
1037 unnamed,
1038 skipped,
1039 skipped_names,
1040 truncated: false,
1041 dataset_dict,
1042 };
1043 let directory = |name: &str| directories.iter().any(|f| last(f) == name);
1048 if directory("_delta_log") {
1049 return (EntryKind::Delta, holds);
1050 }
1051 if directory(".hoodie") {
1052 return (EntryKind::Hudi, holds);
1053 }
1054 let parquet = files
1055 .iter()
1056 .filter(|key| crate::discover::is_parquet_key(key))
1057 .count();
1058 if directory("metadata") && directory("data") && parquet == 0 {
1067 return (EntryKind::Iceberg, holds);
1068 }
1069 let seen = counted.len() + present.len() + orphan_markers.len();
1075 if partitions > 0 && partitions >= files.len() {
1078 return (EntryKind::Hive, holds);
1079 }
1080 let kind = if parquet > 1 && parquet == files.len() && parquet * 2 >= seen {
1081 EntryKind::MultiFile
1082 } else {
1083 EntryKind::Directory
1084 };
1085 (kind, holds)
1086}
1087
1088pub fn split_bucket_url(url: &str) -> Option<(ProviderKind, String, String)> {
1094 let (_, plain) = crate::source::split_source_id(url);
1095 let (scheme, rest) = plain.split_once("://")?;
1096 let kind = match scheme {
1097 "gs" | "gcs" => ProviderKind::Gcs,
1098 "s3" | "s3a" => ProviderKind::S3,
1099 _ => return None,
1100 };
1101 let rest = rest.trim_end_matches('/');
1102 let (bucket, prefix) = match rest.split_once('/') {
1103 Some((bucket, prefix)) => (bucket, prefix),
1104 None => (rest, ""),
1105 };
1106 if bucket.is_empty() {
1107 return None;
1108 }
1109 Some((
1110 kind,
1111 bucket.to_string(),
1112 prefix.trim_matches('/').to_string(),
1113 ))
1114}
1115
1116pub const MAX_LEVEL_ROWS: usize = crate::discover::MAX_ENTRIES_PER_DIR;
1120
1121#[derive(Debug, Clone, Default)]
1123pub struct Level {
1124 pub rows: Vec<crate::discover::Entry>,
1125 pub truncated: bool,
1127 pub cancelled: bool,
1129}
1130
1131pub type Progress = std::sync::Arc<dyn Fn(&[crate::discover::Entry]) + Send + Sync>;
1133
1134#[derive(Clone, Default)]
1136pub struct Watch {
1137 pub progress: Option<Progress>,
1139 pub cancelled: std::sync::Arc<std::sync::atomic::AtomicBool>,
1141 pub names_from: Option<String>,
1143}
1144
1145impl Watch {
1146 fn cancelled(&self) -> bool {
1147 self.cancelled.load(std::sync::atomic::Ordering::Relaxed)
1148 }
1149}
1150
1151pub async fn list_objects(
1158 url: &str,
1159 config: &CloudConfig,
1160) -> Result<Vec<crate::discover::Entry>, String> {
1161 list_objects_watched(url, config, &Watch::default())
1162 .await
1163 .map(|level| level.rows)
1164}
1165
1166pub async fn list_objects_watched(
1169 url: &str,
1170 config: &CloudConfig,
1171 watch: &Watch,
1172) -> Result<Level, String> {
1173 let resolved = {
1176 let (url, config) = (url.to_string(), config.clone());
1177 tokio::task::spawn_blocking(move || crate::cloud_sources::resolve(&url, &config))
1178 .await
1179 .map_err(|e| format!("{e}"))??
1180 };
1181 let signing = resolved.signing;
1182 let place = resolved.place.clone();
1183 let listed = list_level(url, &resolved, watch).await;
1184 let (listed, resolved) = match listed {
1186 Err(refusal)
1187 if resolved.kind == ProviderKind::Azure
1188 && crate::azure::is_permission_mismatch(&refusal)
1189 && resolved.azure.identity.is_some() =>
1190 {
1191 let enabled = config.use_azure_account_keys;
1192 let keyed = {
1193 let (resolved, refusal) = (resolved.clone(), refusal.clone());
1194 tokio::task::spawn_blocking(move || {
1195 let (account, _, _) =
1196 crate::source::azure_parts(&resolved.url).ok_or_else(|| refusal.clone())?;
1197 crate::azure::with_account_key(
1198 &account,
1199 &resolved.azure,
1200 &refusal,
1201 enabled,
1202 &Environment::current(),
1203 )
1204 .map(|azure| crate::cloud_sources::Resolved { azure, ..resolved })
1205 })
1206 .await
1207 .map_err(|e| format!("{e}"))?
1208 };
1209 match keyed {
1210 Ok(keyed) => (list_level(url, &keyed, watch).await, keyed),
1211 Err(why) => (Err(why), resolved),
1212 }
1213 }
1214 other => (other, resolved),
1215 };
1216 if resolved.kind == ProviderKind::Azure
1217 && listed.is_ok()
1218 && matches!(resolved.azure.auth, crate::azure::AzureAuth::Bearer(_))
1219 && let Some((account, _, _)) = crate::source::azure_parts(&resolved.url)
1220 {
1221 crate::azure::remember_token_reads(&account);
1222 }
1223 match listed {
1224 Err(refused) if signing == Signing::Try && is_refusal(&refused) => {
1225 let level = list_level(url, &resolved.unsigned(), watch)
1228 .await
1229 .map_err(|_| refused)?;
1230 crate::cloud_sources::remember_access(&place, true);
1231 Ok(level)
1232 }
1233 Ok(level) => {
1234 if signing == Signing::Try {
1235 crate::cloud_sources::remember_access(&place, false);
1236 }
1237 Ok(level)
1238 }
1239 Err(refused) if is_refusal(&refused) && resolved.login_error.is_some() => {
1241 Err(resolved.login_error.clone().unwrap_or(refused))
1242 }
1243 Err(e) => Err(e),
1244 }
1245}
1246
1247pub fn narrowing_prefix(filter: &str, names: &[&str]) -> Option<String> {
1255 let filter = filter.trim();
1256 if filter.is_empty() || filter.contains('/') {
1257 return None;
1258 }
1259 let first = names.first()?;
1260 let mut common = first.len();
1261 for name in &names[1..] {
1262 common = common.min(
1263 first
1264 .bytes()
1265 .zip(name.bytes())
1266 .take_while(|(a, b)| a == b)
1267 .count(),
1268 );
1269 }
1270 while !first.is_char_boundary(common) {
1271 common -= 1;
1272 }
1273 let shared = first[..common]
1276 .rfind(|c: char| !c.is_alphanumeric())
1277 .map_or("", |at| &first[..=at]);
1278 let typed = match filter.get(..shared.len()) {
1279 Some(head) if !shared.is_empty() && head.eq_ignore_ascii_case(shared) => {
1280 &filter[shared.len()..]
1281 }
1282 _ => filter,
1283 };
1284 let rest = names.iter().flat_map(|n| n[shared.len()..].chars());
1285 let (mut upper, mut lower) = (false, false);
1286 for c in rest {
1287 upper |= c.is_uppercase();
1288 lower |= c.is_lowercase();
1289 }
1290 let typed = match (upper, lower) {
1291 (true, false) => typed.to_uppercase(),
1292 (false, true) => typed.to_lowercase(),
1293 _ => typed.to_string(),
1294 };
1295 Some(format!("{shared}{typed}"))
1296}
1297
1298pub fn is_refusal(error: &str) -> bool {
1300 let lower = error.to_ascii_lowercase();
1301 [
1302 "403",
1303 "401",
1304 "forbidden",
1305 "unauthorized",
1306 "accessdenied",
1307 "access denied",
1308 "permissiondenied",
1309 "authorizationfailure",
1310 "authenticationfailed",
1311 "invalidauthenticationinfo",
1312 "noauthenticationinformation",
1313 ]
1314 .iter()
1315 .any(|word| lower.contains(word))
1316}
1317
1318pub fn is_marker(name: &str) -> bool {
1327 name == "_SUCCESS"
1328 || name.starts_with("_committed_")
1329 || name.starts_with("_started_")
1330 || name.ends_with("_$folder$")
1331}
1332
1333pub fn is_empty_marker(name: &str, size: u64) -> bool {
1337 size == 0 && !name.contains('.')
1338}
1339
1340fn is_listed_object(location: &str, size: u64, prefix: &str, prefixes: &[String]) -> bool {
1348 let name = location.rsplit('/').next().unwrap_or(location);
1349 !(name.is_empty()
1350 || is_marker(name)
1351 || crate::azure::is_folder_marker(location, size, prefixes)
1352 || is_empty_marker(name, size)
1353 || location.trim_end_matches('/') == prefix)
1354}
1355
1356async fn list_level(
1358 url: &str,
1359 resolved: &crate::cloud_sources::Resolved,
1360 watch: &Watch,
1361) -> Result<Level, String> {
1362 if resolved.kind == ProviderKind::Azure {
1363 return list_azure_objects(resolved, watch).await;
1364 }
1365 let (kind, bucket, prefix) =
1366 split_bucket_url(&resolved.url).ok_or_else(|| format!("not an object-store URL: {url}"))?;
1367 let pager = pager_for_bucket(
1368 kind,
1369 &bucket,
1370 &resolved.s3,
1371 resolved.signing == Signing::Unsigned,
1372 resolved.gcloud.as_ref().map(|(_, token)| token.as_str()),
1373 resolved.google_credentials.as_deref(),
1374 )?;
1375
1376 let base = match crate::source::split_source_id(url).0 {
1379 Some(id) => format!("{}://{id}@{bucket}", kind.scheme()),
1380 None => format!("{}://{bucket}", kind.scheme()),
1381 };
1382 list_pages(pager.as_ref(), &prefix, watch, |result| {
1383 let prefixes: Vec<String> = result
1384 .common_prefixes
1385 .iter()
1386 .map(|p| p.as_ref().to_string())
1387 .collect();
1388 let directories = result
1389 .common_prefixes
1390 .iter()
1391 .map(|common| {
1392 let name = common
1393 .as_ref()
1394 .rsplit('/')
1395 .find(|part| !part.is_empty())
1396 .unwrap_or(common.as_ref())
1397 .to_string();
1398 let mut row = crate::discover::Entry::directory(Path::new(&format!(
1399 "{base}/{}",
1400 common.as_ref()
1401 )));
1402 row.name = name;
1403 row
1404 })
1405 .collect();
1406 let objects = result
1407 .objects
1408 .into_iter()
1409 .filter(|object| {
1410 is_listed_object(object.location.as_ref(), object.size, &prefix, &prefixes)
1411 })
1412 .map(|object| {
1413 let location = object.location.as_ref().to_string();
1414 let name = location.rsplit('/').next().unwrap_or(&location).to_string();
1415 let path = PathBuf::from(format!("{base}/{location}"));
1416 let kind = if crate::discover::unreadable_by_name(&path) {
1419 crate::discover::EntryKind::Other
1420 } else {
1421 crate::discover::EntryKind::File
1422 };
1423 object_row(path, kind, name, &object)
1424 })
1425 .collect();
1426 (directories, objects)
1427 })
1428 .await
1429}
1430
1431fn object_row(
1433 path: PathBuf,
1434 kind: crate::discover::EntryKind,
1435 name: String,
1436 object: &object_store::ObjectMeta,
1437) -> crate::discover::Entry {
1438 crate::discover::Entry {
1439 path,
1440 kind,
1441 name,
1442 size: Some(object.size),
1443 modified: Some(object.last_modified.into()),
1444 rows: None,
1445 cols: None,
1446 cols_sampled: false,
1447 columns: Vec::new(),
1448 cost: Default::default(),
1449 holds: Default::default(),
1450 opens_whole_directory: false,
1451 format_spec: None,
1452 table: None,
1453 }
1454}
1455
1456async fn list_pages(
1463 pager: &dyn object_store::list::PaginatedListStore,
1464 prefix: &str,
1465 watch: &Watch,
1466 mut rows_of: impl FnMut(
1467 object_store::ListResult,
1468 ) -> (Vec<crate::discover::Entry>, Vec<crate::discover::Entry>),
1469) -> Result<Level, String> {
1470 let mut key_prefix = if prefix.is_empty() {
1471 String::new()
1472 } else {
1473 format!("{prefix}/")
1474 };
1475 if let Some(names) = &watch.names_from {
1476 key_prefix.push_str(names);
1477 }
1478 let (mut directories, mut objects) = (Vec::new(), Vec::new());
1479 let mut token = None;
1480 let rows = |directories: &[crate::discover::Entry], objects: &[crate::discover::Entry]| {
1482 let mut rows = directories.to_vec();
1483 rows.extend_from_slice(objects);
1484 rows
1485 };
1486 loop {
1487 if watch.cancelled() {
1488 return Ok(Level {
1489 rows: rows(&directories, &objects),
1490 truncated: false,
1491 cancelled: true,
1492 });
1493 }
1494 let page = pager
1495 .list_paginated(
1496 (!key_prefix.is_empty()).then_some(key_prefix.as_str()),
1497 object_store::list::PaginatedListOptions {
1498 delimiter: Some("/".into()),
1499 page_token: token.take(),
1500 ..Default::default()
1501 },
1502 )
1503 .await
1504 .map_err(|e| format!("{e}"))?;
1505 let (more_directories, more_objects) = rows_of(page.result);
1506 directories.extend(more_directories);
1507 objects.extend(more_objects);
1508 let mut listed = rows(&directories, &objects);
1509 if listed.len() > MAX_LEVEL_ROWS {
1510 listed.truncate(MAX_LEVEL_ROWS);
1511 return Ok(Level {
1512 rows: listed,
1513 truncated: true,
1514 cancelled: false,
1515 });
1516 }
1517 match page.page_token {
1518 Some(next) => {
1519 if let Some(progress) = &watch.progress {
1520 progress(&listed);
1521 }
1522 token = Some(next);
1523 }
1524 None => {
1525 return Ok(Level {
1526 rows: listed,
1527 truncated: false,
1528 cancelled: false,
1529 });
1530 }
1531 }
1532 }
1533}
1534
1535async fn list_azure_objects(
1539 resolved: &crate::cloud_sources::Resolved,
1540 watch: &Watch,
1541) -> Result<Level, String> {
1542 let (account, container, prefix) = crate::source::azure_parts(&resolved.url)
1543 .ok_or_else(|| format!("not an Azure URL: {}", resolved.url))?;
1544 let pager = crate::azure::paginated_store(&account, &container, &resolved.azure)?;
1545 let prefix = prefix.trim_matches('/').to_string();
1546 list_pages(pager.as_ref(), &prefix, watch, |result| {
1547 let prefixes: Vec<String> = result
1548 .common_prefixes
1549 .iter()
1550 .map(|p| p.as_ref().to_string())
1551 .collect();
1552 let directories = prefixes
1553 .iter()
1554 .map(|common| {
1555 let name = common
1556 .rsplit('/')
1557 .find(|part| !part.is_empty())
1558 .unwrap_or(common)
1559 .to_string();
1560 let mut row = crate::discover::Entry::directory(Path::new(
1561 &crate::source::azure_url(&account, &container, &format!("{common}/")),
1562 ));
1563 row.name = name;
1564 row
1565 })
1566 .collect();
1567 let objects = result
1568 .objects
1569 .into_iter()
1570 .filter_map(|object| {
1571 let location = object.location.as_ref().to_string();
1572 let name = location.rsplit('/').next().unwrap_or(&location).to_string();
1573 if name.is_empty()
1574 || is_marker(&name)
1575 || crate::azure::is_folder_marker(&location, object.size, &prefixes)
1576 || is_empty_marker(&name, object.size)
1577 || (object.size == 0 && location.trim_end_matches('/') == prefix)
1578 {
1579 return None;
1580 }
1581 let path = PathBuf::from(crate::source::azure_url(&account, &container, &location));
1582 Some(object_row(
1583 path,
1584 crate::discover::EntryKind::File,
1585 name,
1586 &object,
1587 ))
1588 })
1589 .collect();
1590 (directories, objects)
1591 })
1592 .await
1593}
1594
1595#[derive(Debug, Clone, PartialEq, Eq)]
1598pub struct Listed {
1599 pub name: String,
1600 pub place: PathBuf,
1602 pub details: Vec<(String, String)>,
1604}
1605
1606pub async fn list_first_level(source: &Source) -> Result<Vec<Listed>, String> {
1608 if source.kind == ProviderKind::Gcs {
1609 let source = source.clone();
1610 return tokio::task::spawn_blocking(move || {
1611 tokio::runtime::Handle::current().block_on(list_gcs_projects(&source))
1612 })
1613 .await
1614 .map_err(|e| format!("{e}"))?;
1615 }
1616 if source.kind != ProviderKind::Azure {
1617 let blocking = source.clone();
1621 let names = tokio::task::spawn_blocking(move || {
1622 tokio::runtime::Handle::current().block_on(list_buckets(&blocking))
1623 })
1624 .await
1625 .map_err(|e| format!("{e}"))??;
1626 return Ok(names
1627 .into_iter()
1628 .map(|name| Listed {
1629 place: PathBuf::from(source.bucket_url(&name)),
1630 name,
1631 details: Vec::new(),
1632 })
1633 .collect());
1634 }
1635 let source = source.clone();
1636 tokio::task::spawn_blocking(move || {
1637 if let Some(problem) = &source.problem {
1638 return Err(problem.clone());
1639 }
1640 if let Some(account) = &source.azure.account {
1642 return Ok(vec![Listed {
1643 name: account.clone(),
1644 place: PathBuf::from(source.bucket_url(account)),
1645 details: Vec::new(),
1646 }]);
1647 }
1648 let accounts =
1649 crate::azure::discover_accounts(&source.azure.auth, &Environment::current())?;
1650 Ok(accounts
1651 .into_iter()
1652 .map(|account| {
1653 let mut details = Vec::new();
1654 if let Some(subscription) = account.subscription {
1655 details.push(("subscription".to_string(), subscription));
1656 }
1657 if let Some(location) = account.location {
1658 details.push(("region".to_string(), location));
1659 }
1660 let namespace = if account.hierarchical_namespace {
1661 "hierarchical"
1662 } else {
1663 "flat"
1664 };
1665 details.push(("namespace".to_string(), namespace.to_string()));
1666 if account.private_network {
1667 details.push(("network".to_string(), "private".to_string()));
1668 }
1669 if !account.shared_key_access {
1670 details.push(("shared keys".to_string(), "disabled".to_string()));
1671 }
1672 Listed {
1673 place: PathBuf::from(source.bucket_url(&account.name)),
1674 name: account.name,
1675 details,
1676 }
1677 })
1678 .collect())
1679 })
1680 .await
1681 .map_err(|e| format!("{e}"))?
1682}
1683
1684pub fn s3_bucket_region(bucket: &str) -> Option<String> {
1688 static REGIONS: std::sync::OnceLock<
1689 std::sync::Mutex<std::collections::HashMap<String, Option<String>>>,
1690 > = std::sync::OnceLock::new();
1691 let regions = REGIONS.get_or_init(Default::default);
1692 if let Some(known) = regions.lock().ok()?.get(bucket) {
1693 return known.clone();
1694 }
1695 let url = if bucket.contains('.') {
1697 format!("https://s3.amazonaws.com/{bucket}")
1698 } else {
1699 format!("https://{bucket}.s3.amazonaws.com/")
1700 };
1701 let response = probe_agent().head(&url).call();
1702 let region = match response {
1703 Ok(response) => response
1704 .headers()
1705 .get("x-amz-bucket-region")
1706 .and_then(|v| v.to_str().ok())
1707 .map(str::to_string),
1708 Err(_) => return None,
1710 };
1711 if let Ok(mut map) = regions.lock() {
1712 map.insert(bucket.to_string(), region.clone());
1713 }
1714 region
1715}
1716
1717fn probe_agent() -> ureq::Agent {
1720 crate::user_agent::ureq_config()
1721 .timeout_global(Some(std::time::Duration::from_secs(10)))
1722 .http_status_as_error(false)
1723 .max_redirects(0)
1724 .build()
1725 .into()
1726}
1727
1728pub fn probe_unsigned(resolved: &crate::cloud_sources::Resolved) -> Option<bool> {
1734 let url = probe_url(resolved)?;
1735 let agent = probe_agent();
1736 let mut request = if url.contains('?') {
1737 agent.get(&url)
1738 } else {
1739 agent.head(&url)
1740 };
1741 if resolved.kind == ProviderKind::Azure {
1742 request = request.header("x-ms-version", crate::azure::API_VERSION);
1743 }
1744 let response = request.call().ok()?;
1745 match response.status().as_u16() {
1746 200..=299 => Some(true),
1747 401 | 403 => Some(false),
1748 _ => None,
1749 }
1750}
1751
1752fn probe_url(resolved: &crate::cloud_sources::Resolved) -> Option<String> {
1755 let encode = |key: &str| key.split('/').map(urlencode).collect::<Vec<_>>().join("/");
1756 let split = |key: &str| -> (String, bool) {
1758 let before_glob = key.split('*').next().unwrap_or("");
1759 if key.contains('*') || key.is_empty() || key.ends_with('/') {
1760 let directory = match before_glob.rsplit_once('/') {
1761 Some((directory, _)) => format!("{directory}/"),
1762 None => String::new(),
1763 };
1764 (directory, true)
1765 } else {
1766 (key.to_string(), false)
1767 }
1768 };
1769 match resolved.kind {
1770 ProviderKind::S3 => {
1771 if resolved.s3.endpoint.is_some() {
1772 return None;
1773 }
1774 let (_, bucket, _) = split_bucket_url(&resolved.url)?;
1775 let key = resolved.url.split_once("://")?.1;
1776 let key = key.split_once('/').map_or("", |(_, key)| key);
1777 let region = resolved.s3.region.as_deref().unwrap_or("us-east-1");
1778 let base = if bucket.contains('.') {
1779 format!("https://s3.{region}.amazonaws.com/{bucket}")
1780 } else {
1781 format!("https://{bucket}.s3.{region}.amazonaws.com")
1782 };
1783 Some(match split(key) {
1784 (directory, true) => format!(
1785 "{base}/?list-type=2&max-keys=1&prefix={}",
1786 urlencode(&directory)
1787 ),
1788 (object, false) => format!("{base}/{}", encode(&object)),
1789 })
1790 }
1791 ProviderKind::Gcs => {
1792 let (_, bucket, _) = split_bucket_url(&resolved.url)?;
1793 let key = resolved.url.split_once("://")?.1;
1794 let key = key.split_once('/').map_or("", |(_, key)| key);
1795 Some(match split(key) {
1796 (directory, true) => format!(
1797 "https://storage.googleapis.com/storage/v1/b/{bucket}/o?maxResults=1&prefix={}",
1798 urlencode(&directory)
1799 ),
1800 (object, false) => {
1801 format!(
1802 "https://storage.googleapis.com/{bucket}/{}",
1803 encode(&object)
1804 )
1805 }
1806 })
1807 }
1808 ProviderKind::Azure => {
1809 if resolved.azure.blob_endpoint.is_some() || resolved.azure.use_emulator {
1810 return None;
1811 }
1812 let (account, container, key) = crate::source::azure_parts(&resolved.url)?;
1813 let base = format!("https://{account}.blob.core.windows.net/{container}");
1814 Some(match split(&key) {
1815 (directory, true) => format!(
1816 "{base}?restype=container&comp=list&maxresults=1&prefix={}",
1817 urlencode(&directory)
1818 ),
1819 (object, false) => format!("{base}/{}", encode(&object)),
1820 })
1821 }
1822 }
1823}
1824
1825pub async fn list_account(
1827 source_id: &str,
1828 account: &str,
1829 config: &CloudConfig,
1830) -> Result<Vec<crate::discover::Entry>, String> {
1831 let (source_id, account, config) = (source_id.to_string(), account.to_string(), config.clone());
1832 let source = {
1833 let (source_id, config) = (source_id.clone(), config.clone());
1834 tokio::task::spawn_blocking(move || {
1835 crate::cloud_sources::discover(&config, &Environment::current())
1836 .into_iter()
1837 .find(|s| s.id == source_id)
1838 .ok_or_else(|| format!("source not found: {source_id}"))
1839 })
1840 .await
1841 .map_err(|e| format!("{e}"))??
1842 };
1843 if source.kind == ProviderKind::Gcs {
1844 let blocking = Source {
1845 project: Some(account.clone()),
1846 ..source.clone()
1847 };
1848 let buckets = tokio::task::spawn_blocking(move || {
1849 tokio::runtime::Handle::current().block_on(list_gcs_buckets(&blocking))
1850 })
1851 .await
1852 .map_err(|e| format!("{e}"))??;
1853 return Ok(buckets
1854 .into_iter()
1855 .map(|bucket| {
1856 crate::cloud_sources::remember_bucket(&source, &bucket);
1858 let mut entry =
1859 crate::discover::Entry::directory(Path::new(&format!("gs://{bucket}")));
1860 entry.name = bucket;
1861 entry
1862 })
1863 .collect());
1864 }
1865 tokio::task::spawn_blocking(move || {
1866 let env = Environment::current();
1867 if source.kind != ProviderKind::Azure {
1868 return Err(format!("{source_id} has no accounts"));
1869 }
1870 let settings = source.azure.with_token(&env)?;
1871 let containers = crate::azure::list_containers(&account, &settings)?;
1872 Ok(containers
1873 .into_iter()
1874 .map(|container| {
1875 let mut entry = crate::discover::Entry::directory(Path::new(
1876 &crate::source::azure_url(&account, &container, ""),
1877 ));
1878 entry.name = container;
1879 entry
1880 })
1881 .collect())
1882 })
1883 .await
1884 .map_err(|e| format!("{e}"))?
1885}
1886
1887pub async fn list_buckets(source: &Source) -> Result<Vec<String>, String> {
1896 if let Some(problem) = &source.problem {
1897 return Err(problem.clone());
1898 }
1899 let source = {
1900 let source = source.clone();
1901 tokio::task::spawn_blocking(move || source.with_credentials(&Environment::current()))
1902 .await
1903 .map_err(|e| format!("{e}"))??
1904 };
1905 let source = &source;
1906 match source.kind {
1907 ProviderKind::Gcs => list_gcs_buckets(source).await,
1908 ProviderKind::S3 => list_s3_buckets(&source.s3).await,
1909 ProviderKind::Azure => Err("Azure lists storage accounts, not buckets".to_string()),
1910 }
1911}
1912
1913async fn list_gcs_buckets(source: &Source) -> Result<Vec<String>, String> {
1921 let project = source.project.as_deref().ok_or_else(|| {
1922 "no GCP project is set, so there is nothing to list buckets for. Set \
1923 GOOGLE_CLOUD_PROJECT or DATUI_GCP_PROJECT."
1924 .to_string()
1925 })?;
1926 let bearer = google_bearer(source).await?;
1927
1928 let mut buckets = Vec::new();
1929 let mut page_token: Option<String> = None;
1930 for _ in 0..MAX_BUCKET_PAGES {
1933 let mut url = format!(
1934 "https://storage.googleapis.com/storage/v1/b?project={}&maxResults=1000",
1935 urlencode(project)
1936 );
1937 if let Some(token) = &page_token {
1938 url.push_str(&format!("&pageToken={}", urlencode(token)));
1939 }
1940 let mut response = http_agent()
1941 .get(&url)
1942 .config()
1943 .http_status_as_error(false)
1944 .build()
1945 .header("Authorization", &format!("Bearer {bearer}"))
1946 .call()
1947 .map_err(|e| format!("{e}"))?;
1948 let status = response.status().as_u16();
1949 let body = response
1950 .body_mut()
1951 .read_to_string()
1952 .map_err(|e| format!("could not read the response: {e}"))?;
1953 if status != 200 {
1954 return Err(crate::gcloud::describe_error(status, &body));
1955 }
1956
1957 buckets.extend(parse_gcs_buckets(&body)?);
1958 match gcs_next_page_token(&body) {
1959 Some(token) => page_token = Some(token),
1960 None => break,
1961 }
1962 }
1963
1964 buckets.sort();
1965 Ok(buckets)
1966}
1967
1968async fn google_bearer(source: &Source) -> Result<String, String> {
1971 if let Some(problem) = &source.problem {
1972 return Err(problem.clone());
1973 }
1974 if let Some(configuration) = source.gcloud.clone() {
1975 return tokio::task::spawn_blocking(move || {
1976 crate::gcloud::token(&configuration, &Environment::current()).map(|(token, _)| token)
1977 })
1978 .await
1979 .map_err(|e| format!("{e}"))?;
1980 }
1981 let builder = match &source.google_credentials {
1984 Some(file) => object_store::gcp::GoogleCloudStorageBuilder::new()
1985 .with_application_credentials(file.to_string_lossy()),
1986 None => object_store::gcp::GoogleCloudStorageBuilder::from_env(),
1987 };
1988 let store = builder
1989 .with_bucket_name("datui-credential-probe")
1990 .with_config(
1991 object_store::gcp::GoogleConfigKey::Client(crate::user_agent::CLIENT_KEY),
1992 crate::user_agent::get(),
1993 )
1994 .build()
1995 .map_err(|e| format!("Google Cloud Storage is not configured: {e}"))?;
1996 store
1997 .credentials()
1998 .get_credential()
1999 .await
2000 .map(|credential| credential.bearer.clone())
2001 .map_err(|e| format!("could not obtain Google credentials: {e}"))
2002}
2003
2004async fn list_gcs_projects(source: &Source) -> Result<Vec<Listed>, String> {
2009 let bearer = google_bearer(source).await?;
2010 let searched = {
2011 let bearer = bearer.clone();
2012 tokio::task::spawn_blocking(move || crate::gcloud::search_projects(&bearer))
2013 .await
2014 .map_err(|e| format!("{e}"))?
2015 };
2016 let mut projects = match (searched, &source.project) {
2017 (Ok(projects), _) => projects,
2018 (Err(_), Some(project)) => vec![crate::gcloud::Project {
2019 id: project.clone(),
2020 name: None,
2021 }],
2022 (Err(e), None) => return Err(e),
2023 };
2024 if let Some(configured) = &source.project {
2025 match projects.iter().position(|p| &p.id == configured) {
2026 Some(i) => {
2027 let first = projects.remove(i);
2028 projects.insert(0, first);
2029 }
2030 None => projects.insert(
2031 0,
2032 crate::gcloud::Project {
2033 id: configured.clone(),
2034 name: None,
2035 },
2036 ),
2037 }
2038 }
2039 Ok(projects
2040 .into_iter()
2041 .map(|project| {
2042 let mut details = Vec::new();
2043 if let Some(name) = project.name.filter(|n| n != &project.id) {
2044 details.push(("name".to_string(), name));
2045 }
2046 if source.project.as_deref() == Some(project.id.as_str()) {
2047 details.push(("project".to_string(), "configured".to_string()));
2048 }
2049 Listed {
2050 place: PathBuf::from(source.bucket_url(&project.id)),
2051 name: project.id,
2052 details,
2053 }
2054 })
2055 .collect())
2056}
2057
2058async fn list_s3_buckets(settings: &S3Settings) -> Result<Vec<String>, String> {
2064 use object_store::aws::AwsAuthorizer;
2065
2066 let s3 = s3_builder("datui-credential-probe", settings)
2069 .build()
2070 .map_err(|e| format!("S3 is not configured: {e}"))?;
2071 let credential = s3
2072 .credentials()
2073 .get_credential()
2074 .await
2075 .map_err(|e| format!("could not obtain AWS credentials: {e}"))?;
2076
2077 let region = settings
2079 .region
2080 .clone()
2081 .unwrap_or_else(|| "us-east-1".to_string());
2082 let url = s3_list_buckets_url(settings);
2083
2084 let mut signed = http::Request::builder()
2088 .method("GET")
2089 .uri(&url)
2090 .body(object_store::client::HttpRequestBody::empty())
2091 .map_err(|e| format!("could not build the request: {e}"))?;
2092 AwsAuthorizer::new(&credential, "s3", ®ion).authorize(&mut signed, None);
2093
2094 let mut request = http_agent().get(&url);
2095 for (name, value) in signed.headers() {
2096 if let Ok(value) = value.to_str() {
2097 request = request.header(name.as_str(), value);
2098 }
2099 }
2100 let body = request
2101 .call()
2102 .map_err(|e| format!("{e}"))?
2103 .body_mut()
2104 .read_to_string()
2105 .map_err(|e| format!("could not read the response: {e}"))?;
2106
2107 let mut buckets = parse_s3_buckets(&body)?;
2108 buckets.sort();
2109 Ok(buckets)
2110}
2111
2112fn s3_list_buckets_url(settings: &S3Settings) -> String {
2116 let endpoint = settings
2117 .endpoint
2118 .as_deref()
2119 .unwrap_or("https://s3.amazonaws.com");
2120 format!("{}/", endpoint.trim_end_matches('/'))
2121}
2122
2123const MAX_BUCKET_PAGES: usize = 20;
2126
2127pub(crate) fn urlencode(value: &str) -> String {
2133 let mut out = String::with_capacity(value.len());
2134 for byte in value.as_bytes() {
2135 match byte {
2136 b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
2137 out.push(*byte as char)
2138 }
2139 _ => out.push_str(&format!("%{byte:02X}")),
2140 }
2141 }
2142 out
2143}
2144
2145#[cfg(test)]
2146mod tests {
2147 use super::*;
2148
2149 #[test]
2151 fn a_filter_narrows_by_the_prefix_the_names_share() {
2152 let stations = [
2153 "STATION=ACW00011604",
2154 "STATION=AE000041196",
2155 "STATION=AFM00040938",
2156 ];
2157 assert_eq!(
2158 narrowing_prefix("usw", &stations).as_deref(),
2159 Some("STATION=USW")
2160 );
2161 assert_eq!(
2163 narrowing_prefix("station=usw", &stations).as_deref(),
2164 Some("STATION=USW")
2165 );
2166 let years = ["year=2019", "year=2020", "year=2021"];
2167 assert_eq!(narrowing_prefix("202", &years).as_deref(), Some("year=202"));
2168 let files = ["Sales.csv", "returns.csv"];
2170 assert_eq!(narrowing_prefix("Sal", &files).as_deref(), Some("Sal"));
2171 assert_eq!(narrowing_prefix("", &files), None);
2172 assert_eq!(narrowing_prefix("a/b", &files), None);
2173 assert_eq!(narrowing_prefix("x", &[]), None);
2174 }
2175
2176 fn resolved(url: &str, kind: ProviderKind) -> crate::cloud_sources::Resolved {
2177 crate::cloud_sources::Resolved {
2178 url: url.to_string(),
2179 kind,
2180 source_id: String::new(),
2181 s3: S3Settings::default(),
2182 azure: Default::default(),
2183 signing: Signing::Try,
2184 place: crate::cloud_sources::access_key(url).unwrap(),
2185 gcloud: None,
2186 google_credentials: None,
2187 login_error: None,
2188 }
2189 }
2190
2191 #[test]
2192 fn the_unsigned_look_is_one_small_request() {
2193 let mut s3 = resolved("s3://aws-public-blockchain/v1.0/btc/", ProviderKind::S3);
2194 s3.s3.region = Some("us-east-2".to_string());
2195 assert_eq!(
2196 probe_url(&s3).unwrap(),
2197 "https://aws-public-blockchain.s3.us-east-2.amazonaws.com/?list-type=2&max-keys=1&prefix=v1.0%2Fbtc%2F"
2198 );
2199 let object = resolved("s3://my.dotted.bucket/a b/100%.parquet", ProviderKind::S3);
2200 assert_eq!(
2201 probe_url(&object).unwrap(),
2202 "https://s3.us-east-1.amazonaws.com/my.dotted.bucket/a%20b/100%25.parquet"
2203 );
2204 let glob = resolved("gs://bucket/year=*/part-*.parquet", ProviderKind::Gcs);
2205 assert_eq!(
2206 probe_url(&glob).unwrap(),
2207 "https://storage.googleapis.com/storage/v1/b/bucket/o?maxResults=1&prefix="
2208 );
2209 let azure = resolved(
2210 "abfss://release@overturemapswestus2.dfs.core.windows.net/2026-08-19.0/",
2211 ProviderKind::Azure,
2212 );
2213 assert_eq!(
2214 probe_url(&azure).unwrap(),
2215 "https://overturemapswestus2.blob.core.windows.net/release?restype=container&comp=list&maxresults=1&prefix=2026-08-19.0%2F"
2216 );
2217 let mut minio = resolved("s3://data/x.parquet", ProviderKind::S3);
2218 minio.s3.endpoint = Some("http://127.0.0.1:9000".to_string());
2219 assert_eq!(probe_url(&minio), None, "a custom endpoint is not probed");
2220 }
2221
2222 fn directories(names: &[&str]) -> Vec<String> {
2225 names.iter().map(|n| n.to_string()).collect()
2226 }
2227
2228 fn objects(names: &[(&str, u64)]) -> Vec<(String, u64)> {
2229 names.iter().map(|(n, s)| (n.to_string(), *s)).collect()
2230 }
2231
2232 #[test]
2233 fn the_prefix_being_listed_is_not_something_it_holds() {
2234 let holds = look_at_listing(
2239 "out/sub/",
2240 &directories(&[]),
2241 &objects(&[("out/sub/", 0), ("out/sub/a.parquet", 5)]),
2242 )
2243 .1;
2244 assert_eq!(holds.formats, vec![("parquet".to_string(), 1)]);
2245 assert_eq!(holds.not_read, 0, "the prefix is not a file it cannot read");
2246 assert_eq!(holds.skipped, 0);
2247 assert_eq!(holds.directories, 0);
2248
2249 let holds = look_at_listing(
2251 "bigquery/census/",
2252 &directories(&["bigquery/census/data/"]),
2253 &objects(&[("bigquery/census", 11)]),
2254 )
2255 .1;
2256 assert_eq!(holds.not_read, 0);
2257 assert_eq!(holds.label(), "1 dir");
2258 }
2259
2260 #[test]
2263 fn a_hugging_face_prefix_is_arrow() {
2264 let holds = look_at_listing(
2265 "hf/",
2266 &directories(&[]),
2267 &objects(&[
2268 ("hf/data-00000-of-00001.arrow", 5),
2269 ("hf/dataset_info.json", 5),
2270 ("hf/state.json", 5),
2271 ]),
2272 )
2273 .1;
2274 assert_eq!(holds.formats, vec![("arrow".to_string(), 1)]);
2275 assert_eq!(holds.skipped_names, ["dataset_info.json", "state.json"]);
2276 let holds = look_at_listing("j/", &directories(&[]), &objects(&[("j/state.json", 5)])).1;
2277 assert_eq!(holds.formats, vec![("json".to_string(), 1)], "no Arrow");
2278 }
2279
2280 #[test]
2284 fn a_saved_dataset_dict_prefix_is_arrow() {
2285 let (kind, holds) = look_at_listing(
2286 "dd/",
2287 &directories(&["dd/test/", "dd/train/"]),
2288 &objects(&[("dd/dataset_dict.json", 30)]),
2289 );
2290 assert_eq!(kind, crate::discover::EntryKind::Directory);
2291 assert!(holds.dataset_dict);
2292 assert!(holds.formats.is_empty(), "{:?}", holds.formats);
2293 assert_eq!(holds.skipped_names, ["dataset_dict.json"]);
2294 assert_eq!(holds.directories, 2);
2295
2296 let holds = look_at_listing(
2297 "j/",
2298 &directories(&[]),
2299 &objects(&[("j/dataset_dict.json", 30)]),
2300 )
2301 .1;
2302 assert!(!holds.dataset_dict);
2303 assert_eq!(holds.formats, vec![("json".to_string(), 1)]);
2304 }
2305
2306 #[test]
2307 fn the_prefix_being_listed_is_not_a_row_in_it() {
2308 let prefixes = directories(&["bigquery/census/data/"]);
2311 assert!(!is_listed_object(
2312 "bigquery/census",
2313 11,
2314 "bigquery/census",
2315 &prefixes
2316 ));
2317 assert!(!is_listed_object(
2318 "bigquery/census",
2319 0,
2320 "bigquery/census",
2321 &prefixes
2322 ));
2323 assert!(is_listed_object(
2324 "bigquery/census/test.csv",
2325 240,
2326 "bigquery/census",
2327 &prefixes
2328 ));
2329 assert!(is_listed_object("bigquery/census", 11, "bigquery", &[]));
2331 }
2332
2333 #[test]
2334 fn sub_prefixes_count_against_a_prefix_being_one_table() {
2335 let subs: Vec<String> = (0..10).map(|i| format!("out/sub{i}/")).collect();
2340 let (kind, holds) = look_at_listing(
2341 "out/",
2342 &subs,
2343 &objects(&[("out/a.parquet", 5), ("out/b.parquet", 5)]),
2344 );
2345 assert_eq!(kind, crate::discover::EntryKind::Directory);
2346 assert_eq!(holds.label(), "2 parquet");
2347 assert_eq!(holds.directories, 10);
2348 }
2349
2350 #[test]
2351 fn a_consoles_folder_placeholder_is_not_also_a_file_it_cannot_read() {
2352 let holds = look_at_listing(
2355 "out/",
2356 &directories(&["out/sub/"]),
2357 &objects(&[("out/sub", 0)]),
2358 )
2359 .1;
2360 assert_eq!(holds.directories, 1);
2361 assert_eq!(holds.not_read, 0, "the placeholder is the directory itself");
2362 assert_eq!(holds.label(), "1 dir");
2363 }
2364
2365 #[test]
2366 fn an_emr_folder_marker_beside_a_partition_is_a_writers_own_file() {
2367 let holds = look_at_listing(
2372 "out/",
2373 &directories(&["out/year=2024/", "out/year=2025/"]),
2374 &objects(&[("out/year=2024_$folder$", 0), ("out/year=2025_$folder$", 0)]),
2375 )
2376 .1;
2377 assert_eq!(holds.not_read, 0);
2378 assert_eq!(holds.partitions, 2);
2379 assert_eq!(holds.skipped, 2);
2380 }
2381
2382 #[test]
2383 fn partitions_carry_a_prefix_only_while_they_are_the_most_of_it() {
2384 let parts = directories(&["out/year=2024/", "out/year=2025/"]);
2388 let two = objects(&[("out/a.parquet", 5), ("out/b.parquet", 5)]);
2389 assert_eq!(
2390 look_at_listing("out/", &parts, &two).0,
2391 crate::discover::EntryKind::Hive,
2392 "two partitions against two files"
2393 );
2394 let three = objects(&[
2395 ("out/a.parquet", 5),
2396 ("out/b.parquet", 5),
2397 ("out/c.parquet", 5),
2398 ]);
2399 assert_ne!(
2400 look_at_listing("out/", &parts, &three).0,
2401 crate::discover::EntryKind::Hive,
2402 "one more file than partitions"
2403 );
2404 }
2405
2406 #[test]
2407 fn a_stray_it_cannot_read_counts_against_a_prefix_the_way_it_does_on_disk() {
2408 let (kind, holds) = look_at_listing(
2415 "yellow/",
2416 &directories(&[]),
2417 &objects(&[
2418 ("yellow/a.parquet", 5),
2419 ("yellow/b.parquet", 5),
2420 ("yellow/year=2028", 0),
2421 ("yellow/year=2029", 0),
2422 ("yellow/year=2030", 0),
2423 ("yellow/year=2031", 0),
2424 ("yellow/year=2032", 0),
2425 ]),
2426 );
2427 assert_eq!(holds.not_read, 5);
2428 assert_eq!(
2429 kind,
2430 crate::discover::EntryKind::Directory,
2431 "five it cannot read outvote two it can"
2432 );
2433 }
2434
2435 #[test]
2436 fn a_prefix_with_another_page_behind_it_says_so() {
2437 let keys: Vec<(String, u64)> = (0..100)
2438 .map(|i| (format!("out/part-{i:05}.parquet"), 100u64))
2439 .collect();
2440 let holds = look_at_page("out/", &directories(&[]), &keys, Some("next-page-token")).1;
2443 assert!(holds.truncated);
2444 assert_eq!(holds.label(), "100+ parquet");
2445 let holds = look_at_page("out/", &directories(&[]), &keys, None).1;
2448 assert!(!holds.truncated);
2449 assert_eq!(holds.label(), "100 parquet");
2450 }
2451
2452 #[test]
2453 fn a_pane_never_lists_more_skipped_names_than_it_promised() {
2454 let keys: Vec<(String, u64)> = (0..30)
2457 .map(|i| (format!("out/.part-{i:05}.parquet.crc"), 8u64))
2458 .collect();
2459 let holds = look_at_listing("out/", &directories(&[]), &keys).1;
2460 assert_eq!(holds.skipped, 30, "all of them are counted");
2461 assert_eq!(
2462 holds.skipped_names.len(),
2463 crate::discover::SKIPPED_NAMES_SHOWN,
2464 "but only a few are named"
2465 );
2466
2467 let holds = look_at_listing(
2471 "out/",
2472 &directories(&["out/_temporary/"]),
2473 &objects(&[("out/_SUCCESS", 0), ("out/.part.crc", 8)]),
2474 )
2475 .1;
2476 assert_eq!(
2477 holds.skipped_names,
2478 vec![
2479 ".part.crc".to_string(),
2480 "_SUCCESS".to_string(),
2481 "_temporary".to_string()
2482 ]
2483 );
2484 }
2485
2486 #[test]
2487 fn two_formats_that_tie_are_named_in_the_same_order_every_time() {
2488 let holds = look_at_listing(
2492 "out/",
2493 &directories(&[]),
2494 &objects(&[
2495 ("out/a.json", 5),
2496 ("out/b.json", 5),
2497 ("out/y.csv", 5),
2498 ("out/z.csv", 5),
2499 ]),
2500 )
2501 .1;
2502 assert_eq!(holds.line(false).unwrap(), "2 csv · 2 json");
2503 assert_eq!(holds.label(), "mixed");
2504 }
2505
2506 #[test]
2507 fn a_directory_is_classified_by_one_page_of_its_listing() {
2508 use crate::discover::EntryKind;
2509 assert_eq!(
2514 look_at_listing(
2515 "v1.0/btc/blocks/",
2516 &directories(&[
2517 "v1.0/btc/blocks/date=2009-01-03/",
2518 "v1.0/btc/blocks/date=2009-01-09/"
2519 ]),
2520 &objects(&[("v1.0/btc/blocks/_SUCCESS", 0)]),
2521 )
2522 .0,
2523 EntryKind::Hive
2524 );
2525 assert_eq!(
2526 look_at_listing(
2527 "gbif/occurrence.parquet/",
2528 &directories(&[]),
2529 &objects(&[
2530 ("gbif/occurrence.parquet/000001", 10),
2531 ("gbif/occurrence.parquet/000002", 10)
2532 ]),
2533 )
2534 .0,
2535 EntryKind::MultiFile,
2536 "part files with no extension"
2537 );
2538 assert_eq!(
2539 look_at_listing(
2540 "a/",
2541 &directories(&[]),
2542 &objects(&[("a/x.csv", 5), ("a/y.csv", 5)])
2543 )
2544 .0,
2545 EntryKind::Directory,
2546 "CSV cannot be read in place as one table"
2547 );
2548 assert_eq!(
2549 look_at_listing(
2550 "a/",
2551 &directories(&["a/by_year/", "a/by_station/"]),
2552 &objects(&[])
2553 )
2554 .0,
2555 EntryKind::Directory
2556 );
2557 assert_eq!(
2558 look_at_listing(
2559 "a/",
2560 &directories(&["a/b/"]),
2561 &objects(&[("a/one.parquet", 5)])
2562 )
2563 .0,
2564 EntryKind::Directory,
2565 "one file is a file to open, not a dataset"
2566 );
2567 }
2568
2569 #[test]
2572 fn a_lake_table_is_not_a_directory_of_parquet_files() {
2573 use crate::discover::EntryKind;
2574 let parts = objects(&[
2575 ("t/part-00000.parquet", 10),
2576 ("t/part-00001.parquet", 10),
2577 ("t/part-00002.parquet", 10),
2578 ]);
2579
2580 assert_eq!(
2581 look_at_listing("t/", &directories(&["t/_delta_log/"]), &parts).0,
2582 EntryKind::Delta
2583 );
2584 assert_eq!(
2585 look_at_listing("t/", &directories(&["t/.hoodie/"]), &parts).0,
2586 EntryKind::Hudi
2587 );
2588 assert_eq!(
2589 look_at_listing(
2590 "t/",
2591 &directories(&["t/metadata/", "t/data/"]),
2592 &objects(&[])
2593 )
2594 .0,
2595 EntryKind::Iceberg
2596 );
2597
2598 assert_eq!(
2600 look_at_listing("t/", &directories(&["t/metadata/"]), &parts).0,
2601 EntryKind::MultiFile,
2602 "a directory called metadata beside part files is not an Iceberg table"
2603 );
2604 assert_eq!(
2605 look_at_listing("t/", &directories(&["t/metadata/", "t/data/"]), &parts).0,
2606 EntryKind::MultiFile,
2607 "an Iceberg root holds its data under data/, not beside it"
2608 );
2609 assert_eq!(
2610 look_at_listing(
2611 "t/",
2612 &directories(&["t/metadata/", "t/data/"]),
2613 &objects(&[("t/README.md", 20)])
2614 )
2615 .0,
2616 EntryKind::Iceberg,
2617 "but something else beside them does not disqualify it"
2618 );
2619 assert_eq!(
2620 look_at_listing("t/", &directories(&[]), &parts).0,
2621 EntryKind::MultiFile,
2622 "and a directory of part files with no log is still one table"
2623 );
2624 }
2625
2626 #[test]
2627 fn refusals_and_job_files() {
2628 assert!(is_refusal(
2629 "Client error with status 403 Forbidden: <Code>AccessDenied</Code>"
2630 ));
2631 assert!(is_refusal(
2632 "Server returned 401 NoAuthenticationInformation"
2633 ));
2634 assert!(!is_refusal("error sending request: connection refused"));
2635 for name in [
2636 "_SUCCESS",
2637 "_committed_123",
2638 "_started_123",
2639 "yellow_$folder$",
2640 "_metadata.json",
2641 ".crc",
2642 ] {
2643 assert!(
2644 crate::discover::is_bookkeeping(name),
2645 "{name} is a writer's own file"
2646 );
2647 }
2648 assert!(!crate::discover::is_bookkeeping("part-0000.parquet"));
2649
2650 let directories = ["out/_temporary/".to_string()];
2654 let objects: Vec<(String, u64)> = [
2655 ("out/_temporary", 12u64),
2656 ("out/NOTES", 0),
2657 ("out/a.parquet", 100),
2658 ("out/b.parquet", 100),
2659 ]
2660 .iter()
2661 .map(|(k, s)| ((*k).to_string(), *s))
2662 .collect();
2663 let holds = look_at_listing("out", &directories, &objects).1;
2664 assert_eq!(holds.line(true).as_deref(), Some("2 parquet"));
2665
2666 let own_marker: Vec<(String, u64)> = [("out", 0u64), ("out/a.parquet", 100)]
2670 .iter()
2671 .map(|(k, s)| ((*k).to_string(), *s))
2672 .collect();
2673 let holds = look_at_listing("out/", &[], &own_marker).1;
2674 assert_eq!(holds.line(true).as_deref(), Some("1 parquet"));
2675
2676 for (prefix, key) in [("v1.0/", "v1.0"), ("out/_temporary/", "out/_temporary")] {
2680 let objects: Vec<(String, u64)> =
2681 [(key.to_string(), 0u64), (format!("{key}/a.parquet"), 100)]
2682 .into_iter()
2683 .collect();
2684 let holds = look_at_listing(prefix, &[], &objects).1;
2685 assert_eq!(
2686 holds.line(true).as_deref(),
2687 Some("1 parquet"),
2688 "{prefix} counted its own key"
2689 );
2690 }
2691 assert!(crate::discover::is_bookkeeping("_manifest.parquet"));
2695 assert!(!is_marker("_manifest.parquet"));
2696 assert!(!is_marker("_2024_sales.csv"));
2697 for name in ["_SUCCESS", "_committed_1", "_started_1", "yellow_$folder$"] {
2698 assert!(is_marker(name), "{name} stands for no data at all");
2699 }
2700 assert!(is_empty_marker("year=2032", 0));
2701 assert!(!is_empty_marker("year=2032", 10));
2702 assert!(!is_empty_marker("empty.csv", 0));
2703 }
2704
2705 #[test]
2706 fn a_key_is_taken_as_the_service_stores_it() {
2707 assert_eq!(object_path("edge/100%.csv.gz").as_ref(), "edge/100%.csv.gz");
2708 assert_eq!(
2709 object_path("edge/a+b=c&d#e.parquet").as_ref(),
2710 "edge/a+b=c&d#e.parquet"
2711 );
2712 assert_eq!(
2713 object_path("edge/name with spaces").as_ref(),
2714 "edge/name with spaces"
2715 );
2716 assert_eq!(
2718 object_path("a/../b").as_ref(),
2719 object_store::path::Path::from("a/../b").as_ref()
2720 );
2721 }
2722 use std::collections::HashMap;
2723
2724 fn env_of(
2727 vars: &[(&str, &str)],
2728 files: &[&str],
2729 home: Option<&str>,
2730 ) -> (HashMap<String, String>, Vec<PathBuf>, Option<PathBuf>) {
2731 (
2732 vars.iter()
2733 .map(|(k, v)| (k.to_string(), v.to_string()))
2734 .collect(),
2735 files.iter().map(PathBuf::from).collect(),
2736 home.map(PathBuf::from),
2737 )
2738 }
2739
2740 macro_rules! environment {
2741 ($vars:expr_2021, $files:expr_2021, $home:expr_2021) => {
2742 Environment {
2743 var: &|key| $vars.get(key).cloned(),
2744 exists: &|path| $files.iter().any(|f: &PathBuf| f == path),
2745 read: &|_| None,
2746 home: $home.clone(),
2747 windows: false,
2748 run: &|_, _| {
2749 Err(crate::cloud_command::CommandError::Missing(
2750 "test".to_string(),
2751 ))
2752 },
2753 all_vars: &|| Vec::new(),
2754 list: &|_| Vec::new(),
2755 }
2756 };
2757 ($vars:expr_2021, $files:expr_2021, $home:expr_2021, $contents:expr_2021) => {
2758 Environment {
2759 var: &|key| $vars.get(key).cloned(),
2760 exists: &|path| $files.iter().any(|f: &PathBuf| f == path),
2761 read: &|_| Some($contents.to_string()),
2762 home: $home.clone(),
2763 windows: false,
2764 run: &|_, _| {
2765 Err(crate::cloud_command::CommandError::Missing(
2766 "test".to_string(),
2767 ))
2768 },
2769 all_vars: &|| Vec::new(),
2770 list: &|_| Vec::new(),
2771 }
2772 };
2773 }
2774
2775 #[test]
2776 fn nothing_configured_finds_nothing() {
2777 let (vars, files, home) = env_of(&[], &[], Some("/home/u"));
2778 let env = environment!(vars, files, home);
2779 assert!(detect(&CloudConfig::default(), &env).is_empty());
2780 }
2781
2782 #[test]
2783 fn gcloud_default_credentials_are_enough_to_list_gcs() {
2784 let (vars, files, home) = env_of(
2785 &[("GOOGLE_CLOUD_PROJECT", "example-project")],
2786 &["/home/u/.config/gcloud/application_default_credentials.json"],
2787 Some("/home/u"),
2788 );
2789 let env = environment!(vars, files, home);
2790 let found = detect(&CloudConfig::default(), &env);
2791 assert_eq!(found.len(), 1);
2792 assert_eq!(found[0].kind, ProviderKind::Gcs);
2793 assert_eq!(found[0].project.as_deref(), Some("example-project"));
2794 assert!(found[0].can_list_buckets());
2795 assert_eq!(found[0].note, "gcloud");
2796 assert_eq!(
2797 found[0].detail().as_deref(),
2798 Some("project: example-project")
2799 );
2800 }
2801
2802 #[test]
2803 fn an_aws_profile_is_the_detail_for_s3() {
2804 let (vars, files, home) = env_of(&[("AWS_PROFILE", "research")], &[], Some("/home/u"));
2805 let env = environment!(vars, files, home);
2806 let found = detect(&CloudConfig::default(), &env);
2807 assert_eq!(found[0].detail().as_deref(), Some("profile: research"));
2808 }
2809
2810 #[test]
2811 fn gcs_without_a_project_is_shown_but_cannot_enumerate() {
2812 let (vars, files, home) = env_of(
2815 &[],
2816 &["/home/u/.config/gcloud/application_default_credentials.json"],
2817 Some("/home/u"),
2818 );
2819 let env = environment!(vars, files, home);
2820 let found = detect(&CloudConfig::default(), &env);
2821 assert_eq!(found.len(), 1);
2822 assert!(found[0].project.is_none());
2823 assert!(!found[0].can_list_buckets());
2824 }
2825
2826 #[test]
2827 fn a_service_account_outranks_the_developer_login_in_the_note() {
2828 let (vars, files, home) = env_of(
2829 &[("GOOGLE_SERVICE_ACCOUNT", "/keys/sa.json")],
2830 &["/home/u/.config/gcloud/application_default_credentials.json"],
2831 Some("/home/u"),
2832 );
2833 let env = environment!(vars, files, home);
2834 let found = detect(&CloudConfig::default(), &env);
2835 assert_eq!(found[0].note, "service account");
2836 }
2837
2838 #[test]
2839 fn aws_keys_in_the_environment_are_amazon_until_an_endpoint_says_otherwise() {
2840 let (vars, files, home) = env_of(&[("AWS_ACCESS_KEY_ID", "AKIA")], &[], Some("/home/u"));
2841 let env = environment!(vars, files, home);
2842 let found = detect(&CloudConfig::default(), &env);
2843 assert_eq!(found.len(), 1);
2844 assert_eq!(found[0].label, "Amazon S3");
2845 assert_eq!(found[0].note, "AWS_ACCESS_KEY_ID");
2846 }
2847
2848 #[test]
2849 fn a_custom_endpoint_is_named_by_its_host_and_not_guessed_at() {
2850 let config = CloudConfig {
2851 s3_endpoint_url: Some("http://localhost:9000".to_string()),
2852 s3_access_key_id: Some("minioadmin".to_string()),
2853 ..CloudConfig::default()
2854 };
2855 let (vars, files, home) = env_of(&[], &[], Some("/home/u"));
2856 let env = environment!(vars, files, home);
2857 let found = detect(&config, &env);
2858 assert_eq!(found.len(), 1);
2859 assert_eq!(found[0].label, "S3-compatible (localhost:9000)");
2860 assert_eq!(found[0].note, "datui config");
2861 assert_eq!(found[0].endpoint.as_deref(), Some("http://localhost:9000"));
2862 }
2863
2864 fn effective(config: &CloudConfig, env: &Environment<'_>) -> CloudConfig {
2867 let mut merged = config.clone();
2868 merged.overlay(CloudConfig::from_env(env.var));
2869 merged
2870 }
2871
2872 #[test]
2873 fn an_endpoint_from_the_environment_counts_too() {
2874 let (vars, files, home) = env_of(
2875 &[
2876 ("AWS_ACCESS_KEY_ID", "minioadmin"),
2877 ("AWS_ENDPOINT_URL", "https://minio.internal:9000/"),
2878 ],
2879 &[],
2880 Some("/home/u"),
2881 );
2882 let env = environment!(vars, files, home);
2883 let config = effective(&CloudConfig::default(), &env);
2884 let found = detect(&config, &env);
2885 assert_eq!(found[0].label, "S3-compatible (minio.internal:9000)");
2886 assert_eq!(
2889 s3_list_buckets_url(&S3Settings::from_config(&config)),
2890 "https://minio.internal:9000/"
2891 );
2892 }
2893
2894 #[test]
2895 fn the_service_specific_endpoint_variable_outranks_the_general_one() {
2896 let (vars, files, home) = env_of(
2897 &[
2898 ("AWS_ACCESS_KEY_ID", "k"),
2899 ("AWS_ENDPOINT", "http://third:1"),
2900 ("AWS_ENDPOINT_URL", "http://second:2"),
2901 ("AWS_ENDPOINT_URL_S3", "http://first:3"),
2902 ],
2903 &[],
2904 None,
2905 );
2906 let env = environment!(vars, files, home);
2907 let config = effective(&CloudConfig::default(), &env);
2908 assert_eq!(
2909 s3_list_buckets_url(&S3Settings::from_config(&config)),
2910 "http://first:3/"
2911 );
2912 assert_eq!(detect(&config, &env)[0].label, "S3-compatible (first:3)");
2913 }
2914
2915 #[test]
2916 fn a_blank_endpoint_variable_does_not_erase_the_configured_one() {
2917 let (vars, files, home) = env_of(
2918 &[("AWS_ACCESS_KEY_ID", "k"), ("AWS_ENDPOINT_URL", " ")],
2919 &[],
2920 None,
2921 );
2922 let env = environment!(vars, files, home);
2923 let file = CloudConfig {
2924 s3_endpoint_url: Some("http://localhost:9000".to_string()),
2925 ..CloudConfig::default()
2926 };
2927 let config = effective(&file, &env);
2928 assert_eq!(
2929 s3_list_buckets_url(&S3Settings::from_config(&config)),
2930 "http://localhost:9000/"
2931 );
2932 assert_eq!(
2933 detect(&config, &env)[0].label,
2934 "S3-compatible (localhost:9000)"
2935 );
2936 }
2937
2938 #[test]
2939 fn a_key_without_a_secret_still_reaches_the_builder() {
2940 use object_store::aws::AmazonS3ConfigKey;
2943 let config = CloudConfig {
2944 s3_access_key_id: Some("from-config".to_string()),
2945 ..CloudConfig::default()
2946 };
2947 let builder = s3_builder("bucket", &S3Settings::from_config(&config));
2948 assert_eq!(
2949 builder.get_config_value(&AmazonS3ConfigKey::AccessKeyId),
2950 Some("from-config".to_string())
2951 );
2952 let config = CloudConfig {
2953 s3_secret_access_key: Some("from-env".to_string()),
2954 ..CloudConfig::default()
2955 };
2956 let builder = s3_builder("bucket", &S3Settings::from_config(&config));
2957 assert_eq!(
2958 builder.get_config_value(&AmazonS3ConfigKey::SecretAccessKey),
2959 Some("from-env".to_string())
2960 );
2961 }
2962
2963 #[test]
2964 fn a_shared_credentials_file_is_enough() {
2965 let (vars, files, home) = env_of(&[], &["/home/u/.aws/credentials"], Some("/home/u"));
2966 let env = environment!(vars, files, home);
2967 let found = detect(&CloudConfig::default(), &env);
2968 assert_eq!(found.len(), 1);
2969 assert_eq!(found[0].note, "~/.aws");
2970 }
2971
2972 #[test]
2973 fn both_providers_appear_when_both_are_usable() {
2974 let (vars, files, home) = env_of(
2975 &[("AWS_ACCESS_KEY_ID", "AKIA"), ("GOOGLE_CLOUD_PROJECT", "p")],
2976 &["/home/u/.config/gcloud/application_default_credentials.json"],
2977 Some("/home/u"),
2978 );
2979 let env = environment!(vars, files, home);
2980 let found = detect(&CloudConfig::default(), &env);
2981 assert_eq!(found.len(), 2);
2982 assert_eq!(found[0].kind, ProviderKind::Gcs);
2983 assert_eq!(found[1].kind, ProviderKind::S3);
2984 }
2985
2986 #[test]
2987 fn a_blank_project_variable_is_not_a_project() {
2988 let (vars, files, home) = env_of(
2989 &[("GOOGLE_CLOUD_PROJECT", " ")],
2990 &["/home/u/.config/gcloud/application_default_credentials.json"],
2991 Some("/home/u"),
2992 );
2993 let env = environment!(vars, files, home);
2994 let found = detect(&CloudConfig::default(), &env);
2995 assert!(found[0].project.is_none());
2996 }
2997
2998 #[test]
2999 fn the_project_comes_from_the_credentials_file_when_nothing_else_says() {
3000 let adc = r#"{
3004 "type": "authorized_user",
3005 "client_id": "x.apps.googleusercontent.com",
3006 "refresh_token": "secret-and-not-read-here",
3007 "quota_project_id": "example-project"
3008 }"#;
3009 let (vars, files, home) = env_of(
3010 &[],
3011 &["/home/u/.config/gcloud/application_default_credentials.json"],
3012 Some("/home/u"),
3013 );
3014 let env = environment!(vars, files, home, adc);
3015 let found = detect(&CloudConfig::default(), &env);
3016 assert_eq!(found[0].project.as_deref(), Some("example-project"));
3017 assert!(found[0].can_list_buckets());
3018 }
3019
3020 #[test]
3021 fn an_environment_variable_outranks_the_credentials_file() {
3022 let adc = r#"{"quota_project_id": "from-the-file"}"#;
3023 let (vars, files, home) = env_of(
3024 &[("GOOGLE_CLOUD_PROJECT", "from-the-shell")],
3025 &["/home/u/.config/gcloud/application_default_credentials.json"],
3026 Some("/home/u"),
3027 );
3028 let env = environment!(vars, files, home, adc);
3029 let found = detect(&CloudConfig::default(), &env);
3030 assert_eq!(found[0].project.as_deref(), Some("from-the-shell"));
3031 }
3032
3033 #[test]
3034 fn an_unparseable_credentials_file_costs_the_project_and_nothing_else() {
3035 let (vars, files, home) = env_of(
3036 &[],
3037 &["/home/u/.config/gcloud/application_default_credentials.json"],
3038 Some("/home/u"),
3039 );
3040 let env = environment!(vars, files, home, "{ not json");
3041 let found = detect(&CloudConfig::default(), &env);
3042 assert_eq!(found.len(), 1, "the provider is still usable");
3043 assert!(found[0].project.is_none());
3044 }
3045
3046 #[test]
3047 fn gcs_buckets_come_out_of_a_real_shaped_response() {
3048 let body = r#"{
3049 "kind": "storage#buckets",
3050 "items": [
3051 {"kind": "storage#bucket", "name": "example-data", "location": "US-CENTRAL1"},
3052 {"kind": "storage#bucket", "name": "example-backups"}
3053 ]
3054 }"#;
3055 assert_eq!(
3056 parse_gcs_buckets(body).unwrap(),
3057 vec!["example-data", "example-backups"]
3058 );
3059 }
3060
3061 #[test]
3062 fn a_project_with_no_buckets_is_an_answer_not_an_error() {
3063 assert_eq!(
3064 parse_gcs_buckets(r#"{"kind": "storage#buckets"}"#).unwrap(),
3065 Vec::<String>::new()
3066 );
3067 }
3068
3069 #[test]
3070 fn a_gcs_error_body_is_reported_rather_than_read_as_emptiness() {
3071 let body =
3072 r#"{"error": {"code": 403, "message": "does not have storage.buckets.list access"}}"#;
3073 let err = parse_gcs_buckets(body).unwrap_err();
3074 assert!(err.contains("storage.buckets.list"), "{err}");
3075 }
3076
3077 #[test]
3078 fn a_malformed_gcs_entry_does_not_discard_the_page() {
3079 let body = r#"{"items": [{"name": ""}, {"nome": "typo"}, {"name": "good"}]}"#;
3080 assert_eq!(parse_gcs_buckets(body).unwrap(), vec!["good"]);
3081 }
3082
3083 #[test]
3084 fn gcs_pagination_token_is_found_when_present() {
3085 assert_eq!(
3086 gcs_next_page_token(r#"{"nextPageToken": "abc", "items": []}"#).as_deref(),
3087 Some("abc")
3088 );
3089 assert!(gcs_next_page_token(r#"{"items": []}"#).is_none());
3090 }
3091
3092 #[test]
3093 fn the_listing_goes_to_amazon_when_no_endpoint_is_set() {
3094 assert_eq!(
3095 s3_list_buckets_url(&S3Settings::from_config(&CloudConfig::default())),
3096 "https://s3.amazonaws.com/"
3097 );
3098 }
3099
3100 #[test]
3101 fn the_listing_honours_the_endpoint_override() {
3102 let config = CloudConfig {
3106 s3_endpoint_url: Some("http://localhost:9000/".to_string()),
3107 ..CloudConfig::default()
3108 };
3109 let mut effective = config.clone();
3110 effective.overlay(CloudConfig::from_env(&|name| {
3111 (name == "AWS_ENDPOINT_URL").then(|| "http://127.0.0.1:9101".to_string())
3112 }));
3113 assert_eq!(
3114 s3_list_buckets_url(&S3Settings::from_config(&effective)),
3115 "http://127.0.0.1:9101/"
3116 );
3117 let (vars, files, home) = env_of(&[("AWS_ACCESS_KEY_ID", "testing")], &[], None);
3119 let env = environment!(vars, files, home);
3120 let found = detect(&effective, &env);
3121 assert_eq!(found[0].label, "S3-compatible (127.0.0.1:9101)");
3122
3123 let mut effective = config.clone();
3125 effective.overlay(CloudConfig::from_env(&|_| None));
3126 assert_eq!(
3127 s3_list_buckets_url(&S3Settings::from_config(&effective)),
3128 "http://localhost:9000/"
3129 );
3130 }
3131
3132 #[test]
3133 fn the_override_carries_keys_and_region_too() {
3134 let config = CloudConfig {
3135 s3_access_key_id: Some("from-config".to_string()),
3136 s3_region: Some("eu-west-1".to_string()),
3137 ..CloudConfig::default()
3138 };
3139 let mut effective = config.clone();
3140 effective.overlay(CloudConfig::from_env(&|name| match name {
3141 "AWS_ACCESS_KEY_ID" => Some("from-env".to_string()),
3142 "AWS_SECRET_ACCESS_KEY" => Some("secret".to_string()),
3143 _ => None,
3144 }));
3145 assert_eq!(effective.s3_access_key_id.as_deref(), Some("from-env"));
3146 assert_eq!(effective.s3_secret_access_key.as_deref(), Some("secret"));
3147 assert_eq!(effective.s3_region.as_deref(), Some("eu-west-1"));
3148 }
3149
3150 #[test]
3151 fn s3_buckets_come_out_of_a_real_shaped_response() {
3152 let body = r#"<?xml version="1.0" encoding="UTF-8"?>
3153 <ListAllMyBucketsResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">
3154 <Owner><ID>abc</ID><DisplayName>owner</DisplayName></Owner>
3155 <Buckets>
3156 <Bucket><Name>first-bucket</Name><CreationDate>2024-01-01T00:00:00.000Z</CreationDate></Bucket>
3157 <Bucket><Name>second-bucket</Name><CreationDate>2024-02-01T00:00:00.000Z</CreationDate></Bucket>
3158 </Buckets>
3159 </ListAllMyBucketsResult>"#;
3160 assert_eq!(
3161 parse_s3_buckets(body).unwrap(),
3162 vec!["first-bucket", "second-bucket"]
3163 );
3164 }
3165
3166 #[test]
3167 fn the_owner_display_name_is_not_mistaken_for_a_bucket() {
3168 let body = r#"<ListAllMyBucketsResult>
3172 <Owner><ID>x</ID><DisplayName>Name</DisplayName></Owner>
3173 <Buckets><Bucket><Name>only-bucket</Name></Bucket></Buckets>
3174 </ListAllMyBucketsResult>"#;
3175 assert_eq!(parse_s3_buckets(body).unwrap(), vec!["only-bucket"]);
3176 }
3177
3178 #[test]
3179 fn an_s3_error_document_is_reported() {
3180 let body = r#"<Error><Code>InvalidAccessKeyId</Code><Message>The key is not valid</Message></Error>"#;
3181 let err = parse_s3_buckets(body).unwrap_err();
3182 assert!(err.contains("not valid"), "{err}");
3183 }
3184
3185 #[test]
3186 fn no_buckets_is_not_an_error() {
3187 let body = r#"<ListAllMyBucketsResult><Buckets></Buckets></ListAllMyBucketsResult>"#;
3188 assert_eq!(parse_s3_buckets(body).unwrap(), Vec::<String>::new());
3189 }
3190
3191 #[test]
3192 fn implausibly_deep_xml_is_refused_rather_than_followed() {
3193 let body = "<a>".repeat(MAX_XML_DEPTH + 2);
3194 assert!(parse_s3_buckets(&body).is_err());
3195 }
3196
3197 #[test]
3198 fn malformed_xml_is_an_error_and_not_a_panic() {
3199 assert!(parse_s3_buckets("<Buckets><Bucket><Name>x").is_err());
3200 }
3201
3202 #[test]
3203 fn an_endpoint_host_is_extracted_or_declined() {
3204 assert_eq!(
3205 endpoint_host("http://localhost:9000"),
3206 Some("localhost:9000".into())
3207 );
3208 assert_eq!(endpoint_host("https://a.b/c/d"), Some("a.b".into()));
3209 assert_eq!(endpoint_host("minio:9000"), Some("minio:9000".into()));
3210 assert_eq!(endpoint_host(""), None);
3211 assert_eq!(endpoint_host("http://"), None);
3212 }
3213}
3214
3215#[cfg(test)]
3216mod aws_role_tests {
3217 use super::*;
3218 use std::collections::HashMap;
3219
3220 fn detect_with(vars: &[(&str, &str)]) -> Vec<Provider> {
3221 let vars: HashMap<String, String> = vars
3222 .iter()
3223 .map(|(k, v)| (k.to_string(), v.to_string()))
3224 .collect();
3225 let env = Environment {
3226 var: &|key| vars.get(key).cloned(),
3227 exists: &|_| false,
3228 read: &|_| None,
3229 home: Some(PathBuf::from("/home/u")),
3230 windows: false,
3231 run: &|_, _| {
3232 Err(crate::cloud_command::CommandError::Missing(
3233 "test".to_string(),
3234 ))
3235 },
3236 all_vars: &|| Vec::new(),
3237 list: &|_| Vec::new(),
3238 };
3239 detect(&CloudConfig::default(), &env)
3240 }
3241
3242 #[test]
3243 fn a_gcloud_login_on_windows_is_found_under_appdata() {
3244 let vars: HashMap<String, String> = [("APPDATA", r"C:\Users\u\AppData\Roaming")]
3245 .iter()
3246 .map(|(k, v)| (k.to_string(), v.to_string()))
3247 .collect();
3248 let adc = PathBuf::from(r"C:\Users\u\AppData\Roaming")
3249 .join("gcloud")
3250 .join("application_default_credentials.json");
3251 let under_home = PathBuf::from(r"C:\Users\u")
3252 .join(".config")
3253 .join("gcloud")
3254 .join("application_default_credentials.json");
3255 let windows = |exists: PathBuf, windows: bool| {
3256 let env = Environment {
3257 var: &|key| vars.get(key).cloned(),
3258 exists: &|path| path == exists,
3259 read: &|_| None,
3260 home: Some(PathBuf::from(r"C:\Users\u")),
3261 windows,
3262 run: &|_, _| {
3263 Err(crate::cloud_command::CommandError::Missing(
3264 "test".to_string(),
3265 ))
3266 },
3267 all_vars: &|| Vec::new(),
3268 list: &|_| Vec::new(),
3269 };
3270 detect(&CloudConfig::default(), &env)
3271 .iter()
3272 .any(|p| p.kind == ProviderKind::Gcs)
3273 };
3274 assert!(windows(adc.clone(), true));
3275 assert!(!windows(under_home, true));
3277 assert!(!windows(adc, false));
3278 }
3279
3280 #[test]
3281 fn an_ecs_task_role_is_enough() {
3282 for key in [
3286 "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI",
3287 "AWS_CONTAINER_CREDENTIALS_FULL_URI",
3288 ] {
3289 let found = detect_with(&[(key, "/v2/credentials/abc")]);
3290 assert_eq!(found.len(), 1, "{key} should be enough");
3291 assert_eq!(found[0].kind, ProviderKind::S3);
3292 assert_eq!(found[0].label, "Amazon S3");
3293 assert_eq!(found[0].note, "container role");
3294 }
3295 }
3296
3297 #[test]
3298 fn an_eks_web_identity_is_enough() {
3299 let found = detect_with(&[
3300 ("AWS_WEB_IDENTITY_TOKEN_FILE", "/var/run/secrets/token"),
3301 ("AWS_ROLE_ARN", "arn:aws:iam::1:role/r"),
3302 ]);
3303 assert_eq!(found.len(), 1);
3304 assert_eq!(found[0].note, "web identity");
3305 }
3306
3307 #[test]
3308 fn a_region_alone_is_not_a_credential() {
3309 assert!(detect_with(&[("AWS_REGION", "us-east-1")]).is_empty());
3312 assert!(detect_with(&[("AWS_DEFAULT_REGION", "us-east-1")]).is_empty());
3313 }
3314
3315 #[test]
3316 fn a_key_still_outranks_a_role_in_the_note() {
3317 let found = detect_with(&[
3318 ("AWS_ACCESS_KEY_ID", "AKIA"),
3319 ("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", "/v2/creds"),
3320 ]);
3321 assert_eq!(found[0].note, "AWS_ACCESS_KEY_ID");
3322 }
3323}
3324
3325#[cfg(test)]
3326mod one_table_tests {
3327 use super::*;
3328 use object_store::{ObjectStore, ObjectStoreExt, PutPayload, path::Path as OsPath};
3329 use polars::prelude::*;
3330 use std::sync::Arc;
3331
3332 fn parquet(columns: &[&str]) -> Vec<u8> {
3334 let mut frame = DataFrame::new(
3335 1,
3336 columns
3337 .iter()
3338 .map(|c| Column::new((*c).into(), &[1i32]))
3339 .collect::<Vec<_>>(),
3340 )
3341 .unwrap();
3342 let mut bytes = Vec::new();
3343 ParquetWriter::new(&mut bytes).finish(&mut frame).unwrap();
3344 bytes
3345 }
3346
3347 fn kind_of(files: &[(&str, &[&str])]) -> Option<crate::discover::EntryKind> {
3349 let rt = tokio::runtime::Runtime::new().unwrap();
3350 let store: Arc<dyn ObjectStore> = Arc::new(object_store::memory::InMemory::new());
3351 rt.block_on(async {
3352 let mut objects = Vec::new();
3353 for (key, columns) in files {
3354 let bytes = parquet(columns);
3355 objects.push((format!("data/{key}"), bytes.len() as u64));
3356 store
3357 .put(
3358 &OsPath::from(format!("data/{key}")),
3359 PutPayload::from(bytes),
3360 )
3361 .await
3362 .unwrap();
3363 }
3364 kind_from_footers(&store, &objects).await
3365 })
3366 }
3367
3368 #[test]
3370 fn separate_tables_in_a_bucket_are_a_directory() {
3371 let kind = kind_of(&[
3372 ("circuits.parquet", &["circuit_id", "lat", "lng"]),
3373 ("drivers.parquet", &["driver_id", "code", "nationality"]),
3374 ("laps.parquet", &["lap", "position", "time_millis"]),
3375 ]);
3376 assert_eq!(kind, Some(crate::discover::EntryKind::Directory));
3377 }
3378
3379 #[test]
3380 fn parts_of_one_table_stay_one_dataset() {
3381 let kind = kind_of(&[
3382 ("part-00000.parquet", &["id", "ts", "amount"]),
3383 ("part-00001.parquet", &["id", "ts", "amount"]),
3384 ("part-00002.parquet", &["id", "ts", "amount"]),
3385 ]);
3386 assert_eq!(kind, Some(crate::discover::EntryKind::MultiFile));
3387 }
3388
3389 #[test]
3392 fn a_dataset_that_gained_columns_stays_one_dataset() {
3393 let kind = kind_of(&[
3394 ("date=2009-01-03.parquet", &["id", "ts"]),
3395 ("date=2015-06-01.parquet", &["id", "ts", "fee"]),
3396 (
3397 "date=2025-06-01.parquet",
3398 &["id", "ts", "fee", "witness", "address"],
3399 ),
3400 ]);
3401 assert_eq!(kind, Some(crate::discover::EntryKind::MultiFile));
3402 }
3403
3404 #[test]
3406 fn a_directory_that_cannot_be_read_is_left_as_it_was() {
3407 let rt = tokio::runtime::Runtime::new().unwrap();
3408 let store: Arc<dyn ObjectStore> = Arc::new(object_store::memory::InMemory::new());
3409 let objects = vec![
3410 ("data/a.parquet".to_string(), 10),
3411 ("data/b.parquet".to_string(), 10),
3412 ];
3413 assert_eq!(rt.block_on(kind_from_footers(&store, &objects)), None);
3414 }
3415
3416 #[test]
3418 fn a_single_file_decides_nothing() {
3419 assert_eq!(kind_of(&[("only.parquet", &["a", "b"])]), None);
3420 }
3421
3422 #[test]
3425 fn the_files_read_span_the_listing() {
3426 assert_eq!(footers_to_verify(2), vec![0, 1]);
3427 assert_eq!(footers_to_verify(3), vec![0, 1, 2]);
3428 assert_eq!(footers_to_verify(15), vec![0, 7, 14]);
3429 }
3430}