UNPKG

@bytebase/dbhub

Version:

Token-efficient database MCP server for PostgreSQL, MySQL, MariaDB, SQL Server, Oracle, SQLite

313 lines (311 loc) • 9.83 kB
import { MissingDatabaseError, getDefaultSchema, getMySQLFamilyHealthCheck, getSchemas, getStoredProcedureDetail, getStoredProcedures, getTableComment, getTableIndexes, getTableSchema, getTables, getViews, isTiDBVersion, parseQueryResultSets, requireDatabaseInDSN, tableExists, withReadOnlyTransaction } from "./chunk-SJ5BRPQA.js"; import { CANCEL_QUERY_TIMEOUT_MS, clientQueryTimeoutMs, isClientSideTimeout } from "./chunk-5XSBOK3W.js"; import "./chunk-FU2ZJE4E.js"; import "./chunk-X2EC47W4.js"; import { SQLRowLimiter, closeQuietly } from "./chunk-X5OCFRNL.js"; import "./chunk-55X2VY5D.js"; import { ConnectorRegistry, SafeURL, obfuscateDSNPassword, splitSQLStatements } from "./chunk-IXC4NQGC.js"; // src/connectors/mysql/index.ts import mysql from "mysql2/promise"; var MySQLDSNParser = class { async parse(dsn, config) { const connectionTimeoutSeconds = config?.connectionTimeoutSeconds; const timezone = config?.timezone; const charset = config?.charset; const collation = config?.collation; if (!this.isValidDSN(dsn)) { const obfuscatedDSN = obfuscateDSNPassword(dsn); const expectedFormat = this.getSampleDSN(); throw new Error( `Invalid MySQL DSN format. Provided: ${obfuscatedDSN} Expected: ${expectedFormat}` ); } try { const url = new SafeURL(dsn); const database = url.pathname ? url.pathname.substring(1) : ""; requireDatabaseInDSN(database, dsn, "MySQL"); const config2 = { host: url.hostname, port: url.port ? parseInt(url.port) : 3306, database, user: url.username, password: url.password, multipleStatements: true, // Enable native multi-statement support supportBigNumbers: true // Return BIGINT as string when value exceeds Number.MAX_SAFE_INTEGER }; url.forEachSearchParam((value, key) => { if (key === "sslmode") { if (value === "disable") { config2.ssl = void 0; } else if (value === "require") { config2.ssl = { rejectUnauthorized: false }; } else { config2.ssl = {}; } } }); if (connectionTimeoutSeconds !== void 0) { config2.connectTimeout = connectionTimeoutSeconds * 1e3; } if (timezone !== void 0) { config2.timezone = timezone; } const charsetOrCollation = collation ?? charset; if (charsetOrCollation !== void 0) { config2.charset = charsetOrCollation; } if (url.password && url.password.includes("X-Amz-Credential")) { config2.authPlugins = { mysql_clear_password: () => () => { return Buffer.from(url.password + "\0"); } }; if (config2.ssl === void 0) { config2.ssl = { rejectUnauthorized: false }; } } return config2; } catch (error) { if (error instanceof MissingDatabaseError) { throw error; } throw new Error( `Failed to parse MySQL DSN: ${error instanceof Error ? error.message : String(error)}` ); } } getSampleDSN() { return "mysql://root:password@localhost:3306/mysql?sslmode=require"; } isValidDSN(dsn) { try { return dsn.startsWith("mysql://"); } catch (error) { return false; } } }; var MySQLConnector = class _MySQLConnector { constructor() { this.id = "mysql"; this.name = "MySQL"; this.dsnParser = new MySQLDSNParser(); this.pool = null; // Source ID is set by ConnectorManager after cloning this.sourceId = "default"; // TiDB speaks the MySQL protocol but rejects `START TRANSACTION READ ONLY` // unless tidb_enable_noop_functions is on. Detected once at connect time. this.supportsReadOnlyTransaction = true; } getId() { return this.sourceId; } clone() { return new _MySQLConnector(); } async connect(dsn, initScript, config) { try { const connectionOptions = await this.dsnParser.parse(dsn, config); this.pool = mysql.createPool(connectionOptions); if (config?.queryTimeoutSeconds !== void 0) { const queryTimeoutMs = Math.ceil(config.queryTimeoutSeconds * 1e3); this.queryTimeoutMs = queryTimeoutMs; const setServerLimit = `SET SESSION max_execution_time = ${queryTimeoutMs}`; this.pool.on( "connection", (connection) => { connection.query(setServerLimit, () => { }); } ); } const [rows] = await this.pool.query("SELECT VERSION() AS version"); this.supportsReadOnlyTransaction = !isTiDBVersion(rows[0]?.version); } catch (err) { if (this.pool) { const pool = this.pool; this.pool = null; await closeQuietly(() => pool.end()); } console.error("Failed to connect to MySQL database:", err); throw err; } } async disconnect() { if (this.pool) { await this.pool.end(); this.pool = null; } } /** * Row-array query adapter over the pool for the shared MySQL-family * introspection and health-check code. Throws when not connected, so the * shared functions never see a missing pool. mysql2 resolves to * [rows, fields]; only the rows are handed on. */ requireQuery() { if (!this.pool) { throw new Error("Not connected to database"); } const pool = this.pool; return async (sql, params) => { const [rows] = await (params === void 0 ? pool.query(sql) : pool.query(sql, params)); return rows; }; } async getSchemas() { return getSchemas(this.requireQuery()); } async getTables(schema) { return getTables(this.requireQuery(), schema); } async getViews(schema) { return getViews(this.requireQuery(), schema); } async tableExists(tableName, schema) { return tableExists(this.requireQuery(), tableName, schema); } async getTableIndexes(tableName, schema) { return getTableIndexes(this.requireQuery(), tableName, schema); } async getTableSchema(tableName, schema) { return getTableSchema(this.requireQuery(), tableName, schema); } async getTableComment(tableName, schema) { return getTableComment(this.requireQuery(), tableName, schema); } async getHealthCheck() { return getMySQLFamilyHealthCheck(this.requireQuery()); } async getStoredProcedures(schema, routineType) { return getStoredProcedures(this.requireQuery(), schema, routineType); } async getStoredProcedureDetail(procedureName, schema) { return getStoredProcedureDetail(this.requireQuery(), procedureName, schema); } async getDefaultSchema() { return getDefaultSchema(this.requireQuery()); } async executeSQL(sql, options, parameters) { if (!this.pool) { throw new Error("Not connected to database"); } const conn = await this.pool.getConnection(); const threadId = conn.threadId; let isConnectionDiscarded = false; const timeout = this.queryTimeoutMs !== void 0 ? clientQueryTimeoutMs(this.queryTimeoutMs) : void 0; try { return await withReadOnlyTransaction( conn, options.readonly, this.supportsReadOnlyTransaction, async () => { const statements = splitSQLStatements(sql, "mysql"); let processedSQL = sql; let probes = []; if (options.maxRows) { const rewrites = statements.map( (statement) => SQLRowLimiter.applyMaxRowsWithTruncationProbe(statement, options.maxRows, "mysql") ); probes = rewrites.map((rewrite) => rewrite.probeApplied); processedSQL = rewrites.map((rewrite) => rewrite.sql).join("; "); if (sql.trim().endsWith(";")) { processedSQL += ";"; } } let results; if (parameters && parameters.length > 0) { results = await conn.query({ sql: processedSQL, timeout }, parameters); } else { results = await conn.query({ sql: processedSQL, timeout }); } const [firstResult] = results; const resultSets = parseQueryResultSets(firstResult, statements); if (resultSets.length === probes.length) { resultSets.forEach( (set, index) => SQLRowLimiter.flagTruncation(set, options.maxRows, probes[index]) ); } return { resultSets }; } ); } catch (error) { if (isClientSideTimeout(error)) { isConnectionDiscarded = true; conn.destroy(); await this.killQuery(threadId); } throw error; } finally { if (!isConnectionDiscarded) { conn.release(); } } } /** * Best-effort server-side kill for a query abandoned by mysql2's client-side * timeout. Uses a separate connection: the original connection's command * queue is stuck behind the abandoned statement (see isClientSideTimeout) * and cannot itself be used to send KILL QUERY. * * Bounded by its own short timeout (CANCEL_QUERY_TIMEOUT_MS), independent * of the user's query_timeout. */ async killQuery(threadId) { if (!this.pool) return; let killer; let killerPoisoned = false; try { killer = await this.pool.getConnection(); await killer.query({ sql: `KILL QUERY ${threadId}`, timeout: CANCEL_QUERY_TIMEOUT_MS }); } catch (error) { killerPoisoned = isClientSideTimeout(error); } finally { if (killer) { if (killerPoisoned) { killer.destroy(); } else { killer.release(); } } } } }; var mysqlConnector = new MySQLConnector(); ConnectorRegistry.register(mysqlConnector); export { MySQLConnector };