From a3b4401adcbb05392566b0f3b55c619cc9655c56 Mon Sep 17 00:00:00 2001
From: Elvis Pranskevichus <elvis@edgedb.com>
Date: Mon, 4 Dec 2023 09:56:27 -0800
Subject: [PATCH] Restore preservation of session configuration across
 reconnects (#1183)

This regressed in #1156
---
 src/interactive.rs |  2 +-
 src/repl.rs        | 14 ++++++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/interactive.rs b/src/interactive.rs
index 7b5d4c5a3..8cb8076ed 100644
--- a/src/interactive.rs
+++ b/src/interactive.rs
@@ -149,7 +149,7 @@ pub fn main(options: Options, cfg: Config) -> Result<(), anyhow::Error> {
 pub async fn _main(options: Options, mut state: repl::State, cfg: Config)
     -> anyhow::Result<()>
 {
-    state.reconnect().await?;
+    state.connect().await?;
     if let Some(config_path) = &cfg.file_name {
         echo!(
             format_args!("Applied {} configuration file",
diff --git a/src/repl.rs b/src/repl.rs
index 4f793db9e..75f55a324 100644
--- a/src/repl.rs
+++ b/src/repl.rs
@@ -121,9 +121,23 @@ impl PromptRpc {
 }
 
 impl State {
+    pub async fn connect(&mut self) -> anyhow::Result<()> {
+        let db = self.conn_params.get()?.database().to_owned();
+        self.try_connect(&db).await?;
+        Ok(())
+    }
     pub async fn reconnect(&mut self) -> anyhow::Result<()> {
         let db = self.conn_params.get()?.database().to_owned();
+        let cur_state = self.edgeql_state.clone();
+        let cur_state_desc = self.edgeql_state_desc.clone();
         self.try_connect(&db).await?;
+        if let Some(conn) = &mut self.connection {
+            if cur_state_desc == self.edgeql_state_desc {
+                conn.set_state(cur_state);
+            } else {
+                eprintln!("Discarding session configuration because server configuration layout has changed.");
+            }
+        }
         Ok(())
     }
     pub async fn set_idle_transaction_timeout(&mut self) -> anyhow::Result<()> {