KODEGEN Database Tools

Blazing-Fast MCP Database Tools for AI Agents

Part of KODEGEN.ᴀɪ - A Rust-native MCP server providing 7 production-ready database tools for autonomous SQL execution and schema exploration across PostgreSQL, MySQL, MariaDB, and SQLite.

Features

• 🚀 Multi-Database Support - PostgreSQL, MySQL, MariaDB, SQLite with unified interface
• 🔒 SSH Tunnel Support - Secure connections through bastion hosts
• ⚡ Connection Pooling - Efficient connection reuse with configurable pool settings
• 🛡️ Read-Only Mode - Prevent accidental data modifications
• 🎯 Smart Row Limiting - Automatic result set limiting to prevent memory exhaustion
• 🔄 Retry Logic - Exponential backoff with configurable timeouts
• 📊 Pool Health Monitoring - Real-time connection pool statistics
• 🔍 Schema Introspection - Explore databases, tables, columns, indexes, and stored procedures

Installation

As Part of KODEGEN.ᴀɪ

Install the complete KODEGEN.ᴀɪ toolkit (recommended):

curl -fsSL https://kodegen.ai/install | sh

The database tools are automatically included and configured.

Standalone Installation

# Clone repository
git clone https://github.com/cyrup-ai/kodegen-tools-database
cd kodegen-tools-database

# Build and install
cargo install --path .

Running the Server

The database server uses the DATABASE_DSN environment variable to connect to a database. If not provided, it defaults to an in-memory SQLite database (sqlite::memory:):

# Basic usage with PostgreSQL
DATABASE_DSN="postgres://user:pass@localhost:5432/mydb" kodegen-database

# Or use the default in-memory SQLite database (no external database required)
kodegen-database

# With SSH tunnel
SSH_HOST="bastion.example.com" \
SSH_PORT="22" \
SSH_USER="username" \
SSH_AUTH_TYPE="password" \
SSH_PASSWORD="secret" \
SSH_TARGET_HOST="internal.db.host" \
SSH_TARGET_PORT="5432" \
DATABASE_DSN="postgres://user:pass@internal.db.host:5432/db" \
kodegen-database

The 7 Database Tools

1. db_execute_sql

Execute SQL queries with transaction support, retry logic, and automatic row limiting.

Features:

• Multi-statement execution with transactions
• Read-only mode enforcement
• Automatic LIMIT clause injection
• Binary data handling (base64 encoding)
• Configurable timeouts and retries

Example:

db_execute_sql({
  "sql": "SELECT * FROM employees WHERE department_id = 1",
  "readonly": true,
  "max_rows": 100
})

Response:

{
  "rows": [
    {"id": 1, "name": "Alice", "email": "alice@example.com"},
    {"id": 2, "name": "Bob", "email": "bob@example.com"}
  ],
  "row_count": 2
}

2. db_list_schemas

List all databases or schemas available on the server.

Example:

db_list_schemas({})

Response:

{
  "schemas": [
    {"name": "public", "type": "schema"},
    {"name": "information_schema", "type": "schema"}
  ]
}

3. db_list_tables

List all tables within a specific schema/database.

Example:

db_list_tables({
  "schema": "public"
})

Response:

{
  "tables": [
    {"name": "employees", "type": "table"},
    {"name": "departments", "type": "table"}
  ]
}

4. db_table_schema

Get detailed column information for a table.

Example:

db_table_schema({
  "schema": "public",
  "table": "employees"
})

Response:

{
  "columns": [
    {
      "name": "id",
      "data_type": "integer",
      "nullable": false,
      "default": "nextval('employees_id_seq'::regclass)"
    },
    {
      "name": "name",
      "data_type": "varchar(255)",
      "nullable": false,
      "default": null
    }
  ]
}

5. db_table_indexes

Get index information for a table.

Example:

db_table_indexes({
  "schema": "public",
  "table": "employees"
})

Response:

{
  "indexes": [
    {
      "name": "employees_pkey",
      "columns": ["id"],
      "unique": true,
      "primary": true
    },
    {
      "name": "idx_employee_department",
      "columns": ["department_id"],
      "unique": false,
      "primary": false
    }
  ]
}

6. db_stored_procedures

List stored procedures (PostgreSQL and MySQL only).

Example:

db_stored_procedures({
  "schema": "public"
})

Response:

{
  "procedures": [
    {
      "name": "get_department_employee_count",
      "schema": "public",
      "return_type": "integer",
      "language": "plpgsql"
    }
  ]
}

7. db_pool_stats

Monitor connection pool health and performance.

Example:

db_pool_stats({})

Response:

{
  "connections": 5,
  "idle_connections": 3,
  "max_connections": 10,
  "min_connections": 2,
  "wait_queue_size": 0
}

Configuration

Control database tool behavior through ConfigManager settings:

Connection Pool Settings

{
  "db_min_connections": 2,
  "db_max_connections": 10,
  "db_acquire_timeout_secs": 30,
  "db_idle_timeout_secs": 600,
  "db_max_lifetime_secs": 1800
}

Retry Configuration

{
  "db_max_retries": 2,
  "db_retry_backoff_ms": 500,
  "db_max_backoff_ms": 5000
}

• db_max_retries (default: 2) - Maximum retry attempts
• db_retry_backoff_ms (default: 500) - Base backoff duration
• db_max_backoff_ms (default: 5000) - Maximum backoff cap

Backoff progression: 500ms → 1000ms → 2000ms → 4000ms (capped at 5000ms)

Timeout Configuration

{
  "db_query_timeout_secs": 60
}

• db_query_timeout_secs (default: 60) - Per-query timeout in seconds

SSH Tunnel Support

Secure database connections through SSH bastion hosts using environment variables:

Password Authentication

