diff --git a/.gitignore b/.gitignore index ce6f018..1625664 100644 --- a/.gitignore +++ b/.gitignore @@ -25,6 +25,7 @@ venv/ ENV/ config.local.js *.local.json +package-lock.json # Logs logs diff --git a/CLAUDE.md b/CLAUDE.md index af0a823..4606323 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -8,13 +8,15 @@ DBHub is a Universal Database Gateway implementing the Model Context Protocol (M ## Commands -- Build: `pnpm run build` - Compiles TypeScript to JavaScript using tsup -- Start: `pnpm run start` - Runs the compiled server -- Dev: `pnpm run dev` - Runs server with tsx (no compilation needed) -- Test: `pnpm test` - Run all tests -- Test Watch: `pnpm test:watch` - Run tests in watch mode -- Integration Tests: `pnpm test:integration` - Run database integration tests (requires Docker) -- Pre-commit: `./scripts/setup-husky.sh` - Setup git hooks for automated testing +- **Build**: `pnpm run build` - Compiles TypeScript to JavaScript using tsup +- **Start**: `pnpm run start` - Runs the compiled server +- **Dev**: `pnpm run dev` - Runs server with tsx (no compilation needed) +- **Cross-platform Dev**: `pnpm run crossdev` - Cross-platform development with tsx +- **Test**: `pnpm test` - Run all tests with Vitest +- **Test Watch**: `pnpm test:watch` - Run tests in watch mode +- **Integration Tests**: `pnpm test:integration` - Run database integration tests (requires Docker) +- **Pre-commit**: `./scripts/setup-husky.sh` - Setup git hooks for automated testing +- **Pre-commit Hook**: `pnpm run pre-commit` - Run lint-staged checks ## Architecture Overview @@ -49,6 +51,7 @@ Key architectural patterns: - **Connector Registry**: Dynamic registration system for database connectors - **Transport Abstraction**: Support for both stdio (desktop tools) and HTTP (network clients) - **Resource/Tool/Prompt Handlers**: Clean separation of MCP protocol concerns +- **Multi-Database Support**: Simultaneous connections to multiple databases with isolated contexts - **Integration Test Base**: Shared test utilities for consistent connector testing ## Environment @@ -63,6 +66,43 @@ Key architectural patterns: - Demo mode: Use `--demo` flag for bundled SQLite employee database - Read-only mode: Use `--readonly` flag to restrict to read-only SQL operations +## Multi-Database Support + +DBHub supports connecting to multiple databases simultaneously: + +### Configuration +- **Single Database**: Use `DSN` environment variable or `--dsn` command line argument +- **Multiple Databases**: Use `DSN_dev`, `DSN_test`, etc. environment variables + +### Usage Examples + +```bash +# Single database (backward compatible) +export DSN="postgres://user:pass@localhost:5432/mydb" + +# Multiple databases +export DSN_dev="postgres://user:pass@localhost:5432/db1" +export DSN_test="mysql://user:pass@localhost:3306/db2" +export DSN_prod="sqlite:///path/to/database.db" +``` + +### HTTP Transport Endpoints +When using HTTP transport (`--transport=http`), multiple endpoints are available: + +- `http://localhost:8080/message` - Default database (first configured) +- `http://localhost:8080/message/{databaseId}` - Specific database (e.g., `http://localhost:8080/message/db1`) + +### STDIO Transport +- STDIO transport uses the default database +- Available databases are listed in startup messages +- Use HTTP transport for full multi-database access + +### Database Context +All MCP tools, resources, and prompts support database-specific operations: +- Tools: `execute_sql_{databaseId}` +- Resources: Database-specific schema exploration +- Prompts: `generate_sql_{databaseId}`, `explain_db_{databaseId}` + ## Database Connectors - Add new connectors in `src/connectors/{db-type}/index.ts` @@ -78,12 +118,19 @@ Key architectural patterns: ## Testing Approach -- Unit tests for individual components and utilities -- Integration tests using Testcontainers for real database testing -- All connectors have comprehensive integration test coverage -- Pre-commit hooks run related tests automatically -- Test specific databases: `pnpm test src/connectors/__tests__/{db-type}.integration.test.ts` -- SSH tunnel tests: `pnpm test postgres-ssh-simple.integration.test.ts` +- **Unit Tests**: Individual components and utilities using Vitest +- **Integration Tests**: Real database testing using Testcontainers with Docker +- **Test Coverage**: All connectors have comprehensive integration test coverage +- **Pre-commit Hooks**: Automatic test execution via lint-staged +- **Test Specific Databases**: + - PostgreSQL: `pnpm test src/connectors/__tests__/postgres.integration.test.ts` + - MySQL: `pnpm test src/connectors/__tests__/mysql.integration.test.ts` + - MariaDB: `pnpm test src/connectors/__tests__/mariadb.integration.test.ts` + - SQL Server: `pnpm test src/connectors/__tests__/sqlserver.integration.test.ts` + - SQLite: `pnpm test src/connectors/__tests__/sqlite.integration.test.ts` + - SSH Tunnel: `pnpm test src/connectors/__tests__/postgres-ssh.integration.test.ts` + - JSON RPC: `pnpm test src/__tests__/json-rpc-integration.test.ts` +- **Test Utilities**: Shared integration test base in `src/connectors/__tests__/shared/integration-test-base.ts` ## SSH Tunnel Support @@ -99,6 +146,26 @@ DBHub supports SSH tunnels for secure database connections through bastion hosts - Default SSH key detection (tries `~/.ssh/id_rsa`, `~/.ssh/id_ed25519`, etc.) - Tunnel lifecycle managed by `ConnectorManager` +## Development Environment + +- **TypeScript**: Strict mode enabled with ES2020 target +- **Module System**: ES modules with `.js` extension in imports +- **Package Manager**: pnpm for dependency management +- **Build Tool**: tsup for TypeScript compilation +- **Test Framework**: Vitest for unit and integration testing +- **Development Runtime**: tsx for development without compilation + +## Key Architectural Patterns + +- **Connector Registry**: Dynamic registration system for database connectors with automatic DSN detection +- **Transport Abstraction**: Support for both stdio (desktop tools) and HTTP (network clients) with CORS protection +- **Resource/Tool/Prompt Handlers**: Clean separation of MCP protocol concerns +- **Multi-Database Management**: Simultaneous connections to multiple databases with database ID-based routing +- **Database Context Propagation**: Consistent database ID flow through all MCP handlers +- **SSH Tunnel Integration**: Automatic tunnel establishment when SSH config detected +- **Singleton Manager**: `ConnectorManager` provides unified interface across all database operations +- **Integration Test Base**: Shared test utilities for consistent connector testing + ## Code Style - TypeScript with strict mode enabled diff --git a/README.md b/README.md index c9421f0..4204332 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -> [!NOTE] +> [!NOTE] > Brought to you by [Bytebase](https://www.bytebase.com/), open-source database DevSecOps platform.

@@ -118,6 +118,33 @@ dbhub: - database ``` +**Multiple Database Support:** + +DBHub supports connecting to multiple databases simultaneously using .env file with volume mount: + +```bash +# Create .env file with multiple database configurations +cat > .env << EOF +DSN_dev=postgres://user:password@localhost:5432/db1 +DSN_test=mysql://user:password@localhost:3306/db2 +DSN_prod=sqlite:///path/to/database.db +EOF + +# Run container with .env file mounted +docker run --rm --init \ + --name dbhub \ + --publish 8080:8080 \ + --volume "$(pwd)/.env:/app/.env" \ + bytebase/dbhub \ + --transport http \ + --port 8080 +``` + +Available endpoints when using multiple databases: + +- `http://localhost:8080/message` - Default database (first configured) +- `http://localhost:8080/message/{databaseId}` - Specific database (e.g., `http://localhost:8080/message/dev`, `http://localhost:8080/message/test`) + ### NPM ```bash diff --git a/src/config/env.ts b/src/config/env.ts index c3f3dd5..fc411b9 100644 --- a/src/config/env.ts +++ b/src/config/env.ts @@ -235,6 +235,105 @@ export function resolveDSN(): { dsn: string; source: string; isDemo?: boolean } return null; } +/** + * Resolve multiple DSN configurations from environment variables + * Supports both single DSN (backward compatible) and multiple DSN_* formats + * Returns a map of database IDs to DSN strings + */ +export function resolveMultiDSN(): Map { + const multiDSN = new Map(); + + // Get command line arguments + const args = parseCommandLineArgs(); + + // Check for demo mode first (highest priority) + if (isDemoMode()) { + multiDSN.set("default", { + dsn: "sqlite:///:memory:", + source: "demo mode", + isDemo: true, + }); + return multiDSN; + } + + // 1. Check command line arguments for single DSN + if (args.dsn) { + multiDSN.set("default", { dsn: args.dsn, source: "command line argument" }); + return multiDSN; + } + + // 2. Check for multiple DSN configurations from environment variables + const dsnPattern = /^DSN_(\w+)$/; + let foundMultiDSN = false; + + // Check environment variables before loading .env + for (const [key, value] of Object.entries(process.env)) { + if (key === "DSN" && value) { + // Single DSN format (backward compatible) + multiDSN.set("default", { dsn: value, source: "environment variable" }); + } else if (dsnPattern.test(key) && value) { + // Multiple DSN format: DSN_dev=postgres://... + const match = key.match(dsnPattern); + if (match) { + const id = match[1]; + multiDSN.set(id, { dsn: value, source: "environment variable" }); + foundMultiDSN = true; + } + } + } + + // 3. If no multi-DSN found, check for individual DB parameters + if (multiDSN.size === 0) { + const envParamsResult = buildDSNFromEnvParams(); + if (envParamsResult) { + multiDSN.set("default", envParamsResult); + } + } + + // 4. Try loading from .env files if no DSNs found yet + if (multiDSN.size === 0) { + const loadedEnvFile = loadEnvFiles(); + + if (loadedEnvFile) { + // Check for single DSN in .env file + if (process.env.DSN) { + multiDSN.set("default", { dsn: process.env.DSN, source: `${loadedEnvFile} file` }); + } + + // Check for multiple DSN configurations in .env file + for (const [key, value] of Object.entries(process.env)) { + if (dsnPattern.test(key) && value) { + const match = key.match(dsnPattern); + if (match) { + const id = match[1]; + multiDSN.set(id, { dsn: value, source: `${loadedEnvFile} file` }); + foundMultiDSN = true; + } + } + } + + // Check for individual DB parameters from .env file + if (multiDSN.size === 0) { + const envFileParamsResult = buildDSNFromEnvParams(); + if (envFileParamsResult) { + multiDSN.set("default", { + dsn: envFileParamsResult.dsn, + source: `${loadedEnvFile} file (individual parameters)` + }); + } + } + } + } + + // If we found multiple DSN configurations but no default, use the first one as default + if (foundMultiDSN && !multiDSN.has("default") && multiDSN.size > 0) { + const firstEntry = Array.from(multiDSN.entries())[0]; + multiDSN.set("default", { ...firstEntry[1], source: `${firstEntry[1].source} (as default)` }); + } + + return multiDSN; +} + /** * Resolve transport type from command line args or environment variables * Returns 'stdio' or 'http' (streamable HTTP), with 'stdio' as the default diff --git a/src/connectors/interface.ts b/src/connectors/interface.ts index d5f1e37..113b58b 100644 --- a/src/connectors/interface.ts +++ b/src/connectors/interface.ts @@ -149,6 +149,24 @@ export interface Connector { executeSQL(sql: string, options: ExecuteOptions): Promise; } +/** + * Database connection configuration + */ +export interface DatabaseConnection { + /** Unique identifier for this database connection */ + id: string; + /** Database connector instance */ + connector: Connector; + /** Database connection string */ + dsn: string; + /** Source of the DSN configuration */ + source: string; + /** Whether this is a demo database */ + isDemo?: boolean; + /** SSH tunnel configuration if applicable */ + sshConfig?: any; +} + /** * Registry for available database connectors */ diff --git a/src/connectors/manager.ts b/src/connectors/manager.ts index 0eec459..f8579b0 100644 --- a/src/connectors/manager.ts +++ b/src/connectors/manager.ts @@ -1,4 +1,4 @@ -import { Connector, ConnectorType, ConnectorRegistry, ExecuteOptions } from "./interface.js"; +import { Connector, ConnectorType, ConnectorRegistry, ExecuteOptions, DatabaseConnection } from "./interface.js"; import { SSHTunnel } from "../utils/ssh-tunnel.js"; import { resolveSSHConfig, resolveMaxRows } from "../config/env.js"; import type { SSHTunnelConfig } from "../types/ssh.js"; @@ -7,20 +7,19 @@ import type { SSHTunnelConfig } from "../types/ssh.js"; let managerInstance: ConnectorManager | null = null; /** - * Manages database connectors and provides a unified interface to work with them + * Manages multiple database connectors and provides a unified interface to work with them */ export class ConnectorManager { - private activeConnector: Connector | null = null; - private connected = false; - private sshTunnel: SSHTunnel | null = null; - private originalDSN: string | null = null; + private connections: Map = new Map(); + private activeConnectionId: string = "default"; + private sshTunnels: Map = new Map(); private maxRows: number | null = null; constructor() { if (!managerInstance) { managerInstance = this; } - + // Initialize maxRows from command line arguments const maxRowsData = resolveMaxRows(); if (maxRowsData) { @@ -30,107 +29,194 @@ export class ConnectorManager { } /** - * Initialize and connect to the database using a DSN + * Initialize and connect to a database using a DSN with a specific ID */ - async connectWithDSN(dsn: string, initScript?: string): Promise { - // Store original DSN for reference - this.originalDSN = dsn; - + async connectWithDSN(dsn: string, id: string = "default", initScript?: string): Promise { // Check if SSH tunnel is needed const sshConfig = resolveSSHConfig(); let actualDSN = dsn; - + let sshTunnel: SSHTunnel | null = null; + if (sshConfig) { console.error(`SSH tunnel configuration loaded from ${sshConfig.source}`); - + // Parse DSN to get database host and port const url = new URL(dsn); const targetHost = url.hostname; const targetPort = parseInt(url.port) || this.getDefaultPort(dsn); - + // Create and establish SSH tunnel - this.sshTunnel = new SSHTunnel(); - const tunnelInfo = await this.sshTunnel.establish(sshConfig.config, { + sshTunnel = new SSHTunnel(); + const tunnelInfo = await sshTunnel.establish(sshConfig.config, { targetHost, targetPort, }); - + // Update DSN to use local tunnel endpoint url.hostname = '127.0.0.1'; url.port = tunnelInfo.localPort.toString(); actualDSN = url.toString(); - + console.error(`Database connection will use SSH tunnel through localhost:${tunnelInfo.localPort}`); + + // Store SSH tunnel for this connection + this.sshTunnels.set(id, sshTunnel); } // First try to find a connector that can handle this DSN - let connector = ConnectorRegistry.getConnectorForDSN(actualDSN); + const connectorType = ConnectorRegistry.getConnectorForDSN(actualDSN)?.id; - if (!connector) { + if (!connectorType) { throw new Error(`No connector found that can handle the DSN: ${actualDSN}`); } - this.activeConnector = connector; + // Create a new connector instance for this connection + // This ensures each database connection has its own isolated connector + const connector = await this.createConnectorInstance(connectorType); // Connect to the database through tunnel if applicable - await this.activeConnector.connect(actualDSN, initScript); - this.connected = true; + await connector.connect(actualDSN, initScript); + + // Store the connection + const connection: DatabaseConnection = { + id, + connector, + dsn: dsn, + source: "programmatic", + sshConfig: sshConfig?.config + }; + + this.connections.set(id, connection); + + // Set as active if this is the first connection + if (this.connections.size === 1) { + this.activeConnectionId = id; + } } /** * Initialize and connect to the database using a specific connector type */ - async connectWithType(connectorType: ConnectorType, dsn?: string): Promise { - // Get the connector from the registry - const connector = ConnectorRegistry.getConnector(connectorType); - - if (!connector) { - throw new Error(`Connector "${connectorType}" not found`); - } - - this.activeConnector = connector; + async connectWithType(connectorType: ConnectorType, dsn?: string, id: string = "default"): Promise { + // Create a new connector instance for this connection + const connector = await this.createConnectorInstance(connectorType); // Use provided DSN or get sample DSN const connectionString = dsn || connector.dsnParser.getSampleDSN(); // Connect to the database - await this.activeConnector.connect(connectionString); - this.connected = true; + await connector.connect(connectionString); + + // Store the connection + const connection: DatabaseConnection = { + id, + connector, + dsn: connectionString, + source: "programmatic" + }; + + this.connections.set(id, connection); + + // Set as active if this is the first connection + if (this.connections.size === 1) { + this.activeConnectionId = id; + } } /** - * Close the database connection + * Close all database connections */ async disconnect(): Promise { - if (this.activeConnector && this.connected) { - await this.activeConnector.disconnect(); - this.connected = false; + // Close all database connections + for (const [, connection] of this.connections) { + await connection.connector.disconnect(); + } + this.connections.clear(); + + // Close all SSH tunnels + for (const [, sshTunnel] of this.sshTunnels) { + await sshTunnel.close(); + } + this.sshTunnels.clear(); + + this.activeConnectionId = "default"; + } + + /** + * Close a specific database connection + */ + async disconnectConnection(id: string): Promise { + const connection = this.connections.get(id); + if (connection) { + await connection.connector.disconnect(); + this.connections.delete(id); } - - // Close SSH tunnel if it exists - if (this.sshTunnel) { - await this.sshTunnel.close(); - this.sshTunnel = null; + + // Close SSH tunnel for this connection if it exists + const sshTunnel = this.sshTunnels.get(id); + if (sshTunnel) { + await sshTunnel.close(); + this.sshTunnels.delete(id); + } + + // Update active connection if needed + if (this.activeConnectionId === id && this.connections.size > 0) { + this.activeConnectionId = Array.from(this.connections.keys())[0]; + } else if (this.connections.size === 0) { + this.activeConnectionId = "default"; } - - this.originalDSN = null; } /** - * Get the active connector + * Get a connector by ID */ - getConnector(): Connector { - if (!this.activeConnector) { - throw new Error("No active connector. Call connectWithDSN() or connectWithType() first."); + getConnector(id?: string): Connector { + const connectionId = id || this.activeConnectionId; + const connection = this.connections.get(connectionId); + + if (!connection) { + throw new Error(`No database connection found for ID: ${connectionId}. Available connections: ${Array.from(this.connections.keys()).join(', ')}`); } - return this.activeConnector; + + return connection.connector; } /** - * Check if there's an active connection + * Switch to a different database connection + */ + switchConnection(id: string): void { + if (!this.connections.has(id)) { + throw new Error(`Database connection not found: ${id}. Available connections: ${Array.from(this.connections.keys()).join(', ')}`); + } + this.activeConnectionId = id; + } + + /** + * Get the active connection ID + */ + getActiveConnectionId(): string { + return this.activeConnectionId; + } + + /** + * Get all available connection IDs + */ + getAvailableConnections(): string[] { + return Array.from(this.connections.keys()); + } + + /** + * Check if a specific connection exists + */ + hasConnection(id: string): boolean { + return this.connections.has(id); + } + + /** + * Check if there's any active connection */ isConnected(): boolean { - return this.connected; + return this.connections.size > 0; } /** @@ -148,9 +234,27 @@ export class ConnectorManager { } /** - * Get the current active connector instance + * Get a connector by ID * This is used by resource and tool handlers */ + static getConnector(id?: string): Connector { + // Try global instance first (for HTTP transport) + const globalInstance = (global as any).__dbhubConnectorManager; + if (globalInstance) { + return globalInstance.getConnector(id); + } + + // Fall back to singleton instance + if (!managerInstance) { + throw new Error("ConnectorManager not initialized"); + } + return managerInstance.getConnector(id); + } + + /** + * Get the active connector instance (for backward compatibility) + * This is used by resource and tool handlers that don't specify a database ID + */ static getCurrentConnector(): Connector { if (!managerInstance) { throw new Error("ConnectorManager not initialized"); @@ -179,7 +283,64 @@ export class ConnectorManager { } return managerInstance.getExecuteOptions(); } + + /** + * Get all available connection IDs + */ + static getAvailableConnections(): string[] { + if (!managerInstance) { + throw new Error("ConnectorManager not initialized"); + } + return managerInstance.getAvailableConnections(); + } + + /** + * Get the active connection ID + */ + static getActiveConnectionId(): string { + if (!managerInstance) { + throw new Error("ConnectorManager not initialized"); + } + return managerInstance.getActiveConnectionId(); + } + + /** + * Switch to a different database connection + */ + static switchConnection(id: string): void { + if (!managerInstance) { + throw new Error("ConnectorManager not initialized"); + } + managerInstance.switchConnection(id); + } + /** + * Create a new connector instance for a specific connector type + * This ensures each database connection has its own isolated connector instance + */ + private async createConnectorInstance(connectorType: ConnectorType): Promise { + // Import the connector modules dynamically to avoid circular dependencies + switch (connectorType) { + case "postgres": + const { PostgresConnector } = await import("./postgres/index.js"); + return new PostgresConnector(); + case "mysql": + const { MySQLConnector } = await import("./mysql/index.js"); + return new MySQLConnector(); + case "mariadb": + const { MariaDBConnector } = await import("./mariadb/index.js"); + return new MariaDBConnector(); + case "sqlserver": + const { SQLServerConnector } = await import("./sqlserver/index.js"); + return new SQLServerConnector(); + case "sqlite": + const { SQLiteConnector } = await import("./sqlite/index.js"); + return new SQLiteConnector(); + default: + throw new Error(`Unsupported connector type: ${connectorType}`); + } + } + /** * Get default port for a database based on DSN protocol */ diff --git a/src/prompts/db-explainer.ts b/src/prompts/db-explainer.ts index caf5f98..5583a19 100644 --- a/src/prompts/db-explainer.ts +++ b/src/prompts/db-explainer.ts @@ -23,7 +23,8 @@ export async function dbExplainerPromptHandler( schema?: string; table?: string; }, - _extra: any + _extra: any, + databaseId?: string ): Promise<{ messages: { role: "assistant" | "user"; @@ -39,7 +40,8 @@ export async function dbExplainerPromptHandler( [key: string]: unknown; }> { try { - const connector = ConnectorManager.getCurrentConnector(); + // Get connector based on database ID + const connector = databaseId ? ConnectorManager.getConnector(databaseId) : ConnectorManager.getCurrentConnector(); // Verify schema exists if provided if (schema) { diff --git a/src/prompts/index.ts b/src/prompts/index.ts index ea9e282..14a708a 100644 --- a/src/prompts/index.ts +++ b/src/prompts/index.ts @@ -5,20 +5,24 @@ import { dbExplainerPromptHandler, dbExplainerSchema } from "./db-explainer.js"; /** * Register all prompt handlers with the MCP server */ -export function registerPrompts(server: McpServer): void { +export function registerPrompts(server: McpServer, databaseId?: string): void { + // Build prompt names with optional database ID suffix + const sqlGeneratorName = databaseId ? `generate_sql_${databaseId}` : "generate_sql"; + const dbExplainerName = databaseId ? `explain_db_${databaseId}` : "explain_db"; + // Register SQL Generator prompt server.prompt( - "generate_sql", - "Generate SQL queries from natural language descriptions", + sqlGeneratorName, + `Generate SQL queries from natural language descriptions for the ${databaseId ? databaseId : 'current'} database`, sqlGeneratorSchema, - sqlGeneratorPromptHandler + (args, extra) => sqlGeneratorPromptHandler(args, extra, databaseId) ); // Register Database Explainer prompt server.prompt( - "explain_db", - "Get explanations about database tables, columns, and structures", + dbExplainerName, + `Get explanations about database tables, columns, and structures for the ${databaseId ? databaseId : 'current'} database`, dbExplainerSchema, - dbExplainerPromptHandler + (args, extra) => dbExplainerPromptHandler(args, extra, databaseId) ); } diff --git a/src/prompts/sql-generator.ts b/src/prompts/sql-generator.ts index 4cab520..b7765db 100644 --- a/src/prompts/sql-generator.ts +++ b/src/prompts/sql-generator.ts @@ -24,11 +24,12 @@ export async function sqlGeneratorPromptHandler( description: string; schema?: string; }, - _extra: any + _extra: any, + databaseId?: string ) { try { - // Get current connector to determine dialect - const connector = ConnectorManager.getCurrentConnector(); + // Get connector based on database ID + const connector = databaseId ? ConnectorManager.getConnector(databaseId) : ConnectorManager.getCurrentConnector(); // Determine SQL dialect from connector automatically let sqlDialect: SQLDialect; @@ -130,6 +131,11 @@ export async function sqlGeneratorPromptHandler( "SELECT u.name, COUNT(o.id) FROM users u JOIN orders o ON u.id = o.user_id GROUP BY u.name HAVING COUNT(o.id) > 5", "SELECT product_name, price FROM products WHERE price > (SELECT AVG(price) FROM products)", ], + mariadb: [ + "SELECT * FROM users WHERE created_at > NOW() - INTERVAL 1 DAY", + "SELECT u.name, COUNT(o.id) FROM users u JOIN orders o ON u.id = o.user_id GROUP BY u.name HAVING COUNT(o.id) > 5", + "SELECT product_name, price FROM products WHERE price > (SELECT AVG(price) FROM products)", + ], mssql: [ "SELECT * FROM users WHERE created_at > DATEADD(day, -1, GETDATE())", "SELECT u.name, COUNT(o.id) FROM users u JOIN orders o ON u.id = o.user_id GROUP BY u.name HAVING COUNT(o.id) > 5", @@ -145,19 +151,6 @@ export async function sqlGeneratorPromptHandler( // Build a prompt that would help generate the SQL // In a real implementation, this would call an AI model const schemaInfo = schema ? `in schema '${schema}'` : "across all schemas"; - const prompt = ` -Generate a ${sqlDialect} SQL query based on this description: "${description}" - -${schemaContext} -Working ${schemaInfo} - -The query should: -1. Be written for ${sqlDialect} dialect -2. Use only the available tables and columns -3. Prioritize readability -4. Include appropriate comments -5. Be compatible with ${sqlDialect} syntax -`; // In a real implementation, this would be the result from an AI model call // For this demo, we'll generate a simple SQL query based on the description diff --git a/src/resources/index.ts b/src/resources/index.ts index fe9a534..a5f3e4d 100644 --- a/src/resources/index.ts +++ b/src/resources/index.ts @@ -15,22 +15,26 @@ export { proceduresResourceHandler, procedureDetailResourceHandler } from "./pro /** * Register all resource handlers with the MCP server */ -export function registerResources(server: McpServer): void { +export function registerResources(server: McpServer, databaseId?: string): void { // Resource for listing all schemas - server.resource("schemas", "db://schemas", schemasResourceHandler); + server.resource("schemas", "db://schemas", (uri, variables, extra) => + schemasResourceHandler(uri, databaseId, extra) + ); // Allow listing tables within a specific schema server.resource( "tables_in_schema", new ResourceTemplate("db://schemas/{schemaName}/tables", { list: undefined }), - tablesResourceHandler + (uri, variables, extra) => + tablesResourceHandler(uri, variables, databaseId, extra) ); // Resource for getting table structure within a specific database schema server.resource( "table_structure_in_schema", new ResourceTemplate("db://schemas/{schemaName}/tables/{tableName}", { list: undefined }), - tableStructureResourceHandler + (uri, variables, extra) => + tableStructureResourceHandler(uri, variables, databaseId, extra) ); // Resource for getting indexes for a table within a specific database schema @@ -39,14 +43,16 @@ export function registerResources(server: McpServer): void { new ResourceTemplate("db://schemas/{schemaName}/tables/{tableName}/indexes", { list: undefined, }), - indexesResourceHandler + (uri, variables, extra) => + indexesResourceHandler(uri, variables, databaseId, extra) ); // Resource for listing stored procedures within a schema server.resource( "procedures_in_schema", new ResourceTemplate("db://schemas/{schemaName}/procedures", { list: undefined }), - proceduresResourceHandler + (uri, variables, extra) => + proceduresResourceHandler(uri, variables, databaseId, extra) ); // Resource for getting procedure detail within a schema @@ -55,6 +61,7 @@ export function registerResources(server: McpServer): void { new ResourceTemplate("db://schemas/{schemaName}/procedures/{procedureName}", { list: undefined, }), - procedureDetailResourceHandler + (uri, variables, extra) => + procedureDetailResourceHandler(uri, variables, databaseId, extra) ); } diff --git a/src/resources/indexes.ts b/src/resources/indexes.ts index b5ab080..06229e7 100644 --- a/src/resources/indexes.ts +++ b/src/resources/indexes.ts @@ -8,8 +8,8 @@ import { * Indexes resource handler * Returns information about indexes on a table */ -export async function indexesResourceHandler(uri: URL, variables: any, _extra: any) { - const connector = ConnectorManager.getCurrentConnector(); +export async function indexesResourceHandler(uri: URL, variables: any, databaseId?: string, _extra?: any) { + const connector = databaseId ? ConnectorManager.getConnector(databaseId) : ConnectorManager.getCurrentConnector(); // Extract schema and table names from URL variables const schemaName = @@ -62,6 +62,7 @@ export async function indexesResourceHandler(uri: URL, variables: any, _extra: a schema: schemaName, indexes: indexes, count: indexes.length, + database: databaseId || "default" }; // Use the utility to create a standardized response diff --git a/src/resources/procedures.ts b/src/resources/procedures.ts index 71c212e..c497816 100644 --- a/src/resources/procedures.ts +++ b/src/resources/procedures.ts @@ -8,8 +8,8 @@ import { * Stored procedures/functions resource handler * Returns a list of all stored procedures/functions in the database or within a specific schema */ -export async function proceduresResourceHandler(uri: URL, variables: any, _extra: any) { - const connector = ConnectorManager.getCurrentConnector(); +export async function proceduresResourceHandler(uri: URL, variables: any, databaseId?: string, _extra?: any) { + const connector = databaseId ? ConnectorManager.getConnector(databaseId) : ConnectorManager.getCurrentConnector(); // Extract the schema name from URL variables if present const schemaName = @@ -40,6 +40,7 @@ export async function proceduresResourceHandler(uri: URL, variables: any, _extra procedures: procedureNames, count: procedureNames.length, schema: schemaName, + database: databaseId || "default" }; // Use the utility to create a standardized response @@ -57,8 +58,8 @@ export async function proceduresResourceHandler(uri: URL, variables: any, _extra * Stored procedure/function details resource handler * Returns details for a specific stored procedure/function */ -export async function procedureDetailResourceHandler(uri: URL, variables: any, _extra: any) { - const connector = ConnectorManager.getCurrentConnector(); +export async function procedureDetailResourceHandler(uri: URL, variables: any, databaseId?: string, _extra?: any) { + const connector = databaseId ? ConnectorManager.getConnector(databaseId) : ConnectorManager.getCurrentConnector(); // Extract parameters from URL variables const schemaName = @@ -105,6 +106,7 @@ export async function procedureDetailResourceHandler(uri: URL, variables: any, _ returnType: procedureDetails.return_type, definition: procedureDetails.definition, schema: schemaName, + database: databaseId || "default" }; // Use the utility to create a standardized response diff --git a/src/resources/schema.ts b/src/resources/schema.ts index dc922f7..65766cd 100644 --- a/src/resources/schema.ts +++ b/src/resources/schema.ts @@ -9,8 +9,8 @@ import { * Schema resource handler * Returns schema information for a specific table, optionally within a specific database schema */ -export async function tableStructureResourceHandler(uri: URL, variables: Variables, _extra: any) { - const connector = ConnectorManager.getCurrentConnector(); +export async function tableStructureResourceHandler(uri: URL, variables: Variables, databaseId?: string, _extra?: any) { + const connector = databaseId ? ConnectorManager.getConnector(databaseId) : ConnectorManager.getCurrentConnector(); // Handle tableName which could be a string or string array from URL template const tableName = Array.isArray(variables.tableName) @@ -65,6 +65,7 @@ export async function tableStructureResourceHandler(uri: URL, variables: Variabl schema: schemaName, columns: formattedColumns, count: formattedColumns.length, + database: databaseId || "default" }; // Use the utility to create a standardized response diff --git a/src/resources/schemas.ts b/src/resources/schemas.ts index 938d528..e83ba27 100644 --- a/src/resources/schemas.ts +++ b/src/resources/schemas.ts @@ -8,8 +8,8 @@ import { * Schemas resource handler * Returns a list of all schemas in the database */ -export async function schemasResourceHandler(uri: URL, _extra: any) { - const connector = ConnectorManager.getCurrentConnector(); +export async function schemasResourceHandler(uri: URL, databaseId?: string, _extra?: any) { + const connector = databaseId ? ConnectorManager.getConnector(databaseId) : ConnectorManager.getCurrentConnector(); try { const schemas = await connector.getSchemas(); @@ -18,6 +18,7 @@ export async function schemasResourceHandler(uri: URL, _extra: any) { const responseData = { schemas: schemas, count: schemas.length, + database: databaseId || "default" }; // Use the utility to create a standardized response diff --git a/src/resources/tables.ts b/src/resources/tables.ts index 2811e96..55ef163 100644 --- a/src/resources/tables.ts +++ b/src/resources/tables.ts @@ -8,8 +8,8 @@ import { * Tables resource handler * Returns a list of all tables in the database or within a specific schema */ -export async function tablesResourceHandler(uri: URL, variables: any, _extra: any) { - const connector = ConnectorManager.getCurrentConnector(); +export async function tablesResourceHandler(uri: URL, variables: any, databaseId?: string, _extra?: any) { + const connector = databaseId ? ConnectorManager.getConnector(databaseId) : ConnectorManager.getCurrentConnector(); // Extract the schema name from URL variables if present const schemaName = @@ -40,6 +40,7 @@ export async function tablesResourceHandler(uri: URL, variables: any, _extra: an tables: tableNames, count: tableNames.length, schema: schemaName, + database: databaseId || "default" }; // Use the utility to create a standardized response diff --git a/src/server.ts b/src/server.ts index 5cca648..d4ab3a6 100644 --- a/src/server.ts +++ b/src/server.ts @@ -8,7 +8,7 @@ import { fileURLToPath } from "url"; import { ConnectorManager } from "./connectors/manager.js"; import { ConnectorRegistry } from "./connectors/interface.js"; -import { resolveDSN, resolveTransport, resolvePort, isDemoMode, redactDSN, isReadOnlyMode, resolveId } from "./config/env.js"; +import { resolveMultiDSN, resolveTransport, resolvePort, isDemoMode, redactDSN, isReadOnlyMode, resolveId } from "./config/env.js"; import { getSqliteInMemorySetupSql } from "./config/demo-loader.js"; import { registerResources } from "./resources/index.js"; import { registerTools } from "./tools/index.js"; @@ -54,10 +54,10 @@ export async function main(): Promise { const idData = resolveId(); const id = idData?.id; - // Resolve DSN from command line args, environment variables, or .env files - const dsnData = resolveDSN(); + // Resolve DSNs from command line args, environment variables, or .env files + const dsnData = resolveMultiDSN(); - if (!dsnData) { + if (!dsnData || dsnData.size === 0) { const samples = ConnectorRegistry.getAllSampleDSNs(); const sampleFormats = Object.entries(samples) .map(([id, dsn]) => ` - ${id}: ${dsn}`) @@ -72,6 +72,10 @@ Please provide the DSN in one of these ways (in order of priority): 3. Environment variable: export DSN="your-connection-string" 4. .env file: DSN=your-connection-string +For multiple databases: +1. Environment variables: DSN_dev, DSN_test, etc. +2. .env file: DSN_dev=your-connection-string, DSN_test=another-connection-string + Example formats: ${sampleFormats} @@ -81,15 +85,15 @@ See documentation for more details on configuring database connections. } // Create MCP server factory function for HTTP transport - const createServer = () => { + const createServer = (databaseId?: string) => { const server = new McpServer({ name: SERVER_NAME, version: SERVER_VERSION, }); - // Register resources, tools, and prompts - registerResources(server); - registerTools(server, id); + // Register resources, tools, and prompts with optional database ID + registerResources(server, databaseId); + registerTools(server, databaseId); registerPrompts(server); return server; @@ -97,20 +101,29 @@ See documentation for more details on configuring database connections. // Create server factory function (will be used for both STDIO and HTTP transports) - // Create connector manager and connect to database + // Create connector manager and connect to databases const connectorManager = new ConnectorManager(); - console.error(`Connecting with DSN: ${redactDSN(dsnData.dsn)}`); - console.error(`DSN source: ${dsnData.source}`); - if (idData) { - console.error(`ID: ${idData.id} (from ${idData.source})`); - } - - // If in demo mode, load the employee database - if (dsnData.isDemo) { + + if (isDemoMode()) { + // If in demo mode, load the employee database const initScript = getSqliteInMemorySetupSql(); - await connectorManager.connectWithDSN(dsnData.dsn, initScript); + const databaseId = "default"; + const dsnInfo = dsnData.get(databaseId)!; + await connectorManager.connectWithDSN(dsnInfo.dsn, databaseId, initScript); } else { - await connectorManager.connectWithDSN(dsnData.dsn); + // Connect to all databases + for (const [databaseId, dsnInfo] of dsnData) { + console.error(`Connecting to database '${databaseId}' with DSN: ${redactDSN(dsnInfo.dsn)}`); + console.error(`DSN source: ${dsnInfo.source}`); + await connectorManager.connectWithDSN(dsnInfo.dsn, databaseId); + } + } + + // Store the connector manager instance globally so it can be accessed by all endpoints + (global as any).__dbhubConnectorManager = connectorManager; + + if (idData) { + console.error(`ID: ${idData.id} (from ${idData.source})`); } // Resolve transport type @@ -120,21 +133,23 @@ See documentation for more details on configuring database connections. // Print ASCII art banner with version and slogan const readonly = isReadOnlyMode(); - + // Collect active modes const activeModes: string[] = []; const modeDescriptions: string[] = []; - - if (dsnData.isDemo) { + + // Check if any database is in demo mode + const hasDemoMode = Array.from(dsnData.values()).some(dsnInfo => dsnInfo.isDemo); + if (hasDemoMode) { activeModes.push("DEMO"); modeDescriptions.push("using sample employee database"); } - + if (readonly) { activeModes.push("READ-ONLY"); modeDescriptions.push("only read only queries allowed"); } - + // Output mode information if (activeModes.length > 0) { console.error(`Running in ${activeModes.join(' and ')} mode - ${modeDescriptions.join(', ')}`); @@ -179,27 +194,28 @@ See documentation for more details on configuring database connections. res.status(200).send("OK"); }); - // Main endpoint for streamable HTTP transport - app.post("/message", async (req, res) => { - try { - // In stateless mode, create a new instance of transport and server for each request - // to ensure complete isolation. A single instance would cause request ID collisions - // when multiple clients connect concurrently. - const transport = new StreamableHTTPServerTransport({ - sessionIdGenerator: undefined, // Disable session management for stateless mode - enableJsonResponse: false // Use SSE streaming - }); - const server = createServer(); - - await server.connect(transport); - await transport.handleRequest(req, res, req.body); - } catch (error) { - console.error("Error handling request:", error); - if (!res.headersSent) { - res.status(500).json({ error: 'Internal server error' }); + + // Unified endpoints for all databases + for (const [databaseId] of dsnData) { + const path = databaseId === "default" ? "/message" : `/message/${databaseId}`; + app.post(path, async (req, res) => { + try { + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: undefined, + enableJsonResponse: false + }); + const server = createServer(databaseId); + + await server.connect(transport); + await transport.handleRequest(req, res, req.body); + } catch (error) { + console.error(`Error handling request for database '${databaseId}':`, error); + if (!res.headersSent) { + res.status(500).json({ error: 'Internal server error' }); + } } - } - }); + }); + } // Start the HTTP server @@ -208,13 +224,27 @@ See documentation for more details on configuring database connections. console.error(`Port source: ${portData.source}`); app.listen(port, '0.0.0.0', () => { console.error(`DBHub server listening at http://0.0.0.0:${port}`); - console.error(`Connect to MCP server at http://0.0.0.0:${port}/message`); + console.error(`Available database endpoints:`); + for (const [databaseId] of dsnData) { + const path = databaseId === "default" ? "/message" : `/message/${databaseId}`; + console.error(` - ${databaseId}: http://0.0.0.0:${port}${path}`); + } }); } else { // Set up STDIO transport const server = createServer(); const transport = new StdioServerTransport(); console.error("Starting with STDIO transport"); + + // Show available databases for STDIO mode + if (dsnData.size > 1) { + console.error("Available databases:"); + for (const [databaseId, dsnInfo] of dsnData) { + console.error(` - ${databaseId}: ${redactDSN(dsnInfo.dsn)}`); + } + console.error("Note: STDIO mode uses the default database. Use HTTP transport for multi-database access."); + } + await server.connect(transport); // Listen for SIGINT to gracefully shut down diff --git a/src/tools/execute-sql.ts b/src/tools/execute-sql.ts index 9aac8f5..5c35780 100644 --- a/src/tools/execute-sql.ts +++ b/src/tools/execute-sql.ts @@ -78,8 +78,8 @@ function areAllStatementsReadOnly(sql: string, connectorType: ConnectorType): bo * execute_sql tool handler * Executes a SQL query and returns the results */ -export async function executeSqlToolHandler({ sql }: { sql: string }, _extra: any) { - const connector = ConnectorManager.getCurrentConnector(); +export async function executeSqlToolHandler({ sql }: { sql: string }, _extra: any, databaseId?: string) { + const connector = databaseId ? ConnectorManager.getConnector(databaseId) : ConnectorManager.getCurrentConnector(); const executeOptions = ConnectorManager.getCurrentExecuteOptions(); try { @@ -90,7 +90,7 @@ export async function executeSqlToolHandler({ sql }: { sql: string }, _extra: an "READONLY_VIOLATION" ); } - + // Execute the SQL (single or multiple statements) if validation passed const result = await connector.executeSQL(sql, executeOptions); @@ -98,6 +98,7 @@ export async function executeSqlToolHandler({ sql }: { sql: string }, _extra: an const responseData = { rows: result.rows, count: result.rows.length, + database: databaseId || "default" }; return createToolSuccessResponse(responseData); diff --git a/src/tools/index.ts b/src/tools/index.ts index 5724130..62c3195 100644 --- a/src/tools/index.ts +++ b/src/tools/index.ts @@ -3,18 +3,18 @@ import { executeSqlToolHandler, executeSqlSchema } from "./execute-sql.js"; /** * Register all tool handlers with the MCP server * @param server - The MCP server instance - * @param id - Optional ID to suffix tool names (for Cursor multi-instance support) + * @param id - Optional database ID to suffix tool names (for multi-database support) */ export function registerTools(server: McpServer, id?: string): void { - // Build tool name with optional suffix + // Build tool name with optional database ID suffix const toolName = id ? `execute_sql_${id}` : "execute_sql"; // Tool to run a SQL query (read-only for safety) server.tool( toolName, - "Execute a SQL query on the current database", + `Execute a SQL query on the ${id ? id : 'current'} database`, executeSqlSchema, - executeSqlToolHandler + (args, extra) => executeSqlToolHandler(args, extra, id) ); }