Set correct database on reconnect when applicable. (#1439)

* Set correct database on reconnect when applicable.

* Added test for reconnect to verify previously selected database is still selected

Author: scottnemes

changelog.md

@@ -8,7 +8,7 @@ Internal
 Bug Fixes
 --------
 * Update watch query output to display the correct execution time on all iterations (#763).
-
+* Use correct database (if applicable) when reconnecting after a connection loss (#1437).
 
 1.44.1 (2026/01/10)
 ==============

mycli/main.py

@@ -1144,6 +1144,9 @@ def reconnect(self, database: str = "") -> bool:
             self.logger.debug("Attempting to reconnect.")
             self.echo("Reconnecting...", fg="yellow")
             self.sqlexecute.conn.ping(reconnect=True)
+            # if a database is currently selected, set it on the conn again
+            if self.sqlexecute.dbname:
+                self.sqlexecute.conn.select_db(self.sqlexecute.dbname)
             self.logger.debug("Reconnected successfully.")
             self.echo("Reconnected successfully.", fg="yellow")
             self.sqlexecute.reset_connection_id()

test/conftest.py

@@ -3,13 +3,13 @@
 import pytest
 
 import mycli.sqlexecute
-from test.utils import CHARSET, HOST, PASSWORD, PORT, SSH_HOST, SSH_PORT, SSH_USER, USER, create_db, db_connection
+from test.utils import CHARSET, DATABASE, HOST, PASSWORD, PORT, SSH_HOST, SSH_PORT, SSH_USER, USER, create_db, db_connection
 
 
 @pytest.fixture(scope="function")
 def connection():
-    create_db("mycli_test_db")
-    connection = db_connection("mycli_test_db")
+    create_db(DATABASE)
+    connection = db_connection(DATABASE)
     yield connection
 
     connection.close()
@@ -24,7 +24,7 @@ def cursor(connection):
 @pytest.fixture
 def executor(connection):
     return mycli.sqlexecute.SQLExecute(
-        database="mycli_test_db",
+        database=DATABASE,
         user=USER,
         host=HOST,
         password=PASSWORD,

test/test_main.py

@@ -9,12 +9,13 @@
 
 import click
 from click.testing import CliRunner
+from pymysql.err import OperationalError
 
 from mycli.main import MyCli, cli, thanks_picker
 import mycli.packages.special
 from mycli.packages.special.main import COMMANDS as SPECIAL_COMMANDS
 from mycli.sqlexecute import ServerInfo, SQLExecute
-from test.utils import HOST, PASSWORD, PORT, USER, dbtest, run
+from test.utils import DATABASE, HOST, PASSWORD, PORT, USER, dbtest, run
 
 test_dir = os.path.abspath(os.path.dirname(__file__))
 project_dir = os.path.dirname(test_dir)
@@ -94,6 +95,41 @@ def test_ssl_mode_overrides_no_ssl(executor, capsys):
     assert ssl_cipher
 
 
+@dbtest
+def test_reconnect_database_is_selected(executor, capsys):
+    m = MyCli()
+    m.register_special_commands()
+    m.sqlexecute = SQLExecute(
+        None,
+        USER,
+        PASSWORD,
+        HOST,
+        PORT,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+    )
+    try:
+        next(m.sqlexecute.run(f"use {DATABASE}"))
+        next(m.sqlexecute.run(f"kill {m.sqlexecute.connection_id}"))
+    except OperationalError:
+        pass  # expected as the connection was killed
+    except Exception as e:
+        raise e
+    m.reconnect()
+    try:
+        next(m.sqlexecute.run("show tables")).results.fetchall()
+    except Exception as e:
+        raise e
+
+
 @dbtest
 def test_reconnect_no_database(executor, capsys):
     m = MyCli()

test/utils.py

@@ -11,6 +11,7 @@
 
 from mycli.main import special
 
+DATABASE = "mycli_test_db"
 PASSWORD = os.getenv("PYTEST_PASSWORD")
 USER = os.getenv("PYTEST_USER", "root")
 HOST = os.getenv("PYTEST_HOST", "localhost")