Building a Data Engineering Agent with Claude Opus 4.7 and MCP

python3 -m venv agent_env
source agent_env/bin/activate

pip install anthropic mcp psycopg2-binary pandas pyarrow

import json
import psycopg2
from psycopg2.extras import RealDictCursor
from typing import Any, Dict, List
 
class DataEngineeringMCPServer:
    def __init__(self, db_config: Dict[str, str]):
        self.db_config = db_config
        self.connection = None
    
    def connect(self):
        """Establish database connection."""
        self.connection = psycopg2.connect(**self.db_config)
    
    def execute_query(self, query: str) -> Dict[str, Any]:
        """Execute a SQL query and return results."""
        try:
            with self.connection.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute(query)
                if query.strip().upper().startswith('SELECT'):
                    results = cursor.fetchall()
                    return {
                        'status': 'success',
                        'rows': [dict(row) for row in results],
                        'row_count': len(results)
                    }
                else:
                    self.connection.commit()
                    return {
                        'status': 'success',
                        'rows_affected': cursor.rowcount
                    }
        except Exception as e:
            return {'status': 'error', 'message': str(e)}
    
    def get_schema(self, table_name: str) -> Dict[str, Any]:
        """Fetch schema for a table."""
        try:
            query = f"""
            SELECT column_name, data_type, is_nullable
            FROM information_schema.columns
            WHERE table_name = %s
            ORDER BY ordinal_position
            """
            with self.connection.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute(query, (table_name,))
                columns = [dict(row) for row in cursor.fetchall()]
                return {'status': 'success', 'columns': columns}
        except Exception as e:
            return {'status': 'error', 'message': str(e)}
    
    def validate_data(self, table_name: str) -> Dict[str, Any]:
        """Run basic data quality checks."""
        try:
            validation_results = {}
            schema = self.get_schema(table_name)
            
            if schema['status'] != 'success':
                return schema
            
            columns = schema['columns']
            
            # Check row count
            count_result = self.execute_query(f"SELECT COUNT(*) as cnt FROM {table_name}")
            validation_results['total_rows'] = count_result['rows'][0]['cnt']
            
            # Check nulls per column
            for col in columns:
                col_name = col['column_name']
                null_result = self.execute_query(
                    f"SELECT COUNT(*) as null_count FROM {table_name} WHERE {col_name} IS NULL"
                )
                validation_results[f'{col_name}_nulls'] = null_result['rows'][0]['null_count']
            
            return {'status': 'success', 'validation': validation_results}
        except Exception as e:
            return {'status': 'error', 'message': str(e)}
 
# Initialize server
db_config = {
    'host': 'localhost',
    'database': 'analytics',
    'user': 'postgres',
    'password': 'your_password'
}
 
server = DataEngineeringMCPServer(db_config)
server.connect()

import json
from anthropic import Anthropic
from mcp_server import DataEngineeringMCPServer
 
client = Anthropic()
 
# Initialize your MCP server
db_config = {
    'host': 'localhost',
    'database': 'analytics',
    'user': 'postgres',
    'password': 'your_password'
}
 
mcp_server = DataEngineeringMCPServer(db_config)
mcp_server.connect()
 
# Define tools for Claude
tools = [
    {
        "name": "execute_query",
        "description": "Execute a SQL query against the database.",
        "input_schema": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "The SQL query to execute."
                }
            },
            "required": ["query"]
        }
    },
    {
        "name": "get_schema",
        "description": "Fetch the schema for a table.",
        "input_schema": {
            "type": "object",
            "properties": {
                "table_name": {
                    "type": "string",
                    "description": "The name of the table."
                }
            },
            "required": ["table_name"]
        }
    },
    {
        "name": "validate_data",
        "description": "Run data quality checks on a table.",
        "input_schema": {
            "type": "object",
            "properties": {
                "table_name": {
                    "type": "string",
                    "description": "The name of the table to validate."
                }
            },
            "required": ["table_name"]
        }
    }
]
 
