15分でわかる！AIエージェントでMCPを活用するメリット #ChatGPT – Qiita

from typing import Any, List, Dict, Optional
import aiomysql
import os
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP

# .envファイルを読み込む
load_dotenv()

# FastMCPサーバーの初期化
mcp = FastMCP("mysql_server")

# データベース接続情報を環境変数から取得
DB_CONFIG = {
    "host": os.getenv("DB_HOST", "localhost"),
    "port": int(os.getenv("DB_PORT", "3306")),
    "user": os.getenv("DB_USER", ""),
    "password": os.getenv("DB_PASSWORD", ""),
    "db": os.getenv("DB_NAME", ""),
    "charset": os.getenv("DB_CHARSET", "utf8mb4"),
    "cursorclass": aiomysql.DictCursor
}

# 接続プールを保持する変数
pool = None

async def get_pool():
    """MySQLコネクションプールを取得する"""
    global pool
    if pool is None:
        pool = await aiomysql.create_pool(**DB_CONFIG)
    return pool

@mcp.tool()
async def execute_query(query: str, params: Optional[List[Any]] = None) -> Dict[str, Any]:
    """SQLクエリを実行する

    Args:
        query: 実行するSQLクエリ
        params: クエリパラメータ（オプション）

    Returns:
        クエリ結果とメタデータを含む辞書
    """
    # DELETE操作を禁止
    if query.strip().upper().startswith('DELETE'):
        return {
            "success": False,
            "error": "削除操作は許可されていません"
        }
    pool = await get_pool()
    
    async with pool.acquire() as conn:
        async with conn.cursor() as cursor:
            try:
                # パラメータが提供されている場合はそれを使用
                if params:
                    await cursor.execute(query, params)
                else:
                    await cursor.execute(query)
                
                # SELECTクエリの場合は結果を取得
                if query.strip().upper().startswith('SELECT'):
                    results = await cursor.fetchall()
                    return {
                        "success": True,
                        "rowCount": len(results),
                        "results": results
                    }
                # INSERT, UPDATE等の場合
                else:
                    # 変更をコミット
                    await conn.commit()
                    return {
                        "success": True,
                        "affectedRows": cursor.rowcount,
                        "lastRowId": cursor.lastrowid
                    }
            except Exception as e:
                # エラーが発生した場合はロールバック
                await conn.rollback()
                return {
                    "success": False,
                    "error": str(e)
                }

@mcp.tool()
async def get_tables() -> Dict[str, Any]:
    """データベース内のテーブル一覧を取得する"""
    query = "SHOW TABLES"
    return await execute_query(query)

@mcp.tool()
async def describe_table(table_name: str) -> Dict[str, Any]:
    """指定したテーブルの構造を表示する

    Args:
        table_name: 表示するテーブル名
    """
    query = f"DESCRIBE {table_name}"
    return await execute_query(query)

@mcp.tool()
async def select_from_table(table_name: str, columns: str = "*", 
                           where_clause: str = "", 
                           limit: int = 100,
                           offset: int = 0) -> Dict[str, Any]:
    """テーブルからデータを取得する

    Args:
        table_name: 取得対象のテーブル名
        columns: 取得するカラム (デフォルト: "*")
        where_clause: WHERE句 (オプション)
        limit: 取得する最大行数 (デフォルト: 100)
        offset: 開始位置 (デフォルト: 0)
    """
    query = f"SELECT {columns} FROM {table_name}"
    
    if where_clause:
        query += f" WHERE {where_clause}"
    
    query += f" LIMIT {limit} OFFSET {offset}"
    
    return await execute_query(query)

@mcp.tool()
async def insert_into_table(table_name: str, data: Dict[str, Any]) -> Dict[str, Any]:
    """テーブルに新しい行を挿入する

    Args:
        table_name: 挿入先のテーブル名
        data: 挿入するデータ（カラム名: 値 の辞書）
    """
    columns = ", ".join(data.keys())
    placeholders = ", ".join(["%s"] * len(data))
    values = list(data.values())
    
    query = f"INSERT INTO {table_name} ({columns}) VALUES ({placeholders})"
    
    return await execute_query(query, values)

@mcp.tool()
async def update_table(table_name: str, 
                      data: Dict[str, Any], 
                      where_clause: str) -> Dict[str, Any]:
    """テーブル内の行を更新する

    Args:
        table_name: 更新対象のテーブル名
        data: 更新するデータ（カラム名: 値 の辞書）
        where_clause: 更新対象を絞り込むWHERE句
    """
    set_clause = ", ".join([f"{key} = %s" for key in data.keys()])
    values = list(data.values())
    
    query = f"UPDATE {table_name} SET {set_clause} WHERE {where_clause}"
    
    return await execute_query(query, values)

@mcp.tool()
async def execute_transaction(queries: List[Dict[str, Any]]) -> Dict[str, Any]:
    """複数のクエリをトランザクションとして実行する

    Args:
        queries: 実行するクエリのリスト。各クエリは {'query': 'SQLクエリ', 'params': [パラメータ]} の形式
    """
    # DELETE操作を含むトランザクションを禁止
    for query_data in queries:
        query = query_data.get('query', '')
        if query.strip().upper().startswith('DELETE'):
            return {
                "success": False,
                "error": "トランザクション内の削除操作は許可されていません"
            }
    pool = await get_pool()
    results = []
    
    async with pool.acquire() as conn:
        try:
            # 自動コミットを無効化
            await conn.begin()
            
            async with conn.cursor() as cursor:
                for query_data in queries:
                    query = query_data.get('query')
                    params = query_data.get('params')
                    
                    if params:
                        await cursor.execute(query, params)
                    else:
                        await cursor.execute(query)
                    
                    if query.strip().upper().startswith('SELECT'):
                        query_result = await cursor.fetchall()
                        results.append({
                            "success": True,
                            "rowCount": len(query_result),
                            "results": query_result
                        })
                    else:
                        results.append({
                            "success": True,
                            "affectedRows": cursor.rowcount,
                            "lastRowId": cursor.lastrowid
                        })
            
            # すべてのクエリが成功したらコミット
            await conn.commit()
            return {
                "success": True,
                "results": results
            }
        except Exception as e:
            # エラーが発生したらロールバック
            await conn.rollback()
            return {
                "success": False,
                "error": str(e),
                "completed_results": results
            }

if __name__ == "__main__":
    # サーバーの初期化と実行
    mcp.run(transport='stdio')