feat: preserve session state when replacing its runtime

Cargo.toml

@@ -196,7 +196,7 @@ regex = "1.12"
 rstest = "0.26.1"
 serde_json = "1"
 sha2 = "^0.11.0"
-sqlparser = { git = "https://github.com/Embucket/datafusion-sqlparser-rs.git", rev = "33e5ae265571d9a74887cb5d355a264032a89c61", default-features = false, features = [
+sqlparser = { git = "https://github.com/Embucket/datafusion-sqlparser-rs.git", rev = "9e9559a4e8138d96d80917cd79d9553ba00dd164", default-features = false, features = [
     "std",
     "visitor",
 ] }

datafusion/core/src/execution/session_state.rs

@@ -857,6 +857,18 @@ impl SessionState {
         &self.runtime_env
     }
 
+    /// Replace the runtime environment while preserving all other session state.
+    ///
+    /// Function registrations, aliases, configuration and planners retain their
+    /// existing values and instances. Previously created clones are unchanged.
+    /// This does not rebuild existing plans or catalog/table providers, which may
+    /// retain resources from the previous runtime. The caller must provide the
+    /// object stores and other runtime resources required by the preserved state.
+    pub fn with_runtime_env(mut self, runtime_env: Arc<RuntimeEnv>) -> Self {
+        self.runtime_env = runtime_env;
+        self
+    }
+
     /// Return the execution properties
     pub fn execution_props(&self) -> &ExecutionProps {
         &self.execution_props
@@ -2830,3 +2842,7 @@ mod tests {
         }
     }
 }
+
+#[cfg(test)]
+#[path = "session_state_runtime_env_tests.rs"]
+mod runtime_env_tests;