SSH_HOST="bastion.example.com"
SSH_PORT="22"
SSH_USER="username"
SSH_AUTH_TYPE="password"
SSH_PASSWORD="your-password"
SSH_TARGET_HOST="internal.database.host"
SSH_TARGET_PORT="5432"

SSH Key Authentication

SSH_HOST="bastion.example.com"
SSH_PORT="22"
SSH_USER="username"
SSH_AUTH_TYPE="key"
SSH_KEY_PATH="/path/to/private/key"
SSH_KEY_PASSPHRASE="optional-passphrase"  # Optional
SSH_TARGET_HOST="internal.database.host"
SSH_TARGET_PORT="5432"

The tunnel automatically:

• Establishes SSH connection on server startup
• Creates local port forwarding
• Rewrites DSN to use tunnel endpoint
• Handles graceful shutdown and cleanup

Development & Testing

Docker-Based Testing

Test all 7 tools across 4 database types with Docker:

# Start test databases
docker-compose up -d

# Wait for health checks (20-30 seconds)
docker-compose ps

# Run example
cargo run --example database_demo

# Stop containers
docker-compose down

Test Database Schema

The Docker setup provides a universal schema with 5 tables:

departments (5 records)

• id, name, budget, created_at

employees (15 records)

• id, name, email, department_id, salary, hire_date, active
• Indexes on department_id and name

projects (8 records)

• id, name, department_id, start_date, end_date, status
• Indexes on department_id and (status, start_date, end_date)

employee_projects (20 records - junction table)

• employee_id, project_id, role, assigned_at
• Composite primary key

audit_log (10 records)

• id, table_name, record_id, action, changed_at, changed_by
• Indexes on (table_name, record_id) and changed_at

Stored Procedure (PostgreSQL/MySQL/MariaDB):

• get_department_employee_count(dept_id) - Returns employee count

Connection Strings

# PostgreSQL
postgres://testuser:testpass@localhost:5432/testdb

# MySQL
mysql://testuser:testpass@localhost:3306/testdb

# MariaDB (port 3307 to avoid conflict)
mysql://testuser:testpass@localhost:3307/testdb

# SQLite
sqlite:///tmp/kodegen_test.db

Build Commands

# Build
cargo build

# Build release
cargo build --release

# Run tests
cargo test

# Run clippy
cargo clippy

# Format code
cargo fmt

Troubleshooting

Containers won't start

# Check port conflicts
lsof -i :5432
lsof -i :3306
lsof -i :3307

# View logs
docker-compose logs postgres
docker-compose logs mysql
docker-compose logs mariadb

Schema not loading

# Recreate with fresh data
docker-compose down -v
docker-compose up -d

Connection failures

# Verify health
docker-compose ps

# Test manually
docker exec -it kodegen-test-postgres psql -U testuser -d testdb -c "SELECT COUNT(*) FROM employees;"

Architecture

Core Components

• src/connection.rs - Connection pool setup with SSH tunnel support
• src/ssh_tunnel.rs - SSH port forwarding implementation
• src/dsn.rs - Secure DSN parsing with SecretString
• src/tools/ - 7 tool implementations
• src/schema_queries.rs - Database-specific introspection queries
• src/sql_parser.rs - SQL statement parsing and splitting
• src/sql_limiter.rs - Automatic row limiting
• src/readonly.rs - Read-only SQL validation

Tool Pattern

All tools follow a consistent pattern:

1. Implement Tool trait from kodegen_mcp_tool
2. Accept Arc<AnyPool> for connection pooling
3. Use execute_with_timeout() for retry logic
4. Return structured JSON responses
5. Handle database-specific variations

Example Workflows

Schema Exploration

// 1. List all schemas
db_list_schemas({})

// 2. List tables in a schema
db_list_tables({"schema": "public"})

// 3. Get table structure
db_table_schema({"schema": "public", "table": "employees"})

// 4. Get indexes
db_table_indexes({"schema": "public", "table": "employees"})

Data Analysis

// Execute analytical query
db_execute_sql({
  "sql": `
    SELECT 
      d.name as department,
      COUNT(e.id) as employee_count,
      AVG(e.salary) as avg_salary
    FROM departments d
    LEFT JOIN employees e ON d.id = e.department_id
    WHERE e.active = true
    GROUP BY d.name
    ORDER BY employee_count DESC
  `,
  "readonly": true,
  "max_rows": 50
})

Multi-Statement Transaction

db_execute_sql({
  "sql": `
    INSERT INTO departments (name, budget) VALUES ('Engineering', 500000);
    INSERT INTO employees (name, email, department_id, salary, hire_date, active)
    VALUES ('Charlie', 'charlie@example.com', 1, 85000, '2025-01-15', true);
  `,
  "readonly": false
})

Performance

• Startup Time: ~25ms (with warmup)
• Query Latency: <10ms (simple queries, local DB)
• Memory Usage: ~8MB (idle with pool)
• Connection Pool: Configurable 2-10 connections
• Concurrent Queries: Supports high concurrency via connection pool

Security

• SecretString - Passwords never logged or displayed
• Read-Only Mode - Prevent accidental modifications
• SSH Tunneling - Secure connections through bastions
• SQL Validation - Reject unsafe patterns in read-only mode
• Row Limiting - Prevent memory exhaustion attacks

Contributing

See main KODEGEN.ᴀɪ repository for contribution guidelines.

License

Dual-licensed under Apache-2.0 and MIT. See LICENSE.md for details.

Related Repositories

• kodegen - Main MCP server
• kodegen-tools-filesystem - Filesystem tools
• kodegen-tools-terminal - Terminal tools
• kodegen-tools-github - GitHub integration

Part of KODEGEN.ᴀɪ - The ultimate MCP auto-coding toolset 🚀