MCP Vertica

Instructions

Args:
    ctx: FastMCP context for progress reporting and logging
    query: SQL query to execute
    database: Optional database name to execute the query against

Returns:
    Query results as a string

Implementation Reference

• src/mcp_vertica/mcp.py:200-246 (handler)

  The primary handler for the 'execute_query' MCP tool. Decorated with @mcp.tool() for automatic registration. Executes SQL queries on Vertica, manages connections, checks schema permissions based on operation type, fetches results, and returns them as a string.

  @mcp.tool()
  async def execute_query(ctx: Context, query: str) -> str:
      """Execute a SQL query and return the results.
  
      Args:
          ctx: FastMCP context for progress reporting and logging
          query: SQL query to execute
          database: Optional database name to execute the query against
  
      Returns:
          Query results as a string
      """
      await ctx.info(f"Executing query: {query}")
  
      # Get or create connection manager
      manager = await get_or_create_manager(ctx)
      if not manager:
          return "Error: Failed to initialize database connection. Check configuration."
  
      # Extract schema from query if not provided
      schema = extract_schema_from_query(query)
      # Check operation permissions
      operation = extract_operation_type(query)
      if operation and not manager.is_operation_allowed(schema or "default", operation):
          error_msg = f"Operation {operation.name} not allowed for schema {schema}"
          await ctx.error(error_msg)
          return error_msg
  
      conn = None
      cursor = None
      try:
          conn = manager.get_connection()  # Always use default DB connection
          cursor = conn.cursor()
          cursor.execute(query)
          results = cursor.fetchall()
          await ctx.info(f"Query executed successfully, returned {len(results)} rows")
          return str(results)
      except Exception as e:
          error_msg = f"Error executing query: {str(e)}"
          await ctx.error(error_msg)
          return error_msg
      finally:
          if cursor:
              cursor.close()
          if conn:
              manager.release_connection(conn)

• src/mcp_vertica/mcp.py:61-80 (helper)

  Helper function used by execute_query to lazily obtain or create the VerticaConnectionManager from the MCP context.

  async def get_or_create_manager(ctx: Context) -> VerticaConnectionManager | None:
      """Get connection manager from context or create it lazily.
  
      Args:
          ctx: FastMCP context
  
      Returns:
          VerticaConnectionManager instance or None if creation fails
      """
      manager = ctx.request_context.lifespan_context.get("vertica_manager")
      if not manager:
          try:
              manager = VerticaConnectionManager()
              config = VerticaConfig.from_env()
              manager.initialize_default(config)
              await ctx.info("Connection manager initialized from request config")
          except Exception as e:
              await ctx.error(f"Failed to initialize database connection: {str(e)}")
              return None
      return manager

• src/mcp_vertica/mcp.py:83-96 (helper)

  Helper to determine the SQL operation type (e.g., INSERT, DDL) for permission checks in execute_query.

  def extract_operation_type(query: str) -> OperationType | None:
      """Extract the operation type from a SQL query."""
      query = query.strip().upper()
  
      if query.startswith("INSERT"):
          return OperationType.INSERT
      elif query.startswith("UPDATE"):
          return OperationType.UPDATE
      elif query.startswith("DELETE"):
          return OperationType.DELETE
      elif any(query.startswith(op) for op in ["CREATE", "ALTER", "DROP", "TRUNCATE"]):
          return OperationType.DDL
      return None

• src/mcp_vertica/mcp.py:98-105 (helper)

  Helper to extract schema name from SQL query for permission validation in execute_query.

  def extract_schema_from_query(query: str) -> str | None:
      """Extract schema name from a SQL query."""
      # database.table 또는 schema.table 패턴에서 schema 추출
      match = re.search(r"([a-zA-Z0-9_]+)\.[a-zA-Z0-9_]+", query)
      if match:
          return match.group(1)
      return None