def run_agent(task: str) -> str:
    """Run the data engineering agent for a given task."""
    messages = [
        {
            "role": "user",
            "content": task
        }
    ]
    
    while True:
        # Call Claude with tools
        response = client.messages.create(
            model="claude-opus-4-7",
            max_tokens=4096,
            tools=tools,
            messages=messages
        )
        
        # Check if Claude is done
        if response.stop_reason == "end_turn":
            # Extract final text response
            for block in response.content:
                if hasattr(block, 'text'):
                    return block.text
            break
        
        # Process tool calls
        if response.stop_reason == "tool_use":
            # Add assistant's response to messages
            messages.append({"role": "assistant", "content": response.content})
            
            # Process each tool call
            tool_results = []
            for block in response.content:
                if block.type == "tool_use":
                    tool_name = block.name
                    tool_input = block.input
                    
                    # Execute the tool
                    if tool_name == "execute_query":
                        result = mcp_server.execute_query(tool_input["query"])
                    elif tool_name == "get_schema":
                        result = mcp_server.get_schema(tool_input["table_name"])
                    elif tool_name == "validate_data":
                        result = mcp_server.validate_data(tool_input["table_name"])
                    else:
                        result = {"error": f"Unknown tool: {tool_name}"}
                    
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": json.dumps(result)
                    })
            
            # Add tool results to messages
            messages.append({"role": "user", "content": tool_results})
        else:
            # Unexpected stop reason
            break
    
    return "Agent completed without final response."
 
# Example usage
if __name__ == "__main__":
    task = """
    I have a table called 'users' with columns user_id, email, created_at, and status.
    1. Fetch the schema for the users table.
    2. Run a validation check to count total rows and nulls per column.
    3. Find all users created in the last 7 days.
    4. Summarize the results.
    """
    
    result = run_agent(task)
    print("Agent Result:")
    print(result)

import pandas as pd
import os
from datetime import datetime
 
class ExtendedMCPServer(DataEngineeringMCPServer):
    def load_csv(self, file_path: str, table_name: str) -> Dict[str, Any]:
        """Load a CSV file into a staging table."""
        try:
            df = pd.read_csv(file_path)
            # Create table from DataFrame
            from sqlalchemy import create_engine
            engine = create_engine(
                f"postgresql://{self.db_config['user']}:{self.db_config['password']}"
                f"@{self.db_config['host']}/{self.db_config['database']}"
            )
            df.to_sql(table_name, engine, if_exists='replace', index=False)
            return {
                'status': 'success',
                'rows_loaded': len(df),
                'columns': list(df.columns)
            }
        except Exception as e:
            return {'status': 'error', 'message': str(e)}
    
    def transform_data(self, source_table: str, target_table: str, 
                      transformations: Dict[str, str]) -> Dict[str, Any]:
        """Apply transformations and load into target table."""
        try:
            # Build transformation query
            select_clause = ', '.join(
                f"{expr} as {col}" for col, expr in transformations.items()
            )
            query = f"CREATE TABLE {target_table} AS SELECT {select_clause} FROM {source_table}"
            result = self.execute_query(query)
            return result
        except Exception as e:
            return {'status': 'error', 'message': str(e)}
    
    def get_data_sample(self, table_name: str, limit: int = 5) -> Dict[str, Any]:
        """Fetch sample rows from a table."""
        query = f"SELECT * FROM {table_name} LIMIT {limit}"
        return self.execute_query(query)

task = """
You are a data engineering agent. Your job is to:
 
1. Load the file 'orders.csv' into a staging table called 'orders_staging'.
2. Inspect the schema of the staging table.
3. Create a transformation that:
   - Cleans email addresses (lowercase, trim whitespace)
   - Standardizes dates to ISO format
   - Removes duplicate rows based on order_id
4. Load the transformed data into 'orders_production'.
5. Run validation checks on both tables and compare row counts.
6. Report the results, including any data quality issues.
 
Work step-by-step. If you encounter an error, diagnose it and retry.
"""
 
result = run_agent(task)
print(result)

