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