{
  "cells": [
    {
      "cell_type": "raw",
      "metadata": {
        "vscode": {
          "languageId": "raw"
        }
      },
      "source": [
        "# 🗄️ Database Operations with MCP\n",
        "\n",
        "Learn how to safely integrate databases with your MCP tools. This notebook covers secure database access patterns, connection management, and data validation - essential skills for building data-driven MCP tools.\n",
        "\n",
        "## 🎯 Learning Objectives\n",
        "\n",
        "By the end of this notebook, you will:\n",
        "- Implement secure database access in MCP tools\n",
        "- Handle database connections safely\n",
        "- Validate and sanitize database inputs\n",
        "- Build robust error handling for database operations\n",
        "- Create reusable database tool patterns\n",
        "\n",
        "## 📋 Prerequisites\n",
        "\n",
        "- Completed notebooks 01-07\n",
        "- Basic SQL knowledge\n",
        "- Python virtual environment with required packages\n",
        "- SQLite installed (included with Python)\n",
        "\n",
        "## 🔑 Key Concepts\n",
        "\n",
        "1. **Safe Database Access**\n",
        "   - Connection pooling\n",
        "   - Query parameterization\n",
        "   - Error handling\n",
        "\n",
        "2. **Data Validation**\n",
        "   - Input sanitization\n",
        "   - Schema validation\n",
        "   - Type checking\n",
        "\n",
        "3. **Connection Management**\n",
        "   - Connection lifecycle\n",
        "   - Resource cleanup\n",
        "   - Pool management\n",
        "\n",
        "## 📚 Table of Contents\n",
        "\n",
        "1. [Setting Up Database Tools](#setup)\n",
        "2. [Basic Database Operations](#basics)\n",
        "3. [Advanced Query Patterns](#advanced)\n",
        "4. [Error Handling](#errors)\n",
        "5. [Best Practices](#practices)\n",
        "\n",
        "---\n"
      ]
    },
    {
      "cell_type": "raw",
      "metadata": {
        "vscode": {
          "languageId": "raw"
        }
      },
      "source": [
        "# Setting Up Database Tools\n",
        "\n",
        "First, let's create a basic database tool that follows MCP best practices. We'll use SQLite for this example since it's included with Python and doesn't require external setup.\n",
        "\n",
        "Our database tool will have these features:\n",
        "1. Secure connection management\n",
        "2. Input validation\n",
        "3. Error handling\n",
        "4. Resource cleanup\n",
        "\n",
        "Let's start by importing the required packages:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import sqlite3\n",
        "from contextlib import contextmanager\n",
        "from typing import List, Dict, Any, Optional\n",
        "from dataclasses import dataclass\n",
        "from pathlib import Path\n",
        "\n",
        "import modelcontextprotocol as mcp\n",
        "from pydantic import BaseModel, Field\n",
        "\n",
        "# Schema for our database tool\n",
        "class QueryInput(BaseModel):\n",
        "    query: str = Field(..., description=\"SQL query to execute\")\n",
        "    parameters: Optional[List[Any]] = Field(default=None, description=\"Query parameters\")\n",
        "    \n",
        "class QueryResult(BaseModel):\n",
        "    columns: List[str]\n",
        "    rows: List[List[Any]]\n",
        "    row_count: int\n",
        "    \n",
        "# Database connection manager\n",
        "@contextmanager\n",
        "def get_db_connection(db_path: str):\n",
        "    conn = sqlite3.connect(db_path)\n",
        "    try:\n",
        "        # Enable dictionary cursor\n",
        "        conn.row_factory = sqlite3.Row\n",
        "        yield conn\n",
        "    finally:\n",
        "        conn.close()\n",
        "\n",
        "# Create a test database\n",
        "DB_PATH = \"example.db\"\n",
        "with get_db_connection(DB_PATH) as conn:\n",
        "    conn.execute(\"\"\"\n",
        "        CREATE TABLE IF NOT EXISTS users (\n",
        "            id INTEGER PRIMARY KEY,\n",
        "            name TEXT NOT NULL,\n",
        "            email TEXT UNIQUE NOT NULL\n",
        "        )\n",
        "    \"\"\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "class DatabaseTool:\n",
        "    def __init__(self, db_path: str):\n",
        "        self.db_path = db_path\n",
        "        \n",
        "    def execute_query(self, query: str, parameters: Optional[List[Any]] = None) -> QueryResult:\n",
        "        \"\"\"Execute a SQL query and return the results.\"\"\"\n",
        "        with get_db_connection(self.db_path) as conn:\n",
        "            cursor = conn.cursor()\n",
        "            try:\n",
        "                if parameters:\n",
        "                    cursor.execute(query, parameters)\n",
        "                else:\n",
        "                    cursor.execute(query)\n",
        "                    \n",
        "                # Get column names\n",
        "                columns = [description[0] for description in cursor.description] if cursor.description else []\n",
        "                \n",
        "                # Fetch all rows\n",
        "                rows = cursor.fetchall()\n",
        "                \n",
        "                # Convert rows to lists (from sqlite3.Row objects)\n",
        "                rows = [list(row) for row in rows]\n",
        "                \n",
        "                return QueryResult(\n",
        "                    columns=columns,\n",
        "                    rows=rows,\n",
        "                    row_count=len(rows)\n",
        "                )\n",
        "            except sqlite3.Error as e:\n",
        "                raise ValueError(f\"Database error: {str(e)}\")\n",
        "            \n",
        "    async def handle_query(self, input_data: QueryInput) -> QueryResult:\n",
        "        \"\"\"MCP handler for database queries.\"\"\"\n",
        "        return self.execute_query(input_data.query, input_data.parameters)\n",
        "\n",
        "# Create our MCP tool\n",
        "db_tool = DatabaseTool(DB_PATH)\n",
        "\n",
        "# Create the MCP server\n",
        "server = mcp.Server()\n",
        "server.add_tool(\"database\", db_tool.handle_query, QueryInput, QueryResult)\n",
        "\n",
        "# Example usage\n",
        "example_query = QueryInput(\n",
        "    query=\"INSERT INTO users (name, email) VALUES (?, ?)\",\n",
        "    parameters=[\"John Doe\", \"john@example.com\"]\n",
        ")\n",
        "\n",
        "result = db_tool.execute_query(example_query.query, example_query.parameters)\n",
        "print(\"Query executed successfully!\")\n"
      ]
    },
    {
      "cell_type": "raw",
      "metadata": {
        "vscode": {
          "languageId": "raw"
        }
      },
      "source": [
        "# Testing Our Database Tool\n",
        "\n",
        "Let's test our database tool with some common operations:\n",
        "\n",
        "1. **Insert Data**: Add new records to our users table\n",
        "2. **Query Data**: Retrieve and display user information\n",
        "3. **Update Data**: Modify existing records\n",
        "4. **Delete Data**: Remove records safely\n",
        "\n",
        "Each operation will demonstrate proper:\n",
        "- Parameter handling\n",
        "- Error checking\n",
        "- Result formatting\n",
        "- Resource management\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Let's add some test data\n",
        "test_users = [\n",
        "    (\"Alice Smith\", \"alice@example.com\"),\n",
        "    (\"Bob Johnson\", \"bob@example.com\"),\n",
        "    (\"Carol White\", \"carol@example.com\")\n",
        "]\n",
        "\n",
        "# Insert multiple users\n",
        "insert_query = QueryInput(\n",
        "    query=\"INSERT INTO users (name, email) VALUES (?, ?)\",\n",
        "    parameters=test_users[0]  # Insert Alice\n",
        ")\n",
        "\n",
        "result = db_tool.execute_query(insert_query.query, insert_query.parameters)\n",
        "print(\"Inserted first user successfully!\")\n",
        "\n",
        "# Query the data\n",
        "select_query = QueryInput(\n",
        "    query=\"SELECT * FROM users WHERE email = ?\",\n",
        "    parameters=[\"alice@example.com\"]\n",
        ")\n",
        "\n",
        "result = db_tool.execute_query(select_query.query, select_query.parameters)\n",
        "print(\"\\nQuery Result:\")\n",
        "print(f\"Columns: {result.columns}\")\n",
        "print(f\"Rows: {result.rows}\")\n",
        "print(f\"Row count: {result.row_count}\")\n"
      ]
    },