mssql-agent-mcp 1.0.0__py3-none-any.whl

mssql_agent_mcp/tools/procedures.py

@@ -0,0 +1,349 @@
+"""Stored procedure management tools."""
+
+import re
+from pathlib import Path
+from typing import Annotated
+
+from ..config import READONLY_MODE
+from ..utils import (
+    format_result,
+    get_connection,
+    row_to_dict,
+    rows_to_dicts,
+    sanitize_filename,
+)
+
+
+def list_procedures(
+    schema: Annotated[str | None, "Schema name to filter procedures (optional)"] = None,
+    include_definition: Annotated[bool, "If true, include the full procedure definition"] = False
+) -> str:
+    """List all stored procedures in the current database."""
+    if include_definition:
+        base_query = """
+            SELECT SCHEMA_NAME(p.schema_id) as [schema], p.name as [procedure_name],
+                   p.create_date, p.modify_date,
+                   CASE p.is_auto_executed WHEN 1 THEN 'Yes' ELSE 'No' END as [auto_execute],
+                   CASE WHEN m.definition IS NULL THEN 'Encrypted' ELSE 'Visible' END as [definition_status],
+                   m.definition as [definition]
+            FROM sys.procedures p LEFT JOIN sys.sql_modules m ON p.object_id = m.object_id
+        """
+    else:
+        base_query = """
+            SELECT SCHEMA_NAME(p.schema_id) as [schema], p.name as [procedure_name],
+                   p.create_date, p.modify_date,
+                   CASE p.is_auto_executed WHEN 1 THEN 'Yes' ELSE 'No' END as [auto_execute],
+                   CASE WHEN m.definition IS NULL THEN 'Encrypted' ELSE 'Visible' END as [definition_status]
+            FROM sys.procedures p LEFT JOIN sys.sql_modules m ON p.object_id = m.object_id
+        """
+
+    conn = get_connection()
+    try:
+        cursor = conn.cursor()
+        if schema:
+            query = base_query + " WHERE SCHEMA_NAME(p.schema_id) = ? ORDER BY [schema], [procedure_name]"
+            cursor.execute(query, (schema,))
+        else:
+            query = base_query + " ORDER BY [schema], [procedure_name]"
+            cursor.execute(query)
+        return format_result(rows_to_dicts(cursor))
+    finally:
+        conn.close()
+
+
+def list_all_procedures(
+    include_definition: Annotated[bool, "If true, include the full procedure definition"] = False,
+    database_filter: Annotated[str | None, "Filter databases by name (partial match)"] = None
+) -> str:
+    """List all stored procedures across ALL databases on the SQL Server instance."""
+    db_query = """
+        SELECT name FROM sys.databases
+        WHERE state_desc = 'ONLINE' AND name NOT IN ('master', 'tempdb', 'model', 'msdb')
+        ORDER BY name
+    """
+    conn = get_connection()
+    try:
+        cursor = conn.cursor()
+        cursor.execute(db_query)
+        databases = [row[0] for row in cursor.fetchall()]
+    finally:
+        conn.close()
+
+    if database_filter:
+        databases = [db for db in databases if database_filter.lower() in db.lower()]
+
+    all_procedures = []
+    for db_name in databases:
+        try:
+            safe_db_name = "".join(c for c in db_name if c.isalnum() or c in "_-")
+            if include_definition:
+                query = f"""
+                    SELECT '{safe_db_name}' as [database], SCHEMA_NAME(p.schema_id) as [schema],
+                           p.name as [procedure_name], p.create_date, p.modify_date,
+                           CASE p.is_auto_executed WHEN 1 THEN 'Yes' ELSE 'No' END as [auto_execute],
+                           CASE WHEN m.definition IS NULL THEN 'Encrypted' ELSE 'Visible' END as [definition_status],
+                           m.definition as [definition]
+                    FROM [{safe_db_name}].sys.procedures p
+                    LEFT JOIN [{safe_db_name}].sys.sql_modules m ON p.object_id = m.object_id
+                    ORDER BY SCHEMA_NAME(p.schema_id), p.name
+                """
+            else:
+                query = f"""
+                    SELECT '{safe_db_name}' as [database], SCHEMA_NAME(p.schema_id) as [schema],
+                           p.name as [procedure_name], p.create_date, p.modify_date,
+                           CASE p.is_auto_executed WHEN 1 THEN 'Yes' ELSE 'No' END as [auto_execute],
+                           CASE WHEN m.definition IS NULL THEN 'Encrypted' ELSE 'Visible' END as [definition_status]
+                    FROM [{safe_db_name}].sys.procedures p
+                    LEFT JOIN [{safe_db_name}].sys.sql_modules m ON p.object_id = m.object_id
+                    ORDER BY SCHEMA_NAME(p.schema_id), p.name
+                """
+            conn = get_connection()
+            try:
+                cursor = conn.cursor()
+                cursor.execute(query)
+                all_procedures.extend(rows_to_dicts(cursor))
+            finally:
+                conn.close()
+        except Exception as e:
+            all_procedures.append({"database": db_name, "error": f"Cannot access: {str(e)}"})
+
+    return format_result(all_procedures)
+
+
+def get_procedure_details(
+    procedure: Annotated[str, "The stored procedure name"],
+    schema: Annotated[str, "The schema name (defaults to 'dbo')"] = "dbo"
+) -> str:
+    """Get detailed information about a specific stored procedure including its definition."""
+    query = """
+        SELECT SCHEMA_NAME(p.schema_id) as [schema], p.name as [procedure_name],
+               p.create_date, p.modify_date, p.type_desc as [type],
+               CASE p.is_auto_executed WHEN 1 THEN 'Yes' ELSE 'No' END as [auto_execute],
+               m.definition as [definition], m.uses_ansi_nulls, m.uses_quoted_identifier, m.is_schema_bound
+        FROM sys.procedures p LEFT JOIN sys.sql_modules m ON p.object_id = m.object_id
+        WHERE p.name = ? AND SCHEMA_NAME(p.schema_id) = ?
+    """
+    conn = get_connection()
+    try:
+        cursor = conn.cursor()
+        cursor.execute(query, (procedure, schema))
+        result = row_to_dict(cursor)
+        if result:
+            return format_result(result)
+        return format_result({"error": f"Procedure '{schema}.{procedure}' not found"})
+    finally:
+        conn.close()
+
+
+def get_procedure_parameters(
+    procedure: Annotated[str, "The stored procedure name"],
+    schema: Annotated[str, "The schema name (defaults to 'dbo')"] = "dbo"
+) -> str:
+    """Get parameters for a specific stored procedure."""
+    query = """
+        SELECT par.name as [parameter_name], TYPE_NAME(par.user_type_id) as [data_type],
+               par.max_length, par.precision, par.scale,
+               CASE par.is_output WHEN 1 THEN 'Yes' ELSE 'No' END as [is_output],
+               CASE par.has_default_value WHEN 1 THEN 'Yes' ELSE 'No' END as [has_default],
+               par.default_value
+        FROM sys.procedures p
+        JOIN sys.parameters par ON p.object_id = par.object_id
+        JOIN sys.schemas s ON p.schema_id = s.schema_id
+        WHERE p.name = ? AND s.name = ?
+        ORDER BY par.parameter_id
+    """
+    conn = get_connection()
+    try:
+        cursor = conn.cursor()
+        cursor.execute(query, (procedure, schema))
+        return format_result(rows_to_dicts(cursor))
+    finally:
+        conn.close()
+
+
+def export_procedures_to_files(
+    output_dir: Annotated[str, "The root output directory path where database folders will be created"],
+    databases: Annotated[list[str] | None, "List of database names to export (optional, defaults to all)"] = None
+) -> str:
+    """Export all stored procedures from specified databases to SQL files."""
+    if databases:
+        db_list = databases
+    else:
+        db_query = """
+            SELECT name FROM sys.databases
+            WHERE state_desc = 'ONLINE' AND name NOT IN ('master', 'tempdb', 'model', 'msdb')
+            ORDER BY name
+        """
+        conn = get_connection()
+        try:
+            cursor = conn.cursor()
+            cursor.execute(db_query)
+            db_list = [row[0] for row in cursor.fetchall()]
+        finally:
+            conn.close()
+
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    export_summary = {"output_directory": str(output_path), "databases": []}
+
+    for db_name in db_list:
+        db_summary = {"database": db_name, "schemas": [], "total_procedures": 0, "errors": []}
+
+        try:
+            safe_db_name = "".join(c for c in db_name if c.isalnum() or c in "_-")
+            query = f"""
+                SELECT s.name as [schema_name], p.name as [procedure_name],
+                       p.create_date, p.modify_date, m.definition
+                FROM [{safe_db_name}].sys.procedures p
+                JOIN [{safe_db_name}].sys.schemas s ON p.schema_id = s.schema_id
+                LEFT JOIN [{safe_db_name}].sys.sql_modules m ON p.object_id = m.object_id
+                ORDER BY s.name, p.name
+            """
+
+            conn = get_connection()
+            try:
+                cursor = conn.cursor()
+                cursor.execute(query)
+                procedures = rows_to_dicts(cursor)
+            finally:
+                conn.close()
+
+            if not procedures:
+                db_summary["errors"].append("No procedures found")
+                export_summary["databases"].append(db_summary)
+                continue
+
+            db_dir = output_path / sanitize_filename(db_name)
+            db_dir.mkdir(parents=True, exist_ok=True)
+
+            schemas = {}
+            for proc in procedures:
+                schema_name = proc["schema_name"] or "dbo"
+                if schema_name not in schemas:
+                    schemas[schema_name] = []
+                schemas[schema_name].append(proc)
+
+            for schema_name, procs in schemas.items():
+                schema_summary = {"schema": schema_name, "procedures": []}
+                schema_dir = db_dir / sanitize_filename(schema_name)
+                schema_dir.mkdir(parents=True, exist_ok=True)
+
+                for proc in procs:
+                    proc_name = proc["procedure_name"]
+                    definition = proc["definition"]
+
+                    if definition is None:
+                        file_content = f"""-- Database: {db_name}
+-- Schema: {schema_name}
+-- Procedure: {proc_name}
+-- Created: {proc["create_date"]}
+-- Modified: {proc["modify_date"]}
+-- ============================================
+-- WARNING: This procedure is encrypted and cannot be exported.
+"""
+                        schema_summary["procedures"].append({"name": proc_name, "status": "encrypted"})
+                    else:
+                        file_content = f"""-- Database: {db_name}
+-- Schema: {schema_name}
+-- Procedure: {proc_name}
+-- Created: {proc["create_date"]}
+-- Modified: {proc["modify_date"]}
+-- ============================================
+
+{definition}
+"""
+                        schema_summary["procedures"].append({"name": proc_name, "status": "exported"})
+
+                    file_path = schema_dir / f"{sanitize_filename(proc_name)}.sql"
+                    file_path.write_text(file_content, encoding="utf-8")
+                    db_summary["total_procedures"] += 1
+
+                db_summary["schemas"].append(schema_summary)
+        except Exception as e:
+            db_summary["errors"].append(f"Error accessing database: {str(e)}")
+
+        export_summary["databases"].append(db_summary)
+
+    export_summary["total_databases"] = len(export_summary["databases"])
+    export_summary["total_procedures"] = sum(db["total_procedures"] for db in export_summary["databases"])
+    return format_result(export_summary)
+
+
+def update_procedure_from_file(
+    file_path: Annotated[str, "The absolute path to the SQL file (e.g., /path/to/procedures/materialdb/dbo/usp_get_facility_info_data.sql)"]
+) -> str:
+    """Update a stored procedure from an edited SQL file. Automatically converts CREATE to ALTER."""
+    file_path_obj = Path(file_path)
+
+    if not file_path_obj.exists():
+        return format_result({"success": False, "error": f"File not found: {file_path}"})
+
+    if file_path_obj.suffix.lower() != ".sql":
+        return format_result({"success": False, "error": "File must be a .sql file"})
+
+    procedure_name = file_path_obj.stem
+    schema_name = file_path_obj.parent.name
+    database_name = file_path_obj.parent.parent.name
+
+    file_content = file_path_obj.read_text(encoding="utf-8")
+
+    # Extract metadata from header
+    header_db_match = re.search(r"--\s*Database:\s*(.+)", file_content)
+    header_schema_match = re.search(r"--\s*Schema:\s*(.+)", file_content)
+    header_proc_match = re.search(r"--\s*Procedure:\s*(.+)", file_content)
+
+    if header_db_match:
+        database_name = header_db_match.group(1).strip()
+    if header_schema_match:
+        schema_name = header_schema_match.group(1).strip()
+    if header_proc_match:
+        procedure_name = header_proc_match.group(1).strip()
+
+    # Skip header comments
+    lines = file_content.split("\n")
+    sql_start_index = 0
+    for i, line in enumerate(lines):
+        if line.strip() == "-- ============================================":
+            sql_start_index = i + 1
+            break
+
+    sql_content = "\n".join(lines[sql_start_index:]).strip()
+
+    if not sql_content:
+        return format_result({"success": False, "error": "SQL file is empty"})
+
+    if "This procedure is encrypted" in file_content:
+        return format_result({"success": False, "error": "Cannot update an encrypted procedure"})
+
+    # Convert CREATE to ALTER
+    sql_upper = sql_content.upper().strip()
+    if sql_upper.startswith("CREATE"):
+        sql_content = re.sub(r"^CREATE(\s+)PROC(EDURE)?", r"ALTER\1PROC\2", sql_content, count=1, flags=re.IGNORECASE)
+    elif not sql_upper.startswith("ALTER"):
+        return format_result({"success": False, "error": "SQL must contain CREATE or ALTER PROCEDURE statement"})
+
+    if READONLY_MODE:
+        return format_result({"success": False, "error": "Write operations are disabled in read-only mode."})
+
+    safe_db_name = "".join(c for c in database_name if c.isalnum() or c in "_-")
+
+    try:
+        conn = get_connection()
+        try:
+            cursor = conn.cursor()
+            cursor.execute(f"USE [{safe_db_name}]")
+            cursor.execute(sql_content)
+            conn.commit()
+        finally:
+            conn.close()
+
+        return format_result({
+            "success": True,
+            "message": "Successfully updated procedure",
+            "database": database_name,
+            "schema": schema_name,
+            "procedure": procedure_name,
+        })
+    except Exception as e:
+        return format_result({"success": False, "error": f"Failed to update procedure: {str(e)}"})

