nngo/DBClient
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix three critical bugs in ConnectionProfile and MSSQL driver

Browse Source 
- 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>
This commit is contained in:
nngo
2026-05-21 16:02:04 -04:00
co-authored by Claude Sonnet 4.6
parent b01ad5ea40
commit 724594d3f1
2 changed files with 176 additions and 164 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/drivers/mssql_driver.py
+175 -163
View File
@@ -1,4 +1,5 @@
"""Microsoft SQL Server driver using pyodbc.""" """Microsoft SQL Server driver using pyodbc."""
from contextlib import contextmanager
from typing import Optional from typing import Optional
from app.drivers.base import ( from app.drivers.base import (
BaseDriver, ColumnInfo, IndexInfo, ForeignKeyInfo, TableInfo BaseDriver, ColumnInfo, IndexInfo, ForeignKeyInfo, TableInfo
@@ -24,20 +25,23 @@ class MSSQLDriver(BaseDriver):
db = self.config.get("database", "master") db = self.config.get("database", "master")
user = self.config.get("user", "") user = self.config.get("user", "")
pwd = self.config.get("password", "") pwd = self.config.get("password", "")
# Try drivers in order of preference
for driver in [ preferred = [
"ODBC Driver 18 for SQL Server", "ODBC Driver 18 for SQL Server",
"ODBC Driver 17 for SQL Server", "ODBC Driver 17 for SQL Server",
"SQL Server", "SQL Server",
]: ]
return ( installed = pyodbc.drivers() if PYODBC_AVAILABLE else []
f"DRIVER={{{driver}}};" driver = next((d for d in preferred if d in installed), preferred[0])
f"SERVER={host},{port};"
f"DATABASE={db};" return (
f"UID={user};PWD={pwd};" f"DRIVER={{{driver}}};"
f"TrustServerCertificate=yes;" f"SERVER={host},{port};"
f"Connection Timeout={self.config.get('connection_timeout', 30)};" f"DATABASE={db};"
) f"UID={user};PWD={pwd};"
f"TrustServerCertificate=yes;"
f"Connection Timeout={self.config.get('connection_timeout', 30)};"
)
def connect(self) -> None: def connect(self) -> None:
if not PYODBC_AVAILABLE: if not PYODBC_AVAILABLE:
@@ -63,102 +67,111 @@ class MSSQLDriver(BaseDriver):
except Exception as e: except Exception as e:
return False, str(e) return False, str(e)
@contextmanager
def _cur(self): def _cur(self):
return self._connection.cursor() with self._lock:
cur = self._connection.cursor()
try:
yield cur
finally:
try:
cur.close()
except Exception:
pass
# ── Schema introspection ────────────────────────────────────────────────── # ── Schema introspection ──────────────────────────────────────────────────
def get_databases(self) -> list: def get_databases(self) -> list:
c = self._cur() with self._cur() as c:
c.execute("SELECT name FROM sys.databases ORDER BY name") c.execute("SELECT name FROM sys.databases ORDER BY name")
return [r[0] for r in c.fetchall()] return [r[0] for r in c.fetchall()]
def get_tables(self, database: str) -> list: def get_tables(self, database: str) -> list:
c = self._cur() with self._cur() as c:
c.execute(f""" c.execute(f"""
SELECT t.name, s.name, SELECT t.name, s.name,
COALESCE(p.rows, 0), 0, '', '' COALESCE(p.rows, 0), 0, '', ''
FROM [{database}].sys.tables t FROM [{database}].sys.tables t
JOIN [{database}].sys.schemas s ON t.schema_id = s.schema_id JOIN [{database}].sys.schemas s ON t.schema_id = s.schema_id
LEFT JOIN ( LEFT JOIN (
SELECT object_id, SUM(rows) AS rows SELECT object_id, SUM(rows) AS rows
FROM [{database}].sys.partitions WHERE index_id IN (0,1) FROM [{database}].sys.partitions WHERE index_id IN (0,1)
GROUP BY object_id GROUP BY object_id
) p ON p.object_id = t.object_id ) p ON p.object_id = t.object_id
ORDER BY t.name ORDER BY t.name
""") """)
return [TableInfo(name=r[0], schema=r[1], row_count=r[2], return [TableInfo(name=r[0], schema=r[1], row_count=r[2],
size_bytes=r[3], engine=r[4], comment=r[5]) size_bytes=r[3], engine=r[4], comment=r[5])
for r in c.fetchall()] for r in c.fetchall()]
def get_views(self, database: str) -> list: def get_views(self, database: str) -> list:
c = self._cur() with self._cur() as c:
c.execute(f"SELECT name FROM [{database}].sys.views ORDER BY name") c.execute(f"SELECT name FROM [{database}].sys.views ORDER BY name")
return [r[0] for r in c.fetchall()] return [r[0] for r in c.fetchall()]
def get_columns(self, database: str, table: str) -> list: def get_columns(self, database: str, table: str) -> list:
c = self._cur() with self._cur() as c:
c.execute(f""" c.execute(f"""
SELECT c.name, tp.name, c.is_nullable, dc.definition, 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 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 fk.parent_column_id IS NOT NULL THEN 1 ELSE 0 END,
CASE WHEN c.is_identity = 1 THEN 'auto_increment' ELSE '' END CASE WHEN c.is_identity = 1 THEN 'auto_increment' ELSE '' END
FROM [{database}].sys.columns c FROM [{database}].sys.columns c
JOIN [{database}].sys.types tp ON tp.user_type_id = c.user_type_id 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 JOIN [{database}].sys.tables t ON t.object_id = c.object_id
LEFT JOIN [{database}].sys.default_constraints dc LEFT JOIN [{database}].sys.default_constraints dc
ON dc.parent_object_id = c.object_id AND dc.parent_column_id = c.column_id ON dc.parent_object_id = c.object_id AND dc.parent_column_id = c.column_id
LEFT JOIN ( LEFT JOIN (
SELECT ic.column_id, ic.object_id SELECT ic.column_id, ic.object_id
FROM [{database}].sys.index_columns ic 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 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 WHERE i.is_primary_key = 1
) pk ON pk.object_id = c.object_id AND pk.column_id = c.column_id ) pk ON pk.object_id = c.object_id AND pk.column_id = c.column_id
LEFT JOIN ( LEFT JOIN (
SELECT fkc.parent_column_id, fkc.parent_object_id SELECT fkc.parent_column_id, fkc.parent_object_id
FROM [{database}].sys.foreign_key_columns fkc FROM [{database}].sys.foreign_key_columns fkc
) fk ON fk.parent_object_id = c.object_id AND fk.parent_column_id = c.column_id ) fk ON fk.parent_object_id = c.object_id AND fk.parent_column_id = c.column_id
WHERE t.name = ? WHERE t.name = ?
ORDER BY c.column_id ORDER BY c.column_id
""", (table,)) """, (table,))
return [ColumnInfo(name=r[0], data_type=r[1], nullable=bool(r[2]), return [ColumnInfo(name=r[0], data_type=r[1], nullable=bool(r[2]),
default=r[3], is_primary_key=bool(r[4]), default=r[3], is_primary_key=bool(r[4]),
is_foreign_key=bool(r[5]), extra=r[6] or "") is_foreign_key=bool(r[5]), extra=r[6] or "")
for r in c.fetchall()] for r in c.fetchall()]
def get_indexes(self, database: str, table: str) -> list: def get_indexes(self, database: str, table: str) -> list:
c = self._cur() with self._cur() as c:
c.execute(f""" c.execute(f"""
SELECT i.name, i.is_unique, STRING_AGG(c.name, ',') WITHIN GROUP (ORDER BY ic.key_ordinal) SELECT i.name, i.is_unique, STRING_AGG(c.name, ',') WITHIN GROUP (ORDER BY ic.key_ordinal)
FROM [{database}].sys.indexes i 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.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.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 JOIN [{database}].sys.tables t ON t.object_id = i.object_id
WHERE t.name = ? WHERE t.name = ?
GROUP BY i.name, i.is_unique GROUP BY i.name, i.is_unique
""", (table,)) """, (table,))
return [IndexInfo(name=r[0], columns=r[2].split(','), return [IndexInfo(name=r[0], columns=r[2].split(','),
is_unique=bool(r[1])) is_unique=bool(r[1]))
for r in c.fetchall()] for r in c.fetchall()]
def get_foreign_keys(self, database: str, table: str) -> list: def get_foreign_keys(self, database: str, table: str) -> list:
c = self._cur() with self._cur() as c:
c.execute(f""" c.execute(f"""
SELECT fk.name, pc.name, rt.name, rc.name, SELECT fk.name, pc.name, rt.name, rc.name,
fk.update_referential_action_desc, fk.update_referential_action_desc,
fk.delete_referential_action_desc fk.delete_referential_action_desc
FROM [{database}].sys.foreign_keys fk FROM [{database}].sys.foreign_keys fk
JOIN [{database}].sys.tables pt ON pt.object_id = fk.parent_object_id 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.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.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 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 JOIN [{database}].sys.columns rc ON rc.object_id = fkc.referenced_object_id AND rc.column_id = fkc.referenced_column_id
WHERE pt.name = ? WHERE pt.name = ?
""", (table,)) """, (table,))
return [ForeignKeyInfo(name=r[0], column=r[1], return [ForeignKeyInfo(name=r[0], column=r[1],
ref_table=r[2], ref_column=r[3], ref_table=r[2], ref_column=r[3],
on_update=r[4] or "", on_delete=r[5] or "") on_update=r[4] or "", on_delete=r[5] or "")
for r in c.fetchall()] for r in c.fetchall()]
def get_table_ddl(self, database: str, table: str) -> str: def get_table_ddl(self, database: str, table: str) -> str:
cols = self.get_columns(database, table) cols = self.get_columns(database, table)
@@ -174,57 +187,57 @@ class MSSQLDriver(BaseDriver):
return "\n".join(lines) return "\n".join(lines)
def get_functions(self, database: str) -> list: def get_functions(self, database: str) -> list:
c = self._cur() with self._cur() as c:
c.execute(f""" c.execute(f"""
SELECT name FROM [{database}].sys.objects SELECT name FROM [{database}].sys.objects
WHERE type IN ('FN','IF','TF') ORDER BY name WHERE type IN ('FN','IF','TF') ORDER BY name
""") """)
return [r[0] for r in c.fetchall()] return [r[0] for r in c.fetchall()]
def get_stored_procedures(self, database: str) -> list: def get_stored_procedures(self, database: str) -> list:
c = self._cur() with self._cur() as c:
c.execute(f""" c.execute(f"""
SELECT name FROM [{database}].sys.procedures ORDER BY name SELECT name FROM [{database}].sys.procedures ORDER BY name
""") """)
return [r[0] for r in c.fetchall()] return [r[0] for r in c.fetchall()]
def get_triggers(self, database: str, table: str = "") -> list: def get_triggers(self, database: str, table: str = "") -> list:
c = self._cur() with self._cur() as c:
if table: if table:
c.execute(f""" c.execute(f"""
SELECT t.name FROM [{database}].sys.triggers t SELECT t.name FROM [{database}].sys.triggers t
JOIN [{database}].sys.tables tb ON tb.object_id = t.parent_id JOIN [{database}].sys.tables tb ON tb.object_id = t.parent_id
WHERE tb.name = ? ORDER BY t.name WHERE tb.name = ? ORDER BY t.name
""", (table,)) """, (table,))
else: else:
c.execute(f"SELECT name FROM [{database}].sys.triggers ORDER BY name") c.execute(f"SELECT name FROM [{database}].sys.triggers ORDER BY name")
return [r[0] for r in c.fetchall()] return [r[0] for r in c.fetchall()]
# ── Query execution ─────────────────────────────────────────────────────── # ── Query execution ───────────────────────────────────────────────────────
def execute_query(self, sql: str, params: Optional[tuple] = None) -> tuple: def execute_query(self, sql: str, params: Optional[tuple] = None) -> tuple:
c = self._cur() with self._cur() as c:
c.execute(sql, params or ()) c.execute(sql, params or ())
if c.description: if c.description:
cols = [d[0] for d in c.description] cols = [d[0] for d in c.description]
rows = c.fetchall() rows = c.fetchall()
return cols, [tuple(r) for r in rows], len(rows) return cols, [tuple(r) for r in rows], len(rows)
return [], [], c.rowcount return [], [], c.rowcount
def execute_script(self, sql: str) -> list: def execute_script(self, sql: str) -> list:
results = [] results = []
stmts = [s.strip() for s in sql.split(';') if s.strip()] stmts = [s.strip() for s in sql.split(';') if s.strip()]
for stmt in stmts: for stmt in stmts:
try: try:
c = self._cur() with self._cur() as c:
c.execute(stmt) c.execute(stmt)
if c.description: if c.description:
cols = [d[0] for d in c.description] cols = [d[0] for d in c.description]
rows = [tuple(r) for r in c.fetchall()] rows = [tuple(r) for r in c.fetchall()]
results.append((cols, rows, len(rows), "")) results.append((cols, rows, len(rows), ""))
else: else:
results.append(([], [], c.rowcount, results.append(([], [], c.rowcount,
f"{c.rowcount} row(s) affected")) f"{c.rowcount} row(s) affected"))
except Exception as e: except Exception as e:
results.append(([], [], 0, f"Error: {e}")) results.append(([], [], 0, f"Error: {e}"))
return results return results
@@ -243,16 +256,16 @@ class MSSQLDriver(BaseDriver):
def get_table_row_count(self, database: str, table: str, where: str = "") -> int: def get_table_row_count(self, database: str, table: str, where: str = "") -> int:
sql = f"SELECT COUNT(*) FROM [{database}].[dbo].[{table}]" sql = f"SELECT COUNT(*) FROM [{database}].[dbo].[{table}]"
if where: sql += f" WHERE {where}" if where: sql += f" WHERE {where}"
c = self._cur() with self._cur() as c:
c.execute(sql) c.execute(sql)
return c.fetchone()[0] return c.fetchone()[0]
def insert_row(self, database: str, table: str, data: dict) -> bool: def insert_row(self, database: str, table: str, data: dict) -> bool:
cols = ", ".join(f"[{c}]" for c in data) cols = ", ".join(f"[{c}]" for c in data)
ph = ", ".join(["?"] * len(data)) ph = ", ".join(["?"] * len(data))
c = self._cur() with self._cur() as c:
c.execute(f"INSERT INTO [{database}].[dbo].[{table}] ({cols}) VALUES ({ph})", c.execute(f"INSERT INTO [{database}].[dbo].[{table}] ({cols}) VALUES ({ph})",
tuple(data.values())) tuple(data.values()))
return True return True
@staticmethod @staticmethod
@@ -269,40 +282,40 @@ class MSSQLDriver(BaseDriver):
def update_row(self, database: str, table: str, data: dict, where: dict) -> bool: def update_row(self, database: str, table: str, data: dict, where: dict) -> bool:
set_cl = ", ".join(f"[{c}] = ?" for c in data) set_cl = ", ".join(f"[{c}] = ?" for c in data)
where_cl, where_params = self._where(where) where_cl, where_params = self._where(where)
c = self._cur() with self._cur() as c:
c.execute(f"UPDATE [{database}].[dbo].[{table}] SET {set_cl} WHERE {where_cl}", c.execute(f"UPDATE [{database}].[dbo].[{table}] SET {set_cl} WHERE {where_cl}",
tuple(data.values()) + tuple(where_params)) tuple(data.values()) + tuple(where_params))
return True return True
def delete_row(self, database: str, table: str, where: dict) -> bool: def delete_row(self, database: str, table: str, where: dict) -> bool:
where_cl, where_params = self._where(where) where_cl, where_params = self._where(where)
c = self._cur() with self._cur() as c:
c.execute(f"DELETE FROM [{database}].[dbo].[{table}] WHERE {where_cl}", c.execute(f"DELETE FROM [{database}].[dbo].[{table}] WHERE {where_cl}",
tuple(where_params)) tuple(where_params))
return True return True
# ── Server tools ────────────────────────────────────────────────────────── # ── Server tools ──────────────────────────────────────────────────────────
def explain_query(self, sql: str) -> tuple: def explain_query(self, sql: str) -> tuple:
c = self._cur() with self._cur() as c:
c.execute(f"SET SHOWPLAN_TEXT ON; {sql}; SET SHOWPLAN_TEXT OFF") c.execute(f"SET SHOWPLAN_TEXT ON; {sql}; SET SHOWPLAN_TEXT OFF")
return ["Plan"], c.fetchall() return ["Plan"], c.fetchall()
def get_process_list(self) -> tuple: def get_process_list(self) -> tuple:
c = self._cur() with self._cur() as c:
c.execute(""" c.execute("""
SELECT session_id, login_name, status, host_name, SELECT session_id, login_name, status, host_name,
program_name, cpu_time, text program_name, cpu_time, text
FROM sys.dm_exec_sessions s FROM sys.dm_exec_sessions s
CROSS APPLY sys.dm_exec_sql_text(s.most_recent_sql_handle) t CROSS APPLY sys.dm_exec_sql_text(s.most_recent_sql_handle) t
WHERE s.is_user_process = 1 WHERE s.is_user_process = 1
""") """)
cols = [d[0] for d in c.description] cols = [d[0] for d in c.description]
return cols, [tuple(r) for r in c.fetchall()] return cols, [tuple(r) for r in c.fetchall()]
def kill_process(self, process_id: int) -> bool: def kill_process(self, process_id: int) -> bool:
c = self._cur() with self._cur() as c:
c.execute(f"KILL {process_id}") c.execute(f"KILL {int(process_id)}")
return True return True
# ── Table designer ──────────────────────────────────────────────────────── # ── Table designer ────────────────────────────────────────────────────────
@@ -314,20 +327,19 @@ class MSSQLDriver(BaseDriver):
default_clause = f" DEFAULT {default}" if default is not None and default != "" else "" default_clause = f" DEFAULT {default}" if default is not None and default != "" else ""
sql = (f"ALTER TABLE [{database}].[dbo].[{table}] " sql = (f"ALTER TABLE [{database}].[dbo].[{table}] "
f"ADD [{col_name}] {col_type} {null_clause}{default_clause}") f"ADD [{col_name}] {col_type} {null_clause}{default_clause}")
c = self._cur() with self._cur() as c:
c.execute(sql) c.execute(sql)
return True return True
def drop_column(self, database: str, table: str, col_name: str) -> bool: def drop_column(self, database: str, table: str, col_name: str) -> bool:
sql = f"ALTER TABLE [{database}].[dbo].[{table}] DROP COLUMN [{col_name}]" sql = f"ALTER TABLE [{database}].[dbo].[{table}] DROP COLUMN [{col_name}]"
c = self._cur() with self._cur() as c:
c.execute(sql) c.execute(sql)
return True return True
def rename_column(self, database: str, table: str, def rename_column(self, database: str, table: str,
old_name: str, new_name: str) -> bool: old_name: str, new_name: str) -> bool:
# sp_rename is the standard way in MSSQL
sql = f"EXEC sp_rename '[{database}].[dbo].[{table}].[{old_name}]', '{new_name}', 'COLUMN'" sql = f"EXEC sp_rename '[{database}].[dbo].[{table}].[{old_name}]', '{new_name}', 'COLUMN'"
c = self._cur() with self._cur() as c:
c.execute(sql) c.execute(sql)
return True return True
app/models/connection_model.py
+1 -1
View File
@@ -2,7 +2,6 @@
Connection profile dataclass and registry. Connection profile dataclass and registry.
""" """
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import Optional
import uuid import uuid
@@ -14,6 +13,7 @@ class ConnectionProfile:
port: int = 3306 port: int = 3306
database: str = "" database: str = ""
username: str = "" username: str = ""
password: str = field(default="", repr=False)
color: str = "#89b4fa" color: str = "#89b4fa"
id: str = field(default_factory=lambda: str(uuid.uuid4())) id: str = field(default_factory=lambda: str(uuid.uuid4()))
ssl: bool = False ssl: bool = False