""" Schema browser — left sidebar tree. Node states: • saved_connection — profile saved to disk but not yet connected (grey, dashed) • connection — actively connected (coloured, expandable tree) Tree structure when connected: Connection → Databases → (Tables, Views, Functions, Procedures, Triggers) → Columns / Indexes / Foreign Keys """ from PyQt6.QtWidgets import ( QWidget, QVBoxLayout, QHBoxLayout, QTreeWidget, QTreeWidgetItem, QLineEdit, QMenu, QMessageBox, QInputDialog, ) from PyQt6.QtCore import Qt, pyqtSignal from PyQt6.QtGui import QColor, QBrush from app.utils.icons import get_icon from app.utils.worker import SchemaWorker # ── Node type codes stored in UserRole ──────────────────────────────────────── NT = { # saved but not connected "saved_connection": -1, # connected "connection": 0, "databases_folder": 1, "database": 2, "tables_folder": 3, "views_folder": 4, "functions_folder": 5, "procedures_folder": 6, "triggers_folder": 7, "table": 8, "view": 9, "function": 10, "procedure": 11, "trigger": 12, "columns_folder": 13, "indexes_folder": 14, "fks_folder": 15, "column": 16, "index": 17, "fk": 18, } class SchemaBrowser(QWidget): # ── Signals ─────────────────────────────────────────────────────────────── open_table_viewer = pyqtSignal(object, str, str) # driver, db, table open_table_structure = pyqtSignal(object, str, str) open_sql_editor = pyqtSignal(object, str) # driver, db run_query_requested = pyqtSignal(str) # SQL snippet # Emitted so the main window can act on saved profile management connect_requested = pyqtSignal(str) # profile_id — user wants to connect edit_requested = pyqtSignal(str) # profile_id — user wants to edit delete_requested = pyqtSignal(str) # profile_id — user wants to delete def __init__(self, parent=None): super().__init__(parent) self._drivers: dict = {} # profile_id → driver (only connected) self._profiles: dict = {} # profile_id → ConnectionProfile (all) self._items: dict = {} # profile_id → top-level QTreeWidgetItem self._workers: list = [] self._build_ui() # ── UI construction ─────────────────────────────────────────────────────── def _build_ui(self): layout = QVBoxLayout(self) layout.setContentsMargins(0, 0, 0, 0) layout.setSpacing(0) search_row = QHBoxLayout() search_row.setContentsMargins(6, 6, 6, 4) self._search = QLineEdit() self._search.setPlaceholderText("🔍 Filter objects…") self._search.textChanged.connect(self._filter) search_row.addWidget(self._search) layout.addLayout(search_row) self._tree = QTreeWidget() self._tree.setHeaderHidden(True) self._tree.setAnimated(True) self._tree.setIndentation(16) self._tree.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu) self._tree.customContextMenuRequested.connect(self._context_menu) self._tree.itemDoubleClicked.connect(self._on_double_click) self._tree.itemExpanded.connect(self._on_expanded) layout.addWidget(self._tree) # ── Public API ──────────────────────────────────────────────────────────── def add_saved_profile(self, profile) -> None: """Add a profile that is saved but not yet connected (greyed out).""" pid = profile.id self._profiles[pid] = profile # If it already exists as a tree item (e.g., being refreshed), remove it first self._remove_item(pid) item = QTreeWidgetItem() item.setText(0, f" {profile.name} · {profile.db_type_display}") item.setData(0, Qt.ItemDataRole.UserRole, (NT["saved_connection"], pid, "", "")) # Dim colour — show the profile colour but muted color = QColor(profile.color) color.setAlpha(140) item.setForeground(0, QBrush(QColor("#737994"))) # greyed out item.setIcon(0, get_icon(profile.db_type, 16)) item.setToolTip(0, self._connection_tooltip(profile)) f = item.font(0) f.setItalic(True) item.setFont(0, f) self._tree.addTopLevelItem(item) self._items[pid] = item def add_connection(self, profile, driver) -> None: """ Upgrade a saved (disconnected) profile to a live connected node, or add it fresh if it was never in the tree before. """ pid = profile.id self._drivers[pid] = driver self._profiles[pid] = profile # Reuse existing item if already in tree (upgrade from saved_connection) existing = self._items.get(pid) if existing: self._tree.takeTopLevelItem( self._tree.indexOfTopLevelItem(existing)) item = QTreeWidgetItem() item.setText(0, f" {profile.name}") item.setData(0, Qt.ItemDataRole.UserRole, (NT["connection"], pid, "", "")) item.setForeground(0, QBrush(QColor(profile.color))) f = item.font(0) f.setBold(True) f.setItalic(False) item.setFont(0, f) item.setIcon(0, get_icon(profile.db_type, 16)) item.setToolTip(0, self._connection_tooltip(profile)) item.addChild(QTreeWidgetItem(["Loading…"])) # triggers expand arrow self._tree.addTopLevelItem(item) self._items[pid] = item item.setExpanded(True) def remove_connection(self, profile_id: str, keep_saved: bool = True) -> None: """ Disconnect: revert the node to 'saved' state (keep_saved=True) or remove it entirely (keep_saved=False). """ driver = self._drivers.pop(profile_id, None) if driver: try: driver.disconnect() except Exception: pass if keep_saved and profile_id in self._profiles: profile = self._profiles[profile_id] self._remove_item(profile_id) self.add_saved_profile(profile) else: self._profiles.pop(profile_id, None) self._remove_item(profile_id) def update_saved_profile(self, profile) -> None: """Called after an edit — refresh the sidebar label.""" pid = profile.id self._profiles[pid] = profile if pid in self._drivers: # Currently connected — just update the label item = self._items.get(pid) if item: item.setText(0, f" {profile.name}") item.setForeground(0, QBrush(QColor(profile.color))) else: # Saved only — rebuild the item self._remove_item(pid) self.add_saved_profile(profile) def is_connected(self, profile_id: str) -> bool: return profile_id in self._drivers # ── Helpers ─────────────────────────────────────────────────────────────── def _remove_item(self, profile_id: str) -> None: item = self._items.pop(profile_id, None) if item: idx = self._tree.indexOfTopLevelItem(item) if idx >= 0: self._tree.takeTopLevelItem(idx) @staticmethod def _connection_tooltip(profile) -> str: if profile.db_type == "sqlite": return profile.database return (f"{profile.db_type_display} · " f"{profile.username}@{profile.host}:{profile.port}" + (f"/{profile.database}" if profile.database else "")) # ── Tree population (lazy-load on expand) ───────────────────────────────── def _on_expanded(self, item: QTreeWidgetItem): data = item.data(0, Qt.ItemDataRole.UserRole) if not data: return node_type, pid, db, name = data if node_type == NT["connection"]: self._load_databases(item, pid) elif node_type == NT["database"]: self._build_db_folders(item, pid, db) elif node_type == NT["tables_folder"]: self._load_tables(item, pid, db) elif node_type == NT["views_folder"]: self._load_views(item, pid, db) elif node_type == NT["functions_folder"]: self._load_generic(item, pid, db, "function", lambda d: self._drivers[pid].get_functions(d)) elif node_type == NT["procedures_folder"]: self._load_generic(item, pid, db, "procedure", lambda d: self._drivers[pid].get_stored_procedures(d)) elif node_type == NT["triggers_folder"]: self._load_generic(item, pid, db, "trigger", lambda d: self._drivers[pid].get_triggers(d)) elif node_type == NT["table"]: self._build_table_folders(item, pid, db, name) elif node_type == NT["columns_folder"]: self._load_columns(item, pid, db, name) elif node_type == NT["indexes_folder"]: self._load_indexes(item, pid, db, name) elif node_type == NT["fks_folder"]: self._load_fks(item, pid, db, name) def _placeholder(self, item: QTreeWidgetItem) -> bool: if item.childCount() == 1 and item.child(0).text(0) in ("Loading…", ""): item.takeChild(0) return True return False # databases def _load_databases(self, item, pid): if not self._placeholder(item): return w = SchemaWorker(self._drivers[pid].get_databases, parent=self) w.result.connect(lambda dbs: self._populate_databases(item, pid, dbs)) w.error.connect(lambda e: self._show_error(item, e)) self._workers.append(w) w.start() def _populate_databases(self, item, pid, dbs): for db in dbs: child = QTreeWidgetItem([f" 🗄 {db}"]) child.setData(0, Qt.ItemDataRole.UserRole, (NT["database"], pid, db, "")) child.addChild(QTreeWidgetItem(["Loading…"])) item.addChild(child) # database folders def _build_db_folders(self, item, pid, db): if not self._placeholder(item): return for label, nt in [ ("📂 Tables", NT["tables_folder"]), ("📂 Views", NT["views_folder"]), ("📂 Functions", NT["functions_folder"]), ("📂 Procedures", NT["procedures_folder"]), ("📂 Triggers", NT["triggers_folder"]), ]: child = QTreeWidgetItem([label]) child.setData(0, Qt.ItemDataRole.UserRole, (nt, pid, db, "")) child.addChild(QTreeWidgetItem(["Loading…"])) item.addChild(child) # tables def _load_tables(self, item, pid, db): if not self._placeholder(item): return w = SchemaWorker(self._drivers[pid].get_tables, db, parent=self) w.result.connect(lambda t: self._populate_tables(item, pid, db, t)) w.error.connect(lambda e: self._show_error(item, e)) self._workers.append(w) w.start() def _populate_tables(self, item, pid, db, tables): for t in tables: label = f" 📋 {t.name}" if t.row_count: label += f" ({t.row_count:,})" child = QTreeWidgetItem([label]) child.setData(0, Qt.ItemDataRole.UserRole, (NT["table"], pid, db, t.name)) child.setToolTip(0, t.comment or t.name) child.addChild(QTreeWidgetItem(["Loading…"])) item.addChild(child) # views def _load_views(self, item, pid, db): if not self._placeholder(item): return w = SchemaWorker(self._drivers[pid].get_views, db, parent=self) w.result.connect( lambda v: self._populate_generic(item, pid, db, v, "view", "👁")) w.error.connect(lambda e: self._show_error(item, e)) self._workers.append(w) w.start() # generic (functions / procedures / triggers) def _load_generic(self, item, pid, db, kind, fn): if not self._placeholder(item): return icons = {"function": "⚡", "procedure": "📦", "trigger": "⚙"} icon = icons.get(kind, "📄") w = SchemaWorker(fn, db, parent=self) w.result.connect( lambda items: self._populate_generic(item, pid, db, items, kind, icon)) w.error.connect(lambda e: self._show_error(item, e)) self._workers.append(w) w.start() def _populate_generic(self, item, pid, db, names, kind, icon): nt = NT.get(kind, NT["view"]) for name in names: child = QTreeWidgetItem([f" {icon} {name}"]) child.setData(0, Qt.ItemDataRole.UserRole, (nt, pid, db, name)) item.addChild(child) # table sub-folders def _build_table_folders(self, item, pid, db, table): if not self._placeholder(item): return for label, nt in [ ("📊 Columns", NT["columns_folder"]), ("🔍 Indexes", NT["indexes_folder"]), ("🔗 Foreign Keys", NT["fks_folder"]), ]: child = QTreeWidgetItem([label]) child.setData(0, Qt.ItemDataRole.UserRole, (nt, pid, db, table)) child.addChild(QTreeWidgetItem(["Loading…"])) item.addChild(child) def _load_columns(self, item, pid, db, table): if not self._placeholder(item): return w = SchemaWorker(self._drivers[pid].get_columns, db, table, parent=self) w.result.connect(lambda c: self._populate_columns(item, c)) w.error.connect(lambda e: self._show_error(item, e)) self._workers.append(w) w.start() def _populate_columns(self, item, cols): for c in cols: icon = "🔑" if c.is_primary_key else ("🔗" if c.is_foreign_key else "·") null = "NULL" if c.nullable else "NOT NULL" label = f" {icon} {c.name} {c.data_type} {null}" child = QTreeWidgetItem([label]) child.setData(0, Qt.ItemDataRole.UserRole, (NT["column"], "", "", c.name)) if c.is_primary_key: child.setForeground(0, QBrush(QColor("#e5c890"))) elif c.is_foreign_key: child.setForeground(0, QBrush(QColor("#99d1db"))) item.addChild(child) def _load_indexes(self, item, pid, db, table): if not self._placeholder(item): return w = SchemaWorker(self._drivers[pid].get_indexes, db, table, parent=self) w.result.connect(lambda i: self._populate_indexes(item, i)) w.error.connect(lambda e: self._show_error(item, e)) self._workers.append(w) w.start() def _populate_indexes(self, item, indexes): for idx in indexes: u = " [UNIQUE]" if idx.is_unique else "" label = f" 🔍 {idx.name}{u} ({', '.join(idx.columns)})" item.addChild(QTreeWidgetItem([label])) def _load_fks(self, item, pid, db, table): if not self._placeholder(item): return w = SchemaWorker(self._drivers[pid].get_foreign_keys, db, table, parent=self) w.result.connect(lambda f: self._populate_fks(item, f)) w.error.connect(lambda e: self._show_error(item, e)) self._workers.append(w) w.start() def _populate_fks(self, item, fks): for fk in fks: label = f" 🔗 {fk.column} → {fk.ref_table}.{fk.ref_column}" item.addChild(QTreeWidgetItem([label])) def _show_error(self, item, msg): err = QTreeWidgetItem([f"❌ {msg}"]) err.setForeground(0, QBrush(QColor("#e78284"))) item.addChild(err) # ── Context menu ────────────────────────────────────────────────────────── def _context_menu(self, pos): item = self._tree.itemAt(pos) if not item: return data = item.data(0, Qt.ItemDataRole.UserRole) if not data: return node_type, pid, db, name = data menu = QMenu(self) # ── Saved (disconnected) connection ─────────────────────────────────── if node_type == NT["saved_connection"]: menu.addAction("🔌 Connect", lambda: self.connect_requested.emit(pid)) menu.addSeparator() menu.addAction("✏️ Edit Connection…", lambda: self.edit_requested.emit(pid)) menu.addAction("🗑️ Delete Connection", lambda: self.delete_requested.emit(pid)) # ── Active connection ───────────────────────────────────────────────── elif node_type == NT["connection"]: menu.addAction("🔄 Refresh", lambda: self._refresh_node(item)) menu.addAction("🗄️ New Database…", lambda: self._create_db(pid)) menu.addSeparator() menu.addAction("✏️ Edit Connection…", lambda: self.edit_requested.emit(pid)) menu.addAction("🗑️ Delete Connection…", lambda: self.delete_requested.emit(pid)) menu.addSeparator() menu.addAction("🔌 Disconnect", lambda: self._disconnect(pid)) # ── Database ────────────────────────────────────────────────────────── elif node_type == NT["database"]: menu.addAction("✏️ New SQL Tab", lambda: self.open_sql_editor.emit(self._drivers[pid], db)) menu.addAction("🔄 Refresh", lambda: self._refresh_node(item)) # ── Table ────────────────────────────────────────────────────────────── elif node_type == NT["table"]: d = self._drivers[pid] menu.addAction("📋 Open Data", lambda: self.open_table_viewer.emit(d, db, name)) menu.addAction("🏗️ Table Structure", lambda: self.open_table_structure.emit(d, db, name)) menu.addSeparator() menu.addAction("📑 Copy SELECT *", lambda: self.run_query_requested.emit( f"SELECT * FROM {name} LIMIT 100;")) menu.addAction("🔄 Refresh", lambda: self._refresh_node(item)) menu.addSeparator() menu.addAction("🗑️ Drop Table…", lambda: self._drop_table(pid, db, name)) elif node_type in (NT["tables_folder"], NT["views_folder"]): menu.addAction("🔄 Refresh", lambda: self._refresh_node(item)) if menu.actions(): menu.exec(self._tree.viewport().mapToGlobal(pos)) # ── Context menu actions ────────────────────────────────────────────────── def _disconnect(self, pid: str): self.remove_connection(pid, keep_saved=True) def _refresh_node(self, item: QTreeWidgetItem): item.takeChildren() item.addChild(QTreeWidgetItem(["Loading…"])) item.setExpanded(False) item.setExpanded(True) def _create_db(self, pid: str): name, ok = QInputDialog.getText(self, "New Database", "Database name:") if ok and name.strip(): try: self._drivers[pid].create_database(name.strip()) QMessageBox.information(self, "Success", f"Database '{name.strip()}' created.") except Exception as e: QMessageBox.critical(self, "Error", str(e)) def _drop_table(self, pid: str, _db: str, table: str): btn = QMessageBox.warning( self, "Drop Table", f"Are you sure you want to DROP TABLE '{table}'?\nThis cannot be undone.", QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, QMessageBox.StandardButton.No, ) if btn == QMessageBox.StandardButton.Yes: try: self._drivers[pid].execute_query(f"DROP TABLE `{table}`") QMessageBox.information(self, "Dropped", f"Table '{table}' dropped.") except Exception as e: QMessageBox.critical(self, "Error", str(e)) # ── Filter ──────────────────────────────────────────────────────────────── def _filter(self, text: str): text = text.lower() for i in range(self._tree.topLevelItemCount()): self._filter_item(self._tree.topLevelItem(i), text) def _filter_item(self, item: QTreeWidgetItem, text: str) -> bool: matches = text in item.text(0).lower() child_match = any( self._filter_item(item.child(i), text) for i in range(item.childCount()) ) visible = matches or child_match item.setHidden(not visible) return visible # ── Double-click ────────────────────────────────────────────────────────── def _on_double_click(self, item: QTreeWidgetItem, _col: int): data = item.data(0, Qt.ItemDataRole.UserRole) if not data: return node_type, pid, db, name = data if node_type == NT["saved_connection"]: # Double-click on disconnected profile → try to connect self.connect_requested.emit(pid) elif node_type == NT["table"]: driver = self._drivers.get(pid) if driver: self.open_table_viewer.emit(driver, db, name) elif node_type == NT["database"]: driver = self._drivers.get(pid) if driver: self.open_sql_editor.emit(driver, db)