Skip to main content

iceberg_storage_opendal/
resolving.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Resolving storage that auto-detects the scheme from a path and delegates
19//! to the appropriate [`OpenDalStorage`] variant.
20
21use 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
40/// Schemes supported by OpenDalResolvingStorage
41pub 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
56/// Parse a URL scheme string.
57fn 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
74/// Extract the scheme from a path URL.
75fn 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
85/// Build an [`OpenDalStorage`] variant for the given scheme and config properties.
86fn 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/// A resolving storage factory that creates [`OpenDalResolvingStorage`] instances.
157///
158/// This factory accepts paths from any supported storage system and dynamically
159/// delegates operations to the appropriate [`OpenDalStorage`] variant based on
160/// the path scheme.
161///
162/// # Serialization
163///
164/// Serialization fails when a custom S3 credential loader is configured because the loader holds
165/// process-local state that cannot be reconstructed in another process. Construct the factory
166/// without a custom loader before serializing it.
167///
168/// # Example
169///
170/// ```rust,ignore
171/// use std::sync::Arc;
172/// use iceberg::io::FileIOBuilder;
173/// use iceberg_storage_opendal::OpenDalResolvingStorageFactory;
174///
175/// let factory = OpenDalResolvingStorageFactory::new();
176/// let file_io = FileIOBuilder::new(Arc::new(factory))
177///     .with_prop("s3.region", "us-east-1")
178///     .build();
179/// ```
180#[derive(Clone, Debug, Serialize, Deserialize)]
181pub struct OpenDalResolvingStorageFactory {
182    /// Custom AWS credential loader for S3 storage.
183    #[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    /// Create a new resolving storage factory.
200    pub fn new() -> Self {
201        Self {
202            #[cfg(feature = "opendal-s3")]
203            customized_credential_load: None,
204        }
205    }
206
207    /// Set a custom AWS credential loader for S3 storage.
208    #[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/// A resolving storage that auto-detects the scheme from a path and delegates
228/// to the appropriate [`OpenDalStorage`] variant.
229///
230/// Sub-storages are lazily created on first use for each scheme and cached
231/// for subsequent operations. Scheme aliases like `s3`/`s3a`/`s3n` map to
232/// the same canonical scheme, so they share a storage instance.
233#[derive(Debug, Serialize, Deserialize)]
234pub struct OpenDalResolvingStorage {
235    /// Configuration properties shared across all backends.
236    props: HashMap<String, String>,
237    /// Cache of canonical scheme to storage mappings.
238    #[serde(skip, default)]
239    storages: RwLock<HashMap<&'static str, Arc<OpenDalStorage>>>,
240    /// Custom AWS credential loader for S3 storage.
241    #[cfg(feature = "opendal-s3")]
242    #[serde(skip)]
243    customized_credential_load: Option<CustomAwsCredentialLoader>,
244}
245
246impl OpenDalResolvingStorage {
247    /// Resolve the storage for the given path by extracting the canonical scheme and
248    /// returning the cached or newly-created [`OpenDalStorage`].
249    fn resolve(&self, path: &str) -> Result<Arc<OpenDalStorage>> {
250        let scheme = extract_scheme(path)?;
251
252        // Fast path: check read lock first.
253        {
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        // Slow path: build and insert under write lock.
264        let mut cache = self
265            .storages
266            .write()
267            .map_err(|_| Error::new(ErrorKind::Unexpected, "Storage cache lock poisoned"))?;
268
269        // Double-check after acquiring write lock.
270        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        // Group paths by canonical scheme so each resolved storage receives a batch,
323        // avoiding repeated operator creation per path.
324        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    /// Builds a resolving storage with empty props, suitable for `resolve()`
392    /// calls that don't actually hit any backend.
393    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        // All three S3-family schemes must collapse to a single cached
441        // `Arc<OpenDalStorage>` so that catalogs handing the resolver a mix
442        // of `s3://`, `s3a://`, `s3n://` paths don't rebuild operators.
443        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        // All Azure schemes collapse onto one cached instance.
488        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}