diff --git a/datafusion-cli/src/object_storage/instrumented.rs b/datafusion-cli/src/object_storage/instrumented.rs index c4bd44011dc7..4c930f1d8631 100644 --- a/datafusion-cli/src/object_storage/instrumented.rs +++ b/datafusion-cli/src/object_storage/instrumented.rs @@ -44,6 +44,13 @@ impl ObjectStoreRegistry for InstrumentedObjectStoreRegistry { self.inner.register_store(url, store) } + fn deregister_store( + &self, + url: &Url, + ) -> datafusion::common::Result> { + self.inner.deregister_store(url) + } + fn get_store(&self, url: &Url) -> datafusion::common::Result> { self.inner.get_store(url) } diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index a8148b80495e..eace5610e117 100644 --- a/datafusion/core/src/execution/context/mod.rs +++ b/datafusion/core/src/execution/context/mod.rs @@ -505,6 +505,13 @@ impl SessionContext { self.runtime_env().register_object_store(url, object_store) } + /// Deregisters an [`ObjectStore`] associated with the specific URL prefix. + /// + /// See [`RuntimeEnv::deregister_object_store`] for more details. + pub fn deregister_object_store(&self, url: &Url) -> Result> { + self.runtime_env().deregister_object_store(url) + } + /// Registers the [`RecordBatch`] as the specified table name pub fn register_batch( &self, diff --git a/datafusion/execution/src/object_store.rs b/datafusion/execution/src/object_store.rs index ef83128ac681..21135cfde745 100644 --- a/datafusion/execution/src/object_store.rs +++ b/datafusion/execution/src/object_store.rs @@ -154,6 +154,10 @@ pub trait ObjectStoreRegistry: Send + Sync + std::fmt::Debug + 'static { store: Arc, ) -> Option>; + /// Deregister the store previously registered with the same key. Returns the + /// deregistered store if it existed. + fn deregister_store(&self, url: &Url) -> Result>; + /// Get a suitable store for the provided URL. For example: /// /// - URL with scheme `file:///` or no scheme will return the default LocalFS store @@ -230,6 +234,17 @@ impl ObjectStoreRegistry for DefaultObjectStoreRegistry { self.object_stores.insert(s, store) } + fn deregister_store(&self, url: &Url) -> Result> { + let s = get_url_key(url); + let (_, object_store) = self.object_stores + .remove(&s) + .ok_or_else(|| { + internal_datafusion_err!("Failed to deregister object store. No suitable object store found for {url}. See `RuntimeEnv::register_object_store`") + })?; + + Ok(object_store) + } + fn get_store(&self, url: &Url) -> Result> { let s = get_url_key(url); self.object_stores diff --git a/datafusion/execution/src/runtime_env.rs b/datafusion/execution/src/runtime_env.rs index db045a8b7e8a..b0d0a966b7a2 100644 --- a/datafusion/execution/src/runtime_env.rs +++ b/datafusion/execution/src/runtime_env.rs @@ -114,8 +114,6 @@ impl RuntimeEnv { /// ``` /// /// # Example: Register remote URL object store like [Github](https://github.com) - /// - /// /// ``` /// # use std::sync::Arc; /// # use url::Url; @@ -141,6 +139,12 @@ impl RuntimeEnv { self.object_store_registry.register_store(url, object_store) } + /// Deregisters a custom `ObjectStore` previously registered for a specific url. + /// See [`ObjectStoreRegistry::deregister_store`] for more details. + pub fn deregister_object_store(&self, url: &Url) -> Result> { + self.object_store_registry.deregister_store(url) + } + /// Retrieves a `ObjectStore` instance for a url by consulting the /// registry. See [`ObjectStoreRegistry::get_store`] for more /// details.