@bytebase/dbhub
Version:
Token-efficient database MCP server for PostgreSQL, MySQL, MariaDB, SQL Server, Oracle, SQLite
313 lines (311 loc) • 9.83 kB
JavaScript
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
};