nngoandClaude Sonnet 4.6
724594d3f1
- Declare `password` as a proper dataclass field (repr=False) so it is visible to type checkers and any code path that constructs a bare ConnectionProfile; drop the now-unused Optional import. - Fix _conn_str() fallback loop: the return was inside the for-body, so only ODBC Driver 18 was ever tried. Now uses pyodbc.drivers() to pick the first installed driver from the preference list. - Make _cur() a thread-safe @contextmanager that holds self._lock for the cursor's lifetime, matching MySQL/PostgreSQL/SQLite. Updated all 16 call sites to `with self._cur() as c:`. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
346 lines
15 KiB
Python
346 lines
15 KiB
Python
"""Microsoft SQL Server driver using pyodbc."""
|
|
from contextlib import contextmanager
|
|
from typing import Optional
|
|
from app.drivers.base import (
|
|
BaseDriver, ColumnInfo, IndexInfo, ForeignKeyInfo, TableInfo
|
|
)
|
|
|
|
try:
|
|
import pyodbc
|
|
PYODBC_AVAILABLE = True
|
|
except ImportError:
|
|
PYODBC_AVAILABLE = False
|
|
|
|
|
|
class MSSQLDriver(BaseDriver):
|
|
"""SQL Server driver via pyodbc (ODBC Driver 17/18 for SQL Server required)."""
|
|
|
|
def __init__(self, config: dict):
|
|
super().__init__(config)
|
|
self.db_type = "mssql"
|
|
|
|
def _conn_str(self) -> str:
|
|
host = self.config.get("host", "localhost")
|
|
port = int(self.config.get("port", 1433))
|
|
db = self.config.get("database", "master")
|
|
user = self.config.get("user", "")
|
|
pwd = self.config.get("password", "")
|
|
|
|
preferred = [
|
|
"ODBC Driver 18 for SQL Server",
|
|
"ODBC Driver 17 for SQL Server",
|
|
"SQL Server",
|
|
]
|
|
installed = pyodbc.drivers() if PYODBC_AVAILABLE else []
|
|
driver = next((d for d in preferred if d in installed), preferred[0])
|
|
|
|
return (
|
|
f"DRIVER={{{driver}}};"
|
|
f"SERVER={host},{port};"
|
|
f"DATABASE={db};"
|
|
f"UID={user};PWD={pwd};"
|
|
f"TrustServerCertificate=yes;"
|
|
f"Connection Timeout={self.config.get('connection_timeout', 30)};"
|
|
)
|
|
|
|
def connect(self) -> None:
|
|
if not PYODBC_AVAILABLE:
|
|
raise ImportError("pyodbc is not installed. Run: pip install pyodbc")
|
|
self._connection = pyodbc.connect(self._conn_str(), autocommit=True)
|
|
|
|
def disconnect(self) -> None:
|
|
if self._connection:
|
|
try:
|
|
self._connection.close()
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
self._connection = None
|
|
|
|
def test_connection(self) -> tuple:
|
|
if not PYODBC_AVAILABLE:
|
|
return False, "pyodbc is not installed. Run: pip install pyodbc"
|
|
try:
|
|
conn = pyodbc.connect(self._conn_str(), autocommit=True)
|
|
conn.close()
|
|
return True, "Connection successful"
|
|
except Exception as e:
|
|
return False, str(e)
|
|
|
|
@contextmanager
|
|
def _cur(self):
|
|
with self._lock:
|
|
cur = self._connection.cursor()
|
|
try:
|
|
yield cur
|
|
finally:
|
|
try:
|
|
cur.close()
|
|
except Exception:
|
|
pass
|
|
|
|
# ── Schema introspection ──────────────────────────────────────────────────
|
|
|
|
def get_databases(self) -> list:
|
|
with self._cur() as c:
|
|
c.execute("SELECT name FROM sys.databases ORDER BY name")
|
|
return [r[0] for r in c.fetchall()]
|
|
|
|
def get_tables(self, database: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute(f"""
|
|
SELECT t.name, s.name,
|
|
COALESCE(p.rows, 0), 0, '', ''
|
|
FROM [{database}].sys.tables t
|
|
JOIN [{database}].sys.schemas s ON t.schema_id = s.schema_id
|
|
LEFT JOIN (
|
|
SELECT object_id, SUM(rows) AS rows
|
|
FROM [{database}].sys.partitions WHERE index_id IN (0,1)
|
|
GROUP BY object_id
|
|
) p ON p.object_id = t.object_id
|
|
ORDER BY t.name
|
|
""")
|
|
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(f"SELECT name FROM [{database}].sys.views ORDER BY name")
|
|
return [r[0] for r in c.fetchall()]
|
|
|
|
def get_columns(self, database: str, table: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute(f"""
|
|
SELECT c.name, tp.name, c.is_nullable, dc.definition,
|
|
CASE WHEN pk.column_id IS NOT NULL THEN 1 ELSE 0 END,
|
|
CASE WHEN fk.parent_column_id IS NOT NULL THEN 1 ELSE 0 END,
|
|
CASE WHEN c.is_identity = 1 THEN 'auto_increment' ELSE '' END
|
|
FROM [{database}].sys.columns c
|
|
JOIN [{database}].sys.types tp ON tp.user_type_id = c.user_type_id
|
|
JOIN [{database}].sys.tables t ON t.object_id = c.object_id
|
|
LEFT JOIN [{database}].sys.default_constraints dc
|
|
ON dc.parent_object_id = c.object_id AND dc.parent_column_id = c.column_id
|
|
LEFT JOIN (
|
|
SELECT ic.column_id, ic.object_id
|
|
FROM [{database}].sys.index_columns ic
|
|
JOIN [{database}].sys.indexes i ON i.object_id = ic.object_id AND i.index_id = ic.index_id
|
|
WHERE i.is_primary_key = 1
|
|
) pk ON pk.object_id = c.object_id AND pk.column_id = c.column_id
|
|
LEFT JOIN (
|
|
SELECT fkc.parent_column_id, fkc.parent_object_id
|
|
FROM [{database}].sys.foreign_key_columns fkc
|
|
) fk ON fk.parent_object_id = c.object_id AND fk.parent_column_id = c.column_id
|
|
WHERE t.name = ?
|
|
ORDER BY c.column_id
|
|
""", (table,))
|
|
return [ColumnInfo(name=r[0], data_type=r[1], nullable=bool(r[2]),
|
|
default=r[3], is_primary_key=bool(r[4]),
|
|
is_foreign_key=bool(r[5]), extra=r[6] or "")
|
|
for r in c.fetchall()]
|
|
|
|
def get_indexes(self, database: str, table: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute(f"""
|
|
SELECT i.name, i.is_unique, STRING_AGG(c.name, ',') WITHIN GROUP (ORDER BY ic.key_ordinal)
|
|
FROM [{database}].sys.indexes i
|
|
JOIN [{database}].sys.index_columns ic ON ic.object_id = i.object_id AND ic.index_id = i.index_id
|
|
JOIN [{database}].sys.columns c ON c.object_id = i.object_id AND c.column_id = ic.column_id
|
|
JOIN [{database}].sys.tables t ON t.object_id = i.object_id
|
|
WHERE t.name = ?
|
|
GROUP BY i.name, i.is_unique
|
|
""", (table,))
|
|
return [IndexInfo(name=r[0], columns=r[2].split(','),
|
|
is_unique=bool(r[1]))
|
|
for r in c.fetchall()]
|
|
|
|
def get_foreign_keys(self, database: str, table: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute(f"""
|
|
SELECT fk.name, pc.name, rt.name, rc.name,
|
|
fk.update_referential_action_desc,
|
|
fk.delete_referential_action_desc
|
|
FROM [{database}].sys.foreign_keys fk
|
|
JOIN [{database}].sys.tables pt ON pt.object_id = fk.parent_object_id
|
|
JOIN [{database}].sys.tables rt ON rt.object_id = fk.referenced_object_id
|
|
JOIN [{database}].sys.foreign_key_columns fkc ON fkc.constraint_object_id = fk.object_id
|
|
JOIN [{database}].sys.columns pc ON pc.object_id = fkc.parent_object_id AND pc.column_id = fkc.parent_column_id
|
|
JOIN [{database}].sys.columns rc ON rc.object_id = fkc.referenced_object_id AND rc.column_id = fkc.referenced_column_id
|
|
WHERE pt.name = ?
|
|
""", (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:
|
|
cols = self.get_columns(database, table)
|
|
lines = [f"CREATE TABLE [{table}] ("]
|
|
col_defs = []
|
|
for col in cols:
|
|
d = f" [{col.name}] {col.data_type}"
|
|
if not col.nullable: d += " NOT NULL"
|
|
if col.default: d += f" DEFAULT {col.default}"
|
|
col_defs.append(d)
|
|
lines.append(",\n".join(col_defs))
|
|
lines.append(");")
|
|
return "\n".join(lines)
|
|
|
|
def get_functions(self, database: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute(f"""
|
|
SELECT name FROM [{database}].sys.objects
|
|
WHERE type IN ('FN','IF','TF') ORDER BY name
|
|
""")
|
|
return [r[0] for r in c.fetchall()]
|
|
|
|
def get_stored_procedures(self, database: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute(f"""
|
|
SELECT name FROM [{database}].sys.procedures ORDER BY name
|
|
""")
|
|
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(f"""
|
|
SELECT t.name FROM [{database}].sys.triggers t
|
|
JOIN [{database}].sys.tables tb ON tb.object_id = t.parent_id
|
|
WHERE tb.name = ? ORDER BY t.name
|
|
""", (table,))
|
|
else:
|
|
c.execute(f"SELECT name FROM [{database}].sys.triggers ORDER BY name")
|
|
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 or ())
|
|
if c.description:
|
|
cols = [d[0] for d in c.description]
|
|
rows = c.fetchall()
|
|
return cols, [tuple(r) for r in 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()]
|
|
for stmt in stmts:
|
|
try:
|
|
with self._cur() as c:
|
|
c.execute(stmt)
|
|
if c.description:
|
|
cols = [d[0] for d in c.description]
|
|
rows = [tuple(r) for r in 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}].[dbo].[{table}]"
|
|
if where: sql += f" WHERE {where}"
|
|
if order_by: sql += f" ORDER BY {order_by}"
|
|
sql += f" OFFSET {offset} ROWS FETCH NEXT {limit} ROWS ONLY"
|
|
return self.execute_query(sql)
|
|
|
|
def get_table_row_count(self, database: str, table: str, where: str = "") -> int:
|
|
sql = f"SELECT COUNT(*) FROM [{database}].[dbo].[{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(["?"] * len(data))
|
|
with self._cur() as c:
|
|
c.execute(f"INSERT INTO [{database}].[dbo].[{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}] = ?")
|
|
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}] = ?" for c in data)
|
|
where_cl, where_params = self._where(where)
|
|
with self._cur() as c:
|
|
c.execute(f"UPDATE [{database}].[dbo].[{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}].[dbo].[{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"SET SHOWPLAN_TEXT ON; {sql}; SET SHOWPLAN_TEXT OFF")
|
|
return ["Plan"], c.fetchall()
|
|
|
|
def get_process_list(self) -> tuple:
|
|
with self._cur() as c:
|
|
c.execute("""
|
|
SELECT session_id, login_name, status, host_name,
|
|
program_name, cpu_time, text
|
|
FROM sys.dm_exec_sessions s
|
|
CROSS APPLY sys.dm_exec_sql_text(s.most_recent_sql_handle) t
|
|
WHERE s.is_user_process = 1
|
|
""")
|
|
cols = [d[0] for d in c.description]
|
|
return cols, [tuple(r) for r in c.fetchall()]
|
|
|
|
def kill_process(self, process_id: int) -> bool:
|
|
with self._cur() as c:
|
|
c.execute(f"KILL {int(process_id)}")
|
|
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}].[dbo].[{table}] "
|
|
f"ADD [{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}].[dbo].[{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"EXEC sp_rename '[{database}].[dbo].[{table}].[{old_name}]', '{new_name}', 'COLUMN'"
|
|
with self._cur() as c:
|
|
c.execute(sql)
|
|
return True
|