From b2eabb32bf913db984841a45fc121641e25fbab4 Mon Sep 17 00:00:00 2001 From: Richie Varghese Date: Thu, 25 Dec 2025 17:21:28 +0530 Subject: [PATCH 1/8] feat: add advanced PostgreSQL connection options including SSL, various timeouts, application name, and raw parameters. --- docs/ROADMAP.md | 90 ++ package.json | 40 + src/common/types.ts | 67 +- src/connectionForm.ts | 664 ++++++---- src/extension.ts | 1741 +++++++++++++------------ src/providers/DatabaseTreeProvider.ts | 1263 +++++++++++------- src/services/ConnectionManager.ts | 231 ++-- 7 files changed, 2432 insertions(+), 1664 deletions(-) create mode 100644 docs/ROADMAP.md diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md new file mode 100644 index 0000000..6bae294 --- /dev/null +++ b/docs/ROADMAP.md @@ -0,0 +1,90 @@ +# PgStudio Improvement Roadmap + +> Last Updated: December 2025 + +--- + +## ✅ Phase 1: Connection Management UX (COMPLETE) + +- [x] SSL mode dropdown (disable, allow, prefer, require, verify-ca, verify-full) +- [x] SSL certificate paths (CA, client cert, client key) +- [x] Connection timeout setting +- [x] Statement timeout setting +- [x] Application name (shown in `pg_stat_activity`) +- [x] Raw options field (`-c search_path=myschema`) + +--- + +## 🎯 Phase 2: UX Enhancements + +### 2A: Tree View Improvements ✅ COMPLETE +- [x] Quick filter input for searching objects (toggle icon, schema filtering) +- [x] Favorites (star frequently-used tables/views) +- [x] ⭐ Favorites section under connection +- [x] Context menu preserved for favorited items +- [x] 🕒 Recent items tracking (max 10 items) +- [x] Object count badges on category nodes (right-aligned, muted) + +### 2B: Notebook Experience +- [ ] Query cancellation button +- [ ] Virtual scrolling for large result sets (1000+ rows) +- [ ] Sticky column headers +- [ ] Column resizing + +### 2C: AI Assistant +- [ ] Schema context caching +- [ ] Query history in AI context +- [ ] "Explain this error" feature +- [ ] Query optimization suggestions + +--- + +## 🏗️ Phase 3: Architecture Refactoring + +### Code Organization +- [ ] Split `extension.ts` (882 lines) → `commands.ts`, `providers.ts`, `views.ts` +- [ ] Split `tables.ts` (51KB) → `operations.ts`, `scripts.ts`, `maintenance.ts` +- [ ] Split `renderer_v2.ts` (144KB) into modules + +### Service Layer +- [ ] Command factory pattern for CRUD operations +- [ ] Query history service +- [ ] Connection pooling + +--- + +## 📚 Phase 4: Documentation + +- [ ] `ARCHITECTURE.md` with system diagrams +- [ ] `CONTRIBUTING.md` with code style guide +- [ ] Troubleshooting section in README +- [ ] Feature comparison vs pgAdmin/DBeaver + +--- + +## 🚀 Phase 5: Future Features + +### Near-term (1-3 months) +- [ ] Query snippets with variables +- [ ] Table structure diff across connections +- [ ] Smart query bookmarks + +### Mid-term (3-6 months) +- [ ] Connection export/import (encrypted) +- [ ] Shared query library (`.pgstudio/` folder) +- [ ] ERD diagram generation + +### Long-term (6+ months) +- [ ] Audit logging +- [ ] Schema migration tracking +- [ ] Role-based access controls + +--- + +## 🔧 Technical Debt + +| Item | Priority | +|------|----------| +| Migrate inline styles to `htmlStyles.ts` | Medium | +| Standardize error handling | Medium | +| Add JSDoc to exported functions | Low | diff --git a/package.json b/package.json index e790e52..3cdc3d2 100644 --- a/package.json +++ b/package.json @@ -61,6 +61,26 @@ "title": "Refresh Connections", "icon": "$(refresh)" }, + { + "command": "postgres-explorer.filterTree", + "title": "Filter Tree", + "icon": "$(filter)" + }, + { + "command": "postgres-explorer.clearFilter", + "title": "Clear Filter", + "icon": "$(search-remove)" + }, + { + "command": "postgres-explorer.addToFavorites", + "title": "Add to Favorites", + "icon": "$(star-empty)" + }, + { + "command": "postgres-explorer.removeFromFavorites", + "title": "Remove from Favorites", + "icon": "$(star-full)" + }, { "command": "postgres-explorer.manageConnections", "title": "Manage Connections", @@ -931,6 +951,16 @@ "when": "view == postgresExplorer", "group": "navigation" }, + { + "command": "postgres-explorer.filterTree", + "when": "view == postgresExplorer && !postgresExplorer.filterActive", + "group": "navigation" + }, + { + "command": "postgres-explorer.clearFilter", + "when": "view == postgresExplorer && postgresExplorer.filterActive", + "group": "navigation" + }, { "command": "postgres-explorer.refreshConnections", "when": "view == postgresExplorer", @@ -963,6 +993,16 @@ "when": "view == postgresExplorer && viewItem == connection", "group": "9_delete" }, + { + "command": "postgres-explorer.addToFavorites", + "when": "view == postgresExplorer && viewItem =~ /^(table|view|function|materialized-view)$/ && !postgresExplorer.isFavorite", + "group": "0_favorites" + }, + { + "command": "postgres-explorer.removeFromFavorites", + "when": "view == postgresExplorer && viewItem =~ /^(table|view|function|materialized-view)$/ && postgresExplorer.isFavorite", + "group": "0_favorites" + }, { "command": "postgres-explorer.disconnectConnection", "when": "view == postgresExplorer && viewItem == connection", diff --git a/src/common/types.ts b/src/common/types.ts index 59dd7af..af88643 100644 --- a/src/common/types.ts +++ b/src/common/types.ts @@ -1,37 +1,46 @@ export interface ConnectionConfig { - id: string; - name?: string; + id: string; + name?: string; + host: string; + port: number; + username?: string; + password?: string; + database?: string; + // Advanced connection options + sslmode?: 'disable' | 'allow' | 'prefer' | 'require' | 'verify-ca' | 'verify-full'; + sslCertPath?: string; // Client certificate path + sslKeyPath?: string; // Client key path + sslRootCertPath?: string; // CA certificate path + statementTimeout?: number; // milliseconds + connectTimeout?: number; // seconds (default: 5) + applicationName?: string; // Shows in pg_stat_activity + options?: string; // Raw options string (e.g., "-c search_path=myschema") + ssh?: { + enabled: boolean; host: string; port: number; - username?: string; - password?: string; - database?: string; - ssh?: { - enabled: boolean; - host: string; - port: number; - username: string; - privateKeyPath?: string; - }; + username: string; + privateKeyPath?: string; + }; } export interface PostgresMetadata { - connectionId: string; - databaseName: string | undefined; - host: string; - port: number; - username?: string; - password?: string; - custom?: { - cells: any[]; - metadata: { - connectionId: string; - databaseName: string | undefined; - host: string; - port: number; - username?: string; - password?: string; - enableScripts: boolean; - }; + connectionId: string; + databaseName: string | undefined; + host: string; + port: number; + username?: string; + password?: string; + custom?: { + cells: any[]; + metadata: { + connectionId: string; + databaseName: string | undefined; + host: string; + port: number; + username?: string; + password?: string; + enableScripts: boolean; }; + }; } diff --git a/src/connectionForm.ts b/src/connectionForm.ts index 20f6dea..abcd4fb 100644 --- a/src/connectionForm.ts +++ b/src/connectionForm.ts @@ -3,212 +3,230 @@ import * as vscode from 'vscode'; import { SSHService } from './services/SSHService'; export interface ConnectionInfo { - id: string; - name: string; + id: string; + name: string; + host: string; + port: number; + username?: string; + password?: string; + database?: string; + // Advanced connection options + sslmode?: 'disable' | 'allow' | 'prefer' | 'require' | 'verify-ca' | 'verify-full'; + sslCertPath?: string; + sslKeyPath?: string; + sslRootCertPath?: string; + statementTimeout?: number; + connectTimeout?: number; + applicationName?: string; + options?: string; + ssh?: { + enabled: boolean; host: string; port: number; - username?: string; - password?: string; - database?: string; - ssh?: { - enabled: boolean; - host: string; - port: number; - username: string; - privateKeyPath?: string; - }; + username: string; + privateKeyPath?: string; + }; } export class ConnectionFormPanel { - public static currentPanel: ConnectionFormPanel | undefined; - private readonly _panel: vscode.WebviewPanel; - private readonly _extensionUri: vscode.Uri; - private _disposables: vscode.Disposable[] = []; - - private constructor(panel: vscode.WebviewPanel, extensionUri: vscode.Uri, private readonly _extensionContext: vscode.ExtensionContext, private readonly _connectionToEdit?: ConnectionInfo) { - this._panel = panel; - this._extensionUri = extensionUri; - - this._panel.onDidDispose(() => this.dispose(), null, this._disposables); - this._initialize(); - - this._panel.webview.onDidReceiveMessage( - async (message) => { - switch (message.command) { - case 'testConnection': - try { - const config: any = { - user: message.connection.username || undefined, - password: message.connection.password || undefined, - database: message.connection.database || 'postgres' - }; - - if (message.connection.ssh && message.connection.ssh.enabled) { - const stream = await SSHService.getInstance().createStream( - message.connection.ssh, - message.connection.host, - message.connection.port - ); - config.stream = stream; - } else { - config.host = message.connection.host; - config.port = message.connection.port; - } - - // First try with specified database - const client = new Client(config); - try { - await client.connect(); - const result = await client.query('SELECT version()'); - await client.end(); - this._panel.webview.postMessage({ - type: 'testSuccess', - version: result.rows[0].version - }); - } catch (err: any) { - if (config.stream) { - // If using stream, we can't easily fallback without creating a new stream - // simpler to just throw for now or re-create stream - throw err; - } - - // If database doesn't exist, try postgres database - if (err.code === '3D000' && message.connection.database !== 'postgres') { - const fallbackClient = new Client({ - host: message.connection.host, - port: message.connection.port, - user: message.connection.username || undefined, - password: message.connection.password || undefined, - database: 'postgres' - }); - await fallbackClient.connect(); - const result = await fallbackClient.query('SELECT version()'); - await fallbackClient.end(); - this._panel.webview.postMessage({ - type: 'testSuccess', - version: result.rows[0].version + ' (connected to postgres database)' - }); - } else { - throw err; - } - } - } catch (err: any) { - this._panel.webview.postMessage({ - type: 'testError', - error: err.message - }); - } - break; - - case 'saveConnection': - try { - const config: any = { - user: message.connection.username || undefined, - password: message.connection.password || undefined, - database: 'postgres' - }; - - if (message.connection.ssh && message.connection.ssh.enabled) { - const stream = await SSHService.getInstance().createStream( - message.connection.ssh, - message.connection.host, - message.connection.port - ); - config.stream = stream; - } else { - config.host = message.connection.host; - config.port = message.connection.port; - } - - const client = new Client(config); - - await client.connect(); - - // Verify we can query - await client.query('SELECT 1'); - await client.end(); - - const connections = this.getStoredConnections(); - const newConnection: ConnectionInfo = { - id: this._connectionToEdit ? this._connectionToEdit.id : Date.now().toString(), - name: message.connection.name, - host: message.connection.host, - port: message.connection.port, - username: message.connection.username || undefined, - password: message.connection.password || undefined, - database: message.connection.database, - ssh: message.connection.ssh - }; - - if (this._connectionToEdit) { - const index = connections.findIndex(c => c.id === this._connectionToEdit!.id); - if (index !== -1) { - connections[index] = newConnection; - } else { - connections.push(newConnection); - } - } else { - connections.push(newConnection); - } - - await this.storeConnections(connections); - - vscode.window.showInformationMessage(`Connection ${this._connectionToEdit ? 'updated' : 'saved'} successfully!`); - vscode.commands.executeCommand('postgres-explorer.refreshConnections'); - this._panel.dispose(); - } catch (err: any) { - const errorMessage = err?.message || 'Unknown error occurred'; - vscode.window.showErrorMessage(`Failed to connect: ${errorMessage}`); - } - break; + public static currentPanel: ConnectionFormPanel | undefined; + private readonly _panel: vscode.WebviewPanel; + private readonly _extensionUri: vscode.Uri; + private _disposables: vscode.Disposable[] = []; + + private constructor(panel: vscode.WebviewPanel, extensionUri: vscode.Uri, private readonly _extensionContext: vscode.ExtensionContext, private readonly _connectionToEdit?: ConnectionInfo) { + this._panel = panel; + this._extensionUri = extensionUri; + + this._panel.onDidDispose(() => this.dispose(), null, this._disposables); + this._initialize(); + + this._panel.webview.onDidReceiveMessage( + async (message) => { + switch (message.command) { + case 'testConnection': + try { + const config: any = { + user: message.connection.username || undefined, + password: message.connection.password || undefined, + database: message.connection.database || 'postgres' + }; + + if (message.connection.ssh && message.connection.ssh.enabled) { + const stream = await SSHService.getInstance().createStream( + message.connection.ssh, + message.connection.host, + message.connection.port + ); + config.stream = stream; + } else { + config.host = message.connection.host; + config.port = message.connection.port; + } + + // First try with specified database + const client = new Client(config); + try { + await client.connect(); + const result = await client.query('SELECT version()'); + await client.end(); + this._panel.webview.postMessage({ + type: 'testSuccess', + version: result.rows[0].version + }); + } catch (err: any) { + if (config.stream) { + // If using stream, we can't easily fallback without creating a new stream + // simpler to just throw for now or re-create stream + throw err; } - }, - undefined, - this._disposables - ); - } - - public static show(extensionUri: vscode.Uri, extensionContext: vscode.ExtensionContext, connectionToEdit?: ConnectionInfo) { - if (ConnectionFormPanel.currentPanel) { - ConnectionFormPanel.currentPanel._panel.reveal(); - return; - } - const panel = vscode.window.createWebviewPanel( - 'connectionForm', - connectionToEdit ? 'Edit Connection' : 'Add PostgreSQL Connection', - vscode.ViewColumn.One, - { - enableScripts: true + // If database doesn't exist, try postgres database + if (err.code === '3D000' && message.connection.database !== 'postgres') { + const fallbackClient = new Client({ + host: message.connection.host, + port: message.connection.port, + user: message.connection.username || undefined, + password: message.connection.password || undefined, + database: 'postgres' + }); + await fallbackClient.connect(); + const result = await fallbackClient.query('SELECT version()'); + await fallbackClient.end(); + this._panel.webview.postMessage({ + type: 'testSuccess', + version: result.rows[0].version + ' (connected to postgres database)' + }); + } else { + throw err; + } + } + } catch (err: any) { + this._panel.webview.postMessage({ + type: 'testError', + error: err.message + }); } - ); - - ConnectionFormPanel.currentPanel = new ConnectionFormPanel(panel, extensionUri, extensionContext, connectionToEdit); - } - - private async _initialize() { - // The message handler is already set up in the constructor - await this._update(); + break; + + case 'saveConnection': + try { + const config: any = { + user: message.connection.username || undefined, + password: message.connection.password || undefined, + database: 'postgres' + }; + + if (message.connection.ssh && message.connection.ssh.enabled) { + const stream = await SSHService.getInstance().createStream( + message.connection.ssh, + message.connection.host, + message.connection.port + ); + config.stream = stream; + } else { + config.host = message.connection.host; + config.port = message.connection.port; + } + + const client = new Client(config); + + await client.connect(); + + // Verify we can query + await client.query('SELECT 1'); + await client.end(); + + const connections = this.getStoredConnections(); + const newConnection: ConnectionInfo = { + id: this._connectionToEdit ? this._connectionToEdit.id : Date.now().toString(), + name: message.connection.name, + host: message.connection.host, + port: message.connection.port, + username: message.connection.username || undefined, + password: message.connection.password || undefined, + database: message.connection.database, + // Advanced options + sslmode: message.connection.sslmode || undefined, + sslCertPath: message.connection.sslCertPath || undefined, + sslKeyPath: message.connection.sslKeyPath || undefined, + sslRootCertPath: message.connection.sslRootCertPath || undefined, + statementTimeout: message.connection.statementTimeout || undefined, + connectTimeout: message.connection.connectTimeout || undefined, + applicationName: message.connection.applicationName || undefined, + options: message.connection.options || undefined, + ssh: message.connection.ssh + }; + + if (this._connectionToEdit) { + const index = connections.findIndex(c => c.id === this._connectionToEdit!.id); + if (index !== -1) { + connections[index] = newConnection; + } else { + connections.push(newConnection); + } + } else { + connections.push(newConnection); + } + + await this.storeConnections(connections); + + vscode.window.showInformationMessage(`Connection ${this._connectionToEdit ? 'updated' : 'saved'} successfully!`); + vscode.commands.executeCommand('postgres-explorer.refreshConnections'); + this._panel.dispose(); + } catch (err: any) { + const errorMessage = err?.message || 'Unknown error occurred'; + vscode.window.showErrorMessage(`Failed to connect: ${errorMessage}`); + } + break; + } + }, + undefined, + this._disposables + ); + } + + public static show(extensionUri: vscode.Uri, extensionContext: vscode.ExtensionContext, connectionToEdit?: ConnectionInfo) { + if (ConnectionFormPanel.currentPanel) { + ConnectionFormPanel.currentPanel._panel.reveal(); + return; } - private async _update() { - this._panel.webview.html = await this._getHtmlForWebview(this._panel.webview); + const panel = vscode.window.createWebviewPanel( + 'connectionForm', + connectionToEdit ? 'Edit Connection' : 'Add PostgreSQL Connection', + vscode.ViewColumn.One, + { + enableScripts: true + } + ); + + ConnectionFormPanel.currentPanel = new ConnectionFormPanel(panel, extensionUri, extensionContext, connectionToEdit); + } + + private async _initialize() { + // The message handler is already set up in the constructor + await this._update(); + } + + private async _update() { + this._panel.webview.html = await this._getHtmlForWebview(this._panel.webview); + } + + private async _getHtmlForWebview(webview: vscode.Webview): Promise { + const logoPath = webview.asWebviewUri(vscode.Uri.joinPath(this._extensionUri, 'resources', 'postgres-vsc-icon.png')); + + let connectionData = null; + if (this._connectionToEdit) { + // Get the password from secret storage + const password = await this._extensionContext.secrets.get(`postgres-password-${this._connectionToEdit.id}`); + connectionData = { + ...this._connectionToEdit, + password + }; } - private async _getHtmlForWebview(webview: vscode.Webview): Promise { - const logoPath = webview.asWebviewUri(vscode.Uri.joinPath(this._extensionUri, 'resources', 'postgres-vsc-icon.png')); - - let connectionData = null; - if (this._connectionToEdit) { - // Get the password from secret storage - const password = await this._extensionContext.secrets.get(`postgres-password-${this._connectionToEdit.id}`); - connectionData = { - ...this._connectionToEdit, - password - }; - } - - return ` + return ` @@ -659,6 +677,96 @@ export class ConnectionFormPanel { + + + + +
+
+
⚙️
+
Advanced Options
+ ▼ +
+ +
@@ -685,6 +793,45 @@ export class ConnectionFormPanel { document.getElementById('username').value = connectionData.username || ''; document.getElementById('password').value = connectionData.password || ''; + // Populate advanced options + if (connectionData.sslmode) { + document.getElementById('sslmode').value = connectionData.sslmode; + } + if (connectionData.sslCertPath) { + document.getElementById('sslCertPath').value = connectionData.sslCertPath; + } + if (connectionData.sslKeyPath) { + document.getElementById('sslKeyPath').value = connectionData.sslKeyPath; + } + if (connectionData.sslRootCertPath) { + document.getElementById('sslRootCertPath').value = connectionData.sslRootCertPath; + } + if (connectionData.statementTimeout) { + document.getElementById('statementTimeout').value = connectionData.statementTimeout; + } + if (connectionData.connectTimeout) { + document.getElementById('connectTimeout').value = connectionData.connectTimeout; + } + if (connectionData.applicationName) { + document.getElementById('applicationName').value = connectionData.applicationName; + } + if (connectionData.options) { + document.getElementById('options').value = connectionData.options; + } + + // Show advanced section if any advanced options are set + const hasAdvancedOptions = connectionData.sslmode || connectionData.statementTimeout || + connectionData.connectTimeout || connectionData.applicationName || connectionData.options; + if (hasAdvancedOptions) { + setTimeout(() => { + const advSection = document.getElementById('advanced-section'); + const advArrow = document.getElementById('advanced-arrow'); + advSection.style.display = 'block'; + advArrow.style.transform = 'rotate(180deg)'; + updateSSLCertFields(); + }, 100); + } + if (connectionData.ssh) { document.getElementById('sshEnabled').checked = connectionData.ssh.enabled; document.getElementById('sshHost').value = connectionData.ssh.host || ''; @@ -735,6 +882,34 @@ export class ConnectionFormPanel { document.getElementById('sshEnabled').addEventListener('change', updateSSHState); + // Advanced Options toggle + function toggleAdvanced() { + const section = document.getElementById('advanced-section'); + const arrow = document.getElementById('advanced-arrow'); + if (section.style.display === 'none') { + section.style.display = 'block'; + arrow.style.transform = 'rotate(180deg)'; + } else { + section.style.display = 'none'; + arrow.style.transform = 'rotate(0deg)'; + } + } + + // SSL mode change handler - show cert fields for verify modes + function updateSSLCertFields() { + const sslmode = document.getElementById('sslmode').value; + const certFields = document.getElementById('ssl-cert-fields'); + if (sslmode === 'verify-ca' || sslmode === 'verify-full') { + certFields.style.display = 'block'; + document.getElementById('sslRootCertPath').required = true; + } else { + certFields.style.display = 'none'; + document.getElementById('sslRootCertPath').required = false; + } + } + + document.getElementById('sslmode').addEventListener('change', updateSSLCertFields); + let isTested = false; function showMessage(text, type = 'info') { @@ -763,7 +938,16 @@ export class ConnectionFormPanel { port: parseInt(document.getElementById('port').value), database: document.getElementById('database').value || 'postgres', username: usernameInput || undefined, - password: passwordInput || undefined + password: passwordInput || undefined, + // Advanced options + sslmode: document.getElementById('sslmode').value || undefined, + sslCertPath: document.getElementById('sslCertPath').value || undefined, + sslKeyPath: document.getElementById('sslKeyPath').value || undefined, + sslRootCertPath: document.getElementById('sslRootCertPath').value || undefined, + statementTimeout: document.getElementById('statementTimeout').value ? parseInt(document.getElementById('statementTimeout').value) : undefined, + connectTimeout: document.getElementById('connectTimeout').value ? parseInt(document.getElementById('connectTimeout').value) : undefined, + applicationName: document.getElementById('applicationName').value || undefined, + options: document.getElementById('options').value || undefined }; if (sshEnabled) { @@ -846,54 +1030,54 @@ export class ConnectionFormPanel { `; - } + } - private _getNonce(): string { - let text = ''; - const possible = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789'; - for (let i = 0; i < 32; i++) { - text += possible.charAt(Math.floor(Math.random() * possible.length)); - } - return text; + private _getNonce(): string { + let text = ''; + const possible = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789'; + for (let i = 0; i < 32; i++) { + text += possible.charAt(Math.floor(Math.random() * possible.length)); } - - private getStoredConnections(): ConnectionInfo[] { - const connections = vscode.workspace.getConfiguration().get('postgresExplorer.connections') || []; - return connections; - } - - private async storeConnections(connections: ConnectionInfo[]): Promise { - try { - // First store the connections without passwords in settings - const connectionsForSettings = connections.map(({ password, ...connWithoutPassword }) => connWithoutPassword); - await vscode.workspace.getConfiguration().update('postgresExplorer.connections', connectionsForSettings, vscode.ConfigurationTarget.Global); - - // Then store passwords in SecretStorage - const secretsStorage = this._extensionContext.secrets; - for (const conn of connections) { - if (conn.password) { - // Removed logging of sensitive connection information for security. - await secretsStorage.store(`postgres-password-${conn.id}`, conn.password); - } - } - } catch (error) { - console.error('Failed to store connections:', error); - // If anything fails, make sure we don't leave passwords in settings - const existingConnections = vscode.workspace.getConfiguration().get('postgresExplorer.connections') || []; - const sanitizedConnections = existingConnections.map(({ password, ...connWithoutPassword }) => connWithoutPassword); - await vscode.workspace.getConfiguration().update('postgresExplorer.connections', sanitizedConnections, vscode.ConfigurationTarget.Global); - throw error; + return text; + } + + private getStoredConnections(): ConnectionInfo[] { + const connections = vscode.workspace.getConfiguration().get('postgresExplorer.connections') || []; + return connections; + } + + private async storeConnections(connections: ConnectionInfo[]): Promise { + try { + // First store the connections without passwords in settings + const connectionsForSettings = connections.map(({ password, ...connWithoutPassword }) => connWithoutPassword); + await vscode.workspace.getConfiguration().update('postgresExplorer.connections', connectionsForSettings, vscode.ConfigurationTarget.Global); + + // Then store passwords in SecretStorage + const secretsStorage = this._extensionContext.secrets; + for (const conn of connections) { + if (conn.password) { + // Removed logging of sensitive connection information for security. + await secretsStorage.store(`postgres-password-${conn.id}`, conn.password); } + } + } catch (error) { + console.error('Failed to store connections:', error); + // If anything fails, make sure we don't leave passwords in settings + const existingConnections = vscode.workspace.getConfiguration().get('postgresExplorer.connections') || []; + const sanitizedConnections = existingConnections.map(({ password, ...connWithoutPassword }) => connWithoutPassword); + await vscode.workspace.getConfiguration().update('postgresExplorer.connections', sanitizedConnections, vscode.ConfigurationTarget.Global); + throw error; } - - private dispose() { - ConnectionFormPanel.currentPanel = undefined; - this._panel.dispose(); - while (this._disposables.length) { - const disposable = this._disposables.pop(); - if (disposable) { - disposable.dispose(); - } - } + } + + private dispose() { + ConnectionFormPanel.currentPanel = undefined; + this._panel.dispose(); + while (this._disposables.length) { + const disposable = this._disposables.pop(); + if (disposable) { + disposable.dispose(); + } } + } } diff --git a/src/extension.ts b/src/extension.ts index 696370f..62ffaee 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -37,845 +37,932 @@ export let outputChannel: vscode.OutputChannel; let chatViewProviderInstance: ChatViewProvider | undefined; export function getChatViewProvider(): ChatViewProvider | undefined { - return chatViewProviderInstance; + return chatViewProviderInstance; } export async function activate(context: vscode.ExtensionContext) { - outputChannel = vscode.window.createOutputChannel('PgStudio'); - outputChannel.appendLine('postgres-explorer: Activating extension'); - console.log('postgres-explorer: Activating extension'); - - // Initialize services - SecretStorageService.getInstance(context); - ConnectionManager.getInstance(); - - // Create database tree provider instance - const databaseTreeProvider = new DatabaseTreeProvider(context); - - // Register tree data provider and create tree view - const treeView = vscode.window.createTreeView('postgresExplorer', { - treeDataProvider: databaseTreeProvider, - showCollapseAll: true - }); - context.subscriptions.push(treeView); - - // Register the chat view provider - chatViewProviderInstance = new ChatViewProvider(context.extensionUri, context); - context.subscriptions.push( - vscode.window.registerWebviewViewProvider( - ChatViewProvider.viewType, - chatViewProviderInstance, - { webviewOptions: { retainContextWhenHidden: true } } - ) - ); - - // Register all commands - const commands = [ - { - command: 'postgres-explorer.addConnection', - callback: (connection?: any) => { - ConnectionFormPanel.show(context.extensionUri, context, connection); - } - }, - { - command: 'postgres-explorer.refreshConnections', - callback: () => { - databaseTreeProvider.refresh(); - } - }, - { - command: 'postgres-explorer.manageConnections', - callback: () => { - ConnectionManagementPanel.show(context.extensionUri, context); - } - }, - { - command: 'postgres-explorer.aiSettings', - callback: () => { - AiSettingsPanel.show(context.extensionUri, context); - } - }, - { - command: 'postgres-explorer.connect', - callback: async (item: any) => await cmdConnectDatabase(item, context, databaseTreeProvider) - }, - { - command: 'postgres-explorer.disconnect', - callback: async () => { - databaseTreeProvider.refresh(); - vscode.window.showInformationMessage('Disconnected from PostgreSQL database'); - } - }, - { - command: 'postgres-explorer.queryTable', - callback: async (item: any) => { - if (!item || !item.schema) { - return; - } - - const query = `SELECT * FROM ${item.schema}.${item.label} LIMIT 100;`; - const notebook = await vscode.workspace.openNotebookDocument('postgres-notebook', new vscode.NotebookData([ - new vscode.NotebookCellData(vscode.NotebookCellKind.Code, query, 'sql') - ])); - await vscode.window.showNotebookDocument(notebook); + outputChannel = vscode.window.createOutputChannel('PgStudio'); + outputChannel.appendLine('postgres-explorer: Activating extension'); + console.log('postgres-explorer: Activating extension'); + + // Initialize services + SecretStorageService.getInstance(context); + ConnectionManager.getInstance(); + + // Create database tree provider instance + const databaseTreeProvider = new DatabaseTreeProvider(context); + + // Register tree data provider and create tree view + const treeView = vscode.window.createTreeView('postgresExplorer', { + treeDataProvider: databaseTreeProvider, + showCollapseAll: true + }); + context.subscriptions.push(treeView); + + // Update context key when selection changes to enable Add/Remove favorites menu switching + treeView.onDidChangeSelection(e => { + if (e.selection.length > 0) { + const item = e.selection[0]; + vscode.commands.executeCommand('setContext', 'postgresExplorer.isFavorite', item.isFavorite === true); + } else { + vscode.commands.executeCommand('setContext', 'postgresExplorer.isFavorite', false); + } + }); + + // Register the chat view provider + chatViewProviderInstance = new ChatViewProvider(context.extensionUri, context); + context.subscriptions.push( + vscode.window.registerWebviewViewProvider( + ChatViewProvider.viewType, + chatViewProviderInstance, + { webviewOptions: { retainContextWhenHidden: true } } + ) + ); + + // Register all commands + const commands = [ + { + command: 'postgres-explorer.addConnection', + callback: (connection?: any) => { + ConnectionFormPanel.show(context.extensionUri, context, connection); + } + }, + { + command: 'postgres-explorer.refreshConnections', + callback: () => { + databaseTreeProvider.refresh(); + } + }, + { + command: 'postgres-explorer.filterTree', + callback: async () => { + const currentFilter = databaseTreeProvider.filterPattern; + + if (currentFilter) { + // Filter is active - show options to modify or clear + const choice = await vscode.window.showQuickPick([ + { label: '$(close) Clear Filter', value: 'clear' }, + { label: '$(edit) Change Filter', value: 'change', description: `Current: "${currentFilter}"` } + ], { placeHolder: `Filter active: "${currentFilter}"` }); + + if (choice?.value === 'clear') { + databaseTreeProvider.clearFilter(); + vscode.commands.executeCommand('setContext', 'postgresExplorer.filterActive', false); + vscode.window.showInformationMessage('Filter cleared'); + } else if (choice?.value === 'change') { + const pattern = await vscode.window.showInputBox({ + prompt: 'Enter filter pattern', + placeHolder: 'e.g., users, product, order', + value: currentFilter + }); + if (pattern !== undefined) { + databaseTreeProvider.setFilter(pattern); + vscode.commands.executeCommand('setContext', 'postgresExplorer.filterActive', pattern.length > 0); + if (pattern) { + vscode.window.showInformationMessage(`Filter applied: "${pattern}"`); + } } - }, - { - command: 'postgres-explorer.newNotebook', - callback: async (item: any) => await cmdNewNotebook(item) - }, - { - command: 'postgres-explorer.refresh', - callback: () => databaseTreeProvider.refresh() - }, - // Add database commands - { - command: 'postgres-explorer.createInDatabase', - callback: async (item: DatabaseTreeItem) => await cmdAddObjectInDatabase(item, context) - }, - { - command: 'postgres-explorer.createDatabase', - callback: async (item: DatabaseTreeItem) => await cmdCreateDatabase(item, context) - }, - { - command: 'postgres-explorer.dropDatabase', - callback: async (item: DatabaseTreeItem) => await cmdDeleteDatabase(item, context) - }, - { - command: 'postgres-explorer.scriptAlterDatabase', - callback: async (item: DatabaseTreeItem) => await cmdScriptAlterDatabase(item, context) - }, - { - command: 'postgres-explorer.databaseOperations', - callback: async (item: DatabaseTreeItem) => await cmdDatabaseOperations(item, context) - }, - { - command: 'postgres-explorer.showDashboard', - callback: async (item: DatabaseTreeItem) => await cmdDatabaseDashboard(item, context) - }, - { - command: 'postgres-explorer.backupDatabase', - callback: async (item: DatabaseTreeItem) => await cmdBackupDatabase(item, context) - }, - { - command: 'postgres-explorer.restoreDatabase', - callback: async (item: DatabaseTreeItem) => await cmdRestoreDatabase(item, context) - }, - { - command: 'postgres-explorer.generateCreateScript', - callback: async (item: DatabaseTreeItem) => await cmdGenerateCreateScript(item, context) - }, - { - command: 'postgres-explorer.disconnectDatabase', - callback: async (item: DatabaseTreeItem) => await cmdDisconnectDatabaseLegacy(item, context) - }, - { - command: 'postgres-explorer.maintenanceDatabase', - callback: async (item: DatabaseTreeItem) => await cmdMaintenanceDatabase(item, context) - }, - { - command: 'postgres-explorer.queryTool', - callback: async (item: DatabaseTreeItem) => await cmdQueryTool(item, context) - }, - { - command: 'postgres-explorer.psqlTool', - callback: async (item: DatabaseTreeItem) => await cmdPsqlTool(item, context) - }, - { - command: 'postgres-explorer.showConfiguration', - callback: async (item: DatabaseTreeItem) => await cmdShowConfiguration(item, context) - }, - // Add schema commands - { - command: 'postgres-explorer.createSchema', - callback: async (item: DatabaseTreeItem) => await cmdCreateSchema(item, context) - }, - { - command: 'postgres-explorer.createInSchema', - callback: async (item: DatabaseTreeItem) => await cmdCreateObjectInSchema(item, context) - }, - { - command: 'postgres-explorer.schemaOperations', - callback: async (item: DatabaseTreeItem) => await cmdSchemaOperations(item, context) - }, - { - command: 'postgres-explorer.showSchemaProperties', - callback: async (item: DatabaseTreeItem) => await cmdShowSchemaProperties(item, context) - }, - // Add table commands - { - command: 'postgres-explorer.editTable', - callback: async (item: DatabaseTreeItem) => await cmdEditTable(item, context) - }, - { - command: 'postgres-explorer.viewTableData', - callback: async (item: DatabaseTreeItem) => await cmdViewTableData(item, context) - }, - { - command: 'postgres-explorer.dropTable', - callback: async (item: DatabaseTreeItem) => await cmdDropTable(item, context) - }, - { - command: 'postgres-explorer.tableOperations', - callback: async (item: DatabaseTreeItem) => await cmdTableOperations(item, context) - }, - { - command: 'postgres-explorer.truncateTable', - callback: async (item: DatabaseTreeItem) => await cmdTruncateTable(item, context) - }, - { - command: 'postgres-explorer.insertData', - callback: async (item: DatabaseTreeItem) => await cmdInsertTable(item, context) - }, - { - command: 'postgres-explorer.updateData', - callback: async (item: DatabaseTreeItem) => await cmdUpdateTable(item, context) - }, - { - command: 'postgres-explorer.showTableProperties', - callback: async (item: DatabaseTreeItem) => await cmdShowTableProperties(item, context) - }, - // Add script commands - { - command: 'postgres-explorer.scriptSelect', - callback: async (item: DatabaseTreeItem) => await cmdScriptSelect(item, context) - }, - { - command: 'postgres-explorer.scriptInsert', - callback: async (item: DatabaseTreeItem) => await cmdScriptInsert(item, context) - }, - { - command: 'postgres-explorer.scriptUpdate', - callback: async (item: DatabaseTreeItem) => await cmdScriptUpdate(item, context) - }, - { - command: 'postgres-explorer.scriptDelete', - callback: async (item: DatabaseTreeItem) => await cmdScriptDelete(item, context) - }, - { - command: 'postgres-explorer.scriptCreate', - callback: async (item: DatabaseTreeItem) => await cmdScriptCreate(item, context) - }, - // Add maintenance commands - { - command: 'postgres-explorer.maintenanceVacuum', - callback: async (item: DatabaseTreeItem) => await cmdMaintenanceVacuum(item, context) - }, - { - command: 'postgres-explorer.maintenanceAnalyze', - callback: async (item: DatabaseTreeItem) => await cmdMaintenanceAnalyze(item, context) - }, - { - command: 'postgres-explorer.maintenanceReindex', - callback: async (item: DatabaseTreeItem) => await cmdMaintenanceReindex(item, context) - }, - - // Add view commands - { - command: 'postgres-explorer.refreshView', - callback: async (item: DatabaseTreeItem) => await cmdRefreshView(item, context, databaseTreeProvider) - }, - { - command: 'postgres-explorer.editViewDefinition', - callback: async (item: DatabaseTreeItem) => await cmdEditView(item, context) - }, - { - command: 'postgres-explorer.viewViewData', - callback: async (item: DatabaseTreeItem) => await cmdViewData(item, context) - }, - { - command: 'postgres-explorer.dropView', - callback: async (item: DatabaseTreeItem) => await cmdDropView(item, context) - }, - { - command: 'postgres-explorer.viewOperations', - callback: async (item: DatabaseTreeItem) => await cmdViewOperations(item, context) - }, - { - command: 'postgres-explorer.showViewProperties', - callback: async (item: DatabaseTreeItem) => await cmdShowViewProperties(item, context) - }, - { - command: 'postgres-explorer.viewScriptSelect', - callback: async (item: DatabaseTreeItem) => await cmdViewScriptSelect(item, context) - }, - { - command: 'postgres-explorer.viewScriptCreate', - callback: async (item: DatabaseTreeItem) => await cmdViewScriptCreate(item, context) - }, - // Add function commands - { - command: 'postgres-explorer.refreshFunction', - callback: async (item: DatabaseTreeItem) => await cmdRefreshFunction(item, context, databaseTreeProvider) - }, - { - command: 'postgres-explorer.showFunctionProperties', - callback: async (item: DatabaseTreeItem) => await cmdShowFunctionProperties(item, context) - }, - { - command: 'postgres-explorer.functionOperations', - callback: async (item: DatabaseTreeItem) => await cmdFunctionOperations(item, context) - }, - { - command: 'postgres-explorer.createReplaceFunction', - callback: async (item: DatabaseTreeItem) => await cmdEditFunction(item, context) - }, - { - command: 'postgres-explorer.callFunction', - callback: async (item: DatabaseTreeItem) => await cmdCallFunction(item, context) - }, - { - command: 'postgres-explorer.dropFunction', - callback: async (item: DatabaseTreeItem) => await cmdDropFunction(item, context) - }, - // Add materialized view commands - { - command: 'postgres-explorer.refreshMaterializedView', - callback: async (item: DatabaseTreeItem) => await cmdRefreshMatView(item, context) - }, - { - command: 'postgres-explorer.editMatView', - callback: async (item: DatabaseTreeItem) => await cmdEditMatView(item, context) - }, - { - command: 'postgres-explorer.viewMaterializedViewData', - callback: async (item: DatabaseTreeItem) => await cmdViewMatViewData(item, context) - }, - { - command: 'postgres-explorer.showMaterializedViewProperties', - callback: async (item: DatabaseTreeItem) => await cmdViewMatViewProperties(item, context) - }, - { - command: 'postgres-explorer.dropMatView', - callback: async (item: DatabaseTreeItem) => await cmdDropMatView(item, context) - }, - { - command: 'postgres-explorer.materializedViewOperations', - callback: async (item: DatabaseTreeItem) => await cmdMatViewOperations(item, context) - }, - // Add type commands - { - command: 'postgres-explorer.refreshType', - callback: async (item: DatabaseTreeItem) => await cmdRefreshType(item, context, databaseTreeProvider) - }, - { - command: 'postgres-explorer.typeOperations', - callback: async (item: DatabaseTreeItem) => await cmdAllOperationsTypes(item, context) - }, - { - command: 'postgres-explorer.editType', - callback: async (item: DatabaseTreeItem) => await cmdEditTypes(item, context) - }, - { - command: 'postgres-explorer.showTypeProperties', - callback: async (item: DatabaseTreeItem) => await cmdShowTypeProperties(item, context) - }, - { - command: 'postgres-explorer.dropType', - callback: async (item: DatabaseTreeItem) => await cmdDropType(item, context) - }, - // Add foreign table commands - { - command: 'postgres-explorer.refreshForeignTable', - callback: async (item: DatabaseTreeItem) => await cmdRefreshForeignTable(item, context, databaseTreeProvider) - }, - { - command: 'postgres-explorer.foreignTableOperations', - callback: async (item: DatabaseTreeItem) => await cmdForeignTableOperations(item, context) - }, - { - command: 'postgres-explorer.editForeignTable', - callback: async (item: DatabaseTreeItem) => await cmdEditForeignTable(item, context) - }, - // Add role/user commands - { - command: 'postgres-explorer.refreshRole', - callback: async (item: DatabaseTreeItem) => await cmdRefreshRole(item, context, databaseTreeProvider) - }, - { - command: 'postgres-explorer.createUser', - callback: async (item: DatabaseTreeItem) => await cmdAddUser(item, context) - }, - { - command: 'postgres-explorer.createRole', - callback: async (item: DatabaseTreeItem) => await cmdAddRole(item, context) - }, - { - command: 'postgres-explorer.editRole', - callback: async (item: DatabaseTreeItem) => await cmdEditRole(item, context) - }, - { - command: 'postgres-explorer.grantRevoke', - callback: async (item: DatabaseTreeItem) => await cmdGrantRevokeRole(item, context) - }, - { - command: 'postgres-explorer.dropRole', - callback: async (item: DatabaseTreeItem) => await cmdDropRole(item, context) - }, - { - command: 'postgres-explorer.roleOperations', - callback: async (item: DatabaseTreeItem) => await cmdRoleOperations(item, context) - }, - { - command: 'postgres-explorer.showRoleProperties', - callback: async (item: DatabaseTreeItem) => await cmdShowRoleProperties(item, context) - }, - // Add extension commands - { - command: 'postgres-explorer.refreshExtension', - callback: async (item: DatabaseTreeItem) => await cmdRefreshExtension(item, context, databaseTreeProvider) - }, - { - command: 'postgres-explorer.enableExtension', - callback: async (item: DatabaseTreeItem) => await cmdEnableExtension(item, context) - }, - { - command: 'postgres-explorer.extensionOperations', - callback: async (item: DatabaseTreeItem) => await cmdExtensionOperations(item, context) - }, - { - command: 'postgres-explorer.dropExtension', - callback: async (item: DatabaseTreeItem) => await cmdDropExtension(item, context) - }, - // Add connection commands - { - command: 'postgres-explorer.disconnectConnection', - callback: async (item: DatabaseTreeItem) => await cmdDisconnectConnection(item, context, databaseTreeProvider) - }, - { - command: 'postgres-explorer.reconnectConnection', - callback: async (item: DatabaseTreeItem) => await cmdReconnectConnection(item, context, databaseTreeProvider) - }, - { - command: 'postgres-explorer.deleteConnection', - callback: async (item: DatabaseTreeItem) => await cmdDisconnectDatabase(item, context, databaseTreeProvider) - }, - - { - command: 'postgres-explorer.createTable', - callback: async (item: DatabaseTreeItem) => await cmdCreateTable(item, context) - }, - { - command: 'postgres-explorer.createView', - callback: async (item: DatabaseTreeItem) => await cmdCreateView(item, context) - }, - { - command: 'postgres-explorer.createFunction', - callback: async (item: DatabaseTreeItem) => await cmdCreateFunction(item, context) - }, - { - command: 'postgres-explorer.createMaterializedView', - callback: async (item: DatabaseTreeItem) => await cmdCreateMaterializedView(item, context) - }, - { - command: 'postgres-explorer.createType', - callback: async (item: DatabaseTreeItem) => await cmdCreateType(item, context) - }, - { - command: 'postgres-explorer.createForeignTable', - callback: async (item: DatabaseTreeItem) => await cmdCreateForeignTable(item, context) - }, - // Foreign Data Wrapper commands - { - command: 'postgres-explorer.foreignDataWrapperOperations', - callback: async (item: DatabaseTreeItem) => await cmdForeignDataWrapperOperations(item, context) - }, - { - command: 'postgres-explorer.showForeignDataWrapperProperties', - callback: async (item: DatabaseTreeItem) => await cmdShowForeignDataWrapperProperties(item, context) - }, - { - command: 'postgres-explorer.refreshForeignDataWrapper', - callback: async (item: DatabaseTreeItem) => await cmdRefreshForeignDataWrapper(item, context, databaseTreeProvider) - }, - // Foreign Server commands - { - command: 'postgres-explorer.createForeignServer', - callback: async (item: DatabaseTreeItem) => await cmdCreateForeignServer(item, context) - }, - { - command: 'postgres-explorer.foreignServerOperations', - callback: async (item: DatabaseTreeItem) => await cmdForeignServerOperations(item, context) - }, - { - command: 'postgres-explorer.showForeignServerProperties', - callback: async (item: DatabaseTreeItem) => await cmdShowForeignServerProperties(item, context) - }, - { - command: 'postgres-explorer.dropForeignServer', - callback: async (item: DatabaseTreeItem) => await cmdDropForeignServer(item, context) - }, - { - command: 'postgres-explorer.refreshForeignServer', - callback: async (item: DatabaseTreeItem) => await cmdRefreshForeignServer(item, context, databaseTreeProvider) - }, - // User Mapping commands - { - command: 'postgres-explorer.createUserMapping', - callback: async (item: DatabaseTreeItem) => await cmdCreateUserMapping(item, context) - }, - { - command: 'postgres-explorer.userMappingOperations', - callback: async (item: DatabaseTreeItem) => await cmdUserMappingOperations(item, context) - }, - { - command: 'postgres-explorer.showUserMappingProperties', - callback: async (item: DatabaseTreeItem) => await cmdShowUserMappingProperties(item, context) - }, - { - command: 'postgres-explorer.dropUserMapping', - callback: async (item: DatabaseTreeItem) => await cmdDropUserMapping(item, context) - }, - { - command: 'postgres-explorer.refreshUserMapping', - callback: async (item: DatabaseTreeItem) => await cmdRefreshUserMapping(item, context, databaseTreeProvider) - }, - { - command: 'postgres-explorer.createRole', - callback: async (item: DatabaseTreeItem) => await cmdAddRole(item, context) - }, - { - command: 'postgres-explorer.enableExtension', - callback: async (item: DatabaseTreeItem) => await cmdEnableExtension(item, context) - }, - - { - command: 'postgres-explorer.aiAssist', - callback: async (cell: vscode.NotebookCell) => await cmdAiAssist(cell, context, outputChannel) - }, - - // Column commands - { - command: 'postgres-explorer.showColumnProperties', - callback: async (item: DatabaseTreeItem) => await showColumnProperties(item) - }, - { - command: 'postgres-explorer.copyColumnName', - callback: async (item: DatabaseTreeItem) => await copyColumnName(item) - }, - { - command: 'postgres-explorer.copyColumnNameQuoted', - callback: async (item: DatabaseTreeItem) => await copyColumnNameQuoted(item) - }, - { - command: 'postgres-explorer.generateSelectStatement', - callback: async (item: DatabaseTreeItem) => await generateSelectStatement(item) - }, - { - command: 'postgres-explorer.generateWhereClause', - callback: async (item: DatabaseTreeItem) => await generateWhereClause(item) - }, - { - command: 'postgres-explorer.generateAlterColumnScript', - callback: async (item: DatabaseTreeItem) => await generateAlterColumnScript(item) - }, - { - command: 'postgres-explorer.generateDropColumnScript', - callback: async (item: DatabaseTreeItem) => await generateDropColumnScript(item) - }, - { - command: 'postgres-explorer.generateRenameColumnScript', - callback: async (item: DatabaseTreeItem) => await generateRenameColumnScript(item) - }, - { - command: 'postgres-explorer.addColumnComment', - callback: async (item: DatabaseTreeItem) => await addColumnComment(item) - }, - { - command: 'postgres-explorer.generateIndexOnColumn', - callback: async (item: DatabaseTreeItem) => await generateIndexOnColumn(item) - }, - { - command: 'postgres-explorer.viewColumnStatistics', - callback: async (item: DatabaseTreeItem) => await viewColumnStatistics(item) - }, - - // Constraint commands - { - command: 'postgres-explorer.showConstraintProperties', - callback: async (item: DatabaseTreeItem) => await showConstraintProperties(item) - }, - { - command: 'postgres-explorer.copyConstraintName', - callback: async (item: DatabaseTreeItem) => await copyConstraintName(item) - }, - { - command: 'postgres-explorer.generateDropConstraintScript', - callback: async (item: DatabaseTreeItem) => await generateDropConstraintScript(item) - }, - { - command: 'postgres-explorer.generateAlterConstraintScript', - callback: async (item: DatabaseTreeItem) => await generateAlterConstraintScript(item) - }, - { - command: 'postgres-explorer.validateConstraint', - callback: async (item: DatabaseTreeItem) => await validateConstraint(item) - }, - { - command: 'postgres-explorer.generateAddConstraintScript', - callback: async (item: DatabaseTreeItem) => await generateAddConstraintScript(item) - }, - { - command: 'postgres-explorer.viewConstraintDependencies', - callback: async (item: DatabaseTreeItem) => await viewConstraintDependencies(item) - }, - { - command: 'postgres-explorer.constraintOperations', - callback: async (item: DatabaseTreeItem) => await cmdConstraintOperations(item, context) - }, - - // Index commands - { - command: 'postgres-explorer.showIndexProperties', - callback: async (item: DatabaseTreeItem) => await showIndexProperties(item) - }, - { - command: 'postgres-explorer.copyIndexName', - callback: async (item: DatabaseTreeItem) => await copyIndexName(item) - }, - { - command: 'postgres-explorer.generateDropIndexScript', - callback: async (item: DatabaseTreeItem) => await generateDropIndexScript(item) - }, - { - command: 'postgres-explorer.generateReindexScript', - callback: async (item: DatabaseTreeItem) => await generateReindexScript(item) - }, - { - command: 'postgres-explorer.generateScriptCreate', - callback: async (item: DatabaseTreeItem) => await generateScriptCreate(item) - }, - { - command: 'postgres-explorer.analyzeIndexUsage', - callback: async (item: DatabaseTreeItem) => await analyzeIndexUsage(item) - }, - { - command: 'postgres-explorer.generateAlterIndexScript', - callback: async (item: DatabaseTreeItem) => await generateAlterIndexScript(item) - }, - { - command: 'postgres-explorer.addIndexComment', - callback: async (item: DatabaseTreeItem) => await addIndexComment(item) - }, - { - command: 'postgres-explorer.indexOperations', - callback: async (item: DatabaseTreeItem) => await cmdIndexOperations(item, context) - }, - { - command: 'postgres-explorer.addColumn', - callback: async (item: DatabaseTreeItem) => await cmdAddColumn(item) - }, - { - command: 'postgres-explorer.addConstraint', - callback: async (item: DatabaseTreeItem) => await cmdAddConstraint(item) - }, - { - command: 'postgres-explorer.addIndex', - callback: async (item: DatabaseTreeItem) => await cmdAddIndex(item) - }, - ]; - - // Register all commands - console.log('Starting command registration...'); - outputChannel.appendLine('Starting command registration...'); - - commands.forEach(({ command, callback }) => { - try { - console.log(`Registering command: ${command}`); - context.subscriptions.push( - vscode.commands.registerCommand(command, callback) - ); - } catch (e) { - console.error(`Failed to register command ${command}:`, e); - outputChannel.appendLine(`Failed to register command ${command}: ${e}`); + } + } else { + // No filter active - show input + const pattern = await vscode.window.showInputBox({ + prompt: 'Enter filter pattern', + placeHolder: 'e.g., users, product, order' + }); + if (pattern !== undefined && pattern.length > 0) { + databaseTreeProvider.setFilter(pattern); + vscode.commands.executeCommand('setContext', 'postgresExplorer.filterActive', true); + vscode.window.showInformationMessage(`Filter applied: "${pattern}"`); + } } - }); - - outputChannel.appendLine('All commands registered successfully.'); - - // Create kernel with message handler - // Create kernel for postgres-notebook - const kernel = new PostgresKernel(context, 'postgres-notebook', async (message: { type: string; command: string; format?: string; content?: string; filename?: string }) => { - console.log('Extension: Received message from kernel:', message); - if (message.type === 'custom' && message.command === 'export') { - console.log('Extension: Handling export command'); - vscode.commands.executeCommand('postgres-explorer.exportData', { - format: message.format, - content: message.content, - filename: message.filename - }); + } + }, + { + command: 'postgres-explorer.clearFilter', + callback: () => { + databaseTreeProvider.clearFilter(); + vscode.commands.executeCommand('setContext', 'postgresExplorer.filterActive', false); + vscode.window.showInformationMessage('Filter cleared'); + } + }, + { + command: 'postgres-explorer.addToFavorites', + callback: async (item: DatabaseTreeItem) => { + if (item) { + await databaseTreeProvider.addToFavorites(item); } - }); - context.subscriptions.push(kernel); - - // Create kernel for postgres-query (SQL files) - const queryKernel = new PostgresKernel(context, 'postgres-query'); - - // Set up renderer messaging to receive messages from the notebook renderer - console.log('Extension: Setting up renderer messaging for postgres-query-renderer'); - outputChannel.appendLine('Setting up renderer messaging for postgres-query-renderer'); - const rendererMessaging = vscode.notebooks.createRendererMessaging('postgres-query-renderer'); - rendererMessaging.onDidReceiveMessage(async (event) => { - console.log('Extension: Received message from renderer:', event.message); - outputChannel.appendLine('Received message from renderer: ' + JSON.stringify(event.message)); - const message = event.message; - const notebook = event.editor.notebook; - - if (message.type === 'execute_update_background') { - console.log('Extension: Processing execute_update_background'); - const { statements } = message; - - try { - // Get connection from notebook metadata - const metadata = notebook.metadata as PostgresMetadata; - if (!metadata?.connectionId) { - await ErrorHandlers.handleCommandError(new Error('No connection found in notebook metadata'), 'execute background update'); - return; - } - - const password = await SecretStorageService.getInstance().getPassword(metadata.connectionId); - - const client = new Client({ - host: metadata.host, - port: metadata.port, - database: metadata.databaseName, - user: metadata.username, - password: password || metadata.password || undefined, - }); - - await client.connect(); - console.log('Extension: Connected to database for background update'); - - // Execute each statement - let successCount = 0; - let errorCount = 0; - for (const stmt of statements) { - try { - console.log('Extension: Executing:', stmt); - await client.query(stmt); - successCount++; - } catch (err: any) { - console.error('Extension: Statement error:', err.message); - errorCount++; - await ErrorHandlers.handleCommandError(err, 'execute update statement'); - } - } - - await client.end(); - - if (successCount > 0) { - vscode.window.showInformationMessage(`Successfully updated ${successCount} row(s)${errorCount > 0 ? `, ${errorCount} failed` : ''}`); - } - } catch (err: any) { - console.error('Extension: Background update error:', err); - await ErrorHandlers.handleCommandError(err, 'execute background updates'); - } - } else if (message.type === 'script_delete') { - console.log('Extension: Processing script_delete from renderer'); - const { schema, table, primaryKeys, rows, cellIndex } = message; - - try { - // Construct DELETE query - let query = ''; - for (const row of rows) { - const conditions: string[] = []; - - for (const pk of primaryKeys) { - const val = row[pk]; - const valStr = typeof val === 'string' ? `'${val.replace(/'/g, "''")}'` : val; - conditions.push(`"${pk}" = ${valStr}`); - } - query += `DELETE FROM "${schema}"."${table}" WHERE ${conditions.join(' AND ')};\n`; - } - - // Insert new cell with the query - const targetIndex = cellIndex + 1; - const newCell = new vscode.NotebookCellData( - vscode.NotebookCellKind.Code, - query, - 'sql' - ); - - const edit = new vscode.NotebookEdit( - new vscode.NotebookRange(targetIndex, targetIndex), - [newCell] - ); - - const workspaceEdit = new vscode.WorkspaceEdit(); - workspaceEdit.set(notebook.uri, [edit]); - await vscode.workspace.applyEdit(workspaceEdit); - } catch (err: any) { - await ErrorHandlers.handleCommandError(err, 'generate delete script'); - console.error('Extension: Script delete error:', err); - } + } + }, + { + command: 'postgres-explorer.removeFromFavorites', + callback: async (item: DatabaseTreeItem) => { + if (item) { + await databaseTreeProvider.removeFromFavorites(item); + } + } + }, + { + command: 'postgres-explorer.manageConnections', + callback: () => { + ConnectionManagementPanel.show(context.extensionUri, context); + } + }, + { + command: 'postgres-explorer.aiSettings', + callback: () => { + AiSettingsPanel.show(context.extensionUri, context); + } + }, + { + command: 'postgres-explorer.connect', + callback: async (item: any) => await cmdConnectDatabase(item, context, databaseTreeProvider) + }, + { + command: 'postgres-explorer.disconnect', + callback: async () => { + databaseTreeProvider.refresh(); + vscode.window.showInformationMessage('Disconnected from PostgreSQL database'); + } + }, + { + command: 'postgres-explorer.queryTable', + callback: async (item: any) => { + if (!item || !item.schema) { + return; } - }); - // Note: rendererMessaging doesn't have dispose method, so we don't add to subscriptions - - // Register notebook providers - const notebookProvider = new PostgresNotebookProvider(); - context.subscriptions.push( - vscode.workspace.registerNotebookSerializer('postgres-notebook', notebookProvider), - vscode.workspace.registerNotebookSerializer('postgres-query', new PostgresNotebookSerializer()) - ); - - // Register SQL completion provider - const { SqlCompletionProvider } = require('./providers/SqlCompletionProvider'); - const sqlCompletionProvider = new SqlCompletionProvider(); - context.subscriptions.push( - vscode.languages.registerCompletionItemProvider( - { language: 'sql' }, - sqlCompletionProvider, - '.' // Trigger on dot for schema.table suggestions - ), - vscode.languages.registerCompletionItemProvider( - { scheme: 'vscode-notebook-cell', language: 'sql' }, - sqlCompletionProvider, - '.' - ) - ); - - // Register CodeLens Provider for both 'postgres' and 'sql' languages - const aiCodeLensProvider = new AiCodeLensProvider(); - context.subscriptions.push( - vscode.languages.registerCodeLensProvider( - { language: 'postgres', scheme: 'vscode-notebook-cell' }, - aiCodeLensProvider - ), - vscode.languages.registerCodeLensProvider( - { language: 'sql', scheme: 'vscode-notebook-cell' }, - aiCodeLensProvider - ) - ); - outputChannel.appendLine('AiCodeLensProvider registered for postgres and sql languages.'); - - // Immediately migrate any existing passwords to SecretStorage - await migrateExistingPasswords(context); -} - -export async function deactivate() { - await ConnectionManager.getInstance().closeAll(); -} -async function migrateExistingPasswords(context: vscode.ExtensionContext) { + const query = `SELECT * FROM ${item.schema}.${item.label} LIMIT 100;`; + const notebook = await vscode.workspace.openNotebookDocument('postgres-notebook', new vscode.NotebookData([ + new vscode.NotebookCellData(vscode.NotebookCellKind.Code, query, 'sql') + ])); + await vscode.window.showNotebookDocument(notebook); + } + }, + { + command: 'postgres-explorer.newNotebook', + callback: async (item: any) => await cmdNewNotebook(item) + }, + { + command: 'postgres-explorer.refresh', + callback: () => databaseTreeProvider.refresh() + }, + // Add database commands + { + command: 'postgres-explorer.createInDatabase', + callback: async (item: DatabaseTreeItem) => await cmdAddObjectInDatabase(item, context) + }, + { + command: 'postgres-explorer.createDatabase', + callback: async (item: DatabaseTreeItem) => await cmdCreateDatabase(item, context) + }, + { + command: 'postgres-explorer.dropDatabase', + callback: async (item: DatabaseTreeItem) => await cmdDeleteDatabase(item, context) + }, + { + command: 'postgres-explorer.scriptAlterDatabase', + callback: async (item: DatabaseTreeItem) => await cmdScriptAlterDatabase(item, context) + }, + { + command: 'postgres-explorer.databaseOperations', + callback: async (item: DatabaseTreeItem) => await cmdDatabaseOperations(item, context) + }, + { + command: 'postgres-explorer.showDashboard', + callback: async (item: DatabaseTreeItem) => await cmdDatabaseDashboard(item, context) + }, + { + command: 'postgres-explorer.backupDatabase', + callback: async (item: DatabaseTreeItem) => await cmdBackupDatabase(item, context) + }, + { + command: 'postgres-explorer.restoreDatabase', + callback: async (item: DatabaseTreeItem) => await cmdRestoreDatabase(item, context) + }, + { + command: 'postgres-explorer.generateCreateScript', + callback: async (item: DatabaseTreeItem) => await cmdGenerateCreateScript(item, context) + }, + { + command: 'postgres-explorer.disconnectDatabase', + callback: async (item: DatabaseTreeItem) => await cmdDisconnectDatabaseLegacy(item, context) + }, + { + command: 'postgres-explorer.maintenanceDatabase', + callback: async (item: DatabaseTreeItem) => await cmdMaintenanceDatabase(item, context) + }, + { + command: 'postgres-explorer.queryTool', + callback: async (item: DatabaseTreeItem) => await cmdQueryTool(item, context) + }, + { + command: 'postgres-explorer.psqlTool', + callback: async (item: DatabaseTreeItem) => await cmdPsqlTool(item, context) + }, + { + command: 'postgres-explorer.showConfiguration', + callback: async (item: DatabaseTreeItem) => await cmdShowConfiguration(item, context) + }, + // Add schema commands + { + command: 'postgres-explorer.createSchema', + callback: async (item: DatabaseTreeItem) => await cmdCreateSchema(item, context) + }, + { + command: 'postgres-explorer.createInSchema', + callback: async (item: DatabaseTreeItem) => await cmdCreateObjectInSchema(item, context) + }, + { + command: 'postgres-explorer.schemaOperations', + callback: async (item: DatabaseTreeItem) => await cmdSchemaOperations(item, context) + }, + { + command: 'postgres-explorer.showSchemaProperties', + callback: async (item: DatabaseTreeItem) => await cmdShowSchemaProperties(item, context) + }, + // Add table commands + { + command: 'postgres-explorer.editTable', + callback: async (item: DatabaseTreeItem) => await cmdEditTable(item, context) + }, + { + command: 'postgres-explorer.viewTableData', + callback: async (item: DatabaseTreeItem) => { + await databaseTreeProvider.addToRecent(item); + await cmdViewTableData(item, context); + } + }, + { + command: 'postgres-explorer.dropTable', + callback: async (item: DatabaseTreeItem) => await cmdDropTable(item, context) + }, + { + command: 'postgres-explorer.tableOperations', + callback: async (item: DatabaseTreeItem) => await cmdTableOperations(item, context) + }, + { + command: 'postgres-explorer.truncateTable', + callback: async (item: DatabaseTreeItem) => await cmdTruncateTable(item, context) + }, + { + command: 'postgres-explorer.insertData', + callback: async (item: DatabaseTreeItem) => await cmdInsertTable(item, context) + }, + { + command: 'postgres-explorer.updateData', + callback: async (item: DatabaseTreeItem) => await cmdUpdateTable(item, context) + }, + { + command: 'postgres-explorer.showTableProperties', + callback: async (item: DatabaseTreeItem) => await cmdShowTableProperties(item, context) + }, + // Add script commands + { + command: 'postgres-explorer.scriptSelect', + callback: async (item: DatabaseTreeItem) => await cmdScriptSelect(item, context) + }, + { + command: 'postgres-explorer.scriptInsert', + callback: async (item: DatabaseTreeItem) => await cmdScriptInsert(item, context) + }, + { + command: 'postgres-explorer.scriptUpdate', + callback: async (item: DatabaseTreeItem) => await cmdScriptUpdate(item, context) + }, + { + command: 'postgres-explorer.scriptDelete', + callback: async (item: DatabaseTreeItem) => await cmdScriptDelete(item, context) + }, + { + command: 'postgres-explorer.scriptCreate', + callback: async (item: DatabaseTreeItem) => await cmdScriptCreate(item, context) + }, + // Add maintenance commands + { + command: 'postgres-explorer.maintenanceVacuum', + callback: async (item: DatabaseTreeItem) => await cmdMaintenanceVacuum(item, context) + }, + { + command: 'postgres-explorer.maintenanceAnalyze', + callback: async (item: DatabaseTreeItem) => await cmdMaintenanceAnalyze(item, context) + }, + { + command: 'postgres-explorer.maintenanceReindex', + callback: async (item: DatabaseTreeItem) => await cmdMaintenanceReindex(item, context) + }, + + // Add view commands + { + command: 'postgres-explorer.refreshView', + callback: async (item: DatabaseTreeItem) => await cmdRefreshView(item, context, databaseTreeProvider) + }, + { + command: 'postgres-explorer.editViewDefinition', + callback: async (item: DatabaseTreeItem) => await cmdEditView(item, context) + }, + { + command: 'postgres-explorer.viewViewData', + callback: async (item: DatabaseTreeItem) => { + await databaseTreeProvider.addToRecent(item); + await cmdViewData(item, context); + } + }, + { + command: 'postgres-explorer.dropView', + callback: async (item: DatabaseTreeItem) => await cmdDropView(item, context) + }, + { + command: 'postgres-explorer.viewOperations', + callback: async (item: DatabaseTreeItem) => await cmdViewOperations(item, context) + }, + { + command: 'postgres-explorer.showViewProperties', + callback: async (item: DatabaseTreeItem) => await cmdShowViewProperties(item, context) + }, + { + command: 'postgres-explorer.viewScriptSelect', + callback: async (item: DatabaseTreeItem) => await cmdViewScriptSelect(item, context) + }, + { + command: 'postgres-explorer.viewScriptCreate', + callback: async (item: DatabaseTreeItem) => await cmdViewScriptCreate(item, context) + }, + // Add function commands + { + command: 'postgres-explorer.refreshFunction', + callback: async (item: DatabaseTreeItem) => await cmdRefreshFunction(item, context, databaseTreeProvider) + }, + { + command: 'postgres-explorer.showFunctionProperties', + callback: async (item: DatabaseTreeItem) => { + await databaseTreeProvider.addToRecent(item); + await cmdShowFunctionProperties(item, context); + } + }, + { + command: 'postgres-explorer.functionOperations', + callback: async (item: DatabaseTreeItem) => await cmdFunctionOperations(item, context) + }, + { + command: 'postgres-explorer.createReplaceFunction', + callback: async (item: DatabaseTreeItem) => await cmdEditFunction(item, context) + }, + { + command: 'postgres-explorer.callFunction', + callback: async (item: DatabaseTreeItem) => await cmdCallFunction(item, context) + }, + { + command: 'postgres-explorer.dropFunction', + callback: async (item: DatabaseTreeItem) => await cmdDropFunction(item, context) + }, + // Add materialized view commands + { + command: 'postgres-explorer.refreshMaterializedView', + callback: async (item: DatabaseTreeItem) => await cmdRefreshMatView(item, context) + }, + { + command: 'postgres-explorer.editMatView', + callback: async (item: DatabaseTreeItem) => await cmdEditMatView(item, context) + }, + { + command: 'postgres-explorer.viewMaterializedViewData', + callback: async (item: DatabaseTreeItem) => await cmdViewMatViewData(item, context) + }, + { + command: 'postgres-explorer.showMaterializedViewProperties', + callback: async (item: DatabaseTreeItem) => await cmdViewMatViewProperties(item, context) + }, + { + command: 'postgres-explorer.dropMatView', + callback: async (item: DatabaseTreeItem) => await cmdDropMatView(item, context) + }, + { + command: 'postgres-explorer.materializedViewOperations', + callback: async (item: DatabaseTreeItem) => await cmdMatViewOperations(item, context) + }, + // Add type commands + { + command: 'postgres-explorer.refreshType', + callback: async (item: DatabaseTreeItem) => await cmdRefreshType(item, context, databaseTreeProvider) + }, + { + command: 'postgres-explorer.typeOperations', + callback: async (item: DatabaseTreeItem) => await cmdAllOperationsTypes(item, context) + }, + { + command: 'postgres-explorer.editType', + callback: async (item: DatabaseTreeItem) => await cmdEditTypes(item, context) + }, + { + command: 'postgres-explorer.showTypeProperties', + callback: async (item: DatabaseTreeItem) => await cmdShowTypeProperties(item, context) + }, + { + command: 'postgres-explorer.dropType', + callback: async (item: DatabaseTreeItem) => await cmdDropType(item, context) + }, + // Add foreign table commands + { + command: 'postgres-explorer.refreshForeignTable', + callback: async (item: DatabaseTreeItem) => await cmdRefreshForeignTable(item, context, databaseTreeProvider) + }, + { + command: 'postgres-explorer.foreignTableOperations', + callback: async (item: DatabaseTreeItem) => await cmdForeignTableOperations(item, context) + }, + { + command: 'postgres-explorer.editForeignTable', + callback: async (item: DatabaseTreeItem) => await cmdEditForeignTable(item, context) + }, + // Add role/user commands + { + command: 'postgres-explorer.refreshRole', + callback: async (item: DatabaseTreeItem) => await cmdRefreshRole(item, context, databaseTreeProvider) + }, + { + command: 'postgres-explorer.createUser', + callback: async (item: DatabaseTreeItem) => await cmdAddUser(item, context) + }, + { + command: 'postgres-explorer.createRole', + callback: async (item: DatabaseTreeItem) => await cmdAddRole(item, context) + }, + { + command: 'postgres-explorer.editRole', + callback: async (item: DatabaseTreeItem) => await cmdEditRole(item, context) + }, + { + command: 'postgres-explorer.grantRevoke', + callback: async (item: DatabaseTreeItem) => await cmdGrantRevokeRole(item, context) + }, + { + command: 'postgres-explorer.dropRole', + callback: async (item: DatabaseTreeItem) => await cmdDropRole(item, context) + }, + { + command: 'postgres-explorer.roleOperations', + callback: async (item: DatabaseTreeItem) => await cmdRoleOperations(item, context) + }, + { + command: 'postgres-explorer.showRoleProperties', + callback: async (item: DatabaseTreeItem) => await cmdShowRoleProperties(item, context) + }, + // Add extension commands + { + command: 'postgres-explorer.refreshExtension', + callback: async (item: DatabaseTreeItem) => await cmdRefreshExtension(item, context, databaseTreeProvider) + }, + { + command: 'postgres-explorer.enableExtension', + callback: async (item: DatabaseTreeItem) => await cmdEnableExtension(item, context) + }, + { + command: 'postgres-explorer.extensionOperations', + callback: async (item: DatabaseTreeItem) => await cmdExtensionOperations(item, context) + }, + { + command: 'postgres-explorer.dropExtension', + callback: async (item: DatabaseTreeItem) => await cmdDropExtension(item, context) + }, + // Add connection commands + { + command: 'postgres-explorer.disconnectConnection', + callback: async (item: DatabaseTreeItem) => await cmdDisconnectConnection(item, context, databaseTreeProvider) + }, + { + command: 'postgres-explorer.reconnectConnection', + callback: async (item: DatabaseTreeItem) => await cmdReconnectConnection(item, context, databaseTreeProvider) + }, + { + command: 'postgres-explorer.deleteConnection', + callback: async (item: DatabaseTreeItem) => await cmdDisconnectDatabase(item, context, databaseTreeProvider) + }, + + { + command: 'postgres-explorer.createTable', + callback: async (item: DatabaseTreeItem) => await cmdCreateTable(item, context) + }, + { + command: 'postgres-explorer.createView', + callback: async (item: DatabaseTreeItem) => await cmdCreateView(item, context) + }, + { + command: 'postgres-explorer.createFunction', + callback: async (item: DatabaseTreeItem) => await cmdCreateFunction(item, context) + }, + { + command: 'postgres-explorer.createMaterializedView', + callback: async (item: DatabaseTreeItem) => await cmdCreateMaterializedView(item, context) + }, + { + command: 'postgres-explorer.createType', + callback: async (item: DatabaseTreeItem) => await cmdCreateType(item, context) + }, + { + command: 'postgres-explorer.createForeignTable', + callback: async (item: DatabaseTreeItem) => await cmdCreateForeignTable(item, context) + }, + // Foreign Data Wrapper commands + { + command: 'postgres-explorer.foreignDataWrapperOperations', + callback: async (item: DatabaseTreeItem) => await cmdForeignDataWrapperOperations(item, context) + }, + { + command: 'postgres-explorer.showForeignDataWrapperProperties', + callback: async (item: DatabaseTreeItem) => await cmdShowForeignDataWrapperProperties(item, context) + }, + { + command: 'postgres-explorer.refreshForeignDataWrapper', + callback: async (item: DatabaseTreeItem) => await cmdRefreshForeignDataWrapper(item, context, databaseTreeProvider) + }, + // Foreign Server commands + { + command: 'postgres-explorer.createForeignServer', + callback: async (item: DatabaseTreeItem) => await cmdCreateForeignServer(item, context) + }, + { + command: 'postgres-explorer.foreignServerOperations', + callback: async (item: DatabaseTreeItem) => await cmdForeignServerOperations(item, context) + }, + { + command: 'postgres-explorer.showForeignServerProperties', + callback: async (item: DatabaseTreeItem) => await cmdShowForeignServerProperties(item, context) + }, + { + command: 'postgres-explorer.dropForeignServer', + callback: async (item: DatabaseTreeItem) => await cmdDropForeignServer(item, context) + }, + { + command: 'postgres-explorer.refreshForeignServer', + callback: async (item: DatabaseTreeItem) => await cmdRefreshForeignServer(item, context, databaseTreeProvider) + }, + // User Mapping commands + { + command: 'postgres-explorer.createUserMapping', + callback: async (item: DatabaseTreeItem) => await cmdCreateUserMapping(item, context) + }, + { + command: 'postgres-explorer.userMappingOperations', + callback: async (item: DatabaseTreeItem) => await cmdUserMappingOperations(item, context) + }, + { + command: 'postgres-explorer.showUserMappingProperties', + callback: async (item: DatabaseTreeItem) => await cmdShowUserMappingProperties(item, context) + }, + { + command: 'postgres-explorer.dropUserMapping', + callback: async (item: DatabaseTreeItem) => await cmdDropUserMapping(item, context) + }, + { + command: 'postgres-explorer.refreshUserMapping', + callback: async (item: DatabaseTreeItem) => await cmdRefreshUserMapping(item, context, databaseTreeProvider) + }, + { + command: 'postgres-explorer.createRole', + callback: async (item: DatabaseTreeItem) => await cmdAddRole(item, context) + }, + { + command: 'postgres-explorer.enableExtension', + callback: async (item: DatabaseTreeItem) => await cmdEnableExtension(item, context) + }, + + { + command: 'postgres-explorer.aiAssist', + callback: async (cell: vscode.NotebookCell) => await cmdAiAssist(cell, context, outputChannel) + }, + + // Column commands + { + command: 'postgres-explorer.showColumnProperties', + callback: async (item: DatabaseTreeItem) => await showColumnProperties(item) + }, + { + command: 'postgres-explorer.copyColumnName', + callback: async (item: DatabaseTreeItem) => await copyColumnName(item) + }, + { + command: 'postgres-explorer.copyColumnNameQuoted', + callback: async (item: DatabaseTreeItem) => await copyColumnNameQuoted(item) + }, + { + command: 'postgres-explorer.generateSelectStatement', + callback: async (item: DatabaseTreeItem) => await generateSelectStatement(item) + }, + { + command: 'postgres-explorer.generateWhereClause', + callback: async (item: DatabaseTreeItem) => await generateWhereClause(item) + }, + { + command: 'postgres-explorer.generateAlterColumnScript', + callback: async (item: DatabaseTreeItem) => await generateAlterColumnScript(item) + }, + { + command: 'postgres-explorer.generateDropColumnScript', + callback: async (item: DatabaseTreeItem) => await generateDropColumnScript(item) + }, + { + command: 'postgres-explorer.generateRenameColumnScript', + callback: async (item: DatabaseTreeItem) => await generateRenameColumnScript(item) + }, + { + command: 'postgres-explorer.addColumnComment', + callback: async (item: DatabaseTreeItem) => await addColumnComment(item) + }, + { + command: 'postgres-explorer.generateIndexOnColumn', + callback: async (item: DatabaseTreeItem) => await generateIndexOnColumn(item) + }, + { + command: 'postgres-explorer.viewColumnStatistics', + callback: async (item: DatabaseTreeItem) => await viewColumnStatistics(item) + }, + + // Constraint commands + { + command: 'postgres-explorer.showConstraintProperties', + callback: async (item: DatabaseTreeItem) => await showConstraintProperties(item) + }, + { + command: 'postgres-explorer.copyConstraintName', + callback: async (item: DatabaseTreeItem) => await copyConstraintName(item) + }, + { + command: 'postgres-explorer.generateDropConstraintScript', + callback: async (item: DatabaseTreeItem) => await generateDropConstraintScript(item) + }, + { + command: 'postgres-explorer.generateAlterConstraintScript', + callback: async (item: DatabaseTreeItem) => await generateAlterConstraintScript(item) + }, + { + command: 'postgres-explorer.validateConstraint', + callback: async (item: DatabaseTreeItem) => await validateConstraint(item) + }, + { + command: 'postgres-explorer.generateAddConstraintScript', + callback: async (item: DatabaseTreeItem) => await generateAddConstraintScript(item) + }, + { + command: 'postgres-explorer.viewConstraintDependencies', + callback: async (item: DatabaseTreeItem) => await viewConstraintDependencies(item) + }, + { + command: 'postgres-explorer.constraintOperations', + callback: async (item: DatabaseTreeItem) => await cmdConstraintOperations(item, context) + }, + + // Index commands + { + command: 'postgres-explorer.showIndexProperties', + callback: async (item: DatabaseTreeItem) => await showIndexProperties(item) + }, + { + command: 'postgres-explorer.copyIndexName', + callback: async (item: DatabaseTreeItem) => await copyIndexName(item) + }, + { + command: 'postgres-explorer.generateDropIndexScript', + callback: async (item: DatabaseTreeItem) => await generateDropIndexScript(item) + }, + { + command: 'postgres-explorer.generateReindexScript', + callback: async (item: DatabaseTreeItem) => await generateReindexScript(item) + }, + { + command: 'postgres-explorer.generateScriptCreate', + callback: async (item: DatabaseTreeItem) => await generateScriptCreate(item) + }, + { + command: 'postgres-explorer.analyzeIndexUsage', + callback: async (item: DatabaseTreeItem) => await analyzeIndexUsage(item) + }, + { + command: 'postgres-explorer.generateAlterIndexScript', + callback: async (item: DatabaseTreeItem) => await generateAlterIndexScript(item) + }, + { + command: 'postgres-explorer.addIndexComment', + callback: async (item: DatabaseTreeItem) => await addIndexComment(item) + }, + { + command: 'postgres-explorer.indexOperations', + callback: async (item: DatabaseTreeItem) => await cmdIndexOperations(item, context) + }, + { + command: 'postgres-explorer.addColumn', + callback: async (item: DatabaseTreeItem) => await cmdAddColumn(item) + }, + { + command: 'postgres-explorer.addConstraint', + callback: async (item: DatabaseTreeItem) => await cmdAddConstraint(item) + }, + { + command: 'postgres-explorer.addIndex', + callback: async (item: DatabaseTreeItem) => await cmdAddIndex(item) + }, + ]; + + // Register all commands + console.log('Starting command registration...'); + outputChannel.appendLine('Starting command registration...'); + + commands.forEach(({ command, callback }) => { try { - const config = vscode.workspace.getConfiguration(); - const connections = config.get('postgresExplorer.connections') || []; + console.log(`Registering command: ${command}`); + context.subscriptions.push( + vscode.commands.registerCommand(command, callback) + ); + } catch (e) { + console.error(`Failed to register command ${command}:`, e); + outputChannel.appendLine(`Failed to register command ${command}: ${e}`); + } + }); + + outputChannel.appendLine('All commands registered successfully.'); + + // Create kernel with message handler + // Create kernel for postgres-notebook + const kernel = new PostgresKernel(context, 'postgres-notebook', async (message: { type: string; command: string; format?: string; content?: string; filename?: string }) => { + console.log('Extension: Received message from kernel:', message); + if (message.type === 'custom' && message.command === 'export') { + console.log('Extension: Handling export command'); + vscode.commands.executeCommand('postgres-explorer.exportData', { + format: message.format, + content: message.content, + filename: message.filename + }); + } + }); + context.subscriptions.push(kernel); + + // Create kernel for postgres-query (SQL files) + const queryKernel = new PostgresKernel(context, 'postgres-query'); + + // Set up renderer messaging to receive messages from the notebook renderer + console.log('Extension: Setting up renderer messaging for postgres-query-renderer'); + outputChannel.appendLine('Setting up renderer messaging for postgres-query-renderer'); + const rendererMessaging = vscode.notebooks.createRendererMessaging('postgres-query-renderer'); + rendererMessaging.onDidReceiveMessage(async (event) => { + console.log('Extension: Received message from renderer:', event.message); + outputChannel.appendLine('Received message from renderer: ' + JSON.stringify(event.message)); + const message = event.message; + const notebook = event.editor.notebook; + + if (message.type === 'execute_update_background') { + console.log('Extension: Processing execute_update_background'); + const { statements } = message; + + try { + // Get connection from notebook metadata + const metadata = notebook.metadata as PostgresMetadata; + if (!metadata?.connectionId) { + await ErrorHandlers.handleCommandError(new Error('No connection found in notebook metadata'), 'execute background update'); + return; + } + + const password = await SecretStorageService.getInstance().getPassword(metadata.connectionId); + + const client = new Client({ + host: metadata.host, + port: metadata.port, + database: metadata.databaseName, + user: metadata.username, + password: password || metadata.password || undefined, + }); + + await client.connect(); + console.log('Extension: Connected to database for background update'); + + // Execute each statement + let successCount = 0; + let errorCount = 0; + for (const stmt of statements) { + try { + console.log('Extension: Executing:', stmt); + await client.query(stmt); + successCount++; + } catch (err: any) { + console.error('Extension: Statement error:', err.message); + errorCount++; + await ErrorHandlers.handleCommandError(err, 'execute update statement'); + } + } - // First remove passwords from settings to ensure they don't persist - const sanitizedConnections = connections.map(({ password, ...connWithoutPassword }) => connWithoutPassword); - await config.update('postgresExplorer.connections', sanitizedConnections, vscode.ConfigurationTarget.Global); + await client.end(); - // Then store passwords in SecretStorage - for (const conn of connections) { - if (conn.password) { - await SecretStorageService.getInstance().setPassword(conn.id, conn.password); - } + if (successCount > 0) { + vscode.window.showInformationMessage(`Successfully updated ${successCount} row(s)${errorCount > 0 ? `, ${errorCount} failed` : ''}`); + } + } catch (err: any) { + console.error('Extension: Background update error:', err); + await ErrorHandlers.handleCommandError(err, 'execute background updates'); + } + } else if (message.type === 'script_delete') { + console.log('Extension: Processing script_delete from renderer'); + const { schema, table, primaryKeys, rows, cellIndex } = message; + + try { + // Construct DELETE query + let query = ''; + for (const row of rows) { + const conditions: string[] = []; + + for (const pk of primaryKeys) { + const val = row[pk]; + const valStr = typeof val === 'string' ? `'${val.replace(/'/g, "''")}'` : val; + conditions.push(`"${pk}" = ${valStr}`); + } + query += `DELETE FROM "${schema}"."${table}" WHERE ${conditions.join(' AND ')};\n`; } - return true; - } catch (error) { - console.error('Failed to migrate passwords:', error); - return false; + // Insert new cell with the query + const targetIndex = cellIndex + 1; + const newCell = new vscode.NotebookCellData( + vscode.NotebookCellKind.Code, + query, + 'sql' + ); + + const edit = new vscode.NotebookEdit( + new vscode.NotebookRange(targetIndex, targetIndex), + [newCell] + ); + + const workspaceEdit = new vscode.WorkspaceEdit(); + workspaceEdit.set(notebook.uri, [edit]); + await vscode.workspace.applyEdit(workspaceEdit); + } catch (err: any) { + await ErrorHandlers.handleCommandError(err, 'generate delete script'); + console.error('Extension: Script delete error:', err); + } } + }); + // Note: rendererMessaging doesn't have dispose method, so we don't add to subscriptions + + // Register notebook providers + const notebookProvider = new PostgresNotebookProvider(); + context.subscriptions.push( + vscode.workspace.registerNotebookSerializer('postgres-notebook', notebookProvider), + vscode.workspace.registerNotebookSerializer('postgres-query', new PostgresNotebookSerializer()) + ); + + // Register SQL completion provider + const { SqlCompletionProvider } = require('./providers/SqlCompletionProvider'); + const sqlCompletionProvider = new SqlCompletionProvider(); + context.subscriptions.push( + vscode.languages.registerCompletionItemProvider( + { language: 'sql' }, + sqlCompletionProvider, + '.' // Trigger on dot for schema.table suggestions + ), + vscode.languages.registerCompletionItemProvider( + { scheme: 'vscode-notebook-cell', language: 'sql' }, + sqlCompletionProvider, + '.' + ) + ); + + // Register CodeLens Provider for both 'postgres' and 'sql' languages + const aiCodeLensProvider = new AiCodeLensProvider(); + context.subscriptions.push( + vscode.languages.registerCodeLensProvider( + { language: 'postgres', scheme: 'vscode-notebook-cell' }, + aiCodeLensProvider + ), + vscode.languages.registerCodeLensProvider( + { language: 'sql', scheme: 'vscode-notebook-cell' }, + aiCodeLensProvider + ) + ); + outputChannel.appendLine('AiCodeLensProvider registered for postgres and sql languages.'); + + // Immediately migrate any existing passwords to SecretStorage + await migrateExistingPasswords(context); +} + +export async function deactivate() { + await ConnectionManager.getInstance().closeAll(); +} + +async function migrateExistingPasswords(context: vscode.ExtensionContext) { + try { + const config = vscode.workspace.getConfiguration(); + const connections = config.get('postgresExplorer.connections') || []; + + // First remove passwords from settings to ensure they don't persist + const sanitizedConnections = connections.map(({ password, ...connWithoutPassword }) => connWithoutPassword); + await config.update('postgresExplorer.connections', sanitizedConnections, vscode.ConfigurationTarget.Global); + + // Then store passwords in SecretStorage + for (const conn of connections) { + if (conn.password) { + await SecretStorageService.getInstance().setPassword(conn.id, conn.password); + } + } + + return true; + } catch (error) { + console.error('Failed to migrate passwords:', error); + return false; + } } diff --git a/src/providers/DatabaseTreeProvider.ts b/src/providers/DatabaseTreeProvider.ts index 5dc794b..3f3aa63 100644 --- a/src/providers/DatabaseTreeProvider.ts +++ b/src/providers/DatabaseTreeProvider.ts @@ -2,175 +2,374 @@ import { Client } from 'pg'; import * as vscode from 'vscode'; import { ConnectionManager } from '../services/ConnectionManager'; -export class DatabaseTreeProvider implements vscode.TreeDataProvider { - private _onDidChangeTreeData: vscode.EventEmitter = new vscode.EventEmitter(); - readonly onDidChangeTreeData: vscode.Event = this._onDidChangeTreeData.event; - private disconnectedConnections: Set = new Set(); - - constructor(private readonly extensionContext: vscode.ExtensionContext) { - // Initialize all connections as disconnected by default - this.initializeDisconnectedState(); - } - - private initializeDisconnectedState(): void { - const connections = vscode.workspace.getConfiguration().get('postgresExplorer.connections') || []; - connections.forEach(conn => { - this.disconnectedConnections.add(conn.id); - }); - } - - markConnectionDisconnected(connectionId: string): void { - this.disconnectedConnections.add(connectionId); - // Fire a full refresh to update tree state and collapse items - this._onDidChangeTreeData.fire(undefined); - } +// Key format for favorites: "type:connectionId:database:schema:name" +function buildItemKey(item: DatabaseTreeItem): string { + const parts = [item.type, item.connectionId || '', item.databaseName || '', item.schema || '', item.label]; + return parts.join(':'); +} - markConnectionConnected(connectionId: string): void { - this.disconnectedConnections.delete(connectionId); - // Fire a full refresh to update tree state - this._onDidChangeTreeData.fire(undefined); +export class DatabaseTreeProvider implements vscode.TreeDataProvider { + private _onDidChangeTreeData: vscode.EventEmitter = new vscode.EventEmitter(); + readonly onDidChangeTreeData: vscode.Event = this._onDidChangeTreeData.event; + private disconnectedConnections: Set = new Set(); + + // Filter, Favorites, and Recent Items + private _filterPattern: string = ''; + private _favorites: Set = new Set(); + private _recentItems: string[] = []; + private static readonly MAX_RECENT_ITEMS = 10; + private static readonly FAVORITES_KEY = 'postgresExplorer.favorites'; + private static readonly RECENT_KEY = 'postgresExplorer.recentItems'; + + constructor(private readonly extensionContext: vscode.ExtensionContext) { + // Initialize all connections as disconnected by default + this.initializeDisconnectedState(); + // Load persisted favorites and recent items + this.loadPersistedData(); + } + + private loadPersistedData(): void { + const favorites = this.extensionContext.globalState.get(DatabaseTreeProvider.FAVORITES_KEY, []); + this._favorites = new Set(favorites); + this._recentItems = this.extensionContext.globalState.get(DatabaseTreeProvider.RECENT_KEY, []); + } + + private async saveFavorites(): Promise { + await this.extensionContext.globalState.update(DatabaseTreeProvider.FAVORITES_KEY, Array.from(this._favorites)); + } + + private async saveRecentItems(): Promise { + await this.extensionContext.globalState.update(DatabaseTreeProvider.RECENT_KEY, this._recentItems); + } + + // Filter methods + get filterPattern(): string { + return this._filterPattern; + } + + setFilter(pattern: string): void { + this._filterPattern = pattern.toLowerCase(); + this.refresh(); + } + + clearFilter(): void { + this._filterPattern = ''; + this.refresh(); + } + + // Favorites methods + isFavorite(item: DatabaseTreeItem): boolean { + return this._favorites.has(buildItemKey(item)); + } + + async addToFavorites(item: DatabaseTreeItem): Promise { + const key = buildItemKey(item); + this._favorites.add(key); + await this.saveFavorites(); + this.refresh(); + vscode.window.showInformationMessage(`Added "${item.label}" to favorites`); + } + + async removeFromFavorites(item: DatabaseTreeItem): Promise { + const key = buildItemKey(item); + this._favorites.delete(key); + await this.saveFavorites(); + this.refresh(); + vscode.window.showInformationMessage(`Removed "${item.label}" from favorites`); + } + + getFavoriteKeys(): string[] { + return Array.from(this._favorites); + } + + // Recent items methods + async addToRecent(item: DatabaseTreeItem): Promise { + const key = buildItemKey(item); + // Remove if already exists (to move to front) + this._recentItems = this._recentItems.filter(k => k !== key); + // Add to front + this._recentItems.unshift(key); + // Trim to max size + if (this._recentItems.length > DatabaseTreeProvider.MAX_RECENT_ITEMS) { + this._recentItems = this._recentItems.slice(0, DatabaseTreeProvider.MAX_RECENT_ITEMS); } - - refresh(element?: DatabaseTreeItem): void { - this._onDidChangeTreeData.fire(element); + await this.saveRecentItems(); + } + + getRecentKeys(): string[] { + return [...this._recentItems]; + } + + private matchesFilter(label: string): boolean { + if (!this._filterPattern) return true; + return label.toLowerCase().includes(this._filterPattern); + } + + private isFavoriteItem(type: string, connectionId?: string, databaseName?: string, schema?: string, name?: string): boolean { + const key = `${type}:${connectionId || ''}:${databaseName || ''}:${schema || ''}:${name || ''}`; + return this._favorites.has(key); + } + + private initializeDisconnectedState(): void { + const connections = vscode.workspace.getConfiguration().get('postgresExplorer.connections') || []; + connections.forEach(conn => { + this.disconnectedConnections.add(conn.id); + }); + } + + markConnectionDisconnected(connectionId: string): void { + this.disconnectedConnections.add(connectionId); + // Fire a full refresh to update tree state and collapse items + this._onDidChangeTreeData.fire(undefined); + } + + markConnectionConnected(connectionId: string): void { + this.disconnectedConnections.delete(connectionId); + // Fire a full refresh to update tree state + this._onDidChangeTreeData.fire(undefined); + } + + refresh(element?: DatabaseTreeItem): void { + this._onDidChangeTreeData.fire(element); + } + + collapseAll(): void { + // This will trigger a refresh of the tree view with all items collapsed + this._onDidChangeTreeData.fire(); + } + + getTreeItem(element: DatabaseTreeItem): vscode.TreeItem { + return element; + } + + async getChildren(element?: DatabaseTreeItem): Promise { + const connections = vscode.workspace.getConfiguration().get('postgresExplorer.connections') || []; + + if (!element) { + // Root level - show connections + return connections.map(conn => new DatabaseTreeItem( + conn.name || `${conn.host}:${conn.port}`, + vscode.TreeItemCollapsibleState.Collapsed, + 'connection', + conn.id, + undefined, // databaseName + undefined, // schema + undefined, // tableName + undefined, // columnName + undefined, // comment + undefined, // isInstalled + undefined, // installedVersion + undefined, // roleAttributes + this.disconnectedConnections.has(conn.id) // isDisconnected + )); } - collapseAll(): void { - // This will trigger a refresh of the tree view with all items collapsed - this._onDidChangeTreeData.fire(); + // Auto-connect on expansion: if connection is disconnected, mark it as connected + if (element.type === 'connection' && element.connectionId && this.disconnectedConnections.has(element.connectionId)) { + console.log(`Connection ${element.connectionId} is being expanded, auto-connecting...`); + this.markConnectionConnected(element.connectionId); } - getTreeItem(element: DatabaseTreeItem): vscode.TreeItem { - return element; + const connection = connections.find(c => c.id === element.connectionId); + if (!connection) { + console.error(`Connection not found for ID: ${element.connectionId}`); + vscode.window.showErrorMessage('Connection configuration not found'); + return []; } - async getChildren(element?: DatabaseTreeItem): Promise { - const connections = vscode.workspace.getConfiguration().get('postgresExplorer.connections') || []; - - if (!element) { - // Root level - show connections - return connections.map(conn => new DatabaseTreeItem( - conn.name || `${conn.host}:${conn.port}`, - vscode.TreeItemCollapsibleState.Collapsed, - 'connection', - conn.id, - undefined, // databaseName - undefined, // schema - undefined, // tableName - undefined, // columnName - undefined, // comment - undefined, // isInstalled - undefined, // installedVersion - undefined, // roleAttributes - this.disconnectedConnections.has(conn.id) // isDisconnected + let client: Client | undefined; + try { + const dbName = element.type === 'connection' ? 'postgres' : element.databaseName; + + console.log(`Attempting to connect to ${connection.name} (${dbName})`); + + // Use ConnectionManager to get a shared connection + client = await ConnectionManager.getInstance().getConnection({ + id: connection.id, + host: connection.host, + port: connection.port, + username: connection.username, + database: dbName, + name: connection.name + }); + + console.log(`Successfully connected to ${connection.name}`); + + switch (element.type) { + case 'connection': + // At connection level, show Favorites (if any), Databases group and Users & Roles + const items: DatabaseTreeItem[] = []; + + // Check if there are favorites for this connection + const connectionFavorites = this.getFavoriteKeys().filter(key => { + const parts = key.split(':'); + return parts[1] === element.connectionId; + }); + if (connectionFavorites.length > 0) { + items.push(new DatabaseTreeItem('Favorites', vscode.TreeItemCollapsibleState.Collapsed, 'favorites-group', element.connectionId)); + } + + // Check if there are recent items for this connection + const connectionRecent = this.getRecentKeys().filter(key => { + const parts = key.split(':'); + return parts[1] === element.connectionId; + }); + if (connectionRecent.length > 0) { + items.push(new DatabaseTreeItem('Recent', vscode.TreeItemCollapsibleState.Collapsed, 'recent-group', element.connectionId)); + } + + items.push(new DatabaseTreeItem('Databases', vscode.TreeItemCollapsibleState.Collapsed, 'databases-group', element.connectionId)); + items.push(new DatabaseTreeItem('Users & Roles', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId)); + return items; + + case 'databases-group': + // Show all databases under the Databases group (including system databases) + const dbResult = await client.query( + "SELECT datname FROM pg_database ORDER BY datname" + ); + return dbResult.rows.map(row => new DatabaseTreeItem( + row.datname, + vscode.TreeItemCollapsibleState.Collapsed, + 'database', + element.connectionId, + row.datname + )); + + case 'favorites-group': + // Show all favorited items for this connection + const favoriteItems: DatabaseTreeItem[] = []; + const favoriteKeys = this.getFavoriteKeys().filter(key => { + const parts = key.split(':'); + return parts[1] === element.connectionId; + }); + + for (const key of favoriteKeys) { + const parts = key.split(':'); + // Key format: type:connectionId:database:schema:name + const itemType = parts[0] as 'table' | 'view' | 'function' | 'materialized-view'; + const dbName = parts[2]; + const schemaName = parts[3]; + const itemName = parts[4]; + + // Determine collapsible state based on type + const collapsible = (itemType === 'table' || itemType === 'view') + ? vscode.TreeItemCollapsibleState.Collapsed + : vscode.TreeItemCollapsibleState.None; + + // Use just the item name as label (for SQL commands), put extra info in description via isFavorite handling + favoriteItems.push(new DatabaseTreeItem( + itemName, // Just the name - SQL commands use label + collapsible, + itemType, + element.connectionId, + dbName, + schemaName, + itemName, // tableName - also just the name + undefined, // columnName + `${schemaName}.${dbName}`, // comment - for tooltip + undefined, // isInstalled + undefined, // installedVersion + undefined, // roleAttributes + undefined, // isDisconnected + true // isFavorite )); - } - - // Auto-connect on expansion: if connection is disconnected, mark it as connected - if (element.type === 'connection' && element.connectionId && this.disconnectedConnections.has(element.connectionId)) { - console.log(`Connection ${element.connectionId} is being expanded, auto-connecting...`); - this.markConnectionConnected(element.connectionId); - } - - const connection = connections.find(c => c.id === element.connectionId); - if (!connection) { - console.error(`Connection not found for ID: ${element.connectionId}`); - vscode.window.showErrorMessage('Connection configuration not found'); - return []; - } - - let client: Client | undefined; - try { - const dbName = element.type === 'connection' ? 'postgres' : element.databaseName; - - console.log(`Attempting to connect to ${connection.name} (${dbName})`); - - // Use ConnectionManager to get a shared connection - client = await ConnectionManager.getInstance().getConnection({ - id: connection.id, - host: connection.host, - port: connection.port, - username: connection.username, - database: dbName, - name: connection.name - }); - - console.log(`Successfully connected to ${connection.name}`); - - switch (element.type) { - case 'connection': - // At connection level, show Databases group and Users & Roles - return [ - new DatabaseTreeItem('Databases', vscode.TreeItemCollapsibleState.Collapsed, 'databases-group', element.connectionId), - new DatabaseTreeItem('Users & Roles', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId) - ]; - - case 'databases-group': - // Show all databases under the Databases group (including system databases) - const dbResult = await client.query( - "SELECT datname FROM pg_database ORDER BY datname" - ); - return dbResult.rows.map(row => new DatabaseTreeItem( - row.datname, - vscode.TreeItemCollapsibleState.Collapsed, - 'database', - element.connectionId, - row.datname - )); - - case 'database': - // Return just the categories at database level - return [ - new DatabaseTreeItem('Schemas', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName), - new DatabaseTreeItem('Extensions', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName), - new DatabaseTreeItem('Foreign Data Wrappers', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName) - ]; - - case 'category': - // Handle table sub-categories - if (element.tableName) { - switch (element.label) { - case 'Columns': - const columnResult = await client.query( - "SELECT column_name, data_type FROM information_schema.columns WHERE table_schema = $1 AND table_name = $2 ORDER BY ordinal_position", - [element.schema, element.tableName] - ); - return columnResult.rows.map(row => new DatabaseTreeItem( - `${row.column_name} (${row.data_type})`, - vscode.TreeItemCollapsibleState.None, - 'column', - element.connectionId, - element.databaseName, - element.schema, - element.tableName, - row.column_name - )); - - case 'Constraints': - const constraintResult = await client.query( - `SELECT + } + return favoriteItems; + + case 'recent-group': + // Show all recent items for this connection (max 10) + const recentItems: DatabaseTreeItem[] = []; + const recentKeys = this.getRecentKeys().filter(key => { + const parts = key.split(':'); + return parts[1] === element.connectionId; + }); + + for (const key of recentKeys) { + const parts = key.split(':'); + // Key format: type:connectionId:database:schema:name + const itemType = parts[0] as 'table' | 'view' | 'function' | 'materialized-view'; + const dbName = parts[2]; + const schemaName = parts[3]; + const itemName = parts[4]; + + // Determine collapsible state based on type + const collapsible = (itemType === 'table' || itemType === 'view') + ? vscode.TreeItemCollapsibleState.Collapsed + : vscode.TreeItemCollapsibleState.None; + + // Use just the item name as label + recentItems.push(new DatabaseTreeItem( + itemName, + collapsible, + itemType, + element.connectionId, + dbName, + schemaName, + itemName, + undefined, // columnName + `${schemaName}.${dbName}`, // comment - for tooltip + undefined, // isInstalled + undefined, // installedVersion + undefined, // roleAttributes + undefined, // isDisconnected + false // isFavorite - these are recent, not favorites + )); + } + return recentItems; + + case 'database': + // Return just the categories at database level + return [ + new DatabaseTreeItem('Schemas', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName), + new DatabaseTreeItem('Extensions', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName), + new DatabaseTreeItem('Foreign Data Wrappers', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName) + ]; + + case 'category': + // Handle table sub-categories + if (element.tableName) { + switch (element.label) { + case 'Columns': + const columnResult = await client.query( + "SELECT column_name, data_type FROM information_schema.columns WHERE table_schema = $1 AND table_name = $2 ORDER BY ordinal_position", + [element.schema, element.tableName] + ); + return columnResult.rows.map(row => new DatabaseTreeItem( + `${row.column_name} (${row.data_type})`, + vscode.TreeItemCollapsibleState.None, + 'column', + element.connectionId, + element.databaseName, + element.schema, + element.tableName, + row.column_name + )); + + case 'Constraints': + const constraintResult = await client.query( + `SELECT tc.constraint_name, tc.constraint_type FROM information_schema.table_constraints tc WHERE tc.table_schema = $1 AND tc.table_name = $2 ORDER BY tc.constraint_type, tc.constraint_name`, - [element.schema, element.tableName] - ); - return constraintResult.rows.map(row => { - return new DatabaseTreeItem( - row.constraint_name, - vscode.TreeItemCollapsibleState.None, - 'constraint', - element.connectionId, - element.databaseName, - element.schema, - element.tableName - ); - }); - - case 'Indexes': - const indexResult = await client.query( - `SELECT + [element.schema, element.tableName] + ); + return constraintResult.rows.map(row => { + return new DatabaseTreeItem( + row.constraint_name, + vscode.TreeItemCollapsibleState.None, + 'constraint', + element.connectionId, + element.databaseName, + element.schema, + element.tableName + ); + }); + + case 'Indexes': + const indexResult = await client.query( + `SELECT i.relname as index_name, ix.indisunique as is_unique, ix.indisprimary as is_primary @@ -180,57 +379,58 @@ export class DatabaseTreeProvider implements vscode.TreeDataProvider { - return new DatabaseTreeItem( - row.index_name, - vscode.TreeItemCollapsibleState.None, - 'index', - element.connectionId, - element.databaseName, - element.schema, - element.tableName - ); - }); - } - } - - // Schema-level categories - switch (element.label) { - case 'Users & Roles': - const roleResult = await client.query( - `SELECT r.rolname, + [element.schema, element.tableName] + ); + return indexResult.rows.map(row => { + return new DatabaseTreeItem( + row.index_name, + vscode.TreeItemCollapsibleState.None, + 'index', + element.connectionId, + element.databaseName, + element.schema, + element.tableName + ); + }); + } + } + + // Schema-level categories - extract base name (handle badge format "Tables • 5") + const categoryName = element.label.split(' • ')[0]; + switch (categoryName) { + case 'Users & Roles': + const roleResult = await client.query( + `SELECT r.rolname, r.rolsuper, r.rolcreatedb, r.rolcreaterole, r.rolcanlogin FROM pg_roles r ORDER BY r.rolname` - ); - return roleResult.rows.map(row => new DatabaseTreeItem( - row.rolname, - vscode.TreeItemCollapsibleState.None, - 'role', - element.connectionId, - element.databaseName, - undefined, - undefined, - undefined, - undefined, - undefined, - undefined, - { - rolsuper: row.rolsuper, - rolcreatedb: row.rolcreatedb, - rolcreaterole: row.rolcreaterole, - rolcanlogin: row.rolcanlogin - } - )); - - case 'Schemas': - const schemaResult = await client.query( - `SELECT nspname as schema_name + ); + return roleResult.rows.map(row => new DatabaseTreeItem( + row.rolname, + vscode.TreeItemCollapsibleState.None, + 'role', + element.connectionId, + element.databaseName, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + { + rolsuper: row.rolsuper, + rolcreatedb: row.rolcreatedb, + rolcreaterole: row.rolcreaterole, + rolcanlogin: row.rolcanlogin + } + )); + + case 'Schemas': + const schemaResult = await client.query( + `SELECT nspname as schema_name FROM pg_namespace WHERE nspname NOT LIKE 'pg_%' AND nspname != 'information_schema' @@ -240,308 +440,429 @@ export class DatabaseTreeProvider implements vscode.TreeDataProvider new DatabaseTreeItem( - row.schema_name, - vscode.TreeItemCollapsibleState.Collapsed, - 'schema', - element.connectionId, - element.databaseName, - row.schema_name - )); - - case 'Extensions': - const extensionResult = await client.query( - `SELECT e.name, + ); + + // If filter is active, only show schemas that have matching items + if (this._filterPattern) { + const filteredSchemas: DatabaseTreeItem[] = []; + for (const row of schemaResult.rows) { + // Check if schema has any matching tables, views, or functions + const matchResult = await client.query( + `SELECT 1 FROM information_schema.tables + WHERE table_schema = $1 AND table_type = 'BASE TABLE' + AND LOWER(table_name) LIKE $2 + UNION ALL + SELECT 1 FROM information_schema.views + WHERE table_schema = $1 AND LOWER(table_name) LIKE $2 + UNION ALL + SELECT 1 FROM information_schema.routines + WHERE routine_schema = $1 AND routine_type = 'FUNCTION' + AND LOWER(routine_name) LIKE $2 + LIMIT 1`, + [row.schema_name, `%${this._filterPattern}%`] + ); + if (matchResult.rows.length > 0) { + filteredSchemas.push(new DatabaseTreeItem( + row.schema_name, + vscode.TreeItemCollapsibleState.Collapsed, + 'schema', + element.connectionId, + element.databaseName, + row.schema_name + )); + } + } + return filteredSchemas; + } + + return schemaResult.rows.map(row => new DatabaseTreeItem( + row.schema_name, + vscode.TreeItemCollapsibleState.Collapsed, + 'schema', + element.connectionId, + element.databaseName, + row.schema_name + )); + + case 'Extensions': + const extensionResult = await client.query( + `SELECT e.name, e.installed_version, e.default_version, e.comment, CASE WHEN e.installed_version IS NOT NULL THEN true ELSE false END as is_installed FROM pg_available_extensions e ORDER BY is_installed DESC, name` - ); - return extensionResult.rows.map(row => new DatabaseTreeItem( - row.installed_version ? `${row.name} (${row.installed_version})` : `${row.name} (${row.default_version})`, - vscode.TreeItemCollapsibleState.None, - 'extension', - element.connectionId, - element.databaseName, - undefined, - undefined, - undefined, - row.comment, - row.is_installed, - row.installed_version - )); - - // Existing category cases for schema level items - case 'Tables': - const tableResult = await client.query( - "SELECT table_name FROM information_schema.tables WHERE table_schema = $1 AND table_type = 'BASE TABLE' ORDER BY table_name", - [element.schema] - ); - return tableResult.rows.map(row => new DatabaseTreeItem( - row.table_name, - vscode.TreeItemCollapsibleState.Collapsed, - 'table', - element.connectionId, - element.databaseName, - element.schema - )); - - case 'Views': - const viewResult = await client.query( - "SELECT table_name FROM information_schema.views WHERE table_schema = $1 ORDER BY table_name", - [element.schema] - ); - return viewResult.rows.map(row => new DatabaseTreeItem( - row.table_name, - vscode.TreeItemCollapsibleState.Collapsed, - 'view', - element.connectionId, - element.databaseName, - element.schema - )); - - case 'Functions': - const functionResult = await client.query( - "SELECT routine_name FROM information_schema.routines WHERE routine_schema = $1 AND routine_type = 'FUNCTION' ORDER BY routine_name", - [element.schema] - ); - return functionResult.rows.map(row => new DatabaseTreeItem( - row.routine_name, - vscode.TreeItemCollapsibleState.None, - 'function', - element.connectionId, - element.databaseName, - element.schema - )); - - case 'Materialized Views': - const materializedViewResult = await client.query( - "SELECT matviewname as name FROM pg_matviews WHERE schemaname = $1 ORDER BY matviewname", - [element.schema] - ); - return materializedViewResult.rows.map(row => new DatabaseTreeItem( - row.name, - vscode.TreeItemCollapsibleState.None, - 'materialized-view', - element.connectionId, - element.databaseName, - element.schema - )); - - case 'Types': - const typeResult = await client.query( - `SELECT t.typname as name + ); + return extensionResult.rows.map(row => new DatabaseTreeItem( + row.installed_version ? `${row.name} (${row.installed_version})` : `${row.name} (${row.default_version})`, + vscode.TreeItemCollapsibleState.None, + 'extension', + element.connectionId, + element.databaseName, + undefined, + undefined, + undefined, + row.comment, + row.is_installed, + row.installed_version + )); + + // Existing category cases for schema level items + case 'Tables': + const tableResult = await client.query( + "SELECT table_name FROM information_schema.tables WHERE table_schema = $1 AND table_type = 'BASE TABLE' ORDER BY table_name", + [element.schema] + ); + return tableResult.rows + .filter(row => this.matchesFilter(row.table_name)) + .map(row => { + const isFav = this.isFavoriteItem('table', element.connectionId, element.databaseName, element.schema, row.table_name); + return new DatabaseTreeItem( + row.table_name, + vscode.TreeItemCollapsibleState.Collapsed, + 'table', + element.connectionId, + element.databaseName, + element.schema, + undefined, // tableName + undefined, // columnName + undefined, // comment + undefined, // isInstalled + undefined, // installedVersion + undefined, // roleAttributes + undefined, // isDisconnected + isFav // isFavorite + ); + }); + + case 'Views': + const viewResult = await client.query( + "SELECT table_name FROM information_schema.views WHERE table_schema = $1 ORDER BY table_name", + [element.schema] + ); + return viewResult.rows + .filter(row => this.matchesFilter(row.table_name)) + .map(row => { + const isFav = this.isFavoriteItem('view', element.connectionId, element.databaseName, element.schema, row.table_name); + return new DatabaseTreeItem( + row.table_name, + vscode.TreeItemCollapsibleState.Collapsed, + 'view', + element.connectionId, + element.databaseName, + element.schema, + undefined, undefined, undefined, undefined, undefined, undefined, undefined, + isFav + ); + }); + + case 'Functions': + const functionResult = await client.query( + "SELECT routine_name FROM information_schema.routines WHERE routine_schema = $1 AND routine_type = 'FUNCTION' ORDER BY routine_name", + [element.schema] + ); + return functionResult.rows + .filter(row => this.matchesFilter(row.routine_name)) + .map(row => { + const isFav = this.isFavoriteItem('function', element.connectionId, element.databaseName, element.schema, row.routine_name); + return new DatabaseTreeItem( + row.routine_name, + vscode.TreeItemCollapsibleState.None, + 'function', + element.connectionId, + element.databaseName, + element.schema, + undefined, undefined, undefined, undefined, undefined, undefined, undefined, + isFav + ); + }); + + case 'Materialized Views': + const materializedViewResult = await client.query( + "SELECT matviewname as name FROM pg_matviews WHERE schemaname = $1 ORDER BY matviewname", + [element.schema] + ); + return materializedViewResult.rows + .filter(row => this.matchesFilter(row.name)) + .map(row => { + const isFav = this.isFavoriteItem('materialized-view', element.connectionId, element.databaseName, element.schema, row.name); + return new DatabaseTreeItem( + row.name, + vscode.TreeItemCollapsibleState.None, + 'materialized-view', + element.connectionId, + element.databaseName, + element.schema, + undefined, undefined, undefined, undefined, undefined, undefined, undefined, + isFav + ); + }); + + case 'Types': + const typeResult = await client.query( + `SELECT t.typname as name FROM pg_type t JOIN pg_namespace n ON t.typnamespace = n.oid WHERE n.nspname = $1 AND t.typtype = 'c' ORDER BY t.typname`, - [element.schema] - ); - return typeResult.rows.map(row => new DatabaseTreeItem( - row.name, - vscode.TreeItemCollapsibleState.None, - 'type', - element.connectionId, - element.databaseName, - element.schema - )); - - case 'Foreign Tables': - const foreignTableResult = await client.query( - `SELECT c.relname as name + [element.schema] + ); + return typeResult.rows.map(row => new DatabaseTreeItem( + row.name, + vscode.TreeItemCollapsibleState.None, + 'type', + element.connectionId, + element.databaseName, + element.schema + )); + + case 'Foreign Tables': + const foreignTableResult = await client.query( + `SELECT c.relname as name FROM pg_foreign_table ft JOIN pg_class c ON ft.ftrelid = c.oid JOIN pg_namespace n ON c.relnamespace = n.oid WHERE n.nspname = $1 ORDER BY c.relname`, - [element.schema] - ); - return foreignTableResult.rows.map(row => new DatabaseTreeItem( - row.name, - vscode.TreeItemCollapsibleState.None, - 'foreign-table', - element.connectionId, - element.databaseName, - element.schema - )); - - case 'Foreign Data Wrappers': - const fdwResult = await client.query( - `SELECT fdwname as name + [element.schema] + ); + return foreignTableResult.rows.map(row => new DatabaseTreeItem( + row.name, + vscode.TreeItemCollapsibleState.None, + 'foreign-table', + element.connectionId, + element.databaseName, + element.schema + )); + + case 'Foreign Data Wrappers': + const fdwResult = await client.query( + `SELECT fdwname as name FROM pg_foreign_data_wrapper ORDER BY fdwname` - ); - return fdwResult.rows.map(row => new DatabaseTreeItem( - row.name, - vscode.TreeItemCollapsibleState.Collapsed, - 'foreign-data-wrapper', - element.connectionId, - element.databaseName - )); - } - return []; - - case 'schema': - return [ - new DatabaseTreeItem('Tables', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema), - new DatabaseTreeItem('Views', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema), - new DatabaseTreeItem('Functions', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema), - new DatabaseTreeItem('Materialized Views', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema), - new DatabaseTreeItem('Types', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema), - new DatabaseTreeItem('Foreign Tables', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema) - ]; - - case 'table': - // Show hierarchical structure for tables - return [ - new DatabaseTreeItem('Columns', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema, element.label), - new DatabaseTreeItem('Constraints', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema, element.label), - new DatabaseTreeItem('Indexes', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema, element.label) - ]; - - case 'view': - // Views only have columns - return [ - new DatabaseTreeItem('Columns', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema, element.label) - ]; - - case 'foreign-data-wrapper': - // FDW node - list all foreign servers using this FDW - const serversResult = await client.query( - `SELECT srv.srvname as name + ); + return fdwResult.rows.map(row => new DatabaseTreeItem( + row.name, + vscode.TreeItemCollapsibleState.Collapsed, + 'foreign-data-wrapper', + element.connectionId, + element.databaseName + )); + } + return []; + + case 'schema': + // Query counts for each category (with filter applied if active) + const filterPattern = this._filterPattern ? `%${this._filterPattern.toLowerCase()}%` : null; + + const tablesCountResult = await client.query( + filterPattern + ? "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = $1 AND table_type = 'BASE TABLE' AND LOWER(table_name) LIKE $2" + : "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = $1 AND table_type = 'BASE TABLE'", + filterPattern ? [element.schema, filterPattern] : [element.schema] + ); + + const viewsCountResult = await client.query( + filterPattern + ? "SELECT COUNT(*) FROM information_schema.views WHERE table_schema = $1 AND LOWER(table_name) LIKE $2" + : "SELECT COUNT(*) FROM information_schema.views WHERE table_schema = $1", + filterPattern ? [element.schema, filterPattern] : [element.schema] + ); + + const functionsCountResult = await client.query( + filterPattern + ? "SELECT COUNT(*) FROM information_schema.routines WHERE routine_schema = $1 AND routine_type = 'FUNCTION' AND LOWER(routine_name) LIKE $2" + : "SELECT COUNT(*) FROM information_schema.routines WHERE routine_schema = $1 AND routine_type = 'FUNCTION'", + filterPattern ? [element.schema, filterPattern] : [element.schema] + ); + + const materializedViewsCountResult = await client.query( + filterPattern + ? "SELECT COUNT(*) FROM pg_matviews WHERE schemaname = $1 AND LOWER(matviewname) LIKE $2" + : "SELECT COUNT(*) FROM pg_matviews WHERE schemaname = $1", + filterPattern ? [element.schema, filterPattern] : [element.schema] + ); + + const typesCountResult = await client.query( + "SELECT COUNT(*) FROM pg_type t JOIN pg_namespace n ON t.typnamespace = n.oid WHERE n.nspname = $1 AND t.typtype = 'c'", + [element.schema] + ); + + const foreignTablesCountResult = await client.query( + "SELECT COUNT(*) FROM information_schema.foreign_tables WHERE foreign_table_schema = $1", + [element.schema] + ); + + return [ + new DatabaseTreeItem(`Tables • ${tablesCountResult.rows[0].count}`, vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema), + new DatabaseTreeItem(`Views • ${viewsCountResult.rows[0].count}`, vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema), + new DatabaseTreeItem(`Functions • ${functionsCountResult.rows[0].count}`, vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema), + new DatabaseTreeItem(`Materialized Views • ${materializedViewsCountResult.rows[0].count}`, vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema), + new DatabaseTreeItem(`Types • ${typesCountResult.rows[0].count}`, vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema), + new DatabaseTreeItem(`Foreign Tables • ${foreignTablesCountResult.rows[0].count}`, vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema) + ]; + + case 'table': + // Show hierarchical structure for tables + return [ + new DatabaseTreeItem('Columns', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema, element.label), + new DatabaseTreeItem('Constraints', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema, element.label), + new DatabaseTreeItem('Indexes', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema, element.label) + ]; + + case 'view': + // Views only have columns + return [ + new DatabaseTreeItem('Columns', vscode.TreeItemCollapsibleState.Collapsed, 'category', element.connectionId, element.databaseName, element.schema, element.label) + ]; + + case 'foreign-data-wrapper': + // FDW node - list all foreign servers using this FDW + const serversResult = await client.query( + `SELECT srv.srvname as name FROM pg_foreign_server srv JOIN pg_foreign_data_wrapper fdw ON srv.srvfdw = fdw.oid WHERE fdw.fdwname = $1 ORDER BY srv.srvname`, - [element.label] - ); - return serversResult.rows.map(row => new DatabaseTreeItem( - row.name, - vscode.TreeItemCollapsibleState.Collapsed, - 'foreign-server', - element.connectionId, - element.databaseName, - element.label // Store FDW name in schema field - )); - - case 'foreign-server': - // Foreign server node - list all user mappings - const mappingsResult = await client.query( - `SELECT um.usename as name + [element.label] + ); + return serversResult.rows.map(row => new DatabaseTreeItem( + row.name, + vscode.TreeItemCollapsibleState.Collapsed, + 'foreign-server', + element.connectionId, + element.databaseName, + element.label // Store FDW name in schema field + )); + + case 'foreign-server': + // Foreign server node - list all user mappings + const mappingsResult = await client.query( + `SELECT um.usename as name FROM pg_user_mappings um WHERE um.srvname = $1 ORDER BY um.usename`, - [element.label] - ); - return mappingsResult.rows.map(row => new DatabaseTreeItem( - row.name, - vscode.TreeItemCollapsibleState.None, - 'user-mapping', - element.connectionId, - element.databaseName, - element.label, // Store server name in schema field - element.label // Store server name in tableName for context - )); - - default: - return []; - } - } catch (err: any) { - const errorMessage = err.message || err.toString() || 'Unknown error'; - const errorCode = err.code || 'NO_CODE'; - const errorDetails = `Error getting tree items for ${element?.type || 'root'}: [${errorCode}] ${errorMessage}`; - - console.error(errorDetails); - console.error('Full error:', err); - - // Only show error message to user if it's not a connection initialization issue - if (element && element.type !== 'connection') { - vscode.window.showErrorMessage(`Failed to get tree items: ${errorMessage}`); - } - - return []; - } - // Do NOT close the client here, as it is managed by ConnectionManager + [element.label] + ); + return mappingsResult.rows.map(row => new DatabaseTreeItem( + row.name, + vscode.TreeItemCollapsibleState.None, + 'user-mapping', + element.connectionId, + element.databaseName, + element.label, // Store server name in schema field + element.label // Store server name in tableName for context + )); + + default: + return []; + } + } catch (err: any) { + const errorMessage = err.message || err.toString() || 'Unknown error'; + const errorCode = err.code || 'NO_CODE'; + const errorDetails = `Error getting tree items for ${element?.type || 'root'}: [${errorCode}] ${errorMessage}`; + + console.error(errorDetails); + console.error('Full error:', err); + + // Only show error message to user if it's not a connection initialization issue + if (element && element.type !== 'connection') { + vscode.window.showErrorMessage(`Failed to get tree items: ${errorMessage}`); + } + + return []; } + // Do NOT close the client here, as it is managed by ConnectionManager + } } export class DatabaseTreeItem extends vscode.TreeItem { - constructor( - public readonly label: string, - public readonly collapsibleState: vscode.TreeItemCollapsibleState, - public readonly type: 'connection' | 'database' | 'schema' | 'table' | 'view' | 'function' | 'column' | 'category' | 'materialized-view' | 'type' | 'foreign-table' | 'extension' | 'role' | 'databases-group' | 'constraint' | 'index' | 'foreign-data-wrapper' | 'foreign-server' | 'user-mapping', - public readonly connectionId?: string, - public readonly databaseName?: string, - public readonly schema?: string, - public readonly tableName?: string, - public readonly columnName?: string, - public readonly comment?: string, - public readonly isInstalled?: boolean, - public readonly installedVersion?: string, - public readonly roleAttributes?: { [key: string]: boolean }, - public readonly isDisconnected?: boolean - ) { - super(label, collapsibleState); - if (type === 'category' && label) { - // Create specific context value for categories (e.g., category-tables, category-views) - const suffix = label.toLowerCase().replace(/\s+&\s+/g, '-').replace(/\s+/g, '-'); - this.contextValue = `category-${suffix}`; - } else if (type === 'connection' && isDisconnected) { - this.contextValue = 'connection-disconnected'; - } else { - this.contextValue = isInstalled ? `${type}-installed` : type; - } - this.tooltip = this.getTooltip(type, comment, roleAttributes); - this.description = this.getDescription(type, isInstalled, installedVersion, roleAttributes); - this.iconPath = { - connection: new vscode.ThemeIcon('plug', isDisconnected ? new vscode.ThemeColor('disabledForeground') : new vscode.ThemeColor('charts.blue')), - database: new vscode.ThemeIcon('database', new vscode.ThemeColor('charts.purple')), - 'databases-group': new vscode.ThemeIcon('database', new vscode.ThemeColor('charts.purple')), - schema: new vscode.ThemeIcon('symbol-namespace', new vscode.ThemeColor('charts.yellow')), - table: new vscode.ThemeIcon('table', new vscode.ThemeColor('charts.blue')), - view: new vscode.ThemeIcon('eye', new vscode.ThemeColor('charts.green')), - function: new vscode.ThemeIcon('symbol-method', new vscode.ThemeColor('charts.orange')), - column: new vscode.ThemeIcon('symbol-field', new vscode.ThemeColor('charts.blue')), - category: new vscode.ThemeIcon('list-tree'), - 'materialized-view': new vscode.ThemeIcon('symbol-structure', new vscode.ThemeColor('charts.green')), - type: new vscode.ThemeIcon('symbol-type-parameter', new vscode.ThemeColor('charts.red')), - 'foreign-table': new vscode.ThemeIcon('symbol-interface', new vscode.ThemeColor('charts.blue')), - extension: new vscode.ThemeIcon(isInstalled ? 'extensions-installed' : 'extensions', isInstalled ? new vscode.ThemeColor('charts.green') : undefined), - role: new vscode.ThemeIcon('person', new vscode.ThemeColor('charts.yellow')), - constraint: new vscode.ThemeIcon('lock', new vscode.ThemeColor('charts.orange')), - index: new vscode.ThemeIcon('search', new vscode.ThemeColor('charts.purple')), - 'foreign-data-wrapper': new vscode.ThemeIcon('extensions', new vscode.ThemeColor('charts.blue')), - 'foreign-server': new vscode.ThemeIcon('server', new vscode.ThemeColor('charts.green')), - 'user-mapping': new vscode.ThemeIcon('account', new vscode.ThemeColor('charts.yellow')) - }[type]; + constructor( + public readonly label: string, + public readonly collapsibleState: vscode.TreeItemCollapsibleState, + public readonly type: 'connection' | 'database' | 'schema' | 'table' | 'view' | 'function' | 'column' | 'category' | 'materialized-view' | 'type' | 'foreign-table' | 'extension' | 'role' | 'databases-group' | 'favorites-group' | 'recent-group' | 'constraint' | 'index' | 'foreign-data-wrapper' | 'foreign-server' | 'user-mapping', + public readonly connectionId?: string, + public readonly databaseName?: string, + public readonly schema?: string, + public readonly tableName?: string, + public readonly columnName?: string, + public readonly comment?: string, + public readonly isInstalled?: boolean, + public readonly installedVersion?: string, + public readonly roleAttributes?: { [key: string]: boolean }, + public readonly isDisconnected?: boolean, + public readonly isFavorite?: boolean, + public readonly count?: number // For category item counts + ) { + super(label, collapsibleState); + if (type === 'category' && label) { + // Create specific context value for categories (e.g., category-tables, category-views) + const suffix = label.toLowerCase().replace(/\s+&\s+/g, '-').replace(/\s+/g, '-'); + this.contextValue = `category-${suffix}`; + } else if (type === 'connection' && isDisconnected) { + this.contextValue = 'connection-disconnected'; + } else { + // Keep original contextValue - isFavorite flag is stored separately for star indicator + // For favorites menu detection, we use description containing ★ + this.contextValue = isInstalled ? `${type}-installed` : type; } - - private getTooltip(type: string, comment?: string, roleAttributes?: { [key: string]: boolean }): string { - if (type === 'role' && roleAttributes) { - const attributes = []; - if (roleAttributes.rolsuper) attributes.push('Superuser'); - if (roleAttributes.rolcreatedb) attributes.push('Create DB'); - if (roleAttributes.rolcreaterole) attributes.push('Create Role'); - if (roleAttributes.rolcanlogin) attributes.push('Can Login'); - return `${this.label}\n\nAttributes:\n${attributes.join('\n')}`; - } - return comment ? `${this.label}\n\n${comment}` : this.label; + this.tooltip = this.getTooltip(type, comment, roleAttributes); + this.description = this.getDescription(type, isInstalled, installedVersion, roleAttributes, isFavorite, count); + this.iconPath = { + connection: new vscode.ThemeIcon('plug', isDisconnected ? new vscode.ThemeColor('disabledForeground') : new vscode.ThemeColor('charts.blue')), + database: new vscode.ThemeIcon('database', new vscode.ThemeColor('charts.purple')), + 'databases-group': new vscode.ThemeIcon('database', new vscode.ThemeColor('charts.purple')), + 'favorites-group': new vscode.ThemeIcon('star-full', new vscode.ThemeColor('charts.yellow')), + 'recent-group': new vscode.ThemeIcon('history', new vscode.ThemeColor('charts.green')), + schema: new vscode.ThemeIcon('symbol-namespace', new vscode.ThemeColor('charts.yellow')), + table: new vscode.ThemeIcon('table', new vscode.ThemeColor('charts.blue')), + view: new vscode.ThemeIcon('eye', new vscode.ThemeColor('charts.green')), + function: new vscode.ThemeIcon('symbol-method', new vscode.ThemeColor('charts.orange')), + column: new vscode.ThemeIcon('symbol-field', new vscode.ThemeColor('charts.blue')), + category: new vscode.ThemeIcon('list-tree'), + 'materialized-view': new vscode.ThemeIcon('symbol-structure', new vscode.ThemeColor('charts.green')), + type: new vscode.ThemeIcon('symbol-type-parameter', new vscode.ThemeColor('charts.red')), + 'foreign-table': new vscode.ThemeIcon('symbol-interface', new vscode.ThemeColor('charts.blue')), + extension: new vscode.ThemeIcon(isInstalled ? 'extensions-installed' : 'extensions', isInstalled ? new vscode.ThemeColor('charts.green') : undefined), + role: new vscode.ThemeIcon('person', new vscode.ThemeColor('charts.yellow')), + constraint: new vscode.ThemeIcon('lock', new vscode.ThemeColor('charts.orange')), + index: new vscode.ThemeIcon('search', new vscode.ThemeColor('charts.purple')), + 'foreign-data-wrapper': new vscode.ThemeIcon('extensions', new vscode.ThemeColor('charts.blue')), + 'foreign-server': new vscode.ThemeIcon('server', new vscode.ThemeColor('charts.green')), + 'user-mapping': new vscode.ThemeIcon('account', new vscode.ThemeColor('charts.yellow')) + }[type]; + } + + private getTooltip(type: string, comment?: string, roleAttributes?: { [key: string]: boolean }): string { + if (type === 'role' && roleAttributes) { + const attributes = []; + if (roleAttributes.rolsuper) attributes.push('Superuser'); + if (roleAttributes.rolcreatedb) attributes.push('Create DB'); + if (roleAttributes.rolcreaterole) attributes.push('Create Role'); + if (roleAttributes.rolcanlogin) attributes.push('Can Login'); + return `${this.label}\n\nAttributes:\n${attributes.join('\n')}`; + } + return comment ? `${this.label}\n\n${comment}` : this.label; + } + + private getDescription(type: string, isInstalled?: boolean, installedVersion?: string, roleAttributes?: { [key: string]: boolean }, isFavorite?: boolean, count?: number): string | undefined { + let desc: string | undefined = undefined; + + if (type === 'extension' && isInstalled) { + desc = `v${installedVersion} (installed)`; + } else if (type === 'role' && roleAttributes) { + const tags = []; + if (roleAttributes.rolsuper) tags.push('superuser'); + if (roleAttributes.rolcanlogin) tags.push('login'); + desc = tags.length > 0 ? `(${tags.join(', ')})` : undefined; } - private getDescription(type: string, isInstalled?: boolean, installedVersion?: string, roleAttributes?: { [key: string]: boolean }): string | undefined { - if (type === 'extension' && isInstalled) { - return `v${installedVersion} (installed)`; - } - if (type === 'role' && roleAttributes) { - const tags = []; - if (roleAttributes.rolsuper) tags.push('superuser'); - if (roleAttributes.rolcanlogin) tags.push('login'); - return tags.length > 0 ? `(${tags.join(', ')})` : undefined; - } - return undefined; + // Append muted star for favorites (★ is more subtle than ⭐) + if (isFavorite) { + return desc ? `${desc} ★` : '★'; } + return desc; + } } diff --git a/src/services/ConnectionManager.ts b/src/services/ConnectionManager.ts index 9d593f6..e1a40e9 100644 --- a/src/services/ConnectionManager.ts +++ b/src/services/ConnectionManager.ts @@ -1,127 +1,164 @@ import { Client } from 'pg'; +import * as fs from 'fs'; import { ConnectionConfig } from '../common/types'; import { SecretStorageService } from './SecretStorageService'; import { SSHService } from './SSHService'; export class ConnectionManager { - private static instance: ConnectionManager; - private connections: Map = new Map(); + private static instance: ConnectionManager; + private connections: Map = new Map(); - private constructor() { } + private constructor() { } - public static getInstance(): ConnectionManager { - if (!ConnectionManager.instance) { - ConnectionManager.instance = new ConnectionManager(); - } - return ConnectionManager.instance; + public static getInstance(): ConnectionManager { + if (!ConnectionManager.instance) { + ConnectionManager.instance = new ConnectionManager(); } + return ConnectionManager.instance; + } - public async getConnection(config: ConnectionConfig): Promise { - const key = this.getConnectionKey(config); + public async getConnection(config: ConnectionConfig): Promise { + const key = this.getConnectionKey(config); - if (this.connections.has(key)) { - const client = this.connections.get(key)!; - // Simple check if connection is still alive (optional, pg client handles some of this) - // For now, we assume it's good or will throw on query, handling reconnection could be added here - return client; - } + if (this.connections.has(key)) { + const client = this.connections.get(key)!; + // Simple check if connection is still alive (optional, pg client handles some of this) + // For now, we assume it's good or will throw on query, handling reconnection could be added here + return client; + } - // Get password from secret storage if username is provided - let password: string | undefined; - if (config.username) { - password = await SecretStorageService.getInstance().getPassword(config.id); - // If username is provided but password is not found in storage, it might still work for some auth methods - } + // Get password from secret storage if username is provided + let password: string | undefined; + if (config.username) { + password = await SecretStorageService.getInstance().getPassword(config.id); + // If username is provided but password is not found in storage, it might still work for some auth methods + } - const clientConfig: any = { - user: config.username || undefined, - password: password || undefined, - database: config.database || 'postgres', - connectionTimeoutMillis: 5000 - }; - - if (config.ssh && config.ssh.enabled) { - try { - const stream = await SSHService.getInstance().createStream( - config.ssh, - config.host, - config.port - ); - clientConfig.stream = stream; - } catch (err: any) { - throw new Error(`SSH Connection failed: ${err.message}`); - } - } else { - clientConfig.host = config.host; - clientConfig.port = config.port; + // Build SSL configuration based on sslmode + let sslConfig: boolean | object = false; + if (config.sslmode && config.sslmode !== 'disable') { + sslConfig = { + rejectUnauthorized: config.sslmode === 'verify-ca' || config.sslmode === 'verify-full', + }; + + // Add certificate paths if provided + if (config.sslRootCertPath) { + try { + (sslConfig as any).ca = fs.readFileSync(config.sslRootCertPath).toString(); + } catch (e) { + console.warn('Failed to read SSL CA certificate:', e); + } + } + if (config.sslCertPath) { + try { + (sslConfig as any).cert = fs.readFileSync(config.sslCertPath).toString(); + } catch (e) { + console.warn('Failed to read SSL client certificate:', e); } + } + if (config.sslKeyPath) { + try { + (sslConfig as any).key = fs.readFileSync(config.sslKeyPath).toString(); + } catch (e) { + console.warn('Failed to read SSL client key:', e); + } + } + } - const client = new Client(clientConfig); + const clientConfig: any = { + user: config.username || undefined, + password: password || undefined, + database: config.database || 'postgres', + connectionTimeoutMillis: (config.connectTimeout || 5) * 1000, + statement_timeout: config.statementTimeout || undefined, + application_name: config.applicationName || 'PgStudio', + ssl: sslConfig || undefined, + // Parse additional options string if provided + ...(config.options ? { options: config.options } : {}) + }; + + if (config.ssh && config.ssh.enabled) { + try { + const stream = await SSHService.getInstance().createStream( + config.ssh, + config.host, + config.port + ); + clientConfig.stream = stream; + } catch (err: any) { + throw new Error(`SSH Connection failed: ${err.message}`); + } + } else { + clientConfig.host = config.host; + clientConfig.port = config.port; + } - await client.connect(); - this.connections.set(key, client); + const client = new Client(clientConfig); - // Remove connection on error/end - client.on('end', () => this.connections.delete(key)); - client.on('error', () => this.connections.delete(key)); + await client.connect(); + this.connections.set(key, client); - return client; - } + // Remove connection on error/end + client.on('end', () => this.connections.delete(key)); + client.on('error', () => this.connections.delete(key)); - public async closeConnection(config: ConnectionConfig): Promise { - const key = this.getConnectionKey(config); - const client = this.connections.get(key); - - if (client) { - try { - await client.end(); - this.connections.delete(key); - } catch (e) { - console.error('Error closing connection:', e); - } - } - } + return client; + } - public async closeAllConnectionsById(connectionId: string): Promise { - const keysToClose: string[] = []; + public async closeConnection(config: ConnectionConfig): Promise { + const key = this.getConnectionKey(config); + const client = this.connections.get(key); - // Find all connections with this ID - for (const key of this.connections.keys()) { - if (key.startsWith(`${connectionId}:`)) { - keysToClose.push(key); - } - } + if (client) { + try { + await client.end(); + this.connections.delete(key); + } catch (e) { + console.error('Error closing connection:', e); + } + } + } - // Close all found connections - for (const key of keysToClose) { - const client = this.connections.get(key); - if (client) { - try { - await client.end(); - this.connections.delete(key); - } catch (e) { - console.error(`Error closing connection ${key}:`, e); - } - } - } + public async closeAllConnectionsById(connectionId: string): Promise { + const keysToClose: string[] = []; - console.log(`Closed ${keysToClose.length} connections for ID: ${connectionId}`); + // Find all connections with this ID + for (const key of this.connections.keys()) { + if (key.startsWith(`${connectionId}:`)) { + keysToClose.push(key); + } } - public async closeAll(): Promise { - for (const client of this.connections.values()) { - try { - await client.end(); - } catch (e) { - console.error('Error closing connection:', e); - } + // Close all found connections + for (const key of keysToClose) { + const client = this.connections.get(key); + if (client) { + try { + await client.end(); + this.connections.delete(key); + } catch (e) { + console.error(`Error closing connection ${key}:`, e); } - this.connections.clear(); + } } - private getConnectionKey(config: ConnectionConfig): string { - // Unique key for connection: ID + Database - // If database is not specified, it connects to default (usually postgres) - return `${config.id}:${config.database || 'postgres'}`; + console.log(`Closed ${keysToClose.length} connections for ID: ${connectionId}`); + } + + public async closeAll(): Promise { + for (const client of this.connections.values()) { + try { + await client.end(); + } catch (e) { + console.error('Error closing connection:', e); + } } + this.connections.clear(); + } + + private getConnectionKey(config: ConnectionConfig): string { + // Unique key for connection: ID + Database + // If database is not specified, it connects to default (usually postgres) + return `${config.id}:${config.database || 'postgres'}`; + } } From ce7459dcb6f3ef77d9a252167732fee922c705d9 Mon Sep 17 00:00:00 2001 From: Richie Varghese Date: Thu, 25 Dec 2025 22:36:37 +0530 Subject: [PATCH 2/8] feat: introduce HistoryService, reformat DbObjectService, and update roadmap, kernel, and renderer files. --- docs/ROADMAP.md | 11 +- src/providers/NotebookKernel.ts | 1603 +++++----- src/renderer_v2.ts | 5312 ++++++++++++++++--------------- 3 files changed, 3552 insertions(+), 3374 deletions(-) diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md index 6bae294..cf95c45 100644 --- a/docs/ROADMAP.md +++ b/docs/ROADMAP.md @@ -25,11 +25,12 @@ - [x] 🕒 Recent items tracking (max 10 items) - [x] Object count badges on category nodes (right-aligned, muted) -### 2B: Notebook Experience -- [ ] Query cancellation button -- [ ] Virtual scrolling for large result sets (1000+ rows) -- [ ] Sticky column headers -- [ ] Column resizing +### 2B: Notebook Experience ✅ MOSTLY COMPLETE +- [x] Sticky headers (already implemented) +- [x] Query cancellation backend infrastructure +- [x] Column resizing +- [ ] Virtual scrolling (deferred - 6-8 hrs) +- [ ] Cancel button UI (deferred - requires major refactor) ### 2C: AI Assistant - [ ] Schema context caching diff --git a/src/providers/NotebookKernel.ts b/src/providers/NotebookKernel.ts index da47989..80f0da3 100644 --- a/src/providers/NotebookKernel.ts +++ b/src/providers/NotebookKernel.ts @@ -5,813 +5,852 @@ import { ConnectionManager } from '../services/ConnectionManager'; import { SecretStorageService } from '../services/SecretStorageService'; export class PostgresKernel implements vscode.Disposable { - readonly id = 'postgres-kernel'; - readonly label = 'PostgreSQL'; - readonly supportedLanguages = ['sql']; - - private readonly _controller: vscode.NotebookController; - private readonly _executionOrder = new WeakMap(); - private readonly _messageHandler?: (message: any) => void; - - constructor(private readonly context: vscode.ExtensionContext, viewType: string = 'postgres-notebook', messageHandler?: (message: any) => void) { - console.log(`PostgresKernel: Initializing for viewType: ${viewType}`); - this._controller = vscode.notebooks.createNotebookController( - this.id + '-' + viewType, - viewType, - this.label - ); - - this._messageHandler = messageHandler; - console.log(`PostgresKernel: Message handler registered for ${viewType}:`, !!messageHandler); - - this._controller.supportedLanguages = this.supportedLanguages; - this._controller.supportsExecutionOrder = true; - this._controller.executeHandler = this._executeAll.bind(this); - - // Disable automatic timestamp parsing (this was in original, but removed in new snippet, so removing it) - // const types = require('pg').types; - // const TIMESTAMPTZ_OID = 1184; - // const TIMESTAMP_OID = 1114; - // types.setTypeParser(TIMESTAMPTZ_OID, (val: string) => val); - // types.setTypeParser(TIMESTAMP_OID, (val: string) => val); - - // this._controller.description = 'PostgreSQL Query Executor'; // Removed as per new snippet - - const getClientFromNotebook = async (document: vscode.TextDocument): Promise => { - const cell = vscode.workspace.notebookDocuments - .find(notebook => notebook.getCells().some(c => c.document === document)) - ?.getCells() - .find(c => c.document === document); - - if (!cell) return undefined; + readonly id = 'postgres-kernel'; + readonly label = 'PostgreSQL'; + readonly supportedLanguages = ['sql']; + + private readonly _controller: vscode.NotebookController; + private readonly _executionOrder = new WeakMap(); + private readonly _messageHandler?: (message: any) => void; + + constructor(private readonly context: vscode.ExtensionContext, viewType: string = 'postgres-notebook', messageHandler?: (message: any) => void) { + console.log(`PostgresKernel: Initializing for viewType: ${viewType}`); + this._controller = vscode.notebooks.createNotebookController( + this.id + '-' + viewType, + viewType, + this.label + ); + + this._messageHandler = messageHandler; + console.log(`PostgresKernel: Message handler registered for ${viewType}:`, !!messageHandler); + + this._controller.supportedLanguages = this.supportedLanguages; + this._controller.supportsExecutionOrder = true; + this._controller.executeHandler = this._executeAll.bind(this); + + // Disable automatic timestamp parsing (this was in original, but removed in new snippet, so removing it) + // const types = require('pg').types; + // const TIMESTAMPTZ_OID = 1184; + // const TIMESTAMP_OID = 1114; + // types.setTypeParser(TIMESTAMPTZ_OID, (val: string) => val); + // types.setTypeParser(TIMESTAMP_OID, (val: string) => val); + + // this._controller.description = 'PostgreSQL Query Executor'; // Removed as per new snippet + + const getClientFromNotebook = async (document: vscode.TextDocument): Promise => { + const cell = vscode.workspace.notebookDocuments + .find(notebook => notebook.getCells().some(c => c.document === document)) + ?.getCells() + .find(c => c.document === document); + + if (!cell) return undefined; + + const metadata = cell.notebook.metadata as PostgresMetadata; + if (!metadata?.connectionId) return undefined; + + const connections = vscode.workspace.getConfiguration().get('postgresExplorer.connections') || []; + const connection = connections.find(c => c.id === metadata.connectionId); + if (!connection) return undefined; + + try { + return await ConnectionManager.getInstance().getConnection({ + id: connection.id, + host: connection.host, + port: connection.port, + username: connection.username, + database: metadata.databaseName || connection.database, + name: connection.name + }); + } catch (err) { + console.error('Error connecting to database:', err); + return undefined; + } + }; + + // Create SQL command completions + const sqlCommands = [ + { label: 'SELECT', description: 'Retrieve data from tables', documentation: 'SELECT [columns] FROM [table] WHERE [condition];' }, + { label: 'INSERT', description: 'Add new records', documentation: 'INSERT INTO [table] (columns) VALUES (values);' }, + { label: 'UPDATE', description: 'Modify existing records', documentation: 'UPDATE [table] SET [column = value] WHERE [condition];' }, + { label: 'DELETE', description: 'Remove records', documentation: 'DELETE FROM [table] WHERE [condition];' }, + { label: 'CREATE TABLE', description: 'Create a new table', documentation: 'CREATE TABLE [name] (column_definitions);' }, + { label: 'ALTER TABLE', description: 'Modify table structure', documentation: 'ALTER TABLE [table] [action];' }, + { label: 'DROP TABLE', description: 'Delete a table', documentation: 'DROP TABLE [table];' }, + { label: 'CREATE INDEX', description: 'Create a new index', documentation: 'CREATE INDEX [name] ON [table] (columns);' }, + { label: 'CREATE VIEW', description: 'Create a view', documentation: 'CREATE VIEW [name] AS SELECT ...;' }, + { label: 'GRANT', description: 'Grant permissions', documentation: 'GRANT [privileges] ON [object] TO [role];' }, + { label: 'REVOKE', description: 'Revoke permissions', documentation: 'REVOKE [privileges] ON [object] FROM [role];' }, + { label: 'BEGIN', description: 'Start a transaction', documentation: 'BEGIN; -- transaction statements -- COMMIT;' }, + { label: 'COMMIT', description: 'Commit a transaction', documentation: 'COMMIT;' }, + { label: 'ROLLBACK', description: 'Rollback a transaction', documentation: 'ROLLBACK;' } + ]; + + // Create SQL keyword completions + const sqlKeywords = [ + // DML Keywords + { label: 'SELECT', detail: 'Query data', documentation: 'SELECT [columns] FROM [table] [WHERE condition]' }, + { label: 'FROM', detail: 'Specify source table', documentation: 'FROM table_name [alias]' }, + { label: 'WHERE', detail: 'Filter conditions', documentation: 'WHERE condition' }, + { label: 'GROUP BY', detail: 'Group results', documentation: 'GROUP BY column1, column2' }, + { label: 'HAVING', detail: 'Filter groups', documentation: 'HAVING aggregate_condition' }, + { label: 'ORDER BY', detail: 'Sort results', documentation: 'ORDER BY column1 [ASC|DESC]' }, + { label: 'LIMIT', detail: 'Limit results', documentation: 'LIMIT number' }, + { label: 'OFFSET', detail: 'Skip results', documentation: 'OFFSET number' }, + { label: 'INSERT INTO', detail: 'Add new records', documentation: 'INSERT INTO table (columns) VALUES (values)' }, + { label: 'UPDATE', detail: 'Modify records', documentation: 'UPDATE table SET column = value [WHERE condition]' }, + { label: 'DELETE FROM', detail: 'Remove records', documentation: 'DELETE FROM table [WHERE condition]' }, + + // Joins + { label: 'INNER JOIN', detail: 'Inner join tables', documentation: 'INNER JOIN table ON condition' }, + { label: 'LEFT JOIN', detail: 'Left outer join', documentation: 'LEFT [OUTER] JOIN table ON condition' }, + { label: 'RIGHT JOIN', detail: 'Right outer join', documentation: 'RIGHT [OUTER] JOIN table ON condition' }, + { label: 'FULL JOIN', detail: 'Full outer join', documentation: 'FULL [OUTER] JOIN table ON condition' }, + { label: 'CROSS JOIN', detail: 'Cross join tables', documentation: 'CROSS JOIN table' }, + + // DDL Keywords + { label: 'CREATE TABLE', detail: 'Create new table', documentation: 'CREATE TABLE name (column_definitions)' }, + { label: 'ALTER TABLE', detail: 'Modify table', documentation: 'ALTER TABLE name [action]' }, + { label: 'DROP TABLE', detail: 'Delete table', documentation: 'DROP TABLE [IF EXISTS] name' }, + { label: 'CREATE INDEX', detail: 'Create index', documentation: 'CREATE INDEX name ON table (columns)' }, + { label: 'CREATE VIEW', detail: 'Create view', documentation: 'CREATE VIEW name AS SELECT ...' }, + + // Functions + { label: 'COUNT', detail: 'Count rows', documentation: 'COUNT(*) or COUNT(column)' }, + { label: 'SUM', detail: 'Sum values', documentation: 'SUM(column)' }, + { label: 'AVG', detail: 'Average value', documentation: 'AVG(column)' }, + { label: 'MAX', detail: 'Maximum value', documentation: 'MAX(column)' }, + { label: 'MIN', detail: 'Minimum value', documentation: 'MIN(column)' }, + + // Clauses + { label: 'AS', detail: 'Alias', documentation: 'column AS alias, table AS alias' }, + { label: 'ON', detail: 'Join condition', documentation: 'ON table1.column = table2.column' }, + { label: 'AND', detail: 'Logical AND', documentation: 'condition1 AND condition2' }, + { label: 'OR', detail: 'Logical OR', documentation: 'condition1 OR condition2' }, + { label: 'IN', detail: 'Value in set', documentation: 'column IN (value1, value2, ...)' }, + { label: 'BETWEEN', detail: 'Value in range', documentation: 'column BETWEEN value1 AND value2' }, + { label: 'LIKE', detail: 'Pattern matching', documentation: 'column LIKE pattern' }, + { label: 'IS NULL', detail: 'Null check', documentation: 'column IS NULL' }, + { label: 'IS NOT NULL', detail: 'Not null check', documentation: 'column IS NOT NULL' }, + + // Transaction Control + { label: 'BEGIN', detail: 'Start transaction', documentation: 'BEGIN [TRANSACTION]' }, + { label: 'COMMIT', detail: 'Commit transaction', documentation: 'COMMIT' }, + { label: 'ROLLBACK', detail: 'Rollback transaction', documentation: 'ROLLBACK' } + ]; + + // Register completion provider for SQL + context.subscriptions.push( + vscode.languages.registerCompletionItemProvider( + { scheme: 'vscode-notebook-cell', language: 'sql' }, + { + async provideCompletionItems(document: vscode.TextDocument, position: vscode.Position) { + const linePrefix = document.lineAt(position).text.substr(0, position.character).toLowerCase(); + const wordRange = document.getWordRangeAtPosition(position); + const word = wordRange ? document.getText(wordRange).toLowerCase() : ''; + + // Always provide SQL keyword suggestions + const keywordItems = sqlKeywords.filter(kw => + !word || kw.label.toLowerCase().includes(word) + ).map(kw => { + const item = new vscode.CompletionItem(kw.label, vscode.CompletionItemKind.Keyword); + item.detail = kw.detail; + item.documentation = new vscode.MarkdownString(kw.documentation); + return item; + }); - const metadata = cell.notebook.metadata as PostgresMetadata; - if (!metadata?.connectionId) return undefined; + // Check for column suggestions after table alias (e.g. "t.") + const aliasMatch = linePrefix.match(/(\w+)\.\s*$/); + if (aliasMatch) { + // Look for table alias in previous part of the query + const fullQuery = document.getText(); + const aliasPattern = new RegExp(`(?:FROM|JOIN)\\s+([\\w\\.]+)\\s+(?:AS\\s+)?${aliasMatch[1]}\\b`, 'i'); + const tableMatch = aliasPattern.exec(fullQuery); - const connections = vscode.workspace.getConfiguration().get('postgresExplorer.connections') || []; - const connection = connections.find(c => c.id === metadata.connectionId); - if (!connection) return undefined; + if (tableMatch) { + const [, tablePath] = tableMatch; + const [schema = 'public', table] = tablePath.split('.'); + const client = await getClientFromNotebook(document); + if (!client) return []; - try { - return await ConnectionManager.getInstance().getConnection({ - id: connection.id, - host: connection.host, - port: connection.port, - username: connection.username, - database: metadata.databaseName || connection.database, - name: connection.name - }); - } catch (err) { - console.error('Error connecting to database:', err); - return undefined; - } - }; - - // Create SQL command completions - const sqlCommands = [ - { label: 'SELECT', description: 'Retrieve data from tables', documentation: 'SELECT [columns] FROM [table] WHERE [condition];' }, - { label: 'INSERT', description: 'Add new records', documentation: 'INSERT INTO [table] (columns) VALUES (values);' }, - { label: 'UPDATE', description: 'Modify existing records', documentation: 'UPDATE [table] SET [column = value] WHERE [condition];' }, - { label: 'DELETE', description: 'Remove records', documentation: 'DELETE FROM [table] WHERE [condition];' }, - { label: 'CREATE TABLE', description: 'Create a new table', documentation: 'CREATE TABLE [name] (column_definitions);' }, - { label: 'ALTER TABLE', description: 'Modify table structure', documentation: 'ALTER TABLE [table] [action];' }, - { label: 'DROP TABLE', description: 'Delete a table', documentation: 'DROP TABLE [table];' }, - { label: 'CREATE INDEX', description: 'Create a new index', documentation: 'CREATE INDEX [name] ON [table] (columns);' }, - { label: 'CREATE VIEW', description: 'Create a view', documentation: 'CREATE VIEW [name] AS SELECT ...;' }, - { label: 'GRANT', description: 'Grant permissions', documentation: 'GRANT [privileges] ON [object] TO [role];' }, - { label: 'REVOKE', description: 'Revoke permissions', documentation: 'REVOKE [privileges] ON [object] FROM [role];' }, - { label: 'BEGIN', description: 'Start a transaction', documentation: 'BEGIN; -- transaction statements -- COMMIT;' }, - { label: 'COMMIT', description: 'Commit a transaction', documentation: 'COMMIT;' }, - { label: 'ROLLBACK', description: 'Rollback a transaction', documentation: 'ROLLBACK;' } - ]; - - // Create SQL keyword completions - const sqlKeywords = [ - // DML Keywords - { label: 'SELECT', detail: 'Query data', documentation: 'SELECT [columns] FROM [table] [WHERE condition]' }, - { label: 'FROM', detail: 'Specify source table', documentation: 'FROM table_name [alias]' }, - { label: 'WHERE', detail: 'Filter conditions', documentation: 'WHERE condition' }, - { label: 'GROUP BY', detail: 'Group results', documentation: 'GROUP BY column1, column2' }, - { label: 'HAVING', detail: 'Filter groups', documentation: 'HAVING aggregate_condition' }, - { label: 'ORDER BY', detail: 'Sort results', documentation: 'ORDER BY column1 [ASC|DESC]' }, - { label: 'LIMIT', detail: 'Limit results', documentation: 'LIMIT number' }, - { label: 'OFFSET', detail: 'Skip results', documentation: 'OFFSET number' }, - { label: 'INSERT INTO', detail: 'Add new records', documentation: 'INSERT INTO table (columns) VALUES (values)' }, - { label: 'UPDATE', detail: 'Modify records', documentation: 'UPDATE table SET column = value [WHERE condition]' }, - { label: 'DELETE FROM', detail: 'Remove records', documentation: 'DELETE FROM table [WHERE condition]' }, - - // Joins - { label: 'INNER JOIN', detail: 'Inner join tables', documentation: 'INNER JOIN table ON condition' }, - { label: 'LEFT JOIN', detail: 'Left outer join', documentation: 'LEFT [OUTER] JOIN table ON condition' }, - { label: 'RIGHT JOIN', detail: 'Right outer join', documentation: 'RIGHT [OUTER] JOIN table ON condition' }, - { label: 'FULL JOIN', detail: 'Full outer join', documentation: 'FULL [OUTER] JOIN table ON condition' }, - { label: 'CROSS JOIN', detail: 'Cross join tables', documentation: 'CROSS JOIN table' }, - - // DDL Keywords - { label: 'CREATE TABLE', detail: 'Create new table', documentation: 'CREATE TABLE name (column_definitions)' }, - { label: 'ALTER TABLE', detail: 'Modify table', documentation: 'ALTER TABLE name [action]' }, - { label: 'DROP TABLE', detail: 'Delete table', documentation: 'DROP TABLE [IF EXISTS] name' }, - { label: 'CREATE INDEX', detail: 'Create index', documentation: 'CREATE INDEX name ON table (columns)' }, - { label: 'CREATE VIEW', detail: 'Create view', documentation: 'CREATE VIEW name AS SELECT ...' }, - - // Functions - { label: 'COUNT', detail: 'Count rows', documentation: 'COUNT(*) or COUNT(column)' }, - { label: 'SUM', detail: 'Sum values', documentation: 'SUM(column)' }, - { label: 'AVG', detail: 'Average value', documentation: 'AVG(column)' }, - { label: 'MAX', detail: 'Maximum value', documentation: 'MAX(column)' }, - { label: 'MIN', detail: 'Minimum value', documentation: 'MIN(column)' }, - - // Clauses - { label: 'AS', detail: 'Alias', documentation: 'column AS alias, table AS alias' }, - { label: 'ON', detail: 'Join condition', documentation: 'ON table1.column = table2.column' }, - { label: 'AND', detail: 'Logical AND', documentation: 'condition1 AND condition2' }, - { label: 'OR', detail: 'Logical OR', documentation: 'condition1 OR condition2' }, - { label: 'IN', detail: 'Value in set', documentation: 'column IN (value1, value2, ...)' }, - { label: 'BETWEEN', detail: 'Value in range', documentation: 'column BETWEEN value1 AND value2' }, - { label: 'LIKE', detail: 'Pattern matching', documentation: 'column LIKE pattern' }, - { label: 'IS NULL', detail: 'Null check', documentation: 'column IS NULL' }, - { label: 'IS NOT NULL', detail: 'Not null check', documentation: 'column IS NOT NULL' }, - - // Transaction Control - { label: 'BEGIN', detail: 'Start transaction', documentation: 'BEGIN [TRANSACTION]' }, - { label: 'COMMIT', detail: 'Commit transaction', documentation: 'COMMIT' }, - { label: 'ROLLBACK', detail: 'Rollback transaction', documentation: 'ROLLBACK' } - ]; - - // Register completion provider for SQL - context.subscriptions.push( - vscode.languages.registerCompletionItemProvider( - { scheme: 'vscode-notebook-cell', language: 'sql' }, - { - async provideCompletionItems(document: vscode.TextDocument, position: vscode.Position) { - const linePrefix = document.lineAt(position).text.substr(0, position.character).toLowerCase(); - const wordRange = document.getWordRangeAtPosition(position); - const word = wordRange ? document.getText(wordRange).toLowerCase() : ''; - - // Always provide SQL keyword suggestions - const keywordItems = sqlKeywords.filter(kw => - !word || kw.label.toLowerCase().includes(word) - ).map(kw => { - const item = new vscode.CompletionItem(kw.label, vscode.CompletionItemKind.Keyword); - item.detail = kw.detail; - item.documentation = new vscode.MarkdownString(kw.documentation); - return item; - }); - - // Check for column suggestions after table alias (e.g. "t.") - const aliasMatch = linePrefix.match(/(\w+)\.\s*$/); - if (aliasMatch) { - // Look for table alias in previous part of the query - const fullQuery = document.getText(); - const aliasPattern = new RegExp(`(?:FROM|JOIN)\\s+([\\w\\.]+)\\s+(?:AS\\s+)?${aliasMatch[1]}\\b`, 'i'); - const tableMatch = aliasPattern.exec(fullQuery); - - if (tableMatch) { - const [, tablePath] = tableMatch; - const [schema = 'public', table] = tablePath.split('.'); - const client = await getClientFromNotebook(document); - if (!client) return []; - - try { - const result = await client.query( - `SELECT column_name, data_type, is_nullable + try { + const result = await client.query( + `SELECT column_name, data_type, is_nullable FROM information_schema.columns WHERE table_schema = $1 AND table_name = $2 ORDER BY ordinal_position`, - [schema, table] - ); - - return result.rows.map((row: { column_name: string; data_type: string; is_nullable: string }) => { - const completion = new vscode.CompletionItem(row.column_name); - completion.kind = vscode.CompletionItemKind.Field; - completion.detail = row.data_type; - completion.documentation = `Type: ${row.data_type}\nNullable: ${row.is_nullable === 'YES' ? 'Yes' : 'No'}`; - return completion; - }); - } catch (err) { - console.error('Error getting column completions:', err); - return []; - } - // Do not close client here, it's managed by ConnectionManager - } - } - - // Check if we're after a schema reference (schema.) - const schemaMatch = linePrefix.match(/(\w+)\.\s*$/); - if (schemaMatch) { - const client = await getClientFromNotebook(document); - if (!client) return []; - - try { - const result = await client.query( - `SELECT table_name + [schema, table] + ); + + return result.rows.map((row: { column_name: string; data_type: string; is_nullable: string }) => { + const completion = new vscode.CompletionItem(row.column_name); + completion.kind = vscode.CompletionItemKind.Field; + completion.detail = row.data_type; + completion.documentation = `Type: ${row.data_type}\nNullable: ${row.is_nullable === 'YES' ? 'Yes' : 'No'}`; + return completion; + }); + } catch (err) { + console.error('Error getting column completions:', err); + return []; + } + // Do not close client here, it's managed by ConnectionManager + } + } + + // Check if we're after a schema reference (schema.) + const schemaMatch = linePrefix.match(/(\w+)\.\s*$/); + if (schemaMatch) { + const client = await getClientFromNotebook(document); + if (!client) return []; + + try { + const result = await client.query( + `SELECT table_name FROM information_schema.tables WHERE table_schema = $1 ORDER BY table_name`, - [schemaMatch[1]] - ); - return result.rows.map((row: { table_name: string }) => { - const completion = new vscode.CompletionItem(row.table_name); - completion.kind = vscode.CompletionItemKind.Value; - return completion; - }); - } catch (err) { - console.error('Error getting table completions:', err); - return []; - } - } - - // Provide schema suggestions after 'FROM' or 'JOIN' - const keywords = /(?:from|join)\s+(\w*)$/i; - const match = linePrefix.match(keywords); - if (match) { - const client = await getClientFromNotebook(document); - if (!client) return []; - - try { - const result = await client.query( - `SELECT schema_name + [schemaMatch[1]] + ); + return result.rows.map((row: { table_name: string }) => { + const completion = new vscode.CompletionItem(row.table_name); + completion.kind = vscode.CompletionItemKind.Value; + return completion; + }); + } catch (err) { + console.error('Error getting table completions:', err); + return []; + } + } + + // Provide schema suggestions after 'FROM' or 'JOIN' + const keywords = /(?:from|join)\s+(\w*)$/i; + const match = linePrefix.match(keywords); + if (match) { + const client = await getClientFromNotebook(document); + if (!client) return []; + + try { + const result = await client.query( + `SELECT schema_name FROM information_schema.schemata WHERE schema_name NOT IN ('information_schema', 'pg_catalog') ORDER BY schema_name` - ); - return result.rows.map((row: { schema_name: string }) => { - const completion = new vscode.CompletionItem(row.schema_name); - completion.kind = vscode.CompletionItemKind.Module; - completion.insertText = row.schema_name + '.'; - return completion; - }); - } catch (err) { - console.error('Error getting schema completions:', err); - return []; - } - } - - return keywordItems; - } - } - ) - ); - - // Register completion provider for SQL - context.subscriptions.push( - vscode.languages.registerCompletionItemProvider( - { scheme: 'vscode-notebook-cell', language: 'sql' }, - { - async provideCompletionItems(document: vscode.TextDocument, position: vscode.Position) { - const linePrefix = document.lineAt(position).text.substr(0, position.character); - - // Return SQL command suggestions at start of line or after semicolon - if (linePrefix.trim() === '' || linePrefix.trim().endsWith(';')) { - return sqlCommands.map(cmd => { - const item = new vscode.CompletionItem(cmd.label, vscode.CompletionItemKind.Keyword); - item.detail = cmd.description; - item.documentation = new vscode.MarkdownString(cmd.documentation); - return item; - }); - } - - return []; - } - }, - ' ', ';' // Trigger on space and semicolon - ) - ); - // Handle messages from renderer (e.g., delete row) - console.log(`PostgresKernel: Subscribing to onDidReceiveMessage for Controller ID: ${this._controller.id}`); - (this._controller as any).onDidReceiveMessage(async (event: any) => { - console.log(`PostgresKernel: Received message on Controller ${this._controller.id}`, event.message); - if (event.message.type === 'script_delete') { - console.log('PostgresKernel: Processing script_delete message'); - const { schema, table, primaryKeys, rows, cellIndex } = event.message; - const notebook = event.editor.notebook; - - try { - // Construct DELETE query - let query = ''; - for (const row of rows) { - const conditions: string[] = []; - const values: any[] = []; - - for (const pk of primaryKeys) { - const val = row[pk]; - // Simple quoting for string values, handle numbers/booleans - const valStr = typeof val === 'string' ? `'${val.replace(/'/g, "''")}'` : val; - conditions.push(`"${pk}" = ${valStr}`); - } - query += `DELETE FROM "${schema}"."${table}" WHERE ${conditions.join(' AND ')};\n`; - } - - // Insert new cell with the query - const targetIndex = cellIndex + 1; - const newCell = new vscode.NotebookCellData( - vscode.NotebookCellKind.Code, - query, - 'sql' - ); - - const edit = new vscode.NotebookEdit( - new vscode.NotebookRange(targetIndex, targetIndex), - [newCell] - ); - - const workspaceEdit = new vscode.WorkspaceEdit(); - workspaceEdit.set(notebook.uri, [edit]); - await vscode.workspace.applyEdit(workspaceEdit); - - // Focus the new cell (optional, but good UX) - // Note: Focusing specific cell via API is limited, but inserting it usually reveals it. - - } catch (err: any) { - vscode.window.showErrorMessage(`Failed to generate delete script: ${err.message}`); - console.error('Script delete error:', err); - } - } else if (event.message.type === 'execute_update') { - console.log('PostgresKernel: Processing execute_update message'); - const { statements, cellIndex } = event.message; - const notebook = event.editor.notebook; - - try { - // Insert new cell with the UPDATE statements - const query = statements.join('\n'); - const targetIndex = cellIndex + 1; - const newCell = new vscode.NotebookCellData( - vscode.NotebookCellKind.Code, - `-- Update statements generated from cell edits\n${query}`, - 'sql' - ); - - const edit = new vscode.NotebookEdit( - new vscode.NotebookRange(targetIndex, targetIndex), - [newCell] - ); - - const workspaceEdit = new vscode.WorkspaceEdit(); - workspaceEdit.set(notebook.uri, [edit]); - await vscode.workspace.applyEdit(workspaceEdit); - - vscode.window.showInformationMessage(`Generated ${statements.length} UPDATE statement(s). Review and execute the new cell.`); - } catch (err: any) { - vscode.window.showErrorMessage(`Failed to generate update script: ${err.message}`); - console.error('Script update error:', err); - } - } else if (event.message.type === 'execute_update_background') { - console.log('PostgresKernel: Processing execute_update_background message'); - console.log('PostgresKernel: Statements to execute:', event.message.statements); - const { statements } = event.message; - const notebook = event.editor.notebook; - - try { - // Get connection from notebook metadata - const metadata = notebook.metadata as PostgresMetadata; - console.log('PostgresKernel: Notebook metadata:', metadata); - if (!metadata?.connectionId) { - throw new Error('No connection found in notebook metadata'); - } - - // Get connection configuration - const connections = vscode.workspace.getConfiguration().get('postgresExplorer.connections') || []; - console.log('PostgresKernel: Found connections:', connections.length); - const savedConnection = connections.find((c: any) => c.id === metadata.connectionId); - - if (!savedConnection) { - throw new Error(`Connection not found for id: ${metadata.connectionId}`); - } - console.log('PostgresKernel: Using connection:', savedConnection.name); - - // Get password from secret storage - const secretService = SecretStorageService.getInstance(this.context); - const password = await secretService.getPassword(savedConnection.id); - - const client = new Client({ - host: savedConnection.host, - port: savedConnection.port, - user: savedConnection.username, - password: password || '', - database: metadata.databaseName || savedConnection.database, - ssl: savedConnection.ssl ? { rejectUnauthorized: false } : false - }); - - console.log('PostgresKernel: Connecting to database:', metadata.databaseName || savedConnection.database); - await client.connect(); - - try { - // Execute all UPDATE statements - const combinedQuery = statements.join('\n'); - console.log('PostgresKernel: Executing query:', combinedQuery); - const result = await client.query(combinedQuery); - console.log('PostgresKernel: Query result:', result); - - vscode.window.showInformationMessage(`✅ Successfully saved ${statements.length} change(s) to database.`); - } finally { - await client.end(); - } - } catch (err: any) { - console.error('PostgresKernel: Background update error:', err); - vscode.window.showErrorMessage(`Failed to save changes: ${err.message}`); - } - } else if (event.message.type === 'export_request') { - console.log('PostgresKernel: Processing export_request message'); - const { rows, columns } = event.message; - - const selection = await vscode.window.showQuickPick( - ['Save as CSV', 'Save as JSON', 'Copy to Clipboard'], - { placeHolder: 'Select export format' } ); - - if (!selection) return; - - if (selection === 'Copy to Clipboard') { - // Convert to CSV for clipboard - const header = columns.map((c: string) => `"${c.replace(/"/g, '""')}"`).join(','); - const body = rows.map((row: any) => { - return columns.map((col: string) => { - const val = row[col]; - if (val === null || val === undefined) return ''; - const str = String(val); - if (str.includes(',') || str.includes('"') || str.includes('\n')) { - return `"${str.replace(/"/g, '""')}"`; - } - return str; - }).join(','); - }).join('\n'); - const csv = `${header}\n${body}`; - - await vscode.env.clipboard.writeText(csv); - vscode.window.showInformationMessage('Data copied to clipboard (CSV format).'); - } else if (selection === 'Save as CSV') { - const header = columns.map((c: string) => `"${c.replace(/"/g, '""')}"`).join(','); - const body = rows.map((row: any) => { - return columns.map((col: string) => { - const val = row[col]; - if (val === null || val === undefined) return ''; - const str = String(val); - if (str.includes(',') || str.includes('"') || str.includes('\n')) { - return `"${str.replace(/"/g, '""')}"`; - } - return str; - }).join(','); - }).join('\n'); - const csv = `${header}\n${body}`; - - const uri = await vscode.window.showSaveDialog({ - filters: { 'CSV': ['csv'] }, - saveLabel: 'Export CSV' - }); - - if (uri) { - await vscode.workspace.fs.writeFile(uri, Buffer.from(csv, 'utf8')); - vscode.window.showInformationMessage('CSV exported successfully.'); - } - } else if (selection === 'Save as JSON') { - const json = JSON.stringify(rows, null, 2); - const uri = await vscode.window.showSaveDialog({ - filters: { 'JSON': ['json'] }, - saveLabel: 'Export JSON' - }); - - if (uri) { - await vscode.workspace.fs.writeFile(uri, Buffer.from(json, 'utf8')); - vscode.window.showInformationMessage('JSON exported successfully.'); - } - } + return result.rows.map((row: { schema_name: string }) => { + const completion = new vscode.CompletionItem(row.schema_name); + completion.kind = vscode.CompletionItemKind.Module; + completion.insertText = row.schema_name + '.'; + return completion; + }); + } catch (err) { + console.error('Error getting schema completions:', err); + return []; + } } - if (event.message.type === 'delete_row') { - const { schema, table, primaryKeys, row } = event.message; - const notebook = event.editor.notebook; - const metadata = notebook.metadata as PostgresMetadata; - - if (!metadata?.connectionId) { - vscode.window.showErrorMessage('No connection found for this notebook.'); - return; - } - try { - const connections = vscode.workspace.getConfiguration().get('postgresExplorer.connections') || []; - const connection = connections.find(c => c.id === metadata.connectionId); - if (!connection) throw new Error('Connection not found'); - - const client = await ConnectionManager.getInstance().getConnection({ - id: connection.id, - host: connection.host, - port: connection.port, - username: connection.username, - database: metadata.databaseName || connection.database, - name: connection.name - }); - - // Construct DELETE query - const conditions: string[] = []; - const values: any[] = []; - let paramIndex = 1; - - for (const pk of primaryKeys) { - conditions.push(`"${pk}" = $${paramIndex}`); - values.push(row[pk]); - paramIndex++; - } - - const query = `DELETE FROM "${schema}"."${table}" WHERE ${conditions.join(' AND ')}`; - await client.query(query, values); - vscode.window.showInformationMessage('Row deleted successfully.'); - } catch (err: any) { - vscode.window.showErrorMessage(`Failed to delete row: ${err.message}`); - console.error('Delete row error:', err); - } + return keywordItems; + } + } + ) + ); + + // Register completion provider for SQL + context.subscriptions.push( + vscode.languages.registerCompletionItemProvider( + { scheme: 'vscode-notebook-cell', language: 'sql' }, + { + async provideCompletionItems(document: vscode.TextDocument, position: vscode.Position) { + const linePrefix = document.lineAt(position).text.substr(0, position.character); + + // Return SQL command suggestions at start of line or after semicolon + if (linePrefix.trim() === '' || linePrefix.trim().endsWith(';')) { + return sqlCommands.map(cmd => { + const item = new vscode.CompletionItem(cmd.label, vscode.CompletionItemKind.Keyword); + item.detail = cmd.description; + item.documentation = new vscode.MarkdownString(cmd.documentation); + return item; + }); } - }); - } - private async _executeAll(cells: vscode.NotebookCell[], _notebook: vscode.NotebookDocument, _controller: vscode.NotebookController): Promise { - for (const cell of cells) { - await this._doExecution(cell); - } - } + return []; + } + }, + ' ', ';' // Trigger on space and semicolon + ) + ); + // Handle messages from renderer (e.g., delete row) + console.log(`PostgresKernel: Subscribing to onDidReceiveMessage for Controller ID: ${this._controller.id}`); + (this._controller as any).onDidReceiveMessage(async (event: any) => { + console.log(`PostgresKernel: Received message on Controller ${this._controller.id}`, event.message); + if (event.message.type === 'cancel_query') { + console.log('PostgresKernel: Processing cancel_query message'); + const { backendPid, connectionId, databaseName } = event.message; - /** - * Split SQL text into individual statements, respecting semicolons but ignoring them inside: - * - String literals (single quotes) - * - Dollar-quoted strings ($$...$$, $tag$...$tag$) - * - Comments (-- and /* *\/) - */ - private splitSqlStatements(sql: string): string[] { - const statements: string[] = []; - let currentStatement = ''; - let i = 0; - let inSingleQuote = false; - let inDollarQuote = false; - let dollarQuoteTag = ''; - let inBlockComment = false; - - while (i < sql.length) { - const char = sql[i]; - const nextChar = i + 1 < sql.length ? sql[i + 1] : ''; - const peek = sql.substring(i, i + 10); - - // Handle block comments /* ... */ - if (!inSingleQuote && !inDollarQuote && char === '/' && nextChar === '*') { - inBlockComment = true; - currentStatement += char + nextChar; - i += 2; - continue; - } + try { + const connections = vscode.workspace.getConfiguration().get('postgresExplorer.connections') || []; + const connection = connections.find(c => c.id === connectionId); + if (!connection) { + throw new Error('Connection not found'); + } + + // Create a new connection to cancel the query + const cancelClient = await ConnectionManager.getInstance().getConnection({ + id: connection.id, + host: connection.host, + port: connection.port, + username: connection.username, + database: databaseName || connection.database, + name: connection.name + }); + + // Cancel the backend process + await cancelClient.query('SELECT pg_cancel_backend($1)', [backendPid]); + vscode.window.showInformationMessage(`Query cancelled (PID: ${backendPid})`); + } catch (err: any) { + vscode.window.showErrorMessage(`Failed to cancel query: ${err.message}`); + console.error('Cancel query error:', err); + } + } else if (event.message.type === 'script_delete') { + console.log('PostgresKernel: Processing script_delete message'); + const { schema, table, primaryKeys, rows, cellIndex } = event.message; + const notebook = event.editor.notebook; - if (inBlockComment && char === '*' && nextChar === '/') { - inBlockComment = false; - currentStatement += char + nextChar; - i += 2; - continue; + try { + // Construct DELETE query + let query = ''; + for (const row of rows) { + const conditions: string[] = []; + const values: any[] = []; + + for (const pk of primaryKeys) { + const val = row[pk]; + // Simple quoting for string values, handle numbers/booleans + const valStr = typeof val === 'string' ? `'${val.replace(/'/g, "''")}'` : val; + conditions.push(`"${pk}" = ${valStr}`); } + query += `DELETE FROM "${schema}"."${table}" WHERE ${conditions.join(' AND ')};\n`; + } - // Handle line comments -- ... - if (!inSingleQuote && !inDollarQuote && !inBlockComment && char === '-' && nextChar === '-') { - // Add rest of line to current statement - const lineEnd = sql.indexOf('\n', i); - if (lineEnd === -1) { - currentStatement += sql.substring(i); - break; - } - currentStatement += sql.substring(i, lineEnd + 1); - i = lineEnd + 1; - continue; - } + // Insert new cell with the query + const targetIndex = cellIndex + 1; + const newCell = new vscode.NotebookCellData( + vscode.NotebookCellKind.Code, + query, + 'sql' + ); - // Handle dollar-quoted strings - if (!inSingleQuote && !inBlockComment) { - const dollarMatch = peek.match(/^(\$[a-zA-Z0-9_]*\$)/); - if (dollarMatch) { - const tag = dollarMatch[1]; - if (!inDollarQuote) { - inDollarQuote = true; - dollarQuoteTag = tag; - currentStatement += tag; - i += tag.length; - continue; - } else if (tag === dollarQuoteTag) { - inDollarQuote = false; - dollarQuoteTag = ''; - currentStatement += tag; - i += tag.length; - continue; - } - } - } + const edit = new vscode.NotebookEdit( + new vscode.NotebookRange(targetIndex, targetIndex), + [newCell] + ); - // Handle single-quoted strings - if (!inDollarQuote && !inBlockComment && char === "'") { - if (inSingleQuote && nextChar === "'") { - // Escaped quote '' - currentStatement += "''"; - i += 2; - continue; - } - inSingleQuote = !inSingleQuote; - } + const workspaceEdit = new vscode.WorkspaceEdit(); + workspaceEdit.set(notebook.uri, [edit]); + await vscode.workspace.applyEdit(workspaceEdit); - // Handle semicolon as statement separator - if (!inSingleQuote && !inDollarQuote && !inBlockComment && char === ';') { - currentStatement += char; - const trimmed = currentStatement.trim(); - if (trimmed) { - statements.push(trimmed); - } - currentStatement = ''; - i++; - continue; - } + // Focus the new cell (optional, but good UX) + // Note: Focusing specific cell via API is limited, but inserting it usually reveals it. - currentStatement += char; - i++; + } catch (err: any) { + vscode.window.showErrorMessage(`Failed to generate delete script: ${err.message}`); + console.error('Script delete error:', err); } + } else if (event.message.type === 'execute_update') { + console.log('PostgresKernel: Processing execute_update message'); + const { statements, cellIndex } = event.message; + const notebook = event.editor.notebook; - // Add remaining statement if any - const trimmed = currentStatement.trim(); - if (trimmed) { - statements.push(trimmed); + try { + // Insert new cell with the UPDATE statements + const query = statements.join('\n'); + const targetIndex = cellIndex + 1; + const newCell = new vscode.NotebookCellData( + vscode.NotebookCellKind.Code, + `-- Update statements generated from cell edits\n${query}`, + 'sql' + ); + + const edit = new vscode.NotebookEdit( + new vscode.NotebookRange(targetIndex, targetIndex), + [newCell] + ); + + const workspaceEdit = new vscode.WorkspaceEdit(); + workspaceEdit.set(notebook.uri, [edit]); + await vscode.workspace.applyEdit(workspaceEdit); + + vscode.window.showInformationMessage(`Generated ${statements.length} UPDATE statement(s). Review and execute the new cell.`); + } catch (err: any) { + vscode.window.showErrorMessage(`Failed to generate update script: ${err.message}`); + console.error('Script update error:', err); } - - return statements.filter(s => s.length > 0); - } - - private async _doExecution(cell: vscode.NotebookCell): Promise { - console.log(`PostgresKernel: Starting cell execution. Controller ID: ${this._controller.id}`); - const execution = this._controller.createNotebookCellExecution(cell); - const startTime = Date.now(); - execution.start(startTime); - execution.clearOutput(); + } else if (event.message.type === 'execute_update_background') { + console.log('PostgresKernel: Processing execute_update_background message'); + console.log('PostgresKernel: Statements to execute:', event.message.statements); + const { statements } = event.message; + const notebook = event.editor.notebook; try { - const metadata = cell.notebook.metadata as PostgresMetadata; - if (!metadata || !metadata.connectionId) { - throw new Error('No connection metadata found'); - } - - // Get connection info and password from SecretStorage - const connections = vscode.workspace.getConfiguration().get('postgresExplorer.connections') || []; - const connection = connections.find(c => c.id === metadata.connectionId); - if (!connection) { - throw new Error('Connection not found'); - } - - const client = await ConnectionManager.getInstance().getConnection({ - id: connection.id, - host: connection.host, - port: connection.port, - username: connection.username, - database: metadata.databaseName || connection.database, - name: connection.name - }); + // Get connection from notebook metadata + const metadata = notebook.metadata as PostgresMetadata; + console.log('PostgresKernel: Notebook metadata:', metadata); + if (!metadata?.connectionId) { + throw new Error('No connection found in notebook metadata'); + } + + // Get connection configuration + const connections = vscode.workspace.getConfiguration().get('postgresExplorer.connections') || []; + console.log('PostgresKernel: Found connections:', connections.length); + const savedConnection = connections.find((c: any) => c.id === metadata.connectionId); + + if (!savedConnection) { + throw new Error(`Connection not found for id: ${metadata.connectionId}`); + } + console.log('PostgresKernel: Using connection:', savedConnection.name); + + // Get password from secret storage + const secretService = SecretStorageService.getInstance(this.context); + const password = await secretService.getPassword(savedConnection.id); + + const client = new Client({ + host: savedConnection.host, + port: savedConnection.port, + user: savedConnection.username, + password: password || '', + database: metadata.databaseName || savedConnection.database, + ssl: savedConnection.ssl ? { rejectUnauthorized: false } : false + }); + + console.log('PostgresKernel: Connecting to database:', metadata.databaseName || savedConnection.database); + await client.connect(); + + try { + // Execute all UPDATE statements + const combinedQuery = statements.join('\n'); + console.log('PostgresKernel: Executing query:', combinedQuery); + const result = await client.query(combinedQuery); + console.log('PostgresKernel: Query result:', result); + + vscode.window.showInformationMessage(`✅ Successfully saved ${statements.length} change(s) to database.`); + } finally { + await client.end(); + } + } catch (err: any) { + console.error('PostgresKernel: Background update error:', err); + vscode.window.showErrorMessage(`Failed to save changes: ${err.message}`); + } + } else if (event.message.type === 'export_request') { + console.log('PostgresKernel: Processing export_request message'); + const { rows, columns } = event.message; - console.log('PostgresKernel: Connected to database'); + const selection = await vscode.window.showQuickPick( + ['Save as CSV', 'Save as JSON', 'Copy to Clipboard'], + { placeHolder: 'Select export format' } + ); - // Capture PostgreSQL NOTICE messages - const notices: string[] = []; - const noticeListener = (msg: any) => { - const message = msg.message || msg.toString(); - notices.push(message); - }; - client.on('notice', noticeListener); + if (!selection) return; + + if (selection === 'Copy to Clipboard') { + // Convert to CSV for clipboard + const header = columns.map((c: string) => `"${c.replace(/"/g, '""')}"`).join(','); + const body = rows.map((row: any) => { + return columns.map((col: string) => { + const val = row[col]; + if (val === null || val === undefined) return ''; + const str = String(val); + if (str.includes(',') || str.includes('"') || str.includes('\n')) { + return `"${str.replace(/"/g, '""')}"`; + } + return str; + }).join(','); + }).join('\n'); + const csv = `${header}\n${body}`; + + await vscode.env.clipboard.writeText(csv); + vscode.window.showInformationMessage('Data copied to clipboard (CSV format).'); + } else if (selection === 'Save as CSV') { + const header = columns.map((c: string) => `"${c.replace(/"/g, '""')}"`).join(','); + const body = rows.map((row: any) => { + return columns.map((col: string) => { + const val = row[col]; + if (val === null || val === undefined) return ''; + const str = String(val); + if (str.includes(',') || str.includes('"') || str.includes('\n')) { + return `"${str.replace(/"/g, '""')}"`; + } + return str; + }).join(','); + }).join('\n'); + const csv = `${header}\n${body}`; + + const uri = await vscode.window.showSaveDialog({ + filters: { 'CSV': ['csv'] }, + saveLabel: 'Export CSV' + }); + + if (uri) { + await vscode.workspace.fs.writeFile(uri, Buffer.from(csv, 'utf8')); + vscode.window.showInformationMessage('CSV exported successfully.'); + } + } else if (selection === 'Save as JSON') { + const json = JSON.stringify(rows, null, 2); + const uri = await vscode.window.showSaveDialog({ + filters: { 'JSON': ['json'] }, + saveLabel: 'Export JSON' + }); + + if (uri) { + await vscode.workspace.fs.writeFile(uri, Buffer.from(json, 'utf8')); + vscode.window.showInformationMessage('JSON exported successfully.'); + } + } + } + if (event.message.type === 'delete_row') { + const { schema, table, primaryKeys, row } = event.message; + const notebook = event.editor.notebook; + const metadata = notebook.metadata as PostgresMetadata; + + if (!metadata?.connectionId) { + vscode.window.showErrorMessage('No connection found for this notebook.'); + return; + } - const queryText = cell.document.getText(); - const statements = this.splitSqlStatements(queryText); + try { + const connections = vscode.workspace.getConfiguration().get('postgresExplorer.connections') || []; + const connection = connections.find(c => c.id === metadata.connectionId); + if (!connection) throw new Error('Connection not found'); + + const client = await ConnectionManager.getInstance().getConnection({ + id: connection.id, + host: connection.host, + port: connection.port, + username: connection.username, + database: metadata.databaseName || connection.database, + name: connection.name + }); + + // Construct DELETE query + const conditions: string[] = []; + const values: any[] = []; + let paramIndex = 1; + + for (const pk of primaryKeys) { + conditions.push(`"${pk}" = $${paramIndex}`); + values.push(row[pk]); + paramIndex++; + } + + const query = `DELETE FROM "${schema}"."${table}" WHERE ${conditions.join(' AND ')}`; + await client.query(query, values); + vscode.window.showInformationMessage('Row deleted successfully.'); + } catch (err: any) { + vscode.window.showErrorMessage(`Failed to delete row: ${err.message}`); + console.error('Delete row error:', err); + } + } + }); + } - console.log('PostgresKernel: Executing', statements.length, 'statement(s)'); + private async _executeAll(cells: vscode.NotebookCell[], _notebook: vscode.NotebookDocument, _controller: vscode.NotebookController): Promise { + for (const cell of cells) { + await this._doExecution(cell); + } + } + + /** + * Split SQL text into individual statements, respecting semicolons but ignoring them inside: + * - String literals (single quotes) + * - Dollar-quoted strings ($$...$$, $tag$...$tag$) + * - Comments (-- and /* *\/) + */ + private splitSqlStatements(sql: string): string[] { + const statements: string[] = []; + let currentStatement = ''; + let i = 0; + let inSingleQuote = false; + let inDollarQuote = false; + let dollarQuoteTag = ''; + let inBlockComment = false; + + while (i < sql.length) { + const char = sql[i]; + const nextChar = i + 1 < sql.length ? sql[i + 1] : ''; + const peek = sql.substring(i, i + 10); + + // Handle block comments /* ... */ + if (!inSingleQuote && !inDollarQuote && char === '/' && nextChar === '*') { + inBlockComment = true; + currentStatement += char + nextChar; + i += 2; + continue; + } + + if (inBlockComment && char === '*' && nextChar === '/') { + inBlockComment = false; + currentStatement += char + nextChar; + i += 2; + continue; + } + + // Handle line comments -- ... + if (!inSingleQuote && !inDollarQuote && !inBlockComment && char === '-' && nextChar === '-') { + // Add rest of line to current statement + const lineEnd = sql.indexOf('\n', i); + if (lineEnd === -1) { + currentStatement += sql.substring(i); + break; + } + currentStatement += sql.substring(i, lineEnd + 1); + i = lineEnd + 1; + continue; + } + + // Handle dollar-quoted strings + if (!inSingleQuote && !inBlockComment) { + const dollarMatch = peek.match(/^(\$[a-zA-Z0-9_]*\$)/); + if (dollarMatch) { + const tag = dollarMatch[1]; + if (!inDollarQuote) { + inDollarQuote = true; + dollarQuoteTag = tag; + currentStatement += tag; + i += tag.length; + continue; + } else if (tag === dollarQuoteTag) { + inDollarQuote = false; + dollarQuoteTag = ''; + currentStatement += tag; + i += tag.length; + continue; + } + } + } + + // Handle single-quoted strings + if (!inDollarQuote && !inBlockComment && char === "'") { + if (inSingleQuote && nextChar === "'") { + // Escaped quote '' + currentStatement += "''"; + i += 2; + continue; + } + inSingleQuote = !inSingleQuote; + } - // Execute each statement and collect outputs - const outputs: vscode.NotebookCellOutput[] = []; - let totalExecutionTime = 0; + // Handle semicolon as statement separator + if (!inSingleQuote && !inDollarQuote && !inBlockComment && char === ';') { + currentStatement += char; + const trimmed = currentStatement.trim(); + if (trimmed) { + statements.push(trimmed); + } + currentStatement = ''; + i++; + continue; + } - for (let stmtIndex = 0; stmtIndex < statements.length; stmtIndex++) { - const query = statements[stmtIndex]; - const stmtStartTime = Date.now(); + currentStatement += char; + i++; + } - console.log(`PostgresKernel: Executing statement ${stmtIndex + 1}/${statements.length}:`, query.substring(0, 100)); + // Add remaining statement if any + const trimmed = currentStatement.trim(); + if (trimmed) { + statements.push(trimmed); + } - let result; - try { - result = await client.query(query); - const stmtEndTime = Date.now(); - const executionTime = (stmtEndTime - stmtStartTime) / 1000; - totalExecutionTime += executionTime; - - console.log(`PostgresKernel: Statement ${stmtIndex + 1} result:`, { - hasFields: !!result.fields, - fieldsLength: result.fields?.length, - rowsLength: result.rows?.length, - command: result.command - }); - - let tableInfo: { schema: string; table: string; primaryKeys: string[]; uniqueKeys: string[] } | undefined; - - // Try to get table metadata for SELECT queries to enable deletion - if (result.command === 'SELECT' && result.fields && result.fields.length > 0) { - const tableId = result.fields[0].tableID; - // Check if all fields come from the same table and tableId is valid - const allSameTable = result.fields.every((f: any) => f.tableID === tableId); - - if (tableId && tableId > 0 && allSameTable) { - try { - // Get table name and schema - const tableRes = await client.query( - `SELECT n.nspname, c.relname + return statements.filter(s => s.length > 0); + } + + private async _doExecution(cell: vscode.NotebookCell): Promise { + console.log(`PostgresKernel: Starting cell execution. Controller ID: ${this._controller.id}`); + const execution = this._controller.createNotebookCellExecution(cell); + const startTime = Date.now(); + execution.start(startTime); + execution.clearOutput(); + + try { + const metadata = cell.notebook.metadata as PostgresMetadata; + if (!metadata || !metadata.connectionId) { + throw new Error('No connection metadata found'); + } + + // Get connection info and password from SecretStorage + const connections = vscode.workspace.getConfiguration().get('postgresExplorer.connections') || []; + const connection = connections.find(c => c.id === metadata.connectionId); + if (!connection) { + throw new Error('Connection not found'); + } + + const client = await ConnectionManager.getInstance().getConnection({ + id: connection.id, + host: connection.host, + port: connection.port, + username: connection.username, + database: metadata.databaseName || connection.database, + name: connection.name + }); + + console.log('PostgresKernel: Connected to database'); + + // Get PostgreSQL backend PID for query cancellation + let backendPid: number | null = null; + try { + const pidResult = await client.query('SELECT pg_backend_pid()'); + backendPid = pidResult.rows[0]?.pg_backend_pid || null; + console.log('PostgresKernel: Backend PID:', backendPid); + } catch (err) { + console.warn('Failed to get backend PID:', err); + } + + // Capture PostgreSQL NOTICE messages + const notices: string[] = []; + const noticeListener = (msg: any) => { + const message = msg.message || msg.toString(); + notices.push(message); + }; + client.on('notice', noticeListener); + + const queryText = cell.document.getText(); + const statements = this.splitSqlStatements(queryText); + + console.log('PostgresKernel: Executing', statements.length, 'statement(s)'); + + // Execute each statement and collect outputs + const outputs: vscode.NotebookCellOutput[] = []; + let totalExecutionTime = 0; + + for (let stmtIndex = 0; stmtIndex < statements.length; stmtIndex++) { + const query = statements[stmtIndex]; + const stmtStartTime = Date.now(); + + console.log(`PostgresKernel: Executing statement ${stmtIndex + 1}/${statements.length}:`, query.substring(0, 100)); + + let result; + try { + result = await client.query(query); + const stmtEndTime = Date.now(); + const executionTime = (stmtEndTime - stmtStartTime) / 1000; + totalExecutionTime += executionTime; + + console.log(`PostgresKernel: Statement ${stmtIndex + 1} result:`, { + hasFields: !!result.fields, + fieldsLength: result.fields?.length, + rowsLength: result.rows?.length, + command: result.command + }); + + let tableInfo: { schema: string; table: string; primaryKeys: string[]; uniqueKeys: string[] } | undefined; + + // Try to get table metadata for SELECT queries to enable deletion + if (result.command === 'SELECT' && result.fields && result.fields.length > 0) { + const tableId = result.fields[0].tableID; + // Check if all fields come from the same table and tableId is valid + const allSameTable = result.fields.every((f: any) => f.tableID === tableId); + + if (tableId && tableId > 0 && allSameTable) { + try { + // Get table name and schema + const tableRes = await client.query( + `SELECT n.nspname, c.relname FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE c.oid = $1`, - [tableId] - ); + [tableId] + ); - if (tableRes.rows.length > 0) { - const { nspname, relname } = tableRes.rows[0]; + if (tableRes.rows.length > 0) { + const { nspname, relname } = tableRes.rows[0]; - // Get primary keys - const pkRes = await client.query( - `SELECT a.attname + // Get primary keys + const pkRes = await client.query( + `SELECT a.attname FROM pg_index i JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey) WHERE i.indrelid = $1 AND i.indisprimary`, - [tableId] - ); + [tableId] + ); - const primaryKeys = pkRes.rows.map((r: any) => r.attname); + const primaryKeys = pkRes.rows.map((r: any) => r.attname); - // Get unique keys (columns with unique constraints, excluding primary keys) - const ukRes = await client.query( - `SELECT DISTINCT a.attname + // Get unique keys (columns with unique constraints, excluding primary keys) + const ukRes = await client.query( + `SELECT DISTINCT a.attname FROM pg_index i JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey) WHERE i.indrelid = $1 AND i.indisunique AND NOT i.indisprimary`, - [tableId] - ); - - const uniqueKeys = ukRes.rows.map((r: any) => r.attname); - - if (primaryKeys.length > 0 || uniqueKeys.length > 0) { - tableInfo = { - schema: nspname, - table: relname, - primaryKeys: primaryKeys, - uniqueKeys: uniqueKeys - }; - } - } - } catch (err) { - console.warn('Failed to fetch table metadata:', err); - } - } - } - - // Get column type names from pg_type - let columnTypes: { [key: string]: string } = {}; - if (result.fields && result.fields.length > 0) { - try { - const typeOids = result.fields.map((f: any) => f.dataTypeID); - const uniqueOids = [...new Set(typeOids)]; - const typeRes = await client.query( - `SELECT oid, typname FROM pg_type WHERE oid = ANY($1::oid[])`, - [uniqueOids] - ); - const typeMap = new Map(typeRes.rows.map((r: any) => [r.oid, r.typname])); - result.fields.forEach((f: any) => { - columnTypes[f.name] = typeMap.get(f.dataTypeID) || 'unknown'; - }); - } catch (err) { - console.warn('Failed to fetch column type names:', err); - } - } - - // Generate output for this statement - const data = { - columns: result.fields ? result.fields.map((f: any) => f.name) : [], - columnTypes: columnTypes, - rows: result.rows || [], - rowCount: result.rowCount, - command: result.command, - notices: [...notices], - executionTime: executionTime, - tableInfo: tableInfo, - cellIndex: cell.index, - success: true - }; + [tableId] + ); - const cellOutput = new vscode.NotebookCellOutput([ - vscode.NotebookCellOutputItem.json(data, 'application/x-postgres-result') - ]); - outputs.push(cellOutput); + const uniqueKeys = ukRes.rows.map((r: any) => r.attname); - console.log(`PostgresKernel: Generated output for statement ${stmtIndex + 1}, outputs count: ${outputs.length}`); + if (primaryKeys.length > 0 || uniqueKeys.length > 0) { + tableInfo = { + schema: nspname, + table: relname, + primaryKeys: primaryKeys, + uniqueKeys: uniqueKeys + }; + } + } + } catch (err) { + console.warn('Failed to fetch table metadata:', err); + } + } + } - // Clear notices for next statement - notices.length = 0; - } catch (err: any) { - const stmtEndTime = Date.now(); - const executionTime = (stmtEndTime - stmtStartTime) / 1000; - totalExecutionTime += executionTime; + // Get column type names from pg_type + let columnTypes: { [key: string]: string } = {}; + if (result.fields && result.fields.length > 0) { + try { + const typeOids = result.fields.map((f: any) => f.dataTypeID); + const uniqueOids = [...new Set(typeOids)]; + const typeRes = await client.query( + `SELECT oid, typname FROM pg_type WHERE oid = ANY($1::oid[])`, + [uniqueOids] + ); + const typeMap = new Map(typeRes.rows.map((r: any) => [r.oid, r.typname])); + result.fields.forEach((f: any) => { + columnTypes[f.name] = typeMap.get(f.dataTypeID) || 'unknown'; + }); + } catch (err) { + console.warn('Failed to fetch column type names:', err); + } + } + + // Generate output for this statement + const data = { + columns: result.fields ? result.fields.map((f: any) => f.name) : [], + columnTypes: columnTypes, + rows: result.rows || [], + rowCount: result.rowCount, + command: result.command, + notices: [...notices], + executionTime: executionTime, + tableInfo: tableInfo, + cellIndex: cell.index, + backendPid: backendPid, + success: true + }; + + const cellOutput = new vscode.NotebookCellOutput([ + vscode.NotebookCellOutputItem.json(data, 'application/x-postgres-result') + ]); + outputs.push(cellOutput); + + console.log(`PostgresKernel: Generated output for statement ${stmtIndex + 1}, outputs count: ${outputs.length}`); + + // Clear notices for next statement + notices.length = 0; + } catch (err: any) { + const stmtEndTime = Date.now(); + const executionTime = (stmtEndTime - stmtStartTime) / 1000; + totalExecutionTime += executionTime; - console.error(`PostgresKernel: Statement ${stmtIndex + 1} error:`, err.message); + console.error(`PostgresKernel: Statement ${stmtIndex + 1} error:`, err.message); - // Show error for this specific statement - const errorHtml = ` + // Show error for this specific statement + const errorHtml = `
Execution time: ${executionTime.toFixed(3)} sec.
`; - const cellOutput = new vscode.NotebookCellOutput([ - vscode.NotebookCellOutputItem.text(errorHtml, 'text/html') - ]); - outputs.push(cellOutput); + const cellOutput = new vscode.NotebookCellOutput([ + vscode.NotebookCellOutputItem.text(errorHtml, 'text/html') + ]); + outputs.push(cellOutput); - // Continue with remaining statements even if one fails - notices.length = 0; - } - } - - // Remove notice listener - client.off('notice', noticeListener); - - // Combine all outputs - console.log(`PostgresKernel: Combining ${outputs.length} output(s)`); - - if (outputs.length > 0) { - execution.replaceOutput(outputs); - execution.end(true); - console.log('PostgresKernel: Cell execution completed successfully'); - } else { - throw new Error('No statements to execute'); - } - } catch (err: any) { - console.error('PostgresKernel: Cell execution failed:', err); - const output = new vscode.NotebookCellOutput([ - vscode.NotebookCellOutputItem.error(err) - ]); - execution.replaceOutput([output]); - execution.end(false); + // Continue with remaining statements even if one fails + notices.length = 0; } + } + + // Remove notice listener + client.off('notice', noticeListener); + + // Combine all outputs + console.log(`PostgresKernel: Combining ${outputs.length} output(s)`); + + if (outputs.length > 0) { + execution.replaceOutput(outputs); + execution.end(true); + console.log('PostgresKernel: Cell execution completed successfully'); + } else { + throw new Error('No statements to execute'); + } + } catch (err: any) { + console.error('PostgresKernel: Cell execution failed:', err); + const output = new vscode.NotebookCellOutput([ + vscode.NotebookCellOutputItem.error(err) + ]); + execution.replaceOutput([output]); + execution.end(false); } + } - dispose() { - this._controller.dispose(); - } + dispose() { + this._controller.dispose(); + } } diff --git a/src/renderer_v2.ts b/src/renderer_v2.ts index cdf7231..e580fd5 100644 --- a/src/renderer_v2.ts +++ b/src/renderer_v2.ts @@ -5,372 +5,373 @@ import { Chart, registerables } from 'chart.js'; Chart.register(...registerables); export const activate: ActivationFunction = context => { - return { - renderOutputItem(data, element) { - const json = data.json(); - - if (!json) { - element.innerText = 'No data'; - return; - } - - const { columns, rows, rowCount, command, notices, executionTime, tableInfo, success, columnTypes } = json; - // Deep copy rows to allow modifications without affecting originals - const originalRows: any[] = rows ? JSON.parse(JSON.stringify(rows)) : []; - let currentRows: any[] = rows ? JSON.parse(JSON.stringify(rows)) : []; - const selectedIndices = new Set(); - - // Track modified cells: Map of "rowIndex-columnName" -> { originalValue, newValue } - const modifiedCells = new Map(); - let currentlyEditingCell: HTMLElement | null = null; - - // Track date/time column display mode: true = local time, false = original value - const dateTimeDisplayMode = new Map(); - - // ... (rest of the code) - - // Main Container (Collapsible Wrapper) - const mainContainer = document.createElement('div'); - mainContainer.style.fontFamily = 'var(--vscode-font-family), "Segoe UI", "Helvetica Neue", sans-serif'; - mainContainer.style.fontSize = '13px'; - mainContainer.style.color = 'var(--vscode-editor-foreground)'; - mainContainer.style.border = '1px solid var(--vscode-widget-border)'; - mainContainer.style.borderRadius = '4px'; - mainContainer.style.overflow = 'hidden'; - mainContainer.style.marginBottom = '8px'; - - // Header - const header = document.createElement('div'); - header.style.padding = '6px 12px'; - // Use green background for successful queries, neutral for others - if (success) { - header.style.background = 'rgba(115, 191, 105, 0.25)'; // Green tint for success - header.style.borderLeft = '4px solid var(--vscode-testing-iconPassed)'; - } else { - header.style.background = 'var(--vscode-editor-background)'; - } - header.style.borderBottom = '1px solid var(--vscode-widget-border)'; - header.style.cursor = 'pointer'; - header.style.display = 'flex'; - header.style.alignItems = 'center'; - header.style.gap = '8px'; - header.style.userSelect = 'none'; - - const chevron = document.createElement('span'); - chevron.textContent = '▼'; // Expanded by default - chevron.style.fontSize = '10px'; - chevron.style.transition = 'transform 0.2s'; - chevron.style.display = 'inline-block'; - - const title = document.createElement('span'); - title.textContent = command || 'QUERY'; - title.style.fontWeight = '600'; - title.style.textTransform = 'uppercase'; - - const summary = document.createElement('span'); - summary.style.marginLeft = 'auto'; - summary.style.opacity = '0.7'; - summary.style.fontSize = '0.9em'; - - let summaryText = ''; - if (rowCount !== undefined && rowCount !== null) { - summaryText += `${rowCount} rows`; - } - if (notices && notices.length > 0) { - summaryText += summaryText ? `, ${notices.length} messages` : `${notices.length} messages`; - } - if (executionTime !== undefined) { - summaryText += summaryText ? `, ${executionTime.toFixed(3)}s` : `${executionTime.toFixed(3)}s`; - } - if (!summaryText) summaryText = 'No results'; - summary.textContent = summaryText; - - header.appendChild(chevron); - header.appendChild(title); - header.appendChild(summary); - mainContainer.appendChild(header); - - // Content Container - const contentContainer = document.createElement('div'); - contentContainer.style.display = 'flex'; // Expanded by default - contentContainer.style.flexDirection = 'column'; - contentContainer.style.height = '100%'; // Added to ensure content takes full height if needed - mainContainer.appendChild(contentContainer); - - // Toggle Logic - let isExpanded = true; - header.addEventListener('click', () => { - isExpanded = !isExpanded; - contentContainer.style.display = isExpanded ? 'flex' : 'none'; - chevron.style.transform = isExpanded ? 'rotate(0deg)' : 'rotate(-90deg)'; - header.style.borderBottom = isExpanded ? '1px solid var(--vscode-widget-border)' : 'none'; - }); - - // Messages Section - if (notices && notices.length > 0) { - const messagesContainer = document.createElement('div'); - messagesContainer.style.padding = '8px 12px'; - messagesContainer.style.background = 'var(--vscode-textBlockQuote-background)'; - messagesContainer.style.borderLeft = '4px solid var(--vscode-textBlockQuote-border)'; - messagesContainer.style.margin = '8px 12px 0 12px'; // Add margin - messagesContainer.style.fontFamily = 'var(--vscode-editor-font-family)'; - messagesContainer.style.whiteSpace = 'pre-wrap'; - messagesContainer.style.fontSize = '12px'; - - const title = document.createElement('div'); - title.textContent = 'Messages'; - title.style.fontWeight = '600'; - title.style.marginBottom = '4px'; - title.style.opacity = '0.8'; - messagesContainer.appendChild(title); - - notices.forEach((msg: string) => { - const msgDiv = document.createElement('div'); - msgDiv.textContent = msg; - msgDiv.style.marginBottom = '2px'; - messagesContainer.appendChild(msgDiv); - }); + return { + renderOutputItem(data, element) { + const json = data.json(); + + if (!json) { + element.innerText = 'No data'; + return; + } + + const { columns, rows, rowCount, command, notices, executionTime, tableInfo, success, columnTypes, backendPid } = json; + // Deep copy rows to allow modifications without affecting originals + const originalRows: any[] = rows ? JSON.parse(JSON.stringify(rows)) : []; + let currentRows: any[] = rows ? JSON.parse(JSON.stringify(rows)) : []; + const selectedIndices = new Set(); + + // Track modified cells: Map of "rowIndex-columnName" -> { originalValue, newValue } + const modifiedCells = new Map(); + let currentlyEditingCell: HTMLElement | null = null; + + // Track date/time column display mode: true = local time, false = original value + const dateTimeDisplayMode = new Map(); + + // ... (rest of the code) + + // Main Container (Collapsible Wrapper) + const mainContainer = document.createElement('div'); + mainContainer.style.fontFamily = 'var(--vscode-font-family), "Segoe UI", "Helvetica Neue", sans-serif'; + mainContainer.style.fontSize = '13px'; + mainContainer.style.color = 'var(--vscode-editor-foreground)'; + mainContainer.style.border = '1px solid var(--vscode-widget-border)'; + mainContainer.style.borderRadius = '4px'; + mainContainer.style.overflow = 'hidden'; + mainContainer.style.marginBottom = '8px'; + + // Header + const header = document.createElement('div'); + header.style.padding = '6px 12px'; + // Use green background for successful queries, neutral for others + if (success) { + header.style.background = 'rgba(115, 191, 105, 0.25)'; // Green tint for success + header.style.borderLeft = '4px solid var(--vscode-testing-iconPassed)'; + } else { + header.style.background = 'var(--vscode-editor-background)'; + } + header.style.borderBottom = '1px solid var(--vscode-widget-border)'; + header.style.cursor = 'pointer'; + header.style.display = 'flex'; + header.style.alignItems = 'center'; + header.style.gap = '8px'; + header.style.userSelect = 'none'; + + const chevron = document.createElement('span'); + chevron.textContent = '▼'; // Expanded by default + chevron.style.fontSize = '10px'; + chevron.style.transition = 'transform 0.2s'; + chevron.style.display = 'inline-block'; + + const title = document.createElement('span'); + title.textContent = command || 'QUERY'; + title.style.fontWeight = '600'; + title.style.textTransform = 'uppercase'; + + const summary = document.createElement('span'); + summary.style.marginLeft = 'auto'; + summary.style.opacity = '0.7'; + summary.style.fontSize = '0.9em'; + + let summaryText = ''; + if (rowCount !== undefined && rowCount !== null) { + summaryText += `${rowCount} rows`; + } + if (notices && notices.length > 0) { + summaryText += summaryText ? `, ${notices.length} messages` : `${notices.length} messages`; + } + if (executionTime !== undefined) { + summaryText += summaryText ? `, ${executionTime.toFixed(3)}s` : `${executionTime.toFixed(3)}s`; + } + if (!summaryText) summaryText = 'No results'; + summary.textContent = summaryText; + + header.appendChild(chevron); + header.appendChild(title); + header.appendChild(summary); + mainContainer.appendChild(header); + + // Content Container + const contentContainer = document.createElement('div'); + contentContainer.style.display = 'flex'; // Expanded by default + contentContainer.style.flexDirection = 'column'; + contentContainer.style.height = '100%'; // Added to ensure content takes full height if needed + mainContainer.appendChild(contentContainer); + + // Toggle Logic + let isExpanded = true; + header.addEventListener('click', () => { + isExpanded = !isExpanded; + contentContainer.style.display = isExpanded ? 'flex' : 'none'; + chevron.style.transform = isExpanded ? 'rotate(0deg)' : 'rotate(-90deg)'; + header.style.borderBottom = isExpanded ? '1px solid var(--vscode-widget-border)' : 'none'; + }); + + // Messages Section + if (notices && notices.length > 0) { + const messagesContainer = document.createElement('div'); + messagesContainer.style.padding = '8px 12px'; + messagesContainer.style.background = 'var(--vscode-textBlockQuote-background)'; + messagesContainer.style.borderLeft = '4px solid var(--vscode-textBlockQuote-border)'; + messagesContainer.style.margin = '8px 12px 0 12px'; // Add margin + messagesContainer.style.fontFamily = 'var(--vscode-editor-font-family)'; + messagesContainer.style.whiteSpace = 'pre-wrap'; + messagesContainer.style.fontSize = '12px'; + + const title = document.createElement('div'); + title.textContent = 'Messages'; + title.style.fontWeight = '600'; + title.style.marginBottom = '4px'; + title.style.opacity = '0.8'; + messagesContainer.appendChild(title); + + notices.forEach((msg: string) => { + const msgDiv = document.createElement('div'); + msgDiv.textContent = msg; + msgDiv.style.marginBottom = '2px'; + messagesContainer.appendChild(msgDiv); + }); + + contentContainer.appendChild(messagesContainer); + } + + // Actions Bar + const actionsBar = document.createElement('div'); + actionsBar.style.display = 'none'; // Hidden by default + actionsBar.style.padding = '8px 12px'; + actionsBar.style.gap = '8px'; + actionsBar.style.alignItems = 'center'; + actionsBar.style.borderBottom = '1px solid var(--vscode-panel-border)'; + actionsBar.style.background = 'var(--vscode-editor-background)'; + + const createButton = (text: string, primary: boolean = false) => { + const btn = document.createElement('button'); + btn.textContent = text; + btn.style.background = primary ? 'var(--vscode-button-background)' : 'var(--vscode-button-secondaryBackground)'; + btn.style.color = primary ? 'var(--vscode-button-foreground)' : 'var(--vscode-button-secondaryForeground)'; + btn.style.border = 'none'; + btn.style.padding = '4px 12px'; + btn.style.cursor = 'pointer'; + btn.style.borderRadius = '2px'; + btn.style.fontSize = '12px'; + btn.style.fontWeight = '500'; + return btn; + }; + const selectAllBtn = createButton('Select All', true); + selectAllBtn.addEventListener('click', () => { + const allSelected = selectedIndices.size === currentRows.length; + + if (allSelected) { + selectedIndices.clear(); + selectAllBtn.innerText = 'Select All'; + } else { + currentRows.forEach((_, i) => selectedIndices.add(i)); + selectAllBtn.innerText = 'Deselect All'; + } - contentContainer.appendChild(messagesContainer); + updateTable(); + updateActionsVisibility(); + }); + actionsBar.appendChild(selectAllBtn); + + const copyBtn = createButton('Copy Selected', true); + copyBtn.addEventListener('click', async () => { + if (selectedIndices.size === 0) return; + + const selectedRows = currentRows.filter((_, i) => selectedIndices.has(i)); + + // Convert to CSV + const header = columns.map((c: string) => `"${c.replace(/"/g, '""')}"`).join(','); + const body = selectedRows.map(row => { + return columns.map((col: string) => { + const val = row[col]; + if (val === null || val === undefined) return ''; + // Use JSON.stringify for objects, String for primitives + const str = typeof val === 'object' ? JSON.stringify(val) : String(val); + // Quote strings if they contain commas, quotes, or newlines + if (str.includes(',') || str.includes('"') || str.includes('\n')) { + return `"${str.replace(/"/g, '""')}"`; } - - // Actions Bar - const actionsBar = document.createElement('div'); - actionsBar.style.display = 'none'; // Hidden by default - actionsBar.style.padding = '8px 12px'; - actionsBar.style.gap = '8px'; - actionsBar.style.alignItems = 'center'; - actionsBar.style.borderBottom = '1px solid var(--vscode-panel-border)'; - actionsBar.style.background = 'var(--vscode-editor-background)'; - - const createButton = (text: string, primary: boolean = false) => { - const btn = document.createElement('button'); - btn.textContent = text; - btn.style.background = primary ? 'var(--vscode-button-background)' : 'var(--vscode-button-secondaryBackground)'; - btn.style.color = primary ? 'var(--vscode-button-foreground)' : 'var(--vscode-button-secondaryForeground)'; - btn.style.border = 'none'; - btn.style.padding = '4px 12px'; - btn.style.cursor = 'pointer'; - btn.style.borderRadius = '2px'; - btn.style.fontSize = '12px'; - btn.style.fontWeight = '500'; - return btn; - }; - const selectAllBtn = createButton('Select All', true); - selectAllBtn.addEventListener('click', () => { - const allSelected = selectedIndices.size === currentRows.length; - - if (allSelected) { - selectedIndices.clear(); - selectAllBtn.innerText = 'Select All'; - } else { - currentRows.forEach((_, i) => selectedIndices.add(i)); - selectAllBtn.innerText = 'Deselect All'; - } - - updateTable(); - updateActionsVisibility(); + return str; + }).join(','); + }).join('\n'); + + const csv = `${header}\n${body}`; + + navigator.clipboard.writeText(csv).then(() => { + const originalText = copyBtn.textContent; + copyBtn.textContent = 'Copied!'; + copyBtn.style.background = 'var(--vscode-debugIcon-startForeground)'; + setTimeout(() => { + copyBtn.textContent = originalText; + copyBtn.style.background = 'var(--vscode-button-background)'; + }, 2000); + }).catch((err: Error) => { + console.error('Failed to copy:', err); + copyBtn.textContent = 'Failed'; + copyBtn.style.background = 'var(--vscode-errorForeground)'; + setTimeout(() => { + copyBtn.textContent = 'Copy Selected'; + copyBtn.style.background = 'var(--vscode-button-background)'; + }, 2000); + }); + }); + + const deleteBtn = createButton(tableInfo ? 'Script Delete' : 'Remove from View', !!tableInfo); + + deleteBtn.addEventListener('click', () => { + if (selectedIndices.size === 0) return; + + if (tableInfo) { + // Send script_delete message to kernel + const selectedRows = currentRows.filter((_, i) => selectedIndices.has(i)); + if (context.postMessage) { + context.postMessage({ + type: 'script_delete', + schema: tableInfo.schema, + table: tableInfo.table, + primaryKeys: tableInfo.primaryKeys, + rows: selectedRows, + cellIndex: (json as any).cellIndex // Access cellIndex from JSON }); - actionsBar.appendChild(selectAllBtn); - - const copyBtn = createButton('Copy Selected', true); - copyBtn.addEventListener('click', async () => { - if (selectedIndices.size === 0) return; - - const selectedRows = currentRows.filter((_, i) => selectedIndices.has(i)); - - // Convert to CSV - const header = columns.map((c: string) => `"${c.replace(/"/g, '""')}"`).join(','); - const body = selectedRows.map(row => { - return columns.map((col: string) => { - const val = row[col]; - if (val === null || val === undefined) return ''; - const str = String(val); - // Quote strings if they contain commas, quotes, or newlines - if (str.includes(',') || str.includes('"') || str.includes('\n')) { - return `"${str.replace(/"/g, '""')}"`; - } - return str; - }).join(','); - }).join('\n'); - - const csv = `${header}\n${body}`; - - navigator.clipboard.writeText(csv).then(() => { - const originalText = copyBtn.textContent; - copyBtn.textContent = 'Copied!'; - copyBtn.style.background = 'var(--vscode-debugIcon-startForeground)'; - setTimeout(() => { - copyBtn.textContent = originalText; - copyBtn.style.background = 'var(--vscode-button-background)'; - }, 2000); - }).catch((err: Error) => { - console.error('Failed to copy:', err); - copyBtn.textContent = 'Failed'; - copyBtn.style.background = 'var(--vscode-errorForeground)'; - setTimeout(() => { - copyBtn.textContent = 'Copy Selected'; - copyBtn.style.background = 'var(--vscode-button-background)'; - }, 2000); - }); - }); - - const deleteBtn = createButton(tableInfo ? 'Script Delete' : 'Remove from View', !!tableInfo); - - deleteBtn.addEventListener('click', () => { - if (selectedIndices.size === 0) return; - - if (tableInfo) { - // Send script_delete message to kernel - const selectedRows = currentRows.filter((_, i) => selectedIndices.has(i)); - if (context.postMessage) { - context.postMessage({ - type: 'script_delete', - schema: tableInfo.schema, - table: tableInfo.table, - primaryKeys: tableInfo.primaryKeys, - rows: selectedRows, - cellIndex: (json as any).cellIndex // Access cellIndex from JSON - }); - } - } else { - // Fallback to remove from view - if (confirm('Remove selected rows from this view?')) { - currentRows = currentRows.filter((_, i) => !selectedIndices.has(i)); - selectedIndices.clear(); - updateTable(); - updateActionsVisibility(); - } - } - }); - - const exportBtn = createButton('Export ▼', true); - exportBtn.style.position = 'relative'; - - exportBtn.addEventListener('click', (e) => { - e.stopPropagation(); - - // Remove existing dropdown if any - const existing = document.querySelector('.export-dropdown'); - if (existing) { - existing.remove(); - return; - } - - const menu = document.createElement('div'); - menu.className = 'export-dropdown'; - menu.style.position = 'absolute'; - menu.style.top = '100%'; - menu.style.left = '0'; - menu.style.background = 'var(--vscode-menu-background)'; - menu.style.border = '1px solid var(--vscode-menu-border)'; - menu.style.boxShadow = '0 2px 8px rgba(0,0,0,0.15)'; - menu.style.zIndex = '100'; - menu.style.minWidth = '150px'; - menu.style.borderRadius = '3px'; - menu.style.padding = '4px 0'; - - const createMenuItem = (label: string, onClick: () => void) => { - const item = document.createElement('div'); - item.textContent = label; - item.style.padding = '6px 12px'; - item.style.cursor = 'pointer'; - item.style.color = 'var(--vscode-menu-foreground)'; - item.style.fontSize = '12px'; - - item.addEventListener('mouseenter', () => { - item.style.background = 'var(--vscode-menu-selectionBackground)'; - item.style.color = 'var(--vscode-menu-selectionForeground)'; - }); - item.addEventListener('mouseleave', () => { - item.style.background = 'transparent'; - item.style.color = 'var(--vscode-menu-foreground)'; - }); - item.addEventListener('click', (e) => { - e.stopPropagation(); - onClick(); - menu.remove(); - }); - return item; - }; - - const downloadFile = (content: string, filename: string, type: string) => { - const blob = new Blob([content], { type }); - const url = URL.createObjectURL(blob); - const a = document.createElement('a'); - a.href = url; - a.download = filename; - document.body.appendChild(a); - a.click(); - document.body.removeChild(a); - URL.revokeObjectURL(url); - }; - - const stringifyValue = (val: any): string => { - if (val === null || val === undefined) return ''; - if (typeof val === 'object') return JSON.stringify(val); - return String(val); - }; + } + } else { + // Fallback to remove from view + if (confirm('Remove selected rows from this view?')) { + currentRows = currentRows.filter((_, i) => !selectedIndices.has(i)); + selectedIndices.clear(); + updateTable(); + updateActionsVisibility(); + } + } + }); - const getCSV = () => { - const header = columns.map((c: string) => `"${c.replace(/"/g, '""')}"`).join(','); - const body = currentRows.map(row => { - return columns.map((col: string) => { - const val = row[col]; - const str = stringifyValue(val); - if (str.includes(',') || str.includes('"') || str.includes('\n')) { - return `"${str.replace(/"/g, '""')}"`; - } - return str; - }).join(','); - }).join('\n'); - return `${header}\n${body}`; - }; + const exportBtn = createButton('Export ▼', true); + exportBtn.style.position = 'relative'; - const getMarkdown = () => { - const header = `| ${columns.join(' | ')} |`; - const separator = `| ${columns.map(() => '---').join(' | ')} |`; - const body = currentRows.map(row => { - return `| ${columns.map((col: string) => { - const val = row[col]; - if (val === null || val === undefined) return 'NULL'; - const str = typeof val === 'object' ? JSON.stringify(val) : String(val); - return str.replace(/\|/g, '\\|').replace(/\n/g, ' '); - }).join(' | ')} |`; - }).join('\n'); - return `${header}\n${separator}\n${body}`; - }; + exportBtn.addEventListener('click', (e) => { + e.stopPropagation(); - const getSQLInsert = () => { - if (!tableInfo) return '-- Table information not available for INSERT script'; - const tableName = `"${tableInfo.schema}"."${tableInfo.table}"`; - const cols = columns.map((c: string) => `"${c}"`).join(', '); - - return currentRows.map((row: any) => { - const values = columns.map((col: string) => { - const val = row[col]; - if (val === null || val === undefined) return 'NULL'; - if (typeof val === 'number') return val; - if (typeof val === 'boolean') return val ? 'TRUE' : 'FALSE'; - const str = typeof val === 'object' ? JSON.stringify(val) : String(val); - return `'${str.replace(/'/g, "''")}'`; - }).join(', '); - return `INSERT INTO ${tableName} (${cols}) VALUES (${values});`; - }).join('\n'); - }; + // Remove existing dropdown if any + const existing = document.querySelector('.export-dropdown'); + if (existing) { + existing.remove(); + return; + } - const getExcel = () => { - // Simple HTML-based Excel format - const header = columns.map((c: string) => `${c}`).join(''); - const body = currentRows.map(row => { - const cells = columns.map((col: string) => { - const val = row[col]; - return `${stringifyValue(val)}`; - }).join(''); - return `${cells}`; - }).join(''); - - return ` + const menu = document.createElement('div'); + menu.className = 'export-dropdown'; + menu.style.position = 'absolute'; + menu.style.top = '100%'; + menu.style.left = '0'; + menu.style.background = 'var(--vscode-menu-background)'; + menu.style.border = '1px solid var(--vscode-menu-border)'; + menu.style.boxShadow = '0 2px 8px rgba(0,0,0,0.15)'; + menu.style.zIndex = '100'; + menu.style.minWidth = '150px'; + menu.style.borderRadius = '3px'; + menu.style.padding = '4px 0'; + + const createMenuItem = (label: string, onClick: () => void) => { + const item = document.createElement('div'); + item.textContent = label; + item.style.padding = '6px 12px'; + item.style.cursor = 'pointer'; + item.style.color = 'var(--vscode-menu-foreground)'; + item.style.fontSize = '12px'; + + item.addEventListener('mouseenter', () => { + item.style.background = 'var(--vscode-menu-selectionBackground)'; + item.style.color = 'var(--vscode-menu-selectionForeground)'; + }); + item.addEventListener('mouseleave', () => { + item.style.background = 'transparent'; + item.style.color = 'var(--vscode-menu-foreground)'; + }); + item.addEventListener('click', (e) => { + e.stopPropagation(); + onClick(); + menu.remove(); + }); + return item; + }; + + const downloadFile = (content: string, filename: string, type: string) => { + const blob = new Blob([content], { type }); + const url = URL.createObjectURL(blob); + const a = document.createElement('a'); + a.href = url; + a.download = filename; + document.body.appendChild(a); + a.click(); + document.body.removeChild(a); + URL.revokeObjectURL(url); + }; + + const stringifyValue = (val: any): string => { + if (val === null || val === undefined) return ''; + if (typeof val === 'object') return JSON.stringify(val); + return String(val); + }; + + const getCSV = () => { + const header = columns.map((c: string) => `"${c.replace(/"/g, '""')}"`).join(','); + const body = currentRows.map(row => { + return columns.map((col: string) => { + const val = row[col]; + const str = stringifyValue(val); + if (str.includes(',') || str.includes('"') || str.includes('\n')) { + return `"${str.replace(/"/g, '""')}"`; + } + return str; + }).join(','); + }).join('\n'); + return `${header}\n${body}`; + }; + + const getMarkdown = () => { + const header = `| ${columns.join(' | ')} |`; + const separator = `| ${columns.map(() => '---').join(' | ')} |`; + const body = currentRows.map(row => { + return `| ${columns.map((col: string) => { + const val = row[col]; + if (val === null || val === undefined) return 'NULL'; + const str = typeof val === 'object' ? JSON.stringify(val) : String(val); + return str.replace(/\|/g, '\\|').replace(/\n/g, ' '); + }).join(' | ')} |`; + }).join('\n'); + return `${header}\n${separator}\n${body}`; + }; + + const getSQLInsert = () => { + if (!tableInfo) return '-- Table information not available for INSERT script'; + const tableName = `"${tableInfo.schema}"."${tableInfo.table}"`; + const cols = columns.map((c: string) => `"${c}"`).join(', '); + + return currentRows.map((row: any) => { + const values = columns.map((col: string) => { + const val = row[col]; + if (val === null || val === undefined) return 'NULL'; + if (typeof val === 'number') return val; + if (typeof val === 'boolean') return val ? 'TRUE' : 'FALSE'; + const str = typeof val === 'object' ? JSON.stringify(val) : String(val); + return `'${str.replace(/'/g, "''")}'`; + }).join(', '); + return `INSERT INTO ${tableName} (${cols}) VALUES (${values});`; + }).join('\n'); + }; + + const getExcel = () => { + // Simple HTML-based Excel format + const header = columns.map((c: string) => `${c}`).join(''); + const body = currentRows.map(row => { + const cells = columns.map((col: string) => { + const val = row[col]; + return `${stringifyValue(val)}`; + }).join(''); + return `${cells}`; + }).join(''); + + return ` CMD + EXT --> PROV + EXT --> SVC + + CMD --> SVC + PROV --> SVC + + SVC --> PG + SVC --> AI + + EXT -.Message Passing.-> DASH + EXT -.Message Passing.-> NB + + DASH -.Postback.-> CMD + NB -.Postback.-> CMD +``` + +--- + +## Component Architecture + +### Extension Host Components + +#### 1. **Extension Entry Point** (`src/extension.ts`) +- Registers commands, providers, and services +- Initializes connection manager and secret storage +- Sets up message passing for webviews +- Manages extension lifecycle + +#### 2. **Commands** (`src/commands/`) +Implements VS Code commands for database operations: + +```mermaid +graph LR + subgraph Commands + CONN[Connection
Management] + TBL[Table
Operations] + VIEW[View
Operations] + FUNC[Function
Operations] + FDW[FDW
Operations] + AI_CMD[AI
Commands] + end + + CONN --> |Uses| CM[ConnectionManager] + TBL --> |Uses| CM + VIEW --> |Uses| CM + FUNC --> |Uses| CM + FDW --> |Uses| CM + AI_CMD --> |Uses| AI_SVC[AIService] +``` + +**Key Command Files**: +- `connections.ts` - Add, edit, delete, test connections +- `tables.ts` - CRUD, VACUUM, ANALYZE, REINDEX, scripts +- `views.ts` - View operations and scripts +- `functions.ts` - Function management +- `fdw.ts` - Foreign Data Wrapper operations +- `ai.ts` - AI-powered query generation and optimization + +#### 3. **Providers** (`src/providers/`) +Implements VS Code provider interfaces: + +- **`DatabaseTreeProvider.ts`** - Tree view for database objects +- **`SqlCompletionProvider.ts`** - SQL autocomplete +- **`NotebookKernel.ts`** - SQL notebook execution kernel +- **`DashboardPanel.ts`** - Real-time monitoring dashboard + +#### 4. **Services** (`src/services/`) +Core business logic and infrastructure: + +```mermaid +graph TB + subgraph Services + CM[ConnectionManager] + SS[SecretStorageService] + SSH[SSHService] + AI[AIService] + HIST[HistoryService] + ERR[ErrorService] + DBO[DbObjectService] + end + + CM --> |Stores passwords| SS + CM --> |Tunnels via| SSH + AI --> |Reads schema| DBO + AI --> |Reads history| HIST + CM --> |Reports errors| ERR +``` + +**Service Responsibilities**: + +| Service | Purpose | +|---------|---------| +| `ConnectionManager` | Hybrid pooling (`pg.Pool` + `pg.Client`), connection lifecycle | +| `SecretStorageService` | Encrypted password storage via VS Code API | +| `SSHService` | SSH tunnel management for remote connections | +| `AIService` | AI provider abstraction (OpenAI, Anthropic, Gemini) | +| `HistoryService` | Query history tracking for AI context | +| `ErrorService` | Centralized error handling and reporting | +| `DbObjectService` | Database schema introspection and caching | + +--- + +### Webview Components + +#### 1. **Dashboard Panel** (`src/panels/DashboardPanel.ts`) +Real-time monitoring interface showing: +- Active connections (`pg_stat_activity`) +- Database size and growth +- Table statistics +- Index health +- Long-running queries + +**Technology**: HTML + CSS + Vanilla JavaScript with Chart.js + +#### 2. **Notebook Renderer** (`src/renderer_v2.ts`) +Renders SQL query results in notebooks with: +- **Infinite Scrolling** (200 rows/chunk via `IntersectionObserver`) +- Interactive table with column resizing +- Inline cell editing (double-click) +- Chart visualization (Chart.js) +- Export to CSV/JSON/Excel +- AI-powered data analysis + +**Modular Structure**: +``` +src/renderer/ +├── components/ +│ └── ui.ts # Reusable UI components (buttons, tabs) +└── features/ + ├── export.ts # Export functionality + └── ai.ts # AI analysis features +``` + +--- + +## Data Flow + +### Query Execution Flow + +```mermaid +sequenceDiagram + participant User + participant Notebook as Notebook Cell + participant Kernel as NotebookKernel + participant CM as ConnectionManager + participant Pool as pg.Pool + participant PG as PostgreSQL + participant Renderer as Notebook Renderer + + User->>Notebook: Execute SQL + Notebook->>Kernel: Execute cell + Kernel->>CM: getSessionClient(config, sessionId) + CM->>Pool: Get/Create pool + Pool->>CM: Return pooled client + CM->>Kernel: Return client + + Kernel->>PG: Execute query + PG->>Kernel: Return result (rows) + + alt Result > 10k rows + Kernel->>Kernel: Truncate to 10k + Kernel->>Kernel: Add warning notice + end + + Kernel->>Renderer: Send result JSON + Renderer->>Renderer: Render first 200 rows + + User->>Renderer: Scroll down + Renderer->>Renderer: Load next 200 rows +``` + +### Connection Pooling Strategy + +PgStudio uses a **hybrid pooling approach**: + +```mermaid +graph TB + subgraph "Connection Manager" + direction TB + API[Public API] + + subgraph "Ephemeral Operations" + POOL[pg.Pool] + PC1[PoolClient 1] + PC2[PoolClient 2] + PCN[PoolClient N] + end + + subgraph "Stateful Sessions" + CLIENT[pg.Client] + SC1[Session Client 1
Notebook A] + SC2[Session Client 2
Notebook B] + end + end + + API -->|getPooledClient| POOL + POOL --> PC1 + POOL --> PC2 + POOL --> PCN + + API -->|getSessionClient| CLIENT + CLIENT --> SC1 + CLIENT --> SC2 + + PC1 -.auto release.-> POOL + PC2 -.auto release.-> POOL + + SC1 -.manual close.-> CLIENT + SC2 -.manual close.-> CLIENT +``` + +**Why Hybrid?** +- **`pg.Pool`** for commands (tree refresh, CRUD operations) - automatic connection reuse +- **`pg.Client`** for notebooks - maintains transaction state across cells + +--- + +## Key Design Decisions + +### 1. **Infinite Scrolling over Virtual Scrolling** + +**Problem**: Rendering 10k+ rows freezes the UI. + +**Solution**: +- Backend truncates results to 10k rows (prevents crashes) +- Frontend renders in 200-row chunks +- `IntersectionObserver` triggers loading on scroll + +**Trade-offs**: +- ✅ Simple implementation +- ✅ Works with existing table structure +- ✅ No complex windowing logic +- ⚠️ All 10k rows in memory (acceptable for max limit) + +### 2. **Hybrid Connection Pooling** + +**Problem**: +- Commands need short-lived connections +- Notebooks need persistent connections (transactions) + +**Solution**: +```typescript +// Ephemeral (auto-released) +const client = await ConnectionManager.getInstance().getPooledClient(config); +try { + await client.query('SELECT ...'); +} finally { + client.release(); // Returns to pool +} + +// Stateful (manual lifecycle) +const client = await ConnectionManager.getInstance() + .getSessionClient(config, notebookUri); +// Client persists across cells +// Closed when notebook closes +``` + +### 3. **AI Context Optimization** + +**Problem**: Sending full schema on every AI request is expensive. + +**Solution**: +- **Schema Caching**: `DbObjectService` caches table/column metadata +- **Selective Context**: Only send relevant tables based on query +- **History Integration**: Include recent queries for better suggestions + +### 4. **Modular Renderer Architecture** + +**Problem**: `renderer_v2.ts` was 144KB monolith. + +**Solution**: +``` +renderer_v2.ts (main) +├── renderer/components/ui.ts # Buttons, tabs, modals +└── renderer/features/ + ├── export.ts # CSV/JSON/Excel export + └── ai.ts # AI analysis integration +``` + +**Benefits**: +- Easier testing +- Clear separation of concerns +- Reusable components + +--- + +## Extension Points + +### Adding a New Command + +1. **Define command** in `package.json`: +```json +{ + "command": "postgres-explorer.myCommand", + "title": "My Command", + "category": "PostgreSQL" +} +``` + +2. **Implement handler** in `src/commands/myCommand.ts`: +```typescript +export async function myCommandHandler(node: TreeNode) { + const client = await ConnectionManager.getInstance() + .getPooledClient(node.connection); + try { + await client.query('...'); + } finally { + client.release(); + } +} +``` + +3. **Register** in `src/extension.ts`: +```typescript +context.subscriptions.push( + vscode.commands.registerCommand( + 'postgres-explorer.myCommand', + myCommandHandler + ) +); +``` + +### Adding a New AI Provider + +Implement the `AIProvider` interface in `src/services/AIService.ts`: + +```typescript +interface AIProvider { + generateQuery(prompt: string, context: SchemaContext): Promise; + optimizeQuery(query: string, context: SchemaContext): Promise; + explainError(error: string, query: string): Promise; +} +``` + +### Extending the Tree View + +Add new node types in `DatabaseTreeProvider.ts`: + +```typescript +class MyCustomNode extends TreeNode { + contextValue = 'myCustomNode'; + + async getChildren(): Promise { + // Return child nodes + } +} +``` + +--- + +## Performance Considerations + +### Connection Pool Sizing + +Default pool configuration: +```typescript +{ + max: 10, // Max connections per pool + idleTimeoutMillis: 30000, + connectionTimeoutMillis: 2000 +} +``` + +### Result Set Limits + +- **Backend**: 10,000 rows (hard limit in `NotebookKernel.ts`) +- **Frontend**: 200 rows/chunk (configurable in `renderer_v2.ts`) + +### Memory Management + +- Pools automatically release idle connections +- Session clients closed on notebook disposal +- Result truncation prevents OOM errors + +--- + +## Security + +### Password Storage + +```mermaid +graph LR + USER[User Input] --> |Plain text| EXT[Extension] + EXT --> |Encrypt| SS[SecretStorage API] + SS --> |Store| KEYCHAIN[OS Keychain] + + KEYCHAIN -.Retrieve.-> SS + SS -.Decrypt.-> EXT + EXT -.Use.-> PG[(PostgreSQL)] +``` + +- Passwords stored via VS Code `SecretStorage` API +- Encrypted by OS keychain (Keychain on macOS, Credential Manager on Windows, libsecret on Linux) +- Never stored in plain text + +### SSH Tunnels + +- Supports SSH key-based authentication +- Tunnels created per connection +- Automatically cleaned up on disconnect + +--- + +## Testing Strategy + +### Unit Tests +- `src/test/unit/` - Service layer tests +- Mock `pg.Pool` and `pg.Client` +- Test connection lifecycle + +### Integration Tests +- Require local PostgreSQL instance +- Test actual query execution +- Verify result rendering + +### Manual Testing +- Use test database with large tables +- Verify infinite scrolling +- Test AI features with real providers + +--- + +## Future Architecture Improvements + +1. **WebSocket for Real-time Updates** - Replace polling in dashboard +2. **Worker Threads for Large Exports** - Offload CSV/Excel generation +3. **Query Result Streaming** - Stream rows instead of loading all at once +4. **Distributed Tracing** - Add telemetry for performance monitoring +5. **Plugin System** - Allow third-party extensions + +--- + +## References + +- [VS Code Extension API](https://code.visualstudio.com/api) +- [node-postgres Documentation](https://node-postgres.com/) +- [Chart.js Documentation](https://www.chartjs.org/) +- [PostgreSQL System Catalogs](https://www.postgresql.org/docs/current/catalogs.html) diff --git a/docs/CONTRIBUTING.md b/docs/CONTRIBUTING.md new file mode 100644 index 0000000..4800a13 --- /dev/null +++ b/docs/CONTRIBUTING.md @@ -0,0 +1,478 @@ +# Contributing to PgStudio + +Thank you for your interest in contributing to PgStudio! This guide will help you get started with development. + +--- + +## Table of Contents + +- [Development Setup](#development-setup) +- [Project Structure](#project-structure) +- [Code Style Guide](#code-style-guide) +- [Testing Guidelines](#testing-guidelines) +- [Pull Request Process](#pull-request-process) +- [Commit Message Format](#commit-message-format) + +--- + +## Development Setup + +### Prerequisites + +- **Node.js** >= 18.x +- **npm** >= 9.x +- **VS Code** >= 1.80.0 +- **PostgreSQL** >= 12.x (for testing) + +### Initial Setup + +```bash +# Clone the repository +git clone https://github.com/dev-asterix/yape.git +cd yape + +# Install dependencies +npm install + +# Compile TypeScript +npm run compile + +# Watch mode for development +npm run watch +``` + +### Running the Extension + +1. Open the project in VS Code +2. Press `F5` to launch Extension Development Host +3. The extension will be loaded in a new VS Code window + +### Project Scripts + +```bash +npm run compile # Compile TypeScript +npm run watch # Watch mode (auto-compile on save) +npm run test # Run unit tests +npm run lint # Run ESLint +npm run esbuild-renderer # Bundle renderer for production +``` + +--- + +## Project Structure + +``` +yape/ +├── src/ +│ ├── extension.ts # Extension entry point +│ ├── commands/ # Command implementations +│ │ ├── connections.ts # Connection management +│ │ ├── tables.ts # Table operations +│ │ ├── views.ts # View operations +│ │ ├── functions.ts # Function operations +│ │ ├── fdw.ts # Foreign Data Wrapper ops +│ │ ├── ai.ts # AI commands +│ │ └── helper.ts # Shared utilities +│ ├── providers/ # VS Code providers +│ │ ├── DatabaseTreeProvider.ts +│ │ ├── SqlCompletionProvider.ts +│ │ ├── NotebookKernel.ts +│ │ └── DashboardPanel.ts +│ ├── services/ # Core services +│ │ ├── ConnectionManager.ts # Connection pooling +│ │ ├── SecretStorageService.ts +│ │ ├── SSHService.ts +│ │ ├── AIService.ts +│ │ ├── HistoryService.ts +│ │ ├── ErrorService.ts +│ │ └── DbObjectService.ts +│ ├── renderer_v2.ts # Notebook renderer +│ ├── renderer/ # Renderer modules +│ │ ├── components/ui.ts +│ │ └── features/ +│ │ ├── export.ts +│ │ └── ai.ts +│ ├── common/ # Shared types +│ │ └── types.ts +│ └── test/ # Tests +│ └── unit/ +├── docs/ # Documentation +├── package.json # Extension manifest +└── tsconfig.json # TypeScript config +``` + +--- + +## Code Style Guide + +### TypeScript Conventions + +#### 1. **Strict Typing** +Always use explicit types. Avoid `any` unless absolutely necessary. + +```typescript +// ✅ Good +async function getTableData( + client: PoolClient, + schema: string, + table: string +): Promise { + return await client.query('SELECT * FROM $1.$2', [schema, table]); +} + +// ❌ Bad +async function getTableData(client: any, schema: any, table: any): Promise { + return await client.query('SELECT * FROM $1.$2', [schema, table]); +} +``` + +#### 2. **Naming Conventions** + +| Type | Convention | Example | +|------|------------|---------| +| Classes | PascalCase | `ConnectionManager` | +| Interfaces | PascalCase with `I` prefix (optional) | `ConnectionConfig` | +| Functions | camelCase | `getPooledClient` | +| Constants | UPPER_SNAKE_CASE | `MAX_ROWS` | +| Private members | camelCase with `_` prefix | `_pools` | + +#### 3. **Async/Await** +Prefer `async/await` over `.then()` chains. + +```typescript +// ✅ Good +async function fetchData() { + try { + const result = await client.query('SELECT ...'); + return result.rows; + } catch (error) { + ErrorService.getInstance().handleError(error); + } +} + +// ❌ Bad +function fetchData() { + return client.query('SELECT ...') + .then(result => result.rows) + .catch(error => ErrorService.getInstance().handleError(error)); +} +``` + +--- + +### Connection Management Best Practices + +#### 1. **Always Use Pooling** + +```typescript +// ✅ Good - Pooled client (auto-released) +const client = await ConnectionManager.getInstance().getPooledClient(config); +try { + await client.query('SELECT ...'); +} finally { + client.release(); // CRITICAL: Always release in finally +} + +// ❌ Bad - Direct client creation +const client = new Client(config); +await client.connect(); +await client.query('SELECT ...'); +await client.end(); // May not execute if error occurs +``` + +#### 2. **Session Clients for Notebooks** + +```typescript +// ✅ Good - Session client for stateful operations +const client = await ConnectionManager.getInstance() + .getSessionClient(config, notebook.uri.toString()); + +// Client persists across cells +// Automatically closed when notebook closes +``` + +#### 3. **Error Handling** + +```typescript +// ✅ Good - Centralized error handling +try { + await client.query('...'); +} catch (error) { + ErrorService.getInstance().handleError(error, { + context: 'Table Operations', + operation: 'INSERT', + table: tableName + }); + throw error; // Re-throw if caller needs to handle +} +``` + +--- + +### SQL Query Patterns + +#### 1. **Parameterized Queries** +Always use parameterized queries to prevent SQL injection. + +```typescript +// ✅ Good +await client.query( + 'SELECT * FROM $1.$2 WHERE id = $3', + [schema, table, id] +); + +// ❌ Bad - SQL injection risk +await client.query( + `SELECT * FROM ${schema}.${table} WHERE id = ${id}` +); +``` + +#### 2. **Identifier Quoting** +Use double quotes for identifiers to handle special characters. + +```typescript +// ✅ Good +const query = `SELECT * FROM "${schema}"."${table}"`; + +// Handles schemas/tables with spaces, uppercase, etc. +``` + +--- + +### Error Handling Patterns + +#### 1. **Service Layer Errors** + +```typescript +export class MyService { + private static instance: MyService; + + public static getInstance(): MyService { + if (!MyService.instance) { + MyService.instance = new MyService(); + } + return MyService.instance; + } + + async performOperation(): Promise { + try { + // Operation logic + } catch (error) { + ErrorService.getInstance().handleError(error, { + context: 'MyService', + operation: 'performOperation' + }); + throw error; + } + } +} +``` + +#### 2. **Command Handler Errors** + +```typescript +export async function myCommandHandler(node: TreeNode) { + try { + const client = await ConnectionManager.getInstance() + .getPooledClient(node.connection); + try { + await client.query('...'); + vscode.window.showInformationMessage('Success!'); + } finally { + client.release(); + } + } catch (error) { + vscode.window.showErrorMessage(`Failed: ${error.message}`); + } +} +``` + +--- + +## Testing Guidelines + +### Unit Tests + +Located in `src/test/unit/`. Use Mocha + Chai. + +```typescript +import { expect } from 'chai'; +import * as sinon from 'sinon'; + +describe('ConnectionManager', () => { + let sandbox: sinon.SinonSandbox; + + beforeEach(() => { + sandbox = sinon.createSandbox(); + }); + + afterEach(() => { + sandbox.restore(); + }); + + it('should create a new pool if one does not exist', async () => { + const manager = ConnectionManager.getInstance(); + const config = { /* ... */ }; + + const poolStub = { + connect: sandbox.stub().resolves({ release: sandbox.stub() }), + on: sandbox.stub(), + end: sandbox.stub().resolves() + }; + + sandbox.stub(require('pg'), 'Pool').returns(poolStub); + + const client = await manager.getPooledClient(config); + + expect(poolStub.connect.calledOnce).to.be.true; + expect(client).to.exist; + }); +}); +``` + +### Integration Tests + +Require a local PostgreSQL instance. Use environment variables for configuration: + +```bash +export PGHOST=localhost +export PGPORT=5432 +export PGUSER=postgres +export PGPASSWORD=postgres +export PGDATABASE=test_db +``` + +--- + +## Pull Request Process + +### 1. **Fork and Branch** + +```bash +# Fork the repository on GitHub +# Clone your fork +git clone https://github.com/YOUR_USERNAME/yape.git + +# Create a feature branch +git checkout -b feature/my-new-feature +``` + +### 2. **Make Changes** + +- Follow the code style guide +- Add tests for new functionality +- Update documentation if needed + +### 3. **Test Your Changes** + +```bash +npm run compile +npm run test +npm run lint +``` + +### 4. **Commit** + +Follow the [commit message format](#commit-message-format). + +### 5. **Push and Create PR** + +```bash +git push origin feature/my-new-feature +``` + +Then create a Pull Request on GitHub with: +- Clear description of changes +- Reference to related issues +- Screenshots/GIFs for UI changes + +### 6. **Code Review** + +- Address reviewer feedback +- Keep commits clean (squash if needed) +- Ensure CI passes + +--- + +## Commit Message Format + +We follow the [Conventional Commits](https://www.conventionalcommits.org/) specification. + +### Format + +``` +(): + + + +