Model Context Protocol (MCP) Integration#

pgmi init --template advanced myproject
cd myproject
pgmi deploy --connection "postgresql://user:pass@localhost:5432/mydb"

cd tools
pip install -r requirements.txt
export DATABASE_URL="postgresql://user:pass@localhost:5432/mydb"
python mcp-gateway.py

# Initialize handshake
curl -X POST http://localhost:8080/mcp \
  -H "Content-Type: application/json" \
  -d '{"jsonrpc":"2.0","id":"1","method":"initialize","params":{"protocolVersion":"2024-11-05"}}'

# List available tools
curl -X POST http://localhost:8080/mcp \
  -H "Content-Type: application/json" \
  -d '{"jsonrpc":"2.0","id":"2","method":"tools/list"}'

# Call a tool
curl -X POST http://localhost:8080/mcp \
  -H "Content-Type: application/json" \
  -d '{"jsonrpc":"2.0","id":"3","method":"tools/call","params":{"name":"database_info","arguments":{}}}'

{
  "mcpServers": {
    "my-database": {
      "url": "http://localhost:8080/mcp"
    }
  }
}

-- Tool with outputSchema and tags, emitting structuredContent
SELECT api.create_or_replace_mcp_handler(
    jsonb_build_object(
        'id',   'a0000001-0001-4000-8000-000000000001'::uuid,
        'type', 'tool',
        'name', 'weather_at',
        'description', 'Get weather for a location',
        'inputSchema', jsonb_build_object(
            'type', 'object',
            'properties', jsonb_build_object('location', jsonb_build_object('type', 'string')),
            'required', jsonb_build_array('location')
        ),
        'outputSchema', jsonb_build_object(
            'type', 'object',
            'properties', jsonb_build_object('temp_c', jsonb_build_object('type', 'number'))
        ),
        'tags', jsonb_build_array('weather', 'read-only')
    ),
    $body$
DECLARE v_structured jsonb;
BEGIN
    v_structured := jsonb_build_object('temp_c', 21.5);
    RETURN api.mcp_tool_result(
        jsonb_build_array(api.mcp_text(v_structured::text)),
        (request).request_id,
        false,
        v_structured
    );
END;
    $body$
);

-- Request type (passed to your handler)
CREATE TYPE api.mcp_request AS (
    arguments jsonb,      -- Tool/prompt arguments
    uri text,             -- Resource URI
    context jsonb,        -- Auth: {"user_id": "...", "tenant_id": "..."}
    request_id text       -- Must echo in response
);

-- Response type (JSON-RPC 2.0 envelope)
CREATE TYPE api.mcp_response AS (
    envelope jsonb        -- {jsonrpc, id, result} or {jsonrpc, id, error}
);

SELECT api.create_or_replace_mcp_handler(
    jsonb_build_object(
        'id', 'e3000001-0001-4000-8000-000000000001',
        'type', 'tool',
        'name', 'database_info',
        'description', 'Get database version and connection info',
        'inputSchema', jsonb_build_object(
            'type', 'object',
            'properties', jsonb_build_object(),
            'required', jsonb_build_array()
        )
    ),
    $body$
BEGIN
    RETURN api.mcp_tool_result(
        jsonb_build_array(api.mcp_text(format(
            'Database: %s, Version: %s, User: %s',
            current_database(),
            version(),
            current_user
        ))),
        (request).request_id
    );
END;
    $body$
);

SELECT api.create_or_replace_mcp_handler(
    jsonb_build_object(
        'id', 'e3000001-0002-4000-8000-000000000001',
        'type', 'resource',
        'name', 'table_schema',
        'description', 'Get table schema information',
        'uriTemplate', 'postgres:///{schema}/{table}',
        'mimeType', 'application/json'
    ),
    $body$
DECLARE
    v_uri_parts text[];
    v_schema text;
    v_table text;
    v_columns jsonb;
BEGIN
    -- Parse URI: postgres:///public/users -> ['public', 'users']
    v_uri_parts := string_to_array(
        regexp_replace((request).uri, '^postgres:///', ''), '/'
    );
    v_schema := v_uri_parts[1];
    v_table := v_uri_parts[2];

    -- Query column metadata
    SELECT jsonb_agg(jsonb_build_object(
        'column_name', column_name,
        'data_type', data_type,
        'is_nullable', is_nullable
    ))
    INTO v_columns
    FROM information_schema.columns
    WHERE table_schema = COALESCE(v_schema, 'public')
      AND table_name = v_table;

    RETURN api.mcp_resource_result(
        jsonb_build_array(jsonb_build_object(
            'uri', (request).uri,
            'mimeType', 'application/json',
            'text', COALESCE(v_columns, '[]'::jsonb)::text
        )),
        (request).request_id
    );
END;
    $body$
);

SELECT api.create_or_replace_mcp_handler(
    jsonb_build_object(
        'id', 'e3000001-0003-4000-8000-000000000001',
        'type', 'prompt',
        'name', 'sql_assistant',
        'description', 'Generate a SQL query assistant prompt',
        'arguments', jsonb_build_array(
            jsonb_build_object('name', 'task', 'description', 'Task description', 'required', true),
            jsonb_build_object('name', 'tables', 'description', 'Relevant tables', 'required', false)
        )
    ),
    $body$
DECLARE
    v_task text;
    v_tables text;
BEGIN
    v_task := (request).arguments->>'task';
    v_tables := COALESCE((request).arguments->>'tables', 'any relevant tables');

    RETURN api.mcp_prompt_result(
        jsonb_build_array(
            jsonb_build_object(
                'role', 'user',
                'content', jsonb_build_object(
                    'type', 'text',
                    'text', format(
                        'You are a PostgreSQL expert. Help me write a SQL query for: %s. Consider using: %s.',
                        v_task, v_tables
                    )
                )
            )
        ),
        (request).request_id
    );
END;
    $body$
);

