import { useEffect, useCallback, useRef, useState } from "react"; import { useConnectionStore } from "../stores/connectionStore"; import { useDbViewerStore } from "../stores/dbViewerStore"; import { useNotificationStore } from "../stores/notificationStore"; import * as cmd from "../lib/commands"; import { pickDefaultSchema } from "../lib/utils"; import type { ConnectionInput, TableInfo } from "../lib/types"; export function useDbConnection(connectionId: string) { const reset = useDbViewerStore((s) => s.reset); const populate = useDbViewerStore((s) => s.populate); const setSchemaTreeLoading = useDbViewerStore((s) => s.setSchemaTreeLoading); const currentDatabase = useDbViewerStore((s) => s.currentDatabase); const setCurrentDatabase = useDbViewerStore((s) => s.setCurrentDatabase); const setCurrentSchema = useDbViewerStore((s) => s.setCurrentSchema); const notify = useNotificationStore((s) => s.notify); const [connectionError, setConnectionError] = useState(null); const [passwordPromptOpen, setPasswordPromptOpen] = useState(false); const inputRef = useRef(null); // The database the pool is currently connected to. Unlike the selected // `currentDatabase`, this lets us reconnect whenever the selection drifts // from the live connection (including switching back to the first DB). const connectedDbRef = useRef(null); const connect = useCallback(async () => { const conn = useConnectionStore .getState() .connections.find((c) => c.id === connectionId); if (!conn) { setConnectionError("Connection not found"); return; } try { const password = await useConnectionStore.getState().getConnectionPassword(conn.id).catch(() => null); // Keychain-off + no session password: prompt the user instead of // connecting with an empty password (early return, no Tauri call). if (conn.use_keychain === false && !password) { setPasswordPromptOpen(true); return; } const sshPassword = conn.ssh_host ? await cmd.getConnectionSshPassword(conn.id).catch(() => null) : null; const sshPassphrase = conn.ssh_host ? await cmd.getConnectionSshPassphrase(conn.id).catch(() => null) : null; const input: ConnectionInput = { name: conn.name, db_type: conn.db_type, host: conn.host, port: conn.port, username: conn.username, password, database: conn.database, folder_id: conn.folder_id, ssh_host: conn.ssh_host, ssh_port: conn.ssh_port, ssh_user: conn.ssh_user, ssh_auth_method: conn.ssh_auth_method as "password" | "key" | null | undefined, ssh_private_key_path: conn.ssh_private_key_path, ssh_password: sshPassword, ssh_passphrase: sshPassphrase, ssl_mode: conn.ssl_mode as | "disable" | "require" | "verify-ca" | "verify-full" | null | undefined, ssl_ca_path: conn.ssl_ca_path, ssl_cert_path: conn.ssl_cert_path, ssl_key_path: conn.ssl_key_path, }; await cmd.dbConnect(connectionId, input); setConnectionError(null); inputRef.current = input; connectedDbRef.current = input.db_type === "sqlite" ? "main" : (input.database ?? "postgres"); // Load initial data, smart-selecting the default schema (e.g. `public`) setSchemaTreeLoading(true); try { const databases = await cmd .getDatabases(connectionId) .catch(() => [] as string[]); const schemas = await cmd .getSchemas(connectionId) .catch(() => [] as string[]); const defaultSchema = pickDefaultSchema(schemas); const tables = await cmd.getTables( connectionId, defaultSchema ?? undefined, ); populate(databases, schemas, tables); if (databases.length > 0) { // Prefer the connection's configured database, fall back to the first // available one so the dropdown matches what the pool is connected to. const preferred = input.database && databases.includes(input.database) ? input.database : databases[0]; setCurrentDatabase(preferred); } setCurrentSchema(defaultSchema); } finally { setSchemaTreeLoading(false); } } catch (e) { const msg = e instanceof Error ? e.message : String(e); setConnectionError(msg); notify(`Failed to connect: ${msg}`, "error"); } }, [connectionId, populate, setCurrentDatabase, setCurrentSchema, setSchemaTreeLoading, notify]); useEffect(() => { connect(); return () => { cmd.dbDisconnect(connectionId).catch(() => {}); const hasPending = useDbViewerStore .getState() .changesQueue.some((c) => c.status === "pending"); if (!hasPending) reset(); }; }, [connectionId, connect, reset]); // Database switch effect: whenever the selected database drifts from the // pool's live connection, reconnect and refresh schemas/tables for that DB. useEffect(() => { if (!currentDatabase || !inputRef.current) return; if (currentDatabase === connectedDbRef.current) return; let cancelled = false; const reconnect = async () => { const input = { ...inputRef.current!, database: currentDatabase }; try { setSchemaTreeLoading(true); await cmd.dbConnect(connectionId, input); if (cancelled) return; connectedDbRef.current = currentDatabase; const schemas = await cmd .getSchemas(connectionId) .catch(() => [] as string[]); if (cancelled) return; const newSchema = pickDefaultSchema(schemas); const tables = await cmd .getTables(connectionId, newSchema ?? undefined) .catch(() => [] as TableInfo[]); if (cancelled) return; populate(useDbViewerStore.getState().databases, schemas, tables); setCurrentSchema(newSchema); } catch (e) { const msg = e instanceof Error ? e.message : String(e); // Revert the selection so the dropdown matches the live connection setCurrentDatabase(connectedDbRef.current); notify(`Failed to switch database: ${msg}`, "error"); } finally { if (!cancelled) setSchemaTreeLoading(false); } }; reconnect(); return () => { cancelled = true; }; }, [currentDatabase, connectionId, populate, setCurrentSchema, setSchemaTreeLoading, notify]); const submitPassword = useCallback( (pw: string) => { useConnectionStore.getState().setSessionPassword(connectionId, pw); setPasswordPromptOpen(false); void connect(); }, [connectionId, connect], ); const cancelPassword = useCallback(() => setPasswordPromptOpen(false), []); return { connectionError, connect, passwordPromptOpen, submitPassword, cancelPassword }; }