'u' to Use database as default - and star currently used db in explorer

Author: Maxteabag

sqlit/app.py

@@ -292,6 +292,7 @@ class SSMSTUI(
         Binding("z", "collapse_tree", "Collapse", show=False),
         Binding("j", "tree_cursor_down", "Down", show=False),
         Binding("k", "tree_cursor_up", "Up", show=False),
+        Binding("u", "use_database", "Use as default", show=False),
         Binding("v", "view_cell", "View cell", show=False),
         Binding("u", "edit_cell", "Update cell", show=False),
         Binding("h", "results_cursor_left", "Left", show=False),

sqlit/state_machine.py

@@ -601,8 +601,44 @@ def is_active(self, app: SSMSTUI) -> bool:
         return node is not None and _get_node_kind(node) in ("table", "view")
 
 
+class TreeOnDatabaseState(State):
+    """Tree focused on a database node (in multi-database servers)."""
+
+    help_category = "Explorer"
+
+    def _setup_actions(self) -> None:
+        self.allows("use_database", key="u", label="Use as default", help="Set as default database")
+
+    def get_display_bindings(self, app: SSMSTUI) -> tuple[list[DisplayBinding], list[DisplayBinding]]:
+        left: list[DisplayBinding] = []
+        seen: set[str] = set()
+
+        left.append(DisplayBinding(key="enter", label="Expand", action="toggle_node"))
+        seen.add("toggle_node")
+        left.append(DisplayBinding(key="u", label="Use as default", action="use_database"))
+        seen.add("use_database")
+        left.append(DisplayBinding(key="f", label="Refresh", action="refresh_tree"))
+        seen.add("refresh_tree")
+
+        right: list[DisplayBinding] = []
+        if self.parent:
+            _, parent_right = self.parent.get_display_bindings(app)
+            for binding in parent_right:
+                if binding.action not in seen:
+                    right.append(binding)
+                    seen.add(binding.action)
+
+        return left, right
+
+    def is_active(self, app: SSMSTUI) -> bool:
+        if not app.object_tree.has_focus:
+            return False
+        node = app.object_tree.cursor_node
+        return node is not None and _get_node_kind(node) == "database"
+
+
 class TreeOnFolderState(State):
-    """Tree focused on a folder, database, or schema node."""
+    """Tree focused on a folder or schema node."""
 
     def _setup_actions(self) -> None:
         pass  # Just inherits from parent
@@ -630,7 +666,7 @@ def is_active(self, app: SSMSTUI) -> bool:
         if not app.object_tree.has_focus:
             return False
         node = app.object_tree.cursor_node
-        return node is not None and _get_node_kind(node) in ("folder", "database", "schema")
+        return node is not None and _get_node_kind(node) in ("folder", "schema")
 
 
 class TreeOnObjectState(State):
@@ -906,6 +942,7 @@ def __init__(self) -> None:
         self.tree_focused = TreeFocusedState(parent=self.main_screen)
         self.tree_filter_active = TreeFilterActiveState(parent=self.main_screen)
         self.tree_on_connection = TreeOnConnectionState(parent=self.tree_focused)
+        self.tree_on_database = TreeOnDatabaseState(parent=self.tree_focused)
         self.tree_on_table = TreeOnTableState(parent=self.tree_focused)
         self.tree_on_folder = TreeOnFolderState(parent=self.tree_focused)
         self.tree_on_object = TreeOnObjectState(parent=self.tree_focused)
@@ -924,6 +961,7 @@ def __init__(self) -> None:
             self.leader_pending,
             self.tree_filter_active,  # Before tree_focused (more specific when filter active)
             self.tree_on_connection,
+            self.tree_on_database,  # For database nodes (multi-database servers)
             self.tree_on_table,
             self.tree_on_folder,
             self.tree_on_object,  # For index/trigger/sequence nodes

sqlit/ui/mixins/query.py

@@ -125,7 +125,6 @@ async def _run_query_async(self: AppProtocol, query: str, keep_insert_mode: bool
         """Run query asynchronously using a cancellable dedicated connection."""
         import asyncio
         import time
-        from dataclasses import replace
 
         from ...services import CancellableQuery, QueryResult, QueryService
         from ...services.query import parse_use_statement
@@ -141,11 +140,9 @@ async def _run_query_async(self: AppProtocol, query: str, keep_insert_mode: bool
         # Handle USE database statements
         db_name = parse_use_statement(query)
         if db_name is not None:
-            self.current_config = replace(config, database=db_name)
             self._stop_query_spinner()
             self._display_non_query_result(0, 0)
-            self.notify(f"Switched to database: {db_name}")
-            self._update_status_bar()
+            self.set_default_database(db_name)  # type: ignore[attr-defined]
             if keep_insert_mode:
                 self._restore_insert_mode()
             return

sqlit/ui/mixins/tree.py

@@ -173,8 +173,14 @@ def get_conn_label(config: Any, connected: Any = False) -> str:
                     dbs_node.data = FolderNode(folder_type="databases")
 
                     databases = self._run_db_call(adapter.get_databases, self.current_connection)
+                    default_db = self.current_config.database if self.current_config else None
                     for db_name in databases:
-                        db_node = dbs_node.add(escape_markup(db_name))
+                        # Show default database with star and green text
+                        if default_db and db_name.lower() == default_db.lower():
+                            db_label = f"[#4ADE80]* {escape_markup(db_name)}[/]"
+                        else:
+                            db_label = escape_markup(db_name)
+                        db_node = dbs_node.add(db_label)
                         db_node.data = DatabaseNode(name=db_name)
                         db_node.allow_expand = True
                         self._add_database_object_nodes(db_node, db_name)
@@ -700,3 +706,67 @@ def _display_object_info(self: AppProtocol, object_type: str, info: dict) -> Non
         definition = info.get("definition")
         if definition:
             self.query_input.text = f"/*\n{definition}\n*/"
+
+    def set_default_database(self: AppProtocol, db_name: str) -> None:
+        """Set the default database for the current connection.
+
+        This is the shared function used by both the USE query handler and
+        the explorer 'Use as default' action.
+
+        Args:
+            db_name: The database name to set as default.
+        """
+        from dataclasses import replace
+
+        if not self.current_config:
+            self.notify("Not connected", severity="error")
+            return
+
+        self.current_config = replace(self.current_config, database=db_name)
+        self.notify(f"Switched to database: {db_name}")
+        self._update_status_bar()
+        self._update_database_labels()
+
+    def _update_database_labels(self: AppProtocol) -> None:
+        """Update database node labels to show the default database with a star."""
+        if not self.current_config:
+            return
+
+        default_db = self.current_config.database
+
+        # Find the Databases folder and update labels
+        for conn_node in self.object_tree.root.children:
+            if self._get_node_kind(conn_node) != "connection":
+                continue
+
+            # Check if this is the active connection
+            if not (conn_node.data and conn_node.data.config.name == self.current_config.name):
+                continue
+
+            # Find Databases folder
+            for child in conn_node.children:
+                if self._get_node_kind(child) == "folder" and child.data.folder_type == "databases":
+                    # Update each database node
+                    for db_node in child.children:
+                        if self._get_node_kind(db_node) == "database":
+                            db_name = db_node.data.name
+                            if default_db and db_name.lower() == default_db.lower():
+                                db_node.set_label(f"[#4ADE80]* {escape_markup(db_name)}[/]")
+                            else:
+                                db_node.set_label(escape_markup(db_name))
+                    break
+            break
+
+    def action_use_database(self: AppProtocol) -> None:
+        """Set the selected database as the default for the current connection."""
+        node = self.object_tree.cursor_node
+
+        if not node or self._get_node_kind(node) != "database":
+            return
+
+        if not self.current_connection:
+            self.notify("Not connected", severity="error")
+            return
+
+        db_name = node.data.name
+        self.set_default_database(db_name)