dbmcp-postgres 0.10.0

//! `PostgreSQL` handler: composes a [`PostgresConnection`] with the MCP tool router.
//!
//! All pool ownership and pool initialization logic lives in the
//! [`PostgresConnection`]. This module wires the connection into
//! the MCP `ServerHandler` surface and exposes a small set of thin
//! delegator methods that the per-tool implementations call.

use dbmcp_config::{Config, DatabaseConfig};
use dbmcp_pii::Redactor;
use dbmcp_server::{Server, server_info};
use rmcp::RoleServer;
use rmcp::handler::server::router::tool::ToolRouter;
use rmcp::handler::server::tool::ToolCallContext;
use rmcp::model::{CallToolRequestParams, CallToolResult, ListToolsResult, PaginatedRequestParams, ServerInfo, Tool};
use rmcp::service::RequestContext;
use rmcp::{ErrorData, ServerHandler};

use crate::connection::PostgresConnection;

use crate::tools::{
    CreateDatabaseTool, DropDatabaseTool, DropTableTool, ExplainQueryTool, ListDatabasesTool, ListFunctionsTool,
    ListMaterializedViewsTool, ListProceduresTool, ListTablesTool, ListTriggersTool, ListViewsTool, ReadQueryTool,
    WriteQueryTool,
};

/// Backend-specific description for `PostgreSQL`.
const DESCRIPTION: &str = "Database MCP Server for PostgreSQL";

/// Backend-specific instructions for `PostgreSQL`.
const INSTRUCTIONS: &str = r"## Workflow

1. Call `listDatabases` to discover available databases.
2. Call `listTables` to see tables. Pass `search` to filter by name (case-insensitive substring). Pass `detailed: true` to get columns, constraints, indexes, and triggers in the same call — this supersedes the legacy `getTableSchema` workflow.
3. Call `listViews` to see views in the `public` schema.
4. Call `listTriggers` to see user-defined triggers in the `public` schema.
5. Call `listFunctions` to see user-defined functions in the `public` schema.
6. Call `listProcedures` to see user-defined procedures in the `public` schema.
7. Call `listMaterializedViews` to see materialized views in the `public` schema.
8. Use `readQuery` for read-only SQL (SELECT, SHOW, EXPLAIN).
9. Use `writeQuery` for data changes (INSERT, UPDATE, DELETE, CREATE, ALTER, DROP).
10. Use `explainQuery` to analyze query execution plans and diagnose slow queries.
11. Use `createDatabase` to create a new database.
12. Use `dropDatabase` to drop an existing database.
13. Use `dropTable` to remove a table from a database (supports `cascade` for foreign key dependencies).

Per-database tools default to the active database; pass `database` to target another.

## Constraints

- The `writeQuery`, `createDatabase`, `dropDatabase`, and `dropTable` tools are hidden when read-only mode is active.
- Multi-statement queries are not supported. Send one statement per request.";

/// `PostgreSQL` database handler.
///
/// Composes one [`PostgresConnection`] (which owns every pool
/// and the pool initialization logic) with the per-backend MCP tool
/// router.
#[derive(Clone)]
pub struct PostgresHandler {
    pub(crate) config: DatabaseConfig,
    pub(crate) connection: PostgresConnection,
    pub(crate) redactor: Option<Redactor>,
    tool_router: ToolRouter<Self>,
}

impl std::fmt::Debug for PostgresHandler {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PostgresHandler")
            .field("read_only", &self.config.read_only)
            .field("redact_pii", &self.redactor.is_some())
            .field("connection", &self.connection)
            .finish_non_exhaustive()
    }
}

impl PostgresHandler {
    /// Creates a new `PostgreSQL` handler.
    ///
    /// Constructs the [`PostgresConnection`] (which builds the
    /// lazy default pool and the per-database pool cache) and the MCP
    /// tool router. No network I/O happens here.
    #[must_use]
    pub fn new(config: &Config) -> Self {
        Self {
            config: config.database.clone(),
            connection: PostgresConnection::new(&config.database),
            redactor: config
                .pii
                .enabled
                .then(|| Redactor::with_operator_config(config.pii.operator.into())),
            tool_router: build_tool_router(config.database.read_only),
        }
    }
}

impl From<PostgresHandler> for Server {
    /// Wraps a [`PostgresHandler`] in the type-erased MCP server.
    fn from(handler: PostgresHandler) -> Self {
        Self::new(handler)
    }
}

