From 4154dc3479459d019af5fb93bbe529a39cf42655 Mon Sep 17 00:00:00 2001
From: Will Jones <willjones127@gmail.com>
Date: Fri, 11 Aug 2023 12:01:17 -0700
Subject: [PATCH] handle params better

---
 rust/src/dataset.rs         | 10 ++++------
 rust/src/io/commit.rs       |  2 ++
 rust/src/io/object_store.rs | 19 +++++++++++++++----
 3 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/rust/src/dataset.rs b/rust/src/dataset.rs
index 64ce6f3ca5..21bf78d692 100644
--- a/rust/src/dataset.rs
+++ b/rust/src/dataset.rs
@@ -434,12 +434,10 @@ impl Dataset {
         };
 
         // Need to include params here because it might include a commit mechanism.
-        let (object_store, _) = ObjectStore::from_uri_and_params(
-            &self.object_store.base_uri(),
-            params.store_params.clone().unwrap_or_default(),
-        )
-        .await?;
-        let object_store = Arc::new(object_store);
+        let object_store = Arc::new(
+            self.object_store()
+                .with_params(&params.store_params.clone().unwrap_or_default()),
+        );
 
         let (stream, schema) = reader_to_stream(batches)?;
 
diff --git a/rust/src/io/commit.rs b/rust/src/io/commit.rs
index 09535b932e..0f770eaf5a 100644
--- a/rust/src/io/commit.rs
+++ b/rust/src/io/commit.rs
@@ -807,6 +807,8 @@ mod tests {
                 }
                 _ => unreachable!(),
             }
+
+            dataset.validate().await.unwrap()
         }
     }
 }
diff --git a/rust/src/io/object_store.rs b/rust/src/io/object_store.rs
index 6ecb90686b..4a00bfed67 100644
--- a/rust/src/io/object_store.rs
+++ b/rust/src/io/object_store.rs
@@ -258,6 +258,21 @@ impl ObjectStore {
         ))
     }
 
+    pub fn with_params(&self, params: &ObjectStoreParams) -> Self {
+        Self {
+            inner: params
+                .object_store_wrapper
+                .as_ref()
+                .map(|w| w.wrap(self.inner.clone()))
+                .unwrap_or_else(|| self.inner.clone()),
+            commit_handler: params
+                .commit_handler
+                .clone()
+                .unwrap_or(self.commit_handler.clone()),
+            ..self.clone()
+        }
+    }
+
     fn new_from_path(str_path: &str, params: &ObjectStoreParams) -> Result<(Self, Path)> {
         let expanded = tilde(str_path).to_string();
         let expanded_path = StdPath::new(&expanded);
@@ -371,10 +386,6 @@ impl ObjectStore {
         &self.base_path
     }
 
-    pub fn base_uri(&self) -> String {
-        format!("{}://{}", self.scheme, self.base_path)
-    }
-
     /// Open a file for path.
     ///
     /// Parameters