378 lines
16 KiB
Python
378 lines
16 KiB
Python
"""MySQL driver implementation using pymysql."""
|
|
import pymysql
|
|
import pymysql.cursors
|
|
from contextlib import contextmanager
|
|
from typing import Optional
|
|
from app.drivers.base import (
|
|
BaseDriver, ColumnInfo, IndexInfo, ForeignKeyInfo, TableInfo
|
|
)
|
|
from app.utils.logger import get_logger
|
|
|
|
_log = get_logger(__name__)
|
|
|
|
|
|
class MySQLDriver(BaseDriver):
|
|
"""MySQL / MariaDB database driver."""
|
|
|
|
def __init__(self, config: dict):
|
|
super().__init__(config)
|
|
self.db_type = "mysql"
|
|
|
|
def _connect_kwargs(self) -> dict:
|
|
kw = {
|
|
"host": self.config.get("host", "localhost"),
|
|
"port": int(self.config.get("port", 3306)),
|
|
"user": self.config.get("user", ""),
|
|
"password": self.config.get("password", ""),
|
|
"connect_timeout": int(self.config.get("connection_timeout", 30)),
|
|
"autocommit": True,
|
|
"charset": "utf8mb4",
|
|
}
|
|
db = self.config.get("database", "")
|
|
if db:
|
|
kw["database"] = db
|
|
return kw
|
|
|
|
def connect(self) -> None:
|
|
host = self.config.get("host", "localhost")
|
|
port = self.config.get("port", 3306)
|
|
user = self.config.get("user", "")
|
|
_log.info("MySQL connecting host=%s:%s user=%s", host, port, user)
|
|
try:
|
|
self._connection = pymysql.connect(**self._connect_kwargs())
|
|
_log.info("MySQL connected host=%s:%s user=%s", host, port, user)
|
|
except Exception:
|
|
_log.error("MySQL connection failed host=%s:%s user=%s",
|
|
host, port, user, exc_info=True)
|
|
raise
|
|
|
|
def disconnect(self) -> None:
|
|
if self._connection:
|
|
try:
|
|
self._connection.close()
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
self._connection = None
|
|
|
|
def test_connection(self) -> tuple:
|
|
try:
|
|
conn = pymysql.connect(**self._connect_kwargs())
|
|
conn.close()
|
|
return True, "Connection successful"
|
|
except Exception as e:
|
|
return False, str(e)
|
|
|
|
def _ensure_alive(self) -> None:
|
|
"""Ping the server and silently reconnect if the connection has gone away."""
|
|
if self._connection is None:
|
|
_log.warning("MySQL connection is None — connecting now")
|
|
self.connect()
|
|
return
|
|
try:
|
|
self._connection.ping(reconnect=True)
|
|
except Exception:
|
|
_log.warning("MySQL ping failed — attempting full reconnect", exc_info=True)
|
|
try:
|
|
self.connect()
|
|
_log.info("MySQL reconnected successfully")
|
|
except Exception:
|
|
_log.error("MySQL reconnect failed", exc_info=True)
|
|
raise
|
|
|
|
@contextmanager
|
|
def _cur(self):
|
|
"""Yield a cursor while holding the driver lock.
|
|
|
|
Using a contextmanager means the lock is held for the *entire*
|
|
``with self._cur() as c: c.execute(); c.fetchall()`` block, which
|
|
prevents concurrent SchemaWorker threads from interleaving on the same
|
|
TCP socket (pymysql is not thread-safe).
|
|
"""
|
|
with self._lock:
|
|
self._ensure_alive()
|
|
cursor = self._connection.cursor(pymysql.cursors.Cursor)
|
|
try:
|
|
yield cursor
|
|
finally:
|
|
try:
|
|
cursor.close()
|
|
except Exception:
|
|
pass
|
|
|
|
@staticmethod
|
|
def _fmt_err(e: Exception) -> str:
|
|
"""Return a readable string for a pymysql exception.
|
|
|
|
pymysql errors carry (error_code, message) as args, so ``str(e)``
|
|
prints something like ``(0, '')``. This helper unwraps that.
|
|
"""
|
|
args = getattr(e, "args", ())
|
|
if args and isinstance(args[0], int):
|
|
code, msg = args[0], args[1] if len(args) > 1 else ""
|
|
if msg:
|
|
return f"MySQL error {code}: {msg}"
|
|
if code == 0:
|
|
return "Lost connection to MySQL server (connection timed out or was reset)."
|
|
return f"MySQL error {code}"
|
|
return str(e)
|
|
|
|
# ── Schema introspection ──────────────────────────────────────────────────
|
|
|
|
def get_databases(self) -> list:
|
|
with self._cur() as c:
|
|
c.execute("SHOW DATABASES")
|
|
return [r[0] for r in c.fetchall()]
|
|
|
|
def get_tables(self, database: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute("""
|
|
SELECT TABLE_NAME, TABLE_SCHEMA,
|
|
COALESCE(TABLE_ROWS, 0),
|
|
COALESCE(DATA_LENGTH + INDEX_LENGTH, 0),
|
|
COALESCE(ENGINE, ''),
|
|
COALESCE(TABLE_COMMENT, '')
|
|
FROM information_schema.TABLES
|
|
WHERE TABLE_SCHEMA = %s AND TABLE_TYPE = 'BASE TABLE'
|
|
ORDER BY TABLE_NAME
|
|
""", (database,))
|
|
return [TableInfo(name=r[0], schema=r[1], row_count=r[2],
|
|
size_bytes=r[3], engine=r[4], comment=r[5])
|
|
for r in c.fetchall()]
|
|
|
|
def get_views(self, database: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute("""
|
|
SELECT TABLE_NAME FROM information_schema.VIEWS
|
|
WHERE TABLE_SCHEMA = %s ORDER BY TABLE_NAME
|
|
""", (database,))
|
|
return [r[0] for r in c.fetchall()]
|
|
|
|
def get_columns(self, database: str, table: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute("""
|
|
SELECT COLUMN_NAME, COLUMN_TYPE, IS_NULLABLE,
|
|
COLUMN_DEFAULT, COLUMN_KEY, EXTRA
|
|
FROM information_schema.COLUMNS
|
|
WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
|
|
ORDER BY ORDINAL_POSITION
|
|
""", (database, table))
|
|
return [ColumnInfo(name=r[0], data_type=r[1],
|
|
nullable=(r[2] == "YES"), default=r[3],
|
|
is_primary_key=(r[4] == "PRI"),
|
|
is_foreign_key=(r[4] == "MUL"),
|
|
extra=r[5] or "")
|
|
for r in c.fetchall()]
|
|
|
|
def get_indexes(self, database: str, table: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute(f"SHOW INDEX FROM `{database}`.`{table}`")
|
|
idx_map = {}
|
|
for r in c.fetchall():
|
|
name, non_unique, col, idx_type = r[2], r[1], r[4], r[10]
|
|
if name not in idx_map:
|
|
idx_map[name] = IndexInfo(name=name, columns=[col],
|
|
is_unique=(non_unique == 0),
|
|
index_type=idx_type)
|
|
else:
|
|
idx_map[name].columns.append(col)
|
|
return list(idx_map.values())
|
|
|
|
def get_foreign_keys(self, database: str, table: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute("""
|
|
SELECT kcu.CONSTRAINT_NAME, kcu.COLUMN_NAME,
|
|
kcu.REFERENCED_TABLE_NAME, kcu.REFERENCED_COLUMN_NAME,
|
|
rc.UPDATE_RULE, rc.DELETE_RULE
|
|
FROM information_schema.KEY_COLUMN_USAGE kcu
|
|
JOIN information_schema.REFERENTIAL_CONSTRAINTS rc
|
|
ON rc.CONSTRAINT_NAME = kcu.CONSTRAINT_NAME
|
|
AND rc.CONSTRAINT_SCHEMA = kcu.CONSTRAINT_SCHEMA
|
|
WHERE kcu.TABLE_SCHEMA = %s AND kcu.TABLE_NAME = %s
|
|
AND kcu.REFERENCED_TABLE_NAME IS NOT NULL
|
|
""", (database, table))
|
|
return [ForeignKeyInfo(name=r[0], column=r[1],
|
|
ref_table=r[2], ref_column=r[3],
|
|
on_update=r[4] or "", on_delete=r[5] or "")
|
|
for r in c.fetchall()]
|
|
|
|
def get_table_ddl(self, database: str, table: str) -> str:
|
|
with self._cur() as c:
|
|
c.execute(f"SHOW CREATE TABLE `{database}`.`{table}`")
|
|
row = c.fetchone()
|
|
return row[1] if row else ""
|
|
|
|
def get_functions(self, database: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute("""
|
|
SELECT ROUTINE_NAME FROM information_schema.ROUTINES
|
|
WHERE ROUTINE_SCHEMA = %s AND ROUTINE_TYPE = 'FUNCTION'
|
|
ORDER BY ROUTINE_NAME
|
|
""", (database,))
|
|
return [r[0] for r in c.fetchall()]
|
|
|
|
def get_stored_procedures(self, database: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute("""
|
|
SELECT ROUTINE_NAME FROM information_schema.ROUTINES
|
|
WHERE ROUTINE_SCHEMA = %s AND ROUTINE_TYPE = 'PROCEDURE'
|
|
ORDER BY ROUTINE_NAME
|
|
""", (database,))
|
|
return [r[0] for r in c.fetchall()]
|
|
|
|
def get_triggers(self, database: str, table: str = "") -> list:
|
|
with self._cur() as c:
|
|
if table:
|
|
c.execute("""
|
|
SELECT TRIGGER_NAME FROM information_schema.TRIGGERS
|
|
WHERE TRIGGER_SCHEMA = %s AND EVENT_OBJECT_TABLE = %s
|
|
ORDER BY TRIGGER_NAME
|
|
""", (database, table))
|
|
else:
|
|
c.execute("""
|
|
SELECT TRIGGER_NAME FROM information_schema.TRIGGERS
|
|
WHERE TRIGGER_SCHEMA = %s ORDER BY TRIGGER_NAME
|
|
""", (database,))
|
|
return [r[0] for r in c.fetchall()]
|
|
|
|
# ── Query execution ───────────────────────────────────────────────────────
|
|
|
|
def execute_query(self, sql: str, params: Optional[tuple] = None) -> tuple:
|
|
with self._cur() as c:
|
|
c.execute(sql, params)
|
|
if c.description:
|
|
cols = [d[0] for d in c.description]
|
|
rows = c.fetchall()
|
|
return cols, rows, len(rows)
|
|
return [], [], c.rowcount
|
|
|
|
def execute_script(self, sql: str) -> list:
|
|
results = []
|
|
stmts = [s.strip() for s in sql.split(';') if s.strip()]
|
|
with self._cur() as c:
|
|
for stmt in stmts:
|
|
try:
|
|
c.execute(stmt)
|
|
if c.description:
|
|
cols = [d[0] for d in c.description]
|
|
rows = c.fetchall()
|
|
results.append((cols, rows, len(rows), ""))
|
|
else:
|
|
results.append(([], [], c.rowcount,
|
|
f"{c.rowcount} row(s) affected"))
|
|
except Exception as e:
|
|
results.append(([], [], 0, f"Error: {e}"))
|
|
return results
|
|
|
|
# ── Table data CRUD ───────────────────────────────────────────────────────
|
|
|
|
def get_table_data(self, database: str, table: str,
|
|
where: str = "", order_by: str = "",
|
|
limit: int = 1000, offset: int = 0) -> tuple:
|
|
sql = f"SELECT * FROM `{database}`.`{table}`"
|
|
if where: sql += f" WHERE {where}"
|
|
if order_by: sql += f" ORDER BY {order_by}"
|
|
sql += f" LIMIT {limit} OFFSET {offset}"
|
|
return self.execute_query(sql)
|
|
|
|
def get_table_row_count(self, database: str, table: str, where: str = "") -> int:
|
|
sql = f"SELECT COUNT(*) FROM `{database}`.`{table}`"
|
|
if where: sql += f" WHERE {where}"
|
|
with self._cur() as c:
|
|
c.execute(sql)
|
|
return c.fetchone()[0]
|
|
|
|
def insert_row(self, database: str, table: str, data: dict) -> bool:
|
|
cols = ", ".join(f"`{c}`" for c in data)
|
|
ph = ", ".join(["%s"] * len(data))
|
|
with self._cur() as c:
|
|
c.execute(f"INSERT INTO `{database}`.`{table}` ({cols}) VALUES ({ph})",
|
|
tuple(data.values()))
|
|
return True
|
|
|
|
@staticmethod
|
|
def _where(where: dict) -> tuple:
|
|
parts, params = [], []
|
|
for col, val in where.items():
|
|
if val is None:
|
|
parts.append(f"`{col}` IS NULL")
|
|
else:
|
|
parts.append(f"`{col}` = %s")
|
|
params.append(val)
|
|
return " AND ".join(parts), params
|
|
|
|
def update_row(self, database: str, table: str, data: dict, where: dict) -> bool:
|
|
set_cl = ", ".join(f"`{c}` = %s" for c in data)
|
|
where_cl, where_params = self._where(where)
|
|
with self._cur() as c:
|
|
c.execute(f"UPDATE `{database}`.`{table}` SET {set_cl} WHERE {where_cl}",
|
|
tuple(data.values()) + tuple(where_params))
|
|
return True
|
|
|
|
def delete_row(self, database: str, table: str, where: dict) -> bool:
|
|
where_cl, where_params = self._where(where)
|
|
with self._cur() as c:
|
|
c.execute(f"DELETE FROM `{database}`.`{table}` WHERE {where_cl}",
|
|
tuple(where_params))
|
|
return True
|
|
|
|
# ── Server tools ──────────────────────────────────────────────────────────
|
|
|
|
def explain_query(self, sql: str) -> tuple:
|
|
with self._cur() as c:
|
|
c.execute(f"EXPLAIN {sql}")
|
|
return [d[0] for d in c.description], c.fetchall()
|
|
|
|
def get_process_list(self) -> tuple:
|
|
with self._cur() as c:
|
|
c.execute("SHOW FULL PROCESSLIST")
|
|
return [d[0] for d in c.description], c.fetchall()
|
|
|
|
def kill_process(self, process_id: int) -> bool:
|
|
# MUST use a dedicated connection, not self._cur().
|
|
#
|
|
# If a QueryWorker is running a long query it holds self._lock via
|
|
# _cur(). kill_process is called from a *different* SchemaWorker
|
|
# thread; using self._cur() here would block waiting for that lock,
|
|
# so the KILL command would never reach MySQL and the server would
|
|
# eventually raise error 1317 ("Query execution was interrupted") on
|
|
# its own. A fresh connection bypasses the lock entirely, which is
|
|
# exactly how MySQL's KILL is intended to work.
|
|
conn = pymysql.connect(**self._connect_kwargs())
|
|
try:
|
|
with conn.cursor() as c:
|
|
c.execute(f"KILL {process_id}")
|
|
finally:
|
|
try:
|
|
conn.close()
|
|
except Exception:
|
|
pass
|
|
return True
|
|
|
|
# ── Table designer ────────────────────────────────────────────────────────
|
|
|
|
def add_column(self, database: str, table: str, col_name: str,
|
|
col_type: str, nullable: bool = True,
|
|
default=None) -> bool:
|
|
null_clause = "NULL" if nullable else "NOT NULL"
|
|
default_clause = f" DEFAULT {default}" if default is not None and default != "" else ""
|
|
sql = (f"ALTER TABLE `{database}`.`{table}` "
|
|
f"ADD COLUMN `{col_name}` {col_type} {null_clause}{default_clause}")
|
|
with self._cur() as c:
|
|
c.execute(sql)
|
|
return True
|
|
|
|
def drop_column(self, database: str, table: str, col_name: str) -> bool:
|
|
sql = f"ALTER TABLE `{database}`.`{table}` DROP COLUMN `{col_name}`"
|
|
with self._cur() as c:
|
|
c.execute(sql)
|
|
return True
|
|
|
|
def rename_column(self, database: str, table: str,
|
|
old_name: str, new_name: str) -> bool:
|
|
sql = (f"ALTER TABLE `{database}`.`{table}` "
|
|
f"RENAME COLUMN `{old_name}` TO `{new_name}`")
|
|
with self._cur() as c:
|
|
c.execute(sql)
|
|
return True
|