def run_agent_with_retry(task: str, max_retries: int = 3) -> str:
    """Run the agent with retry logic for transient failures."""
    for attempt in range(max_retries):
        try:
            return run_agent(task)
        except Exception as e:
            if attempt < max_retries - 1:
                print(f"Attempt {attempt + 1} failed: {e}. Retrying...")
                continue
            else:
                return f"Agent failed after {max_retries} attempts: {e}"

import logging
 
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
 
def run_agent_with_logging(task: str) -> str:
    """Run the agent with detailed logging."""
    logger.info(f"Starting agent task: {task}")
    messages = [{"role": "user", "content": task}]
    
    while True:
        response = client.messages.create(
            model="claude-opus-4-7",
            max_tokens=4096,
            tools=tools,
            messages=messages
        )
        
        logger.info(f"Claude response: {response.stop_reason}")
        
        # ... rest of agent loop ...

import signal
 
def timeout_handler(signum, frame):
    raise TimeoutError("Tool execution timed out")
 
def execute_tool_with_timeout(tool_name: str, tool_input: Dict, timeout: int = 30) -> Dict:
    """Execute a tool with a timeout."""
    signal.signal(signal.SIGALRM, timeout_handler)
    signal.alarm(timeout)
    try:
        if tool_name == "execute_query":
            result = mcp_server.execute_query(tool_input["query"])
        # ... other tools ...
        signal.alarm(0)  # Cancel the alarm
        return result
    except TimeoutError:
        return {"status": "error", "message": "Tool execution timed out"}

def run_agent_with_reflection(task: str) -> str:
    """Run the agent and have Claude reflect on the results."""
    # First, run the agent normally
    result = run_agent(task)
    
    # Then, ask Claude to reflect
    reflection_task = f"""
    You just completed a data engineering task. Here's what you did:
    
    {result}
    
    Now, reflect on the work:
    1. Did you accomplish the goal?
    2. Were there any issues or assumptions you made?
    3. What would you do differently next time?
    4. Are there any data quality concerns you should flag?
    """
    
    reflection = run_agent(reflection_task)
    return f"Initial Result:\n{result}\n\nReflection:\n{reflection}"

import time
from collections import defaultdict
 
class AgentMetrics:
    def __init__(self):
        self.tool_calls = defaultdict(list)
        self.agent_runs = []
    
    def record_tool_call(self, tool_name: str, duration: float, success: bool):
        self.tool_calls[tool_name].append({
            'duration': duration,
            'success': success,
            'timestamp': datetime.now()
        })
    
    def record_agent_run(self, task: str, duration: float, success: bool, tokens_used: int):
        self.agent_runs.append({
            'task': task,
            'duration': duration,
            'success': success,
            'tokens_used': tokens_used,
            'timestamp': datetime.now()
        })
    
    def get_summary(self):
        return {
            'total_runs': len(self.agent_runs),
            'success_rate': sum(1 for r in self.agent_runs if r['success']) / len(self.agent_runs),
            'avg_duration': sum(r['duration'] for r in self.agent_runs) / len(self.agent_runs),
            'total_tokens': sum(r['tokens_used'] for r in self.agent_runs),
            'tool_stats': {
                name: {
                    'calls': len(calls),
                    'success_rate': sum(1 for c in calls if c['success']) / len(calls),
                    'avg_duration': sum(c['duration'] for c in calls) / len(calls)
                }
                for name, calls in self.tool_calls.items()
            }
        }
 
metrics = AgentMetrics()

import concurrent.futures
 
def run_agents_parallel(tasks: List[str]) -> List[str]:
    """Run multiple agent tasks in parallel."""
    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
        futures = [executor.submit(run_agent, task) for task in tasks]
        return [f.result() for f in concurrent.futures.as_completed(futures)]

from functools import lru_cache
 
@lru_cache(maxsize=128)
def cached_get_schema(table_name: str):
    return mcp_server.get_schema(table_name)

def run_batch_validation(table_names: List[str]) -> Dict[str, Any]:
    """Validate multiple tables in a single agent run."""
    task = f"Validate these tables and report issues: {', '.join(table_names)}"
    return run_agent(task)