Skip to content

Commit

Permalink
refactor: construct once
Browse files Browse the repository at this point in the history
  • Loading branch information
de-sh committed Dec 8, 2024
1 parent 6372269 commit 2ca1557
Show file tree
Hide file tree
Showing 4 changed files with 10 additions and 4 deletions.
2 changes: 1 addition & 1 deletion src/storage/azure_blob.rs
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,7 @@ impl ObjectStorageProvider for AzureBlobConfig {
RuntimeConfig::new().with_object_store_registry(Arc::new(object_store_registry))
}

fn get_object_store(&self) -> Arc<dyn super::ObjectStorage> {
fn construct_client(&self) -> Arc<dyn super::ObjectStorage> {
let azure = self.get_default_builder().build().unwrap();
// limit objectstore to a concurrent request limit
let azure = LimitStore::new(azure, super::MAX_OBJECT_STORE_REQUESTS);
Expand Down
2 changes: 1 addition & 1 deletion src/storage/localfs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ impl ObjectStorageProvider for FSConfig {
RuntimeConfig::new()
}

fn get_object_store(&self) -> Arc<dyn ObjectStorage> {
fn construct_client(&self) -> Arc<dyn ObjectStorage> {
Arc::new(LocalFS::new(self.root.clone()))
}

Expand Down
8 changes: 7 additions & 1 deletion src/storage/object_storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ use bytes::Bytes;
use chrono::Local;
use datafusion::{datasource::listing::ListingTableUrl, execution::runtime_env::RuntimeConfig};
use itertools::Itertools;
use once_cell::sync::OnceCell;
use relative_path::RelativePath;
use relative_path::RelativePathBuf;

Expand All @@ -60,7 +61,12 @@ use std::{

pub trait ObjectStorageProvider: StorageMetrics + std::fmt::Debug + Send + Sync {
fn get_datafusion_runtime(&self) -> RuntimeConfig;
fn get_object_store(&self) -> Arc<dyn ObjectStorage>;
fn construct_client(&self) -> Arc<dyn ObjectStorage>;
fn get_object_store(&self) -> Arc<dyn ObjectStorage> {
static STORE: OnceCell<Arc<dyn ObjectStorage>> = OnceCell::new();

STORE.get_or_init(|| self.construct_client()).clone()
}
fn get_endpoint(&self) -> String;
fn register_store_metrics(&self, handler: &PrometheusMetrics);
}
Expand Down
2 changes: 1 addition & 1 deletion src/storage/s3.rs
Original file line number Diff line number Diff line change
Expand Up @@ -298,7 +298,7 @@ impl ObjectStorageProvider for S3Config {
RuntimeConfig::new().with_object_store_registry(Arc::new(object_store_registry))
}

fn get_object_store(&self) -> Arc<dyn ObjectStorage> {
fn construct_client(&self) -> Arc<dyn ObjectStorage> {
let s3 = self.get_default_builder().build().unwrap();

// limit objectstore to a concurrent request limit
Expand Down

0 comments on commit 2ca1557

Please sign in to comment.