1use std::collections::HashMap;
22use std::sync::{Arc, RwLock};
23
24use async_trait::async_trait;
25use bytes::Bytes;
26use futures::StreamExt;
27use futures::stream::BoxStream;
28use iceberg::io::{
29 FileMetadata, FileRead, FileWrite, InputFile, OutputFile, Storage, StorageConfig,
30 StorageFactory,
31};
32use iceberg::{Error, ErrorKind, Result};
33use serde::{Deserialize, Serialize};
34use url::Url;
35
36#[cfg(feature = "opendal-s3")]
37use crate::s3::CustomAwsCredentialLoader;
38use crate::{OpenDalClientConfig, OpenDalStorage};
39
40pub const SCHEME_MEMORY: &str = "memory";
42pub const SCHEME_FILE: &str = "file";
43pub const SCHEME_S3: &str = "s3";
44pub const SCHEME_S3A: &str = "s3a";
45pub const SCHEME_S3N: &str = "s3n";
46pub const SCHEME_GS: &str = "gs";
47pub const SCHEME_GCS: &str = "gcs";
48pub const SCHEME_OSS: &str = "oss";
49pub const SCHEME_ABFSS: &str = "abfss";
50pub const SCHEME_ABFS: &str = "abfs";
51pub const SCHEME_WASBS: &str = "wasbs";
52pub const SCHEME_WASB: &str = "wasb";
53pub const SCHEME_HF: &str = "hf";
54pub const SCHEME_HDFS: &str = "hdfs";
55
56fn parse_scheme(scheme: &str) -> Result<&'static str> {
58 match scheme {
59 SCHEME_MEMORY => Ok("memory"),
60 SCHEME_FILE | "" => Ok("file"),
61 SCHEME_S3 | SCHEME_S3A | SCHEME_S3N => Ok("s3"),
62 SCHEME_GS | SCHEME_GCS => Ok("gcs"),
63 SCHEME_OSS => Ok("oss"),
64 SCHEME_ABFSS | SCHEME_ABFS | SCHEME_WASBS | SCHEME_WASB => Ok("azdls"),
65 SCHEME_HF => Ok("hf"),
66 SCHEME_HDFS => Ok("hdfs"),
67 s => Err(Error::new(
68 ErrorKind::FeatureUnsupported,
69 format!("Unsupported storage scheme: {s}"),
70 )),
71 }
72}
73
74fn extract_scheme(path: &str) -> Result<&'static str> {
76 let url = Url::parse(path).map_err(|e| {
77 Error::new(
78 ErrorKind::DataInvalid,
79 format!("Invalid path: {path}, failed to parse URL: {e}"),
80 )
81 })?;
82 parse_scheme(url.scheme())
83}
84
85fn build_storage_for_scheme(
87 scheme: &'static str,
88 props: &HashMap<String, String>,
89 #[cfg(feature = "opendal-s3")] customized_credential_load: &Option<CustomAwsCredentialLoader>,
90) -> Result<OpenDalStorage> {
91 let client_config = OpenDalClientConfig::from_properties(props)?;
92 match scheme {
93 #[cfg(feature = "opendal-s3")]
94 "s3" => {
95 let config = crate::s3::s3_config_parse(props.clone())?;
96 Ok(OpenDalStorage::S3 {
97 config: Arc::new(config),
98 customized_credential_load: customized_credential_load.clone(),
99 client_config,
100 })
101 }
102 #[cfg(feature = "opendal-gcs")]
103 "gcs" => {
104 let config = crate::gcs::gcs_config_parse(props.clone())?;
105 Ok(OpenDalStorage::Gcs {
106 config: Arc::new(config),
107 client_config,
108 })
109 }
110 #[cfg(feature = "opendal-oss")]
111 "oss" => {
112 let config = crate::oss::oss_config_parse(props.clone())?;
113 Ok(OpenDalStorage::Oss {
114 config: Arc::new(config),
115 client_config,
116 })
117 }
118 #[cfg(feature = "opendal-azdls")]
119 "azdls" => {
120 let config = crate::azdls::azdls_config_parse(props.clone())?;
121 Ok(OpenDalStorage::Azdls {
122 config: Arc::new(config),
123 client_config,
124 })
125 }
126 #[cfg(feature = "opendal-fs")]
127 "file" => Ok(OpenDalStorage::LocalFs { client_config }),
128 #[cfg(feature = "opendal-memory")]
129 "memory" => Ok(OpenDalStorage::Memory {
130 operator: crate::memory::memory_config_build()?,
131 client_config,
132 }),
133 #[cfg(feature = "opendal-hf")]
134 "hf" => {
135 let config = crate::hf::hf_config_parse(props.clone())?;
136 Ok(OpenDalStorage::Hf {
137 config: Arc::new(config),
138 client_config,
139 })
140 }
141 #[cfg(feature = "opendal-hdfs-native")]
142 "hdfs" => {
143 let config = crate::hdfs_native::hdfs_native_config_parse(props.clone())?;
144 Ok(OpenDalStorage::HdfsNative(crate::HdfsNativeStorage::new(
145 config,
146 client_config,
147 )))
148 }
149 unsupported => Err(Error::new(
150 ErrorKind::FeatureUnsupported,
151 format!("Unsupported storage scheme: {unsupported}"),
152 )),
153 }
154}
155
156#[derive(Clone, Debug, Serialize, Deserialize)]
181pub struct OpenDalResolvingStorageFactory {
182 #[cfg(feature = "opendal-s3")]
184 #[serde(
185 skip_deserializing,
186 skip_serializing_if = "Option::is_none",
187 serialize_with = "crate::serialize_custom_credential_loader"
188 )]
189 customized_credential_load: Option<CustomAwsCredentialLoader>,
190}
191
192impl Default for OpenDalResolvingStorageFactory {
193 fn default() -> Self {
194 Self::new()
195 }
196}
197
198impl OpenDalResolvingStorageFactory {
199 pub fn new() -> Self {
201 Self {
202 #[cfg(feature = "opendal-s3")]
203 customized_credential_load: None,
204 }
205 }
206
207 #[cfg(feature = "opendal-s3")]
209 pub fn with_s3_credential_loader(mut self, loader: CustomAwsCredentialLoader) -> Self {
210 self.customized_credential_load = Some(loader);
211 self
212 }
213}
214
215#[typetag::serde]
216impl StorageFactory for OpenDalResolvingStorageFactory {
217 fn build(&self, config: &StorageConfig) -> Result<Arc<dyn Storage>> {
218 Ok(Arc::new(OpenDalResolvingStorage {
219 props: config.props().clone(),
220 storages: RwLock::new(HashMap::new()),
221 #[cfg(feature = "opendal-s3")]
222 customized_credential_load: self.customized_credential_load.clone(),
223 }))
224 }
225}
226
227#[derive(Debug, Serialize, Deserialize)]
234pub struct OpenDalResolvingStorage {
235 props: HashMap<String, String>,
237 #[serde(skip, default)]
239 storages: RwLock<HashMap<&'static str, Arc<OpenDalStorage>>>,
240 #[cfg(feature = "opendal-s3")]
242 #[serde(skip)]
243 customized_credential_load: Option<CustomAwsCredentialLoader>,
244}
245
246impl OpenDalResolvingStorage {
247 fn resolve(&self, path: &str) -> Result<Arc<OpenDalStorage>> {
250 let scheme = extract_scheme(path)?;
251
252 {
254 let cache = self
255 .storages
256 .read()
257 .map_err(|_| Error::new(ErrorKind::Unexpected, "Storage cache lock poisoned"))?;
258 if let Some(storage) = cache.get(&scheme) {
259 return Ok(storage.clone());
260 }
261 }
262
263 let mut cache = self
265 .storages
266 .write()
267 .map_err(|_| Error::new(ErrorKind::Unexpected, "Storage cache lock poisoned"))?;
268
269 if let Some(storage) = cache.get(&scheme) {
271 return Ok(storage.clone());
272 }
273
274 let storage = build_storage_for_scheme(
275 scheme,
276 &self.props,
277 #[cfg(feature = "opendal-s3")]
278 &self.customized_credential_load,
279 )?;
280 let storage = Arc::new(storage);
281 cache.insert(scheme, storage.clone());
282 Ok(storage)
283 }
284}
285
286#[async_trait]
287#[typetag::serde]
288impl Storage for OpenDalResolvingStorage {
289 async fn exists(&self, path: &str) -> Result<bool> {
290 self.resolve(path)?.exists(path).await
291 }
292
293 async fn metadata(&self, path: &str) -> Result<FileMetadata> {
294 self.resolve(path)?.metadata(path).await
295 }
296
297 async fn read(&self, path: &str) -> Result<Bytes> {
298 self.resolve(path)?.read(path).await
299 }
300
301 async fn reader(&self, path: &str) -> Result<Box<dyn FileRead>> {
302 self.resolve(path)?.reader(path).await
303 }
304
305 async fn write(&self, path: &str, bs: Bytes) -> Result<()> {
306 self.resolve(path)?.write(path, bs).await
307 }
308
309 async fn writer(&self, path: &str) -> Result<Box<dyn FileWrite>> {
310 self.resolve(path)?.writer(path).await
311 }
312
313 async fn delete(&self, path: &str) -> Result<()> {
314 self.resolve(path)?.delete(path).await
315 }
316
317 async fn delete_prefix(&self, path: &str) -> Result<()> {
318 self.resolve(path)?.delete_prefix(path).await
319 }
320
321 async fn delete_stream(&self, mut paths: BoxStream<'static, String>) -> Result<()> {
322 let mut grouped: HashMap<&'static str, Vec<String>> = HashMap::new();
325 while let Some(path) = paths.next().await {
326 let scheme = extract_scheme(&path)?;
327 grouped.entry(scheme).or_default().push(path);
328 }
329
330 for (_, paths) in grouped {
331 let storage = self.resolve(&paths[0])?;
332 storage
333 .delete_stream(futures::stream::iter(paths).boxed())
334 .await?;
335 }
336 Ok(())
337 }
338
339 fn new_input(&self, path: &str) -> Result<InputFile> {
340 Ok(InputFile::new(
341 Arc::new(self.resolve(path)?.as_ref().clone()),
342 path.to_string(),
343 ))
344 }
345
346 fn new_output(&self, path: &str) -> Result<OutputFile> {
347 Ok(OutputFile::new(
348 Arc::new(self.resolve(path)?.as_ref().clone()),
349 path.to_string(),
350 ))
351 }
352}
353
354#[cfg(test)]
355mod tests {
356 use std::time::Duration;
357
358 use super::*;
359 use crate::OPENDAL_IO_TIMEOUT_MS;
360
361 #[cfg(feature = "opendal-s3")]
362 #[derive(Debug)]
363 struct EmptyCredentialLoader;
364
365 #[cfg(feature = "opendal-s3")]
366 impl crate::s3::ProvideCredential for EmptyCredentialLoader {
367 type Credential = crate::s3::AwsCredential;
368
369 async fn provide_credential(
370 &self,
371 _ctx: &reqsign_core::Context,
372 ) -> reqsign_core::Result<Option<Self::Credential>> {
373 Ok(None)
374 }
375 }
376
377 #[cfg(feature = "opendal-s3")]
378 #[test]
379 fn test_custom_credential_loader_serialization_fails() {
380 let factory = OpenDalResolvingStorageFactory::new()
381 .with_s3_credential_loader(CustomAwsCredentialLoader::new(EmptyCredentialLoader));
382 let file_io = iceberg::io::FileIOBuilder::new(Arc::new(factory)).build();
383
384 let err = file_io.serialize_all().unwrap_err();
385 assert!(
386 err.to_string()
387 .contains("custom AWS credential loaders cannot be serialized")
388 );
389 }
390
391 fn empty_resolving_storage() -> OpenDalResolvingStorage {
394 OpenDalResolvingStorage {
395 props: HashMap::new(),
396 storages: RwLock::new(HashMap::new()),
397 #[cfg(feature = "opendal-s3")]
398 customized_credential_load: None,
399 }
400 }
401
402 #[test]
403 fn test_resolve_propagates_io_timeout() {
404 let mut storage = empty_resolving_storage();
405 storage
406 .props
407 .insert(OPENDAL_IO_TIMEOUT_MS.to_string(), "45000".to_string());
408
409 let paths: &[&str] = &[
410 #[cfg(feature = "opendal-memory")]
411 "memory:/key",
412 #[cfg(feature = "opendal-fs")]
413 "file:/key",
414 #[cfg(feature = "opendal-s3")]
415 "s3://bucket/key",
416 #[cfg(feature = "opendal-gcs")]
417 "gs://bucket/key",
418 #[cfg(feature = "opendal-oss")]
419 "oss://bucket/key",
420 #[cfg(feature = "opendal-azdls")]
421 "abfss://myfs@myaccount.dfs.core.windows.net/key",
422 #[cfg(feature = "opendal-hf")]
423 "hf://datasets/user/repo/key",
424 ];
425 for path in paths {
426 let resolved = storage.resolve(path).unwrap();
427 assert_eq!(
428 resolved.client_config().io_timeout(),
429 Duration::from_secs(45),
430 "{path}"
431 );
432 }
433 }
434
435 #[cfg(feature = "opendal-s3")]
436 #[test]
437 fn test_resolve_s3_aliases_share_instance() {
438 let storage = empty_resolving_storage();
439
440 let a = storage.resolve("s3://bucket/key").unwrap();
444 let b = storage.resolve("s3a://bucket/key").unwrap();
445 let c = storage.resolve("s3n://bucket/key").unwrap();
446
447 assert!(Arc::ptr_eq(&a, &b), "s3 and s3a should share one instance");
448 assert!(Arc::ptr_eq(&a, &c), "s3 and s3n should share one instance");
449 }
450
451 #[cfg(feature = "opendal-hdfs-native")]
452 #[test]
453 fn test_resolve_hdfs_native_returns_hdfs_variant() {
454 let storage = empty_resolving_storage();
455
456 let resolved = storage.resolve("hdfs://nameservice1/a/b").unwrap();
457
458 assert!(
459 matches!(&*resolved, OpenDalStorage::HdfsNative(_)),
460 "expected HdfsNative variant, got {resolved:?}"
461 );
462 }
463
464 #[cfg(feature = "opendal-hdfs-native")]
465 #[test]
466 fn test_resolve_hdfs_native_distinct_authorities_share_instance() {
467 let storage = empty_resolving_storage();
468
469 let a = storage.resolve("hdfs://ns1/a").unwrap();
470 let b = storage.resolve("hdfs://ns2/b").unwrap();
471
472 assert!(
473 Arc::ptr_eq(&a, &b),
474 "different authorities should share the OpenDalStorage::HdfsNative instance (operator cache is internal)"
475 );
476 }
477
478 #[cfg(feature = "opendal-azdls")]
479 #[test]
480 fn test_resolve_azdls_aliases_share_instance() {
481 let storage = empty_resolving_storage();
482
483 let path_for = |scheme: &str| {
484 format!("{scheme}://myfs@myaccount.dfs.core.windows.net/path/to/file.parquet")
485 };
486
487 let abfss = storage.resolve(&path_for("abfss")).unwrap();
489 let abfs = storage.resolve(&path_for("abfs")).unwrap();
490
491 assert!(
492 Arc::ptr_eq(&abfss, &abfs),
493 "abfss and abfs should share one instance"
494 );
495 }
496}