mssql_agent_mcp/utils.py

@@ -0,0 +1,127 @@
+"""Database connection and utility functions for MSSQL MCP Server."""
+
+import json
+import re
+import struct
+
+import pyodbc
+import sqlparse
+
+from .config import BLOCKED_STATEMENT_TYPES, DB_CONFIG, READONLY_MODE
+
+
+def validate_query(sql: str) -> tuple[bool, str]:
+    """Validate SQL query against readonly restrictions."""
+    if not READONLY_MODE:
+        return True, ""
+
+    try:
+        parsed = sqlparse.parse(sql)
+        for statement in parsed:
+            stmt_type = statement.get_type()
+            if stmt_type and stmt_type.upper() in BLOCKED_STATEMENT_TYPES:
+                return False, f"Blocked: {stmt_type} statements are not allowed in read-only mode."
+
+            tokens = [t for t in statement.tokens if not t.is_whitespace]
+            if tokens:
+                first_token = str(tokens[0]).upper().strip()
+                for blocked in BLOCKED_STATEMENT_TYPES:
+                    if first_token.startswith(blocked):
+                        return False, f"Blocked: {blocked} statements are not allowed in read-only mode."
+    except Exception as e:
+        return False, f"Query validation failed: {str(e)}"
+
+    return True, ""
+
+
+def get_azure_token() -> bytes:
+    """Get Azure AD token for database authentication."""
+    try:
+        from azure.identity import DefaultAzureCredential
+    except ImportError:
+        raise ImportError(
+            "Azure AD authentication requires azure-identity. "
+            "Install with: pip install mssql-agent-mcp[azure]"
+        )
+
+    credential = DefaultAzureCredential()
+    token = credential.get_token("https://database.windows.net/.default")
+    token_bytes = token.token.encode("utf-16-le")
+    return struct.pack(f"<I{len(token_bytes)}s", len(token_bytes), token_bytes)
+
+
+def get_connection():
+    """Create and return a database connection."""
+    auth_mode = DB_CONFIG["auth_mode"].lower()
+
+    conn_str_parts = [
+        f"DRIVER={{{DB_CONFIG['driver']}}}",
+        f"SERVER={DB_CONFIG['server']},{DB_CONFIG['port']}",
+        f"DATABASE={DB_CONFIG['database']}",
+        f"Encrypt={DB_CONFIG['encrypt']}",
+        f"TrustServerCertificate={DB_CONFIG['trust_server_certificate']}",
+    ]
+
+    if auth_mode == "windows":
+        conn_str_parts.append("Trusted_Connection=yes")
+        return pyodbc.connect(";".join(conn_str_parts))
+    elif auth_mode == "azure":
+        token_struct = get_azure_token()
+        return pyodbc.connect(";".join(conn_str_parts), attrs_before={1256: token_struct})
+    else:
+        conn_str_parts.extend([
+            f"UID={DB_CONFIG['user']}",
+            f"PWD={DB_CONFIG['password']}",
+        ])
+        return pyodbc.connect(";".join(conn_str_parts))
+
+
+def rows_to_dicts(cursor) -> list[dict]:
+    """Convert pyodbc cursor results to list of dictionaries."""
+    columns = [column[0] for column in cursor.description] if cursor.description else []
+    return [dict(zip(columns, row)) for row in cursor.fetchall()]
+
+
+def row_to_dict(cursor) -> dict | None:
+    """Convert single pyodbc cursor result to dictionary."""
+    if not cursor.description:
+        return None
+    columns = [column[0] for column in cursor.description]
+    row = cursor.fetchone()
+    return dict(zip(columns, row)) if row else None
+
+
+def format_result(data) -> str:
+    """Format query results as JSON string."""
+    return json.dumps(data, indent=2, default=str)
+
+
+def sanitize_filename(name: str) -> str:
+    """Sanitize a string for use as a filename."""
+    sanitized = re.sub(r'[<>:"/\\|?*]', "_", name)
+    sanitized = re.sub(r"[\s_]+", "_", sanitized)
+    return sanitized.strip("_ ")
+
+
+def validate_sql_syntax(sql: str) -> dict:
+    """Validate SQL syntax using SET PARSEONLY."""
+    sql_stripped = sql.strip()
+    if not sql_stripped:
+        return {"valid": False, "error": "SQL command is empty"}
+
+    try:
+        conn = get_connection()
+        try:
+            cursor = conn.cursor()
+            cursor.execute("SET PARSEONLY ON")
+            try:
+                cursor.execute(sql)
+            except Exception as parse_error:
+                return {"valid": False, "error": f"SQL syntax error: {str(parse_error)}"}
+            finally:
+                cursor.execute("SET PARSEONLY OFF")
+        finally:
+            conn.close()
+        return {"valid": True, "error": None}
+    except Exception as e:
+        return {"valid": True, "error": None, "warning": f"Could not validate: {str(e)}"}