From 724594d3f193ec3eeba6b2df9c89389ff3fce239 Mon Sep 17 00:00:00 2001 From: NguyenND Date: Thu, 21 May 2026 16:02:04 -0400 Subject: [PATCH] Fix three critical bugs in ConnectionProfile and MSSQL driver - 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 --- app/drivers/mssql_driver.py | 338 +++++++++++++++++---------------- app/models/connection_model.py | 2 +- 2 files changed, 176 insertions(+), 164 deletions(-) diff --git a/app/drivers/mssql_driver.py b/app/drivers/mssql_driver.py index 9aa9a9a..d32a468 100644 --- a/app/drivers/mssql_driver.py +++ b/app/drivers/mssql_driver.py @@ -1,4 +1,5 @@ """Microsoft SQL Server driver using pyodbc.""" +from contextlib import contextmanager from typing import Optional from app.drivers.base import ( BaseDriver, ColumnInfo, IndexInfo, ForeignKeyInfo, TableInfo @@ -24,20 +25,23 @@ class MSSQLDriver(BaseDriver): db = self.config.get("database", "master") user = self.config.get("user", "") pwd = self.config.get("password", "") - # Try drivers in order of preference - for driver in [ + + preferred = [ "ODBC Driver 18 for SQL Server", "ODBC Driver 17 for SQL Server", "SQL Server", - ]: - 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)};" - ) + ] + 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: @@ -63,102 +67,111 @@ class MSSQLDriver(BaseDriver): except Exception as e: return False, str(e) + @contextmanager 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 ────────────────────────────────────────────────── def get_databases(self) -> list: - c = self._cur() - c.execute("SELECT name FROM sys.databases ORDER BY name") - return [r[0] for r in c.fetchall()] + 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: - c = self._cur() - 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()] + 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: - c = self._cur() - c.execute(f"SELECT name FROM [{database}].sys.views ORDER BY name") - return [r[0] for r in c.fetchall()] + 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: - c = self._cur() - 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()] + 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: - c = self._cur() - 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()] + 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: - c = self._cur() - 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()] + 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) @@ -174,57 +187,57 @@ class MSSQLDriver(BaseDriver): return "\n".join(lines) def get_functions(self, database: str) -> list: - c = self._cur() - 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()] + 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: - c = self._cur() - c.execute(f""" - SELECT name FROM [{database}].sys.procedures ORDER BY name - """) - return [r[0] for r in c.fetchall()] + 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: - c = self._cur() - 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()] + 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: - c = self._cur() - 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 + 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: - c = self._cur() - 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")) + 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 @@ -243,16 +256,16 @@ class MSSQLDriver(BaseDriver): 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}" - c = self._cur() - c.execute(sql) - return c.fetchone()[0] + 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)) - c = self._cur() - c.execute(f"INSERT INTO [{database}].[dbo].[{table}] ({cols}) VALUES ({ph})", - tuple(data.values())) + with self._cur() as c: + c.execute(f"INSERT INTO [{database}].[dbo].[{table}] ({cols}) VALUES ({ph})", + tuple(data.values())) return True @staticmethod @@ -269,40 +282,40 @@ class MSSQLDriver(BaseDriver): 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) - c = self._cur() - c.execute(f"UPDATE [{database}].[dbo].[{table}] SET {set_cl} WHERE {where_cl}", - tuple(data.values()) + tuple(where_params)) + 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) - c = self._cur() - c.execute(f"DELETE FROM [{database}].[dbo].[{table}] WHERE {where_cl}", - tuple(where_params)) + 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: - c = self._cur() - c.execute(f"SET SHOWPLAN_TEXT ON; {sql}; SET SHOWPLAN_TEXT OFF") - return ["Plan"], c.fetchall() + 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: - c = self._cur() - 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()] + 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: - c = self._cur() - c.execute(f"KILL {process_id}") + with self._cur() as c: + c.execute(f"KILL {int(process_id)}") return True # ── Table designer ──────────────────────────────────────────────────────── @@ -314,20 +327,19 @@ class MSSQLDriver(BaseDriver): 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}") - c = self._cur() - c.execute(sql) + 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}]" - c = self._cur() - c.execute(sql) + 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: - # sp_rename is the standard way in MSSQL sql = f"EXEC sp_rename '[{database}].[dbo].[{table}].[{old_name}]', '{new_name}', 'COLUMN'" - c = self._cur() - c.execute(sql) + with self._cur() as c: + c.execute(sql) return True diff --git a/app/models/connection_model.py b/app/models/connection_model.py index 57068eb..58177ad 100644 --- a/app/models/connection_model.py +++ b/app/models/connection_model.py @@ -2,7 +2,6 @@ Connection profile dataclass and registry. """ from dataclasses import dataclass, field -from typing import Optional import uuid @@ -14,6 +13,7 @@ class ConnectionProfile: port: int = 3306 database: str = "" username: str = "" + password: str = field(default="", repr=False) color: str = "#89b4fa" id: str = field(default_factory=lambda: str(uuid.uuid4())) ssl: bool = False