Skip to content

Commit 5378981

Browse files
author
yangzhong
committed
Allow ObjectStoreProvider to return None (return Result<Option> rather than Result)
1 parent abcb39d commit 5378981

File tree

1 file changed

+33
-18
lines changed

1 file changed

+33
-18
lines changed

datafusion/core/src/datasource/object_store.rs

Lines changed: 33 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ impl std::fmt::Display for ObjectStoreUrl {
8484
/// Object store provider can detector an object store based on the url
8585
pub trait ObjectStoreProvider: Send + Sync + 'static {
8686
/// Return an ObjectStore for the provided url based on its scheme and authority
87-
fn get_by_url(&self, url: &Url) -> Result<Arc<dyn ObjectStore>>;
87+
fn get_by_url(&self, url: &Url) -> Result<Option<Arc<dyn ObjectStore>>>;
8888
}
8989

9090
/// Object store registry
@@ -150,29 +150,44 @@ impl ObjectStoreRegistry {
150150
pub fn get_by_url(&self, url: impl AsRef<Url>) -> Result<Arc<dyn ObjectStore>> {
151151
let url = url.as_ref();
152152
// First check whether can get object store from registry
153-
let store = {
153+
let mut store = {
154154
let stores = self.object_stores.read();
155155
let s = &url[url::Position::BeforeScheme..url::Position::BeforePath];
156156
stores.get(s).cloned()
157157
};
158158

159-
match store {
160-
Some(store) => Ok(store),
161-
None => match &self.provider {
162-
Some(provider) => {
163-
let store = provider.get_by_url(url)?;
164-
let mut stores = self.object_stores.write();
165-
let key =
166-
&url[url::Position::BeforeScheme..url::Position::BeforePath];
167-
stores.insert(key.to_owned(), store.clone());
168-
Ok(store)
169-
}
170-
None => Err(DataFusionError::Internal(format!(
171-
"No suitable object store found for {}",
172-
url
173-
))),
174-
},
159+
// If no store found, then try to detector based on its url from provider.
160+
if store.is_none() {
161+
store = self.get_and_set_by_provider(url)?;
175162
}
163+
164+
if let Some(store) = store {
165+
Ok(store)
166+
} else {
167+
Err(DataFusionError::Internal(format!(
168+
"No suitable object store found for {}",
169+
url
170+
)))
171+
}
172+
}
173+
174+
/// Try to get a suitable store for the provided URL from provider.
175+
/// If find one, then will register it automatically
176+
fn get_and_set_by_provider(&self, url: &Url) -> Result<Option<Arc<dyn ObjectStore>>> {
177+
let store = if let Some(provider) = &self.provider {
178+
if let Some(store) = provider.get_by_url(url)? {
179+
let mut stores = self.object_stores.write();
180+
let key = &url[url::Position::BeforeScheme..url::Position::BeforePath];
181+
stores.insert(key.to_owned(), store.clone());
182+
Some(store)
183+
} else {
184+
None
185+
}
186+
} else {
187+
None
188+
};
189+
190+
Ok(store)
176191
}
177192
}
178193

0 commit comments

Comments
 (0)