Add SessionContext::register_object_store (apache#10621)

Author: alamb

datafusion-cli/src/exec.rs

@@ -379,7 +379,7 @@ pub(crate) async fn register_object_store_and_config_extensions(
     let store = get_object_store(&ctx.state(), scheme, url, &table_options).await?;
 
     // Register the retrieved object store in the session context's runtime environment
-    ctx.runtime_env().register_object_store(url, store);
+    ctx.register_object_store(url, store);
 
     Ok(())
 }

datafusion-examples/examples/external_dependency/dataframe-to-s3.rs

@@ -49,8 +49,7 @@ async fn main() -> Result<()> {
     let path = format!("s3://{bucket_name}");
     let s3_url = Url::parse(&path).unwrap();
     let arc_s3 = Arc::new(s3);
-    ctx.runtime_env()
-        .register_object_store(&s3_url, arc_s3.clone());
+    ctx.register_object_store(&s3_url, arc_s3.clone());
 
     let path = format!("s3://{bucket_name}/test_data/");
     let file_format = ParquetFormat::default().with_enable_pruning(true);

datafusion-examples/examples/external_dependency/query-aws-s3.rs

@@ -48,8 +48,7 @@ async fn main() -> Result<()> {
 
     let path = format!("s3://{bucket_name}");
     let s3_url = Url::parse(&path).unwrap();
-    ctx.runtime_env()
-        .register_object_store(&s3_url, Arc::new(s3));
+    ctx.register_object_store(&s3_url, Arc::new(s3));
 
     // cannot query the parquet files from this bucket because the path contains a whitespace
     // and we don't support that yet

datafusion-examples/examples/parquet_sql_multiple_files.rs

@@ -80,7 +80,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let local_fs = Arc::new(LocalFileSystem::default());
 
     let u = url::Url::parse("file://./")?;
-    ctx.runtime_env().register_object_store(&u, local_fs);
+    ctx.register_object_store(&u, local_fs);
 
     // Register a listing table - this will use all files in the directory as data sources
     // for the query

datafusion-examples/examples/query-http-csv.rs

@@ -34,8 +34,7 @@ async fn main() -> Result<()> {
         .with_url(base_url.clone())
         .build()
         .unwrap();
-    ctx.runtime_env()
-        .register_object_store(&base_url, Arc::new(http_store));
+    ctx.register_object_store(&base_url, Arc::new(http_store));
 
     // register csv file with the execution context
     ctx.register_csv(

datafusion/core/src/datasource/physical_plan/avro.rs

@@ -261,9 +261,7 @@ mod tests {
         let state = session_ctx.state();
 
         let url = Url::parse("file://").unwrap();
-        state
-            .runtime_env()
-            .register_object_store(&url, store.clone());
+        session_ctx.register_object_store(&url, store.clone());
 
         let testdata = crate::test_util::arrow_test_data();
         let filename = format!("{testdata}/avro/alltypes_plain.avro");

datafusion/core/src/datasource/physical_plan/csv.rs

@@ -895,7 +895,7 @@ mod tests {
     ) -> Result<()> {
         let ctx = SessionContext::new();
         let url = Url::parse("file://").unwrap();
-        ctx.runtime_env().register_object_store(&url, store.clone());
+        ctx.register_object_store(&url, store.clone());
 
         let task_ctx = ctx.task_ctx();
 
@@ -968,9 +968,7 @@ mod tests {
         store.put(&path, data).await.unwrap();
 
         let url = Url::parse("memory://").unwrap();
-        session_ctx
-            .runtime_env()
-            .register_object_store(&url, Arc::new(store));
+        session_ctx.register_object_store(&url, Arc::new(store));
 
         let df = session_ctx
             .read_csv("memory:///", CsvReadOptions::new())
@@ -999,7 +997,7 @@ mod tests {
         let tmp_dir = TempDir::new()?;
         let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?);
         let local_url = Url::parse("file://local").unwrap();
-        ctx.runtime_env().register_object_store(&local_url, local);
+        ctx.register_object_store(&local_url, local);
         let options = CsvReadOptions::default()
             .schema_infer_max_records(2)
             .has_header(true);
@@ -1039,7 +1037,7 @@ mod tests {
         let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?);
         let local_url = Url::parse("file://local").unwrap();
 
-        ctx.runtime_env().register_object_store(&local_url, local);
+        ctx.register_object_store(&local_url, local);
 
         // execute a simple query and write the results to CSV
         let out_dir = tmp_dir.as_ref().to_str().unwrap().to_string() + "/out/";

datafusion/core/src/datasource/physical_plan/json.rs

@@ -446,7 +446,7 @@ mod tests {
     ) -> Result<()> {
         let ctx = SessionContext::new();
         let url = Url::parse("file://").unwrap();
-        ctx.runtime_env().register_object_store(&url, store.clone());
+        ctx.register_object_store(&url, store.clone());
         let filename = "1.json";
         let tmp_dir = TempDir::new()?;
         let file_groups = partitioned_file_groups(
@@ -752,7 +752,7 @@ mod tests {
         let tmp_dir = TempDir::new()?;
         let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?);
         let local_url = Url::parse("file://local").unwrap();
-        ctx.runtime_env().register_object_store(&local_url, local);
+        ctx.register_object_store(&local_url, local);
 
         // execute a simple query and write the results to CSV
         let out_dir = tmp_dir.as_ref().to_str().unwrap().to_string() + "/out/";
@@ -845,7 +845,7 @@ mod tests {
         let tmp_dir = TempDir::new()?;
         let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?);
         let local_url = Url::parse("file://local").unwrap();
-        ctx.runtime_env().register_object_store(&local_url, local);
+        ctx.register_object_store(&local_url, local);
         let options = CsvReadOptions::default()
             .schema_infer_max_records(2)
             .has_header(true);

datafusion/core/src/datasource/physical_plan/parquet/mod.rs

@@ -996,7 +996,7 @@ mod tests {
         let tmp_dir = TempDir::new()?;
         let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?);
         let local_url = Url::parse("file://local").unwrap();
-        ctx.runtime_env().register_object_store(&local_url, local);
+        ctx.register_object_store(&local_url, local);
 
         let options = CsvReadOptions::default()
             .schema_infer_max_records(2)
@@ -2052,7 +2052,7 @@ mod tests {
         // register a local file system object store for /tmp directory
         let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?);
         let local_url = Url::parse("file://local").unwrap();
-        ctx.runtime_env().register_object_store(&local_url, local);
+        ctx.register_object_store(&local_url, local);
 
         // Configure listing options
         let file_format = ParquetFormat::default().with_enable_pruning(true);

datafusion/core/src/execution/context/mod.rs

@@ -91,6 +91,7 @@ use sqlparser::dialect::dialect_from_str;
 
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
+use object_store::ObjectStore;
 use parking_lot::RwLock;
 use url::Url;
 use uuid::Uuid;
@@ -354,6 +355,29 @@ impl SessionContext {
         self
     }
 
+    /// Registers an [`ObjectStore`] to be used with a specific URL prefix.
+    ///
+    /// See [`RuntimeEnv::register_object_store`] for more details.
+    ///
+    /// # Example: register a local object store for the "file://" URL prefix
+    /// ```
+    /// # use std::sync::Arc;
+    /// # use datafusion::prelude::SessionContext;
+    /// # use datafusion_execution::object_store::ObjectStoreUrl;
+    /// let object_store_url = ObjectStoreUrl::parse("file://").unwrap();
+    /// let object_store = object_store::local::LocalFileSystem::new();
+    /// let mut ctx = SessionContext::new();
+    /// // All files with the file:// url prefix will be read from the local file system
+    /// ctx.register_object_store(object_store_url.as_ref(), Arc::new(object_store));
+    /// ```
+    pub fn register_object_store(
+        &self,
+        url: &Url,
+        object_store: Arc<dyn ObjectStore>,
+    ) -> Option<Arc<dyn ObjectStore>> {
+        self.runtime_env().register_object_store(url, object_store)
+    }
+
     /// Registers the [`RecordBatch`] as the specified table name
     pub fn register_batch(
         &self,