UNPKG

@bytebase/dbhub

Version:

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

552 lines (550 loc) • 18.1 kB
import { MissingDatabaseError, getMySQLFamilyHealthCheck, isTiDBVersion, parseQueryResultSets, requireDatabaseInDSN, withReadOnlyTransaction } from "./chunk-6SZKW6GD.js"; import "./chunk-FU2ZJE4E.js"; import { quoteIdentifier } from "./chunk-X2EC47W4.js"; import { SQLRowLimiter, closeQuietly } from "./chunk-YIGJHEU6.js"; import "./chunk-ABTVMHD6.js"; import { ConnectorRegistry, SafeURL, obfuscateDSNPassword, splitSQLStatements } from "./chunk-IXC4NQGC.js"; // src/connectors/mariadb/index.ts import * as mariadb from "mariadb"; var MariadbDSNParser = class { async parse(dsn, config) { const connectionTimeoutSeconds = config?.connectionTimeoutSeconds; const queryTimeoutSeconds = config?.queryTimeoutSeconds; if (!this.isValidDSN(dsn)) { const obfuscatedDSN = obfuscateDSNPassword(dsn); const expectedFormat = this.getSampleDSN(); throw new Error( `Invalid MariaDB DSN format. Provided: ${obfuscatedDSN} Expected: ${expectedFormat}` ); } try { const url = new SafeURL(dsn); const database = url.pathname ? url.pathname.substring(1) : ""; requireDatabaseInDSN(database, dsn, "MariaDB"); const connectionConfig = { host: url.hostname, port: url.port ? parseInt(url.port) : 3306, database, user: url.username, password: url.password, multipleStatements: true, // Enable native multi-statement support ...connectionTimeoutSeconds !== void 0 && { connectTimeout: connectionTimeoutSeconds * 1e3 }, ...queryTimeoutSeconds !== void 0 && { queryTimeout: queryTimeoutSeconds * 1e3 }, // Controls how the driver interprets DATETIME values ("Z", "local", or "±HH:MM"). ...config?.timezone !== void 0 && { timezone: config.timezone }, // Connection character set (e.g. "utf8mb4") / collation (e.g. // "utf8mb4_general_ci"). The mariadb driver exposes both as distinct // options, but a configured collation is authoritative (it implies its // character set) and the driver ignores `collation` when `charset` is also // passed. So forward the collation when present, otherwise the charset // (which uses that character set's default collation). ...config?.collation !== void 0 ? { collation: config.collation } : config?.charset !== void 0 ? { charset: config.charset } : {} }; url.forEachSearchParam((value, key) => { if (key === "sslmode") { if (value === "disable") { connectionConfig.ssl = void 0; } else if (value === "require") { connectionConfig.ssl = { rejectUnauthorized: false }; } else { connectionConfig.ssl = {}; } } }); if (url.password && url.password.includes("X-Amz-Credential")) { if (connectionConfig.ssl === void 0) { connectionConfig.ssl = { rejectUnauthorized: false }; } } return connectionConfig; } catch (error) { if (error instanceof MissingDatabaseError) { throw error; } throw new Error( `Failed to parse MariaDB DSN: ${error instanceof Error ? error.message : String(error)}` ); } } getSampleDSN() { return "mariadb://root:password@localhost:3306/db?sslmode=require"; } isValidDSN(dsn) { try { return dsn.startsWith("mariadb://"); } catch (error) { return false; } } }; var MariaDBConnector = class _MariaDBConnector { constructor() { this.id = "mariadb"; this.name = "MariaDB"; this.dsnParser = new MariadbDSNParser(); 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 _MariaDBConnector(); } async connect(dsn, initScript, config) { try { const connectionConfig = await this.dsnParser.parse(dsn, config); this.pool = mariadb.createPool(connectionConfig); this.pool.on("error", (err) => { console.error( `MariaDB pool (source "${this.sourceId}"): background connection error, pool will retry:`, err.message ); }); 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 MariaDB database:", err); throw err; } } async disconnect() { if (this.pool) { await this.pool.end(); this.pool = null; } } async getSchemas() { if (!this.pool) { throw new Error("Not connected to database"); } try { const rows = await this.pool.query(` SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME NOT IN ('information_schema', 'performance_schema', 'mysql', 'sys') ORDER BY SCHEMA_NAME `); return rows.map((row) => row.SCHEMA_NAME); } catch (error) { console.error("Error getting schemas:", error); throw error; } } async getTables(schema) { if (!this.pool) { throw new Error("Not connected to database"); } try { const schemaClause = schema ? "WHERE TABLE_SCHEMA = ?" : "WHERE TABLE_SCHEMA = DATABASE()"; const queryParams = schema ? [schema] : []; const rows = await this.pool.query( ` SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES ${schemaClause} AND TABLE_TYPE = 'BASE TABLE' ORDER BY TABLE_NAME `, queryParams ); return rows.map((row) => row.TABLE_NAME); } catch (error) { console.error("Error getting tables:", error); throw error; } } async getViews(schema) { if (!this.pool) { throw new Error("Not connected to database"); } try { const schemaClause = schema ? "WHERE TABLE_SCHEMA = ?" : "WHERE TABLE_SCHEMA = DATABASE()"; const queryParams = schema ? [schema] : []; const rows = await this.pool.query( ` SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES ${schemaClause} AND TABLE_TYPE = 'VIEW' ORDER BY TABLE_NAME `, queryParams ); return rows.map((row) => row.TABLE_NAME); } catch (error) { console.error("Error getting views:", error); throw error; } } async tableExists(tableName, schema) { if (!this.pool) { throw new Error("Not connected to database"); } try { const schemaClause = schema ? "WHERE TABLE_SCHEMA = ?" : "WHERE TABLE_SCHEMA = DATABASE()"; const queryParams = schema ? [schema, tableName] : [tableName]; const rows = await this.pool.query( ` SELECT COUNT(*) AS COUNT FROM INFORMATION_SCHEMA.TABLES ${schemaClause} AND TABLE_NAME = ? `, queryParams ); return rows[0].COUNT > 0; } catch (error) { console.error("Error checking if table exists:", error); throw error; } } async getTableIndexes(tableName, schema) { if (!this.pool) { throw new Error("Not connected to database"); } try { const schemaClause = schema ? "TABLE_SCHEMA = ?" : "TABLE_SCHEMA = DATABASE()"; const queryParams = schema ? [schema, tableName] : [tableName]; const indexRows = await this.pool.query( ` SELECT INDEX_NAME, COLUMN_NAME, NON_UNIQUE, SEQ_IN_INDEX FROM INFORMATION_SCHEMA.STATISTICS WHERE ${schemaClause} AND TABLE_NAME = ? ORDER BY INDEX_NAME, SEQ_IN_INDEX `, queryParams ); const indexMap = /* @__PURE__ */ new Map(); for (const row of indexRows) { const indexName = row.INDEX_NAME; const columnName = row.COLUMN_NAME; const isUnique = row.NON_UNIQUE === 0; const isPrimary = indexName === "PRIMARY"; if (!indexMap.has(indexName)) { indexMap.set(indexName, { columns: [], is_unique: isUnique, is_primary: isPrimary }); } const indexInfo = indexMap.get(indexName); indexInfo.columns.push(columnName); } const results = []; indexMap.forEach((indexInfo, indexName) => { results.push({ index_name: indexName, column_names: indexInfo.columns, is_unique: indexInfo.is_unique, is_primary: indexInfo.is_primary }); }); return results; } catch (error) { console.error("Error getting table indexes:", error); throw error; } } async getTableSchema(tableName, schema) { if (!this.pool) { throw new Error("Not connected to database"); } try { const schemaClause = schema ? "WHERE TABLE_SCHEMA = ?" : "WHERE TABLE_SCHEMA = DATABASE()"; const queryParams = schema ? [schema, tableName] : [tableName]; const rows = await this.pool.query( ` SELECT COLUMN_NAME as column_name, DATA_TYPE as data_type, IS_NULLABLE as is_nullable, COLUMN_DEFAULT as column_default, COLUMN_COMMENT as description FROM INFORMATION_SCHEMA.COLUMNS ${schemaClause} AND TABLE_NAME = ? ORDER BY ORDINAL_POSITION `, queryParams ); return rows.map((row) => ({ ...row, description: row.description || null })); } catch (error) { console.error("Error getting table schema:", error); throw error; } } async getTableComment(tableName, schema) { if (!this.pool) { throw new Error("Not connected to database"); } try { const schemaClause = schema ? "WHERE TABLE_SCHEMA = ?" : "WHERE TABLE_SCHEMA = DATABASE()"; const queryParams = schema ? [schema, tableName] : [tableName]; const rows = await this.pool.query( ` SELECT TABLE_COMMENT FROM INFORMATION_SCHEMA.TABLES ${schemaClause} AND TABLE_NAME = ? `, queryParams ); if (rows.length > 0) { return rows[0].TABLE_COMMENT || null; } return null; } catch (error) { return null; } } async getHealthCheck() { if (!this.pool) { throw new Error("Not connected to database"); } const pool = this.pool; return getMySQLFamilyHealthCheck((sql) => pool.query(sql)); } async getStoredProcedures(schema, routineType) { if (!this.pool) { throw new Error("Not connected to database"); } try { const schemaClause = schema ? "WHERE ROUTINE_SCHEMA = ?" : "WHERE ROUTINE_SCHEMA = DATABASE()"; const queryParams = schema ? [schema] : []; let typeFilter = ""; if (routineType === "function") { typeFilter = " AND ROUTINE_TYPE = 'FUNCTION'"; } else if (routineType === "procedure") { typeFilter = " AND ROUTINE_TYPE = 'PROCEDURE'"; } const rows = await this.pool.query( ` SELECT ROUTINE_NAME FROM INFORMATION_SCHEMA.ROUTINES ${schemaClause}${typeFilter} ORDER BY ROUTINE_NAME `, queryParams ); return rows.map((row) => row.ROUTINE_NAME); } catch (error) { console.error("Error getting stored procedures:", error); throw error; } } async getStoredProcedureDetail(procedureName, schema) { if (!this.pool) { throw new Error("Not connected to database"); } try { const schemaClause = schema ? "WHERE r.ROUTINE_SCHEMA = ?" : "WHERE r.ROUTINE_SCHEMA = DATABASE()"; const queryParams = schema ? [schema, procedureName] : [procedureName]; const rows = await this.pool.query( ` SELECT r.ROUTINE_NAME AS procedure_name, CASE WHEN r.ROUTINE_TYPE = 'PROCEDURE' THEN 'procedure' ELSE 'function' END AS procedure_type, LOWER(r.ROUTINE_TYPE) AS routine_type, r.ROUTINE_DEFINITION, r.DTD_IDENTIFIER AS return_type, ( SELECT GROUP_CONCAT( CONCAT(p.PARAMETER_NAME, ' ', p.PARAMETER_MODE, ' ', p.DATA_TYPE) ORDER BY p.ORDINAL_POSITION SEPARATOR ', ' ) FROM INFORMATION_SCHEMA.PARAMETERS p WHERE p.SPECIFIC_SCHEMA = r.ROUTINE_SCHEMA AND p.SPECIFIC_NAME = r.ROUTINE_NAME AND p.PARAMETER_NAME IS NOT NULL ) AS parameter_list FROM INFORMATION_SCHEMA.ROUTINES r ${schemaClause} AND r.ROUTINE_NAME = ? `, queryParams ); if (rows.length === 0) { const schemaName = schema || "current schema"; throw new Error(`Stored procedure '${procedureName}' not found in ${schemaName}`); } const procedure = rows[0]; let definition = procedure.ROUTINE_DEFINITION; try { const schemaValue = schema || await this.getCurrentSchema(); const quotedSchema = quoteIdentifier(schemaValue, "mariadb"); const quotedProcName = quoteIdentifier(procedureName, "mariadb"); if (procedure.procedure_type === "procedure") { try { const defRows = await this.pool.query(` SHOW CREATE PROCEDURE ${quotedSchema}.${quotedProcName} `); if (defRows && defRows.length > 0) { definition = defRows[0]["Create Procedure"]; } } catch (err) { console.error(`Error getting procedure definition with SHOW CREATE: ${err}`); } } else { try { const defRows = await this.pool.query(` SHOW CREATE FUNCTION ${quotedSchema}.${quotedProcName} `); if (defRows && defRows.length > 0) { definition = defRows[0]["Create Function"]; } } catch (innerErr) { console.error(`Error getting function definition with SHOW CREATE: ${innerErr}`); } } if (!definition) { const bodyRows = await this.pool.query( ` SELECT ROUTINE_DEFINITION, ROUTINE_BODY FROM INFORMATION_SCHEMA.ROUTINES WHERE ROUTINE_SCHEMA = ? AND ROUTINE_NAME = ? `, [schemaValue, procedureName] ); if (bodyRows && bodyRows.length > 0) { if (bodyRows[0].ROUTINE_DEFINITION) { definition = bodyRows[0].ROUTINE_DEFINITION; } else if (bodyRows[0].ROUTINE_BODY) { definition = bodyRows[0].ROUTINE_BODY; } } } } catch (error) { console.error(`Error getting procedure/function details: ${error}`); } return { procedure_name: procedure.procedure_name, procedure_type: procedure.procedure_type, language: "sql", // MariaDB procedures are generally in SQL parameter_list: procedure.parameter_list || "", return_type: procedure.routine_type === "function" ? procedure.return_type : void 0, definition: definition || void 0 }; } catch (error) { console.error("Error getting stored procedure detail:", error); throw error; } } // Helper method to get current schema (database) name async getCurrentSchema() { const rows = await this.pool.query("SELECT DATABASE() AS DB"); return rows[0].DB; } /** * Default search scope = the database named in the DSN. DATABASE() returns * null when the connection was opened without a database, in which case * callers fall back to the full server-wide schema list. */ async getDefaultSchema() { if (!this.pool) { throw new Error("Not connected to database"); } const rows = await this.pool.query("SELECT DATABASE() AS DB"); return rows[0]?.DB ?? null; } async executeSQL(sql, options, parameters) { if (!this.pool) { throw new Error("Not connected to database"); } const conn = await this.pool.getConnection(); try { return await withReadOnlyTransaction( conn, options.readonly, this.supportsReadOnlyTransaction, async () => { const statements = splitSQLStatements(sql, "mariadb"); let processedSQL = sql; let probes = []; if (options.maxRows) { const rewrites = statements.map( (statement) => SQLRowLimiter.applyMaxRowsWithTruncationProbe(statement, options.maxRows, "mariadb") ); 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(processedSQL, parameters); } else { results = await conn.query(processedSQL); } const resultSets = parseQueryResultSets(results, statements); if (resultSets.length === probes.length) { resultSets.forEach( (set, index) => SQLRowLimiter.flagTruncation(set, options.maxRows, probes[index]) ); } return { resultSets }; } ); } finally { conn.release(); } } }; var mariadbConnector = new MariaDBConnector(); ConnectorRegistry.register(mariadbConnector); export { MariaDBConnector };