/// Builds the tool router, including write tools only when not in read-only mode.
fn build_tool_router(read_only: bool) -> ToolRouter<PostgresHandler> {
    let mut router = ToolRouter::new()
        .with_async_tool::<ListDatabasesTool>()
        .with_async_tool::<ListTablesTool>()
        .with_async_tool::<ListViewsTool>()
        .with_async_tool::<ListTriggersTool>()
        .with_async_tool::<ListFunctionsTool>()
        .with_async_tool::<ListProceduresTool>()
        .with_async_tool::<ListMaterializedViewsTool>()
        .with_async_tool::<ReadQueryTool>()
        .with_async_tool::<ExplainQueryTool>();

    if !read_only {
        router = router
            .with_async_tool::<CreateDatabaseTool>()
            .with_async_tool::<DropDatabaseTool>()
            .with_async_tool::<DropTableTool>()
            .with_async_tool::<WriteQueryTool>();
    }
    router
}

impl ServerHandler for PostgresHandler {
    fn get_info(&self) -> ServerInfo {
        let mut info = server_info();
        info.server_info.description = Some(DESCRIPTION.into());
        info.instructions = Some(INSTRUCTIONS.into());
        info
    }

    async fn call_tool(
        &self,
        request: CallToolRequestParams,
        context: RequestContext<RoleServer>,
    ) -> Result<CallToolResult, ErrorData> {
        let tcc = ToolCallContext::new(self, request, context);
        self.tool_router.call(tcc).await
    }

    async fn list_tools(
        &self,
        _request: Option<PaginatedRequestParams>,
        _context: RequestContext<RoleServer>,
    ) -> Result<ListToolsResult, ErrorData> {
        Ok(ListToolsResult {
            tools: self.tool_router.list_all(),
            next_cursor: None,
            meta: None,
        })
    }

    fn get_tool(&self, name: &str) -> Option<Tool> {
        self.tool_router.get(name).cloned()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use dbmcp_config::DatabaseBackend;

    fn base_config() -> DatabaseConfig {
        DatabaseConfig {
            backend: DatabaseBackend::Postgres,
            host: "pg.example.com".into(),
            port: 5433,
            user: "pgadmin".into(),
            password: Some("pgpass".into()),
            name: Some("mydb".into()),
            ..DatabaseConfig::default()
        }
    }

    fn handler(read_only: bool) -> PostgresHandler {
        PostgresHandler::new(&Config {
            database: DatabaseConfig {
                read_only,
                ..base_config()
            },
            http: None,
            pii: dbmcp_config::PiiConfig::default(),
        })
    }

    #[tokio::test]
    async fn handler_exposes_connection_default_db() {
        let handler = PostgresHandler::new(&Config {
            database: base_config(),
            http: None,
            pii: dbmcp_config::PiiConfig::default(),
        });
        assert_eq!(handler.connection.default_database_name(), "mydb");
    }

    #[tokio::test]
    async fn router_exposes_all_thirteen_tools_in_read_write_mode() {
        let router = handler(false).tool_router;
        for name in [
            "listDatabases",
            "listTables",
            "listViews",
            "listTriggers",
            "listFunctions",
            "listProcedures",
            "listMaterializedViews",
            "readQuery",
            "explainQuery",
            "createDatabase",
            "dropDatabase",
            "dropTable",
            "writeQuery",
        ] {
            assert!(router.has_route(name), "missing tool: {name}");
        }
    }

    #[tokio::test]
    async fn router_hides_write_tools_in_read_only_mode() {
        let router = handler(true).tool_router;
        assert!(router.has_route("listDatabases"));
        assert!(router.has_route("listTables"));
        assert!(router.has_route("listViews"));
        assert!(router.has_route("listTriggers"));
        assert!(router.has_route("listFunctions"));
        assert!(router.has_route("listProcedures"));
        assert!(router.has_route("listMaterializedViews"));
        assert!(router.has_route("readQuery"));
        assert!(router.has_route("explainQuery"));
        assert!(!router.has_route("writeQuery"));
        assert!(!router.has_route("createDatabase"));
        assert!(!router.has_route("dropDatabase"));
        assert!(!router.has_route("dropTable"));
    }

    #[tokio::test]
    async fn router_does_not_expose_get_table_schema() {
        // Postgres drops `getTableSchema` in favour of `listTables({detailed: true})`.
        // MySQL and SQLite still expose it; their routers are tested separately.
        let rw = handler(false).tool_router;
        let ro = handler(true).tool_router;
        assert!(
            !rw.has_route("getTableSchema"),
            "read-write router must not expose getTableSchema"
        );
        assert!(
            !ro.has_route("getTableSchema"),
            "read-only router must not expose getTableSchema"
        );
    }
}