Initial Codes
This commit is contained in:
@@ -0,0 +1,377 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user