From 1207b93211ce95ba1dfd10cda8292b2ddbc9f19f Mon Sep 17 00:00:00 2001
From: "Michael J. Sullivan" <sully@msully.net>
Date: Thu, 19 Dec 2024 11:12:38 -0800
Subject: [PATCH] Fix switching branch after current branch is deleted (#1435)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

There was code that allegedly handled this case, but there wasn't a
test and it didn't work at all.

Fixes #1359.

Co-authored-by: Aljaž Mur Eržen <aljazerzen@users.noreply.github.com>
---
 src/branch/connections.rs     | 23 +++++++++++------------
 tests/func/non_interactive.rs | 22 ++++++++++++++++++++++
 2 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/src/branch/connections.rs b/src/branch/connections.rs
index e66ba569a..eee198612 100644
--- a/src/branch/connections.rs
+++ b/src/branch/connections.rs
@@ -1,8 +1,9 @@
 use std::ops::Deref;
 
 use crate::commands::Options;
-use crate::connect::{Connection, Connector};
+use crate::connect::{Connection, ConnectionError, Connector};
 use crate::print;
+use edgedb_errors::UnknownDatabaseError;
 use uuid::Uuid;
 
 pub struct BranchConnection<'a> {
@@ -40,23 +41,21 @@ impl BranchConnection<'_> {
 
 /// Attempts to connect a provided connection, returning `Ok(Some)` if the connection was
 /// established, `Ok(None)` if the connection couldn't be established because of a
-/// `ClientConnectionFailedError`, and `Err` if any other type of error occurred.
+/// `UnknownDatabaseError`, and `Err` if any other type of error occurred.
 pub async fn connect_if_branch_exists(connector: &Connector) -> anyhow::Result<Option<Connection>> {
     match connector.connect().await {
         Ok(c) => Ok(Some(c)),
-        Err(e) => {
-            match e.downcast::<edgedb_tokio::Error>() {
-                Ok(e) => {
-                    if e.code() == 0x_FF_01_01_00 {
-                        // 0x_FF_01_01_00: ClientConnectionFailedError | https://www.edgedb.com/docs/reference/protocol/errors
-                        return Ok(None);
-                    }
-
+        Err(e) => match e.downcast::<ConnectionError>() {
+            Ok(ConnectionError::Error(e)) => {
+                if e.is::<UnknownDatabaseError>() {
+                    Ok(None)
+                } else {
                     Err(e.into())
                 }
-                Err(e) => Err(e),
             }
-        }
+            Ok(e) => Err(e.into()),
+            Err(e) => Err(e),
+        },
     }
 }
 
diff --git a/tests/func/non_interactive.rs b/tests/func/non_interactive.rs
index eda56185a..ee9963c4e 100644
--- a/tests/func/non_interactive.rs
+++ b/tests/func/non_interactive.rs
@@ -240,6 +240,28 @@ fn branch_commands() {
         .success()
         .stdout(predicates::str::contains("test_branch_2"));
 
+    // Drop test_branch_2 from main
+    crate::edgedb_cli_cmd()
+        .arg("--instance")
+        .arg(instance_name)
+        .arg("-b")
+        .arg("main")
+        .arg("query")
+        .arg("drop branch test_branch_2")
+        .assert()
+        .context("drop", "drop current branch")
+        .success();
+
+    crate::edgedb_cli_cmd()
+        .arg("--instance")
+        .arg(instance_name)
+        .arg("branch")
+        .arg("switch")
+        .arg("main")
+        .assert()
+        .context("switch", "when current branch is destroyed")
+        .success();
+
     // TODO: test how this works in projects
 }