
    {
      "cell_type": "raw",
      "metadata": {
        "vscode": {
          "languageId": "raw"
        }
      },
      "source": [
        "# Error Handling\n",
        "\n",
        "Our database tool should handle various error conditions gracefully:\n",
        "\n",
        "1. **Invalid SQL**: Malformed queries\n",
        "2. **Constraint Violations**: Duplicate unique values\n",
        "3. **Type Mismatches**: Wrong parameter types\n",
        "4. **Missing Tables**: References to non-existent tables\n",
        "\n",
        "Let's test each error case:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# 1. Invalid SQL\n",
        "try:\n",
        "    bad_query = QueryInput(\n",
        "        query=\"SELECT * FROMM users\"  # Intentional typo\n",
        "    )\n",
        "    result = db_tool.execute_query(bad_query.query)\n",
        "except ValueError as e:\n",
        "    print(\"✓ Caught invalid SQL:\", e)\n",
        "\n",
        "# 2. Constraint Violation (duplicate email)\n",
        "try:\n",
        "    duplicate_query = QueryInput(\n",
        "        query=\"INSERT INTO users (name, email) VALUES (?, ?)\",\n",
        "        parameters=[\"Another Alice\", \"alice@example.com\"]  # Email already exists\n",
        "    )\n",
        "    result = db_tool.execute_query(duplicate_query.query, duplicate_query.parameters)\n",
        "except ValueError as e:\n",
        "    print(\"✓ Caught duplicate email:\", e)\n",
        "\n",
        "# 3. Type Mismatch\n",
        "try:\n",
        "    type_mismatch_query = QueryInput(\n",
        "        query=\"SELECT * FROM users WHERE id = ?\",\n",
        "        parameters=[\"not_a_number\"]  # id should be integer\n",
        "    )\n",
        "    result = db_tool.execute_query(type_mismatch_query.query, type_mismatch_query.parameters)\n",
        "except ValueError as e:\n",
        "    print(\"✓ Caught type mismatch:\", e)\n",
        "\n",
        "# 4. Missing Table\n",
        "try:\n",
        "    missing_table_query = QueryInput(\n",
        "        query=\"SELECT * FROM nonexistent_table\"\n",
        "    )\n",
        "    result = db_tool.execute_query(missing_table_query.query)\n",
        "except ValueError as e:\n",
        "    print(\"✓ Caught missing table:\", e)\n"
      ]
    },
    {
      "cell_type": "raw",
      "metadata": {
        "vscode": {
          "languageId": "raw"
        }
      },
      "source": [
        "# Best Practices for Database Tools in MCP\n",
        "\n",
        "When building database tools for MCP, follow these best practices:\n",
        "\n",
        "1. **Security**\n",
        "   - Always use parameterized queries to prevent SQL injection\n",
        "   - Never construct queries by string concatenation\n",
        "   - Validate all inputs before executing queries\n",
        "   - Use connection pooling for better performance\n",
        "\n",
        "2. **Error Handling**\n",
        "   - Catch and handle specific database exceptions\n",
        "   - Provide clear error messages\n",
        "   - Log errors with appropriate context\n",
        "   - Clean up resources in case of errors\n",
        "\n",
        "3. **Resource Management**\n",
        "   - Use context managers for connections\n",
        "   - Close connections properly\n",
        "   - Implement connection pooling for production\n",
        "   - Monitor connection usage\n",
        "\n",
        "4. **Data Validation**\n",
        "   - Validate input types\n",
        "   - Check data constraints\n",
        "   - Sanitize inputs\n",
        "   - Use schema validation\n",
        "\n",
        "5. **Performance**\n",
        "   - Use connection pooling\n",
        "   - Index frequently queried columns\n",
        "   - Optimize queries\n",
        "   - Batch operations when possible\n",
        "\n",
        "## Exercise\n",
        "\n",
        "Try implementing these improvements to our database tool:\n",
        "\n",
        "1. Add connection pooling\n",
        "2. Implement query logging\n",
        "3. Add input validation\n",
        "4. Create a transaction manager\n"
      ]
    }
  ],
  "metadata": {
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}