## Coverage Report

#### Created: 2025-11-22 06:00

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /Users/andrewlamb/Software/datafusion/datafusion/execution/src/object_store.rs

    Line
    Count
    Source
    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
    //! ObjectStoreRegistry holds all the object stores at Runtime with a scheme for each store.
    19
    //! This allows the user to extend DataFusion with different storage systems such as S3 or HDFS
    20
    //! and query data inside these systems.
    21
    22
    use dashmap::DashMap;
    23
    use datafusion_common::{
    24
        exec_err, internal_datafusion_err, not_impl_err, DataFusionError, Result,
    25
    };
    26
    #[cfg(not(target_arch = "wasm32"))]
    27
    use object_store::local::LocalFileSystem;
    28
    use object_store::ObjectStore;
    29
    use std::sync::Arc;
    30
    use url::Url;
    31
    32
    /// A parsed URL identifying a particular [`ObjectStore`] instance
    33
    ///
    34
    /// For example:
    35
    /// * `file://` for local file system
    36
    /// * `s3://bucket` for AWS S3 bucket
    37
    /// * `oss://bucket` for Aliyun OSS bucket
    38
    #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
    39
    pub struct ObjectStoreUrl {
    40
        url: Url,
    41
    }
    42
    43
    impl ObjectStoreUrl {
    44
        /// Parse an [`ObjectStoreUrl`] from a string
    45
        ///
    46
        /// # Example
    47
        /// ```
    48
        /// # use url::Url;
    49
        /// # use datafusion_execution::object_store::ObjectStoreUrl;
    50
        /// let object_store_url = ObjectStoreUrl::parse("s3://bucket").unwrap();
    51
        /// assert_eq!(object_store_url.as_str(), "s3://bucket/");
    52
        /// // can also access the underlying `Url`
    53
        /// let url: &Url = object_store_url.as_ref();
    54
        /// assert_eq!(url.scheme(), "s3");
    55
        /// assert_eq!(url.host_str(), Some("bucket"));
    56
        /// assert_eq!(url.path(), "/");
    57
        /// ```
    58
    0
        pub fn parse(s: impl AsRef<str>) -> Result<Self> {
    59
    0
            let mut parsed =
    60
    0
                Url::parse(s.as_ref()).map_err(|e| DataFusionError::External(Box::new(e)))?;
    61
    62
    0
            let remaining = &parsed[url::Position::BeforePath..];
    63
    0
            if !remaining.is_empty() && remaining != "/" {
    64
    0
                return exec_err!(
    65
    0
                    "ObjectStoreUrl must only contain scheme and authority, got: {remaining}"
    66
                );
    67
    0
            }
    68
    69
            // Always set path for consistency
    70
    0
            parsed.set_path("/");
    71
    0
            Ok(Self { url: parsed })
    72
    0
        }
    73
    74
        /// An [`ObjectStoreUrl`] for the local filesystem (`file://`)
    75
        ///
    76
        /// # Example
    77
        /// ```
    78
        /// # use datafusion_execution::object_store::ObjectStoreUrl;
    79
        /// let local_fs = ObjectStoreUrl::parse("file://").unwrap();
    80
        /// assert_eq!(local_fs, ObjectStoreUrl::local_filesystem())
    81
        /// ```
    82
        pub fn local_filesystem() -> Self {
    83
            Self::parse("file://").unwrap()
    84
        }
    85
    86
        /// Returns this [`ObjectStoreUrl`] as a string
    87
        pub fn as_str(&self) -> &str {
    88
            self.as_ref()
    89
        }
    90
    }
    91
    92
    impl AsRef<str> for ObjectStoreUrl {
    93
        fn as_ref(&self) -> &str {
    94
            self.url.as_ref()
    95
        }
    96
    }
    97
    98
    impl AsRef<Url> for ObjectStoreUrl {
    99
        fn as_ref(&self) -> &Url {
    100
            &self.url
    101
        }
    102
    }
    103
    104
    impl std::fmt::Display for ObjectStoreUrl {
    105
        fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
    106
            self.as_str().fmt(f)
    107
        }
    108
    }
    109
    110
    /// [`ObjectStoreRegistry`] maps a URL to an [`ObjectStore`] instance,
    111
    /// and allows DataFusion to read from different [`ObjectStore`]
    112
    /// instances. For example DataFusion might be configured so that
    113
    ///
    114
    /// 1. `s3://my_bucket/lineitem/` mapped to the `/lineitem` path on an
    115
    ///    AWS S3 object store bound to `my_bucket`
    116
    ///
    117
    /// 2. `s3://my_other_bucket/lineitem/` mapped to the (same)
    118
    ///    `/lineitem` path on a *different* AWS S3 object store bound to
    119
    ///    `my_other_bucket`
    120
    ///
    121
    /// When given a [`ListingTableUrl`], DataFusion tries to find an
    122
    /// appropriate [`ObjectStore`]. For example
    123
    ///
    124
    /// ```sql
    125
    /// create external table unicorns stored as parquet location 's3://my_bucket/lineitem/';
    126
    /// ```
    127
    ///
    128
    /// In this particular case, the url `s3://my_bucket/lineitem/` will be provided to
    129
    /// [`ObjectStoreRegistry::get_store`] and one of three things will happen:
    130
    ///
    131
    /// - If an [`ObjectStore`] has been registered with [`ObjectStoreRegistry::register_store`] with
    132
    ///   `s3://my_bucket`, that [`ObjectStore`] will be returned
    133
    ///
    134
    /// - If an AWS S3 object store can be ad-hoc discovered by the url `s3://my_bucket/lineitem/`, this
    135
    ///   object store will be registered with key `s3://my_bucket` and returned.
    136
    ///
    137
    /// - Otherwise an error will be returned, indicating that no suitable [`ObjectStore`] could
    138
    ///   be found
    139
    ///
    140
    /// This allows for two different use-cases:
    141
    ///
    142
    /// 1. Systems where object store buckets are explicitly created using DDL, can register these
    143
    ///    buckets using [`ObjectStoreRegistry::register_store`]
    144
    ///
    145
    /// 2. Systems relying on ad-hoc discovery, without corresponding DDL, can create [`ObjectStore`]
    146
    ///    lazily by providing a custom implementation of [`ObjectStoreRegistry`]
    147
    ///
    148
    /// <!-- is in a different crate so normal rustdoc links don't work -->
    149
    /// [`ListingTableUrl`]: https://docs.rs/datafusion/latest/datafusion/datasource/listing/struct.ListingTableUrl.html
    150
    /// [`ObjectStore`]: object_store::ObjectStore
    151
    pub trait ObjectStoreRegistry: Send + Sync + std::fmt::Debug + 'static {
    152
        /// If a store with the same key existed before, it is replaced and returned
    153
        fn register_store(
    154
            &self,
    155
            url: &Url,
    156
            store: Arc<dyn ObjectStore>,
    157
        ) -> Option<Arc<dyn ObjectStore>>;
    158
    159
        /// Deregister the store previously registered with the same key. Returns the
    160
        /// deregistered store if it existed.
    161
        #[allow(unused_variables)]
    162
    0
        fn deregister_store(&self, url: &Url) -> Result<Arc<dyn ObjectStore>> {
    163
    0
            not_impl_err!("ObjectStoreRegistry::deregister_store is not implemented for this ObjectStoreRegistry")
    164
    0
        }
    165
    166
        /// Get a suitable store for the provided URL. For example:
    167
        ///
    168
        /// - URL with scheme `file:///` or no scheme will return the default LocalFS store
    169
        /// - URL with scheme `s3://bucket/` will return the S3 store
    170
        /// - URL with scheme `hdfs://hostname:port/` will return the hdfs store
    171
        ///
    172
        /// If no [`ObjectStore`] found for the `url`, ad-hoc discovery may be executed depending on
    173
        /// the `url` and [`ObjectStoreRegistry`] implementation. An [`ObjectStore`] may be lazily
    174
        /// created and registered.
    175
        fn get_store(&self, url: &Url) -> Result<Arc<dyn ObjectStore>>;
    176
    }
    177
    178
    /// The default [`ObjectStoreRegistry`]
    179
    pub struct DefaultObjectStoreRegistry {
    180
        /// A map from scheme to object store that serve list / read operations for the store
    181
        object_stores: DashMap<String, Arc<dyn ObjectStore>>,
    182
    }
    183
    184
    impl std::fmt::Debug for DefaultObjectStoreRegistry {
    185
        fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
    186
            f.debug_struct("DefaultObjectStoreRegistry")
    187
                .field(
    188
                    "schemes",
    189
                    &self
    190
                        .object_stores
    191
                        .iter()
    192
    0
                        .map(|o| o.key().clone())
    193
                        .collect::<Vec<_>>(),
    194
                )
    195
                .finish()
    196
        }
    197
    }
    198
    199
    impl Default for DefaultObjectStoreRegistry {
    200
        fn default() -> Self {
    201
            Self::new()
    202
        }
    203
    }
    204
    205
    impl DefaultObjectStoreRegistry {
    206
        /// This will register [`LocalFileSystem`] to handle `file://` paths
    207
        #[cfg(not(target_arch = "wasm32"))]
    208
        pub fn new() -> Self {
    209
            let object_stores: DashMap<String, Arc<dyn ObjectStore>> = DashMap::new();
    210
            object_stores.insert("file://".to_string(), Arc::new(LocalFileSystem::new()));
    211
            Self { object_stores }
    212
        }
    213
    214
        /// Default without any backend registered.
    215
        #[cfg(target_arch = "wasm32")]
    216
        pub fn new() -> Self {
    217
            let object_stores: DashMap<String, Arc<dyn ObjectStore>> = DashMap::new();
    218
            Self { object_stores }
    219
        }
    220
    }
    221
    222
    ///
    223
    /// Stores are registered based on the scheme, host and port of the provided URL
    224
    /// with a [`LocalFileSystem::new`] automatically registered for `file://` (if the
    225
    /// target arch is not `wasm32`).
    226
    ///
    227
    /// For example:
    228
    ///
    229
    /// - `file:///my_path` will return the default LocalFS store
    230
    /// - `s3://bucket/path` will return a store registered with `s3://bucket` if any
    231
    /// - `hdfs://host:port/path` will return a store registered with `hdfs://host:port` if any
    232
    impl ObjectStoreRegistry for DefaultObjectStoreRegistry {
    233
        fn register_store(
    234
            &self,
    235
            url: &Url,
    236
            store: Arc<dyn ObjectStore>,
    237
        ) -> Option<Arc<dyn ObjectStore>> {
    238
            let s = get_url_key(url);
    239
            self.object_stores.insert(s, store)
    240
        }
    241
    242
        fn deregister_store(&self, url: &Url) -> Result<Arc<dyn ObjectStore>> {
    243
            let s = get_url_key(url);
    244
            let (_, object_store) = self.object_stores
    245
                .remove(&s)
    246
    0
                .ok_or_else(|| {
    247
    0
                    internal_datafusion_err!("Failed to deregister object store. No suitable object store found for {url}. See `RuntimeEnv::register_object_store`")
    248
    0
                })?;
    249
    250
            Ok(object_store)
    251
        }
    252
    253
        fn get_store(&self, url: &Url) -> Result<Arc<dyn ObjectStore>> {
    254
            let s = get_url_key(url);
    255
            self.object_stores
    256
                .get(&s)
    257
    0
                .map(|o| Arc::clone(o.value()))
    258
    0
                .ok_or_else(|| {
    259
    0
                    internal_datafusion_err!("No suitable object store found for {url}. See `RuntimeEnv::register_object_store`")
    260
    0
                })
    261
        }
    262
    }
    263
    264
    /// Get the key of a url for object store registration.
    265
    /// The credential info will be removed
    266
    fn get_url_key(url: &Url) -> String {
    267
        format!(
    268
            "{}://{}",
    269
            url.scheme(),
    270
            &url[url::Position::BeforeHost..url::Position::AfterPort],
    271
        )
    272
    }
    273
    274
    #[cfg(test)]
    275
    mod tests {
    276
        use super::*;
    277
    278
        #[test]
    279
        fn test_object_store_url() {
    280
            let file = ObjectStoreUrl::parse("file://").unwrap();
    281
            assert_eq!(file.as_str(), "file:///");
    282
    283
            let url = ObjectStoreUrl::parse("s3://bucket").unwrap();
    284
            assert_eq!(url.as_str(), "s3://bucket/");
    285
    286
            let url = ObjectStoreUrl::parse("s3://username:password@host:123").unwrap();
    287
            assert_eq!(url.as_str(), "s3://username:password@host:123/");
    288
    289
            let err = ObjectStoreUrl::parse("s3://bucket:invalid").unwrap_err();
    290
            assert_eq!(err.strip_backtrace(), "External error: invalid port number");
    291
    292
            let err = ObjectStoreUrl::parse("s3://bucket?").unwrap_err();
    293
            assert_eq!(err.strip_backtrace(), "Execution error: ObjectStoreUrl must only contain scheme and authority, got: ?");
    294
    295
            let err = ObjectStoreUrl::parse("s3://bucket?foo=bar").unwrap_err();
    296
            assert_eq!(err.strip_backtrace(), "Execution error: ObjectStoreUrl must only contain scheme and authority, got: ?foo=bar");
    297
    298
            let err = ObjectStoreUrl::parse("s3://host:123/foo").unwrap_err();
    299
            assert_eq!(err.strip_backtrace(), "Execution error: ObjectStoreUrl must only contain scheme and authority, got: /foo");
    300
    301
            let err =
    302
                ObjectStoreUrl::parse("s3://username:password@host:123/foo").unwrap_err();
    303
            assert_eq!(err.strip_backtrace(), "Execution error: ObjectStoreUrl must only contain scheme and authority, got: /foo");
    304
        }
    305
    306
        #[test]
    307
        fn test_get_url_key() {
    308
            let file = ObjectStoreUrl::parse("file://").unwrap();
    309
            let key = get_url_key(&file.url);
    310
            assert_eq!(key.as_str(), "file://");
    311
    312
            let url = ObjectStoreUrl::parse("s3://bucket").unwrap();
    313
            let key = get_url_key(&url.url);
    314
            assert_eq!(key.as_str(), "s3://bucket");
    315
    316
            let url = ObjectStoreUrl::parse("s3://username:password@host:123").unwrap();
    317
            let key = get_url_key(&url.url);
    318
            assert_eq!(key.as_str(), "s3://host:123");
    319
        }
    320
    }