-- Tool result (content array)
api.mcp_tool_result(content jsonb, request_id text) RETURNS api.mcp_response

-- Resource result (contents array)
api.mcp_resource_result(contents jsonb, request_id text) RETURNS api.mcp_response

-- Prompt result (messages array)
api.mcp_prompt_result(messages jsonb, request_id text) RETURNS api.mcp_response

-- Generic success
api.mcp_success(result jsonb, request_id text) RETURNS api.mcp_response

-- Generic JSON-RPC error
api.mcp_error(code integer, message text, request_id text) RETURNS api.mcp_response

-- Convenience wrappers (use -32603 Internal Error)
api.mcp_tool_error(message text, request_id text) RETURNS api.mcp_response
api.mcp_resource_error(message text, request_id text) RETURNS api.mcp_response
api.mcp_prompt_error(message text, request_id text) RETURNS api.mcp_response

-- Text content block
api.mcp_text('Hello, world!')
-- Returns: {"type": "text", "text": "Hello, world!"}

-- Gateway extracts from HTTP headers and passes as context:
SELECT api.mcp_handle_request(
    '{"jsonrpc":"2.0","id":"1","method":"tools/call",...}'::jsonb,
    '{"user_id": "auth0|12345", "tenant_id": "org_abc"}'::jsonb
);

-- Inside handlers, access via session variables:
DECLARE
    v_user_id text := current_setting('auth.user_id', true);
    v_tenant_id text := current_setting('auth.tenant_id', true);
BEGIN
    -- Use for RLS, audit logging, etc.
END;

SELECT api.create_or_replace_mcp_handler(
    jsonb_build_object(
        ...
        'requiresAuth', true  -- Default: true
    ),
    $body$...
);

{"jsonrpc": "2.0", "id": "...", "error": {"code": -32001, "message": "Authentication required"}}

{
  "jsonrpc": "2.0",
  "id": "req-1",
  "error": {
    "code": -32602,
    "message": "Tool not found: unknown_tool"
  }
}

-- Test initialize
SELECT (api.mcp_handle_request(
    '{"jsonrpc":"2.0","id":"1","method":"initialize","params":{"protocolVersion":"2024-11-05"}}'::jsonb
)).envelope;

-- Test tools/list
SELECT (api.mcp_handle_request(
    '{"jsonrpc":"2.0","id":"2","method":"tools/list"}'::jsonb
)).envelope;

-- Test tools/call
SELECT (api.mcp_handle_request(
    '{"jsonrpc":"2.0","id":"3","method":"tools/call","params":{"name":"database_info","arguments":{}}}'::jsonb
)).envelope;

-- Test with authentication context
SELECT (api.mcp_handle_request(
    '{"jsonrpc":"2.0","id":"4","method":"tools/call","params":{"name":"execute_query","arguments":{"query":"SELECT 1"}}}'::jsonb,
    '{"user_id":"test|123"}'::jsonb
)).envelope;

DO $$
DECLARE
    v_response api.mcp_response;
    v_envelope jsonb;
BEGIN
    -- Register test handler
    PERFORM api.create_or_replace_mcp_handler(
        jsonb_build_object(
            'id', 'ffffffff-test-4000-8000-000000000001',
            'type', 'tool',
            'name', 'test_tool',
            'description', 'Test',
            'inputSchema', '{}'::jsonb,
            'requiresAuth', false
        ),
        $body$
BEGIN
    RETURN api.mcp_tool_result(
        jsonb_build_array(api.mcp_text('success')),
        (request).request_id
    );
END;
        $body$
    );

    -- Invoke
    v_response := api.mcp_call_tool('test_tool', '{}'::jsonb, NULL, 'test-1');
    v_envelope := (v_response).envelope;

    -- Verify
    IF v_envelope->>'jsonrpc' != '2.0' THEN
        RAISE EXCEPTION 'Missing jsonrpc 2.0';
    END IF;
    IF v_envelope->>'id' != 'test-1' THEN
        RAISE EXCEPTION 'request_id not preserved';
    END IF;
    IF v_envelope->'error' IS NOT NULL THEN
        RAISE EXCEPTION 'Unexpected error';
    END IF;

    RAISE NOTICE 'Test passed';
END $$;

cd myproject/tools
pip install -r requirements.txt
export DATABASE_URL="postgresql://user:pass@localhost:5432/mydb"
python mcp-gateway.py

SET mcp.server_name = 'my-database-server';
SET mcp.server_version = '2.0.0';

SET client_min_messages = DEBUG;
SELECT (api.mcp_handle_request('{"jsonrpc":"2.0","id":"1","method":"tools/call","params":{"name":"my_tool","arguments":{}}}'::jsonb)).envelope;

SELECT * FROM api.mcp_exchange ORDER BY created_at DESC LIMIT 10;