dbmcp-postgres 0.11.1

//! MCP tool: `listFunctions`.

use std::borrow::Cow;

use dbmcp_server::pagination::Pager;
use dbmcp_sql::Connection as _;
use rmcp::handler::server::router::tool::{AsyncTool, ToolBase};
use rmcp::model::{ErrorData, ToolAnnotations};

use crate::PostgresHandler;
use crate::types::{ListFunctionsRequest, ListFunctionsResponse};

/// Marker type for the `listFunctions` MCP tool.
pub(crate) struct ListFunctionsTool;

impl ListFunctionsTool {
    const NAME: &'static str = "listFunctions";
    const TITLE: &'static str = "List Functions";
    const DESCRIPTION: &'static str = include_str!("../../assets/tools/list_functions.md");
}

impl ToolBase for ListFunctionsTool {
    type Parameter = ListFunctionsRequest;
    type Output = ListFunctionsResponse;
    type Error = ErrorData;

    fn name() -> Cow<'static, str> {
        Self::NAME.into()
    }

    fn title() -> Option<String> {
        Some(Self::TITLE.into())
    }

    fn description() -> Option<Cow<'static, str>> {
        Some(Self::DESCRIPTION.into())
    }

    fn annotations() -> Option<ToolAnnotations> {
        Some(
            ToolAnnotations::new()
                .read_only(true)
                .destructive(false)
                .idempotent(true)
                .open_world(false),
        )
    }
}

impl AsyncTool<PostgresHandler> for ListFunctionsTool {
    async fn invoke(handler: &PostgresHandler, params: Self::Parameter) -> Result<Self::Output, Self::Error> {
        handler.list_functions(params).await
    }
}

/// Brief-mode SQL: `pg_proc` scan with `ILIKE` filter on function name.
///
/// `n.nspname = 'public'` plus `p.prokind = 'f'` keep aggregates / window
/// functions / procedures out of the result. The `($1::text IS NULL OR ...)`
/// trinary lets one statement cover both filtered and unfiltered cases.
/// `(p.proname, p.oid)` is the sort key — `oid` is the unique tiebreaker
/// across overloaded names so `OFFSET` pagination is deterministic.
const BRIEF_SQL: &str = r"
    SELECT p.proname
    FROM pg_proc p
    JOIN pg_namespace n ON n.oid = p.pronamespace
    WHERE n.nspname = 'public'
      AND p.prokind = 'f'
      AND ($1::text IS NULL OR p.proname ILIKE '%' || $1 || '%')
    ORDER BY p.proname, p.oid
    LIMIT $2 OFFSET $3";

/// Detailed-mode SQL: per-function `json_build_object` projection.
///
/// `n.nspname = 'public'` and `p.prokind = 'f'` filter to user-defined
/// functions in the `public` schema. Three small lookup joins (`pg_namespace`,
/// `pg_language`, `pg_roles`) supply the language and owner names. Postgres
/// defers SELECT-list evaluation past `LIMIT`, so the expensive `pg_get_*`
/// projections (`pg_get_functiondef`, `pg_get_function_arguments`,
/// `pg_get_function_result`) and `obj_description` only run for the page's
/// rows — never the full schema.
///
/// A `CROSS JOIN LATERAL` materialises `pg_get_function_arguments(p.oid)`
/// into `args.text` so it is computed once per row and reused both in the
/// keyed signature `name(args)` and in the JSON `arguments` field — no
/// double-call, no Rust-side json lookup.
///
/// `provolatile`, `proisstrict`, `prosecdef`, `proparallel` are stable
/// `pg_proc` columns. `provolatile` is `'i'`/`'s'`/`'v'`; `proparallel` is
/// `'s'`/`'r'`/`'u'`; `prosecdef` is a boolean.
const DETAILED_SQL: &str = r"
    SELECT
        p.proname || '(' || args.text || ')' AS name,
        json_build_object(
            'schema',         'public',
            'name',           p.proname,
            'language',       l.lanname,
            'arguments',      args.text,
            'returnType',     pg_get_function_result(p.oid),
            'volatility',     CASE p.provolatile
                                  WHEN 'i' THEN 'IMMUTABLE'
                                  WHEN 's' THEN 'STABLE'
                                  WHEN 'v' THEN 'VOLATILE'
                              END,
            'strict',         p.proisstrict,
            'security',       CASE WHEN p.prosecdef THEN 'DEFINER' ELSE 'INVOKER' END,
            'parallelSafety', CASE p.proparallel
                                  WHEN 's' THEN 'SAFE'
                                  WHEN 'r' THEN 'RESTRICTED'
                                  WHEN 'u' THEN 'UNSAFE'
                              END,
            'owner',          r.rolname,
            'description',    pg_catalog.obj_description(p.oid, 'pg_proc'),
            'definition',     pg_get_functiondef(p.oid)
        ) AS entry
    FROM pg_proc p
    JOIN pg_namespace n ON n.oid = p.pronamespace
    JOIN pg_language  l ON l.oid = p.prolang
    JOIN pg_roles     r ON r.oid = p.proowner
    CROSS JOIN LATERAL (SELECT pg_get_function_arguments(p.oid) AS text) args
    WHERE n.nspname = 'public'
      AND p.prokind = 'f'
      AND ($1::text IS NULL OR p.proname ILIKE '%' || $1 || '%')
    ORDER BY p.proname, p.oid
    LIMIT $2 OFFSET $3";

impl PostgresHandler {
    /// Lists one page of user-defined functions, optionally filtered and/or detailed.
    ///
    /// # Errors
    ///
    /// Returns [`ErrorData`] with code `-32602` if `cursor` is malformed,
    /// or an internal-error [`ErrorData`] if `database` is invalid
    /// or the underlying query fails.
    pub async fn list_functions(
        &self,
        ListFunctionsRequest {
            database,
            cursor,
            search,
            detailed,
        }: ListFunctionsRequest,
    ) -> Result<ListFunctionsResponse, ErrorData> {
        let database = database.as_deref().map(str::trim).filter(|s| !s.is_empty());
        let pattern = search.as_deref().map(str::trim).filter(|s| !s.is_empty());
        let pager = Pager::new(cursor, self.config.page_size);

        if detailed {
            let rows: Vec<(String, sqlx::types::Json<serde_json::Value>)> = self
                .connection
                .fetch(
                    sqlx::query(DETAILED_SQL)
                        .bind(pattern)
                        .bind(pager.limit())
                        .bind(pager.offset()),
                    database,
                )
                .await?;
            let (rows, next_cursor) = pager.paginate(rows);
            return Ok(ListFunctionsResponse::detailed(
                rows.into_iter().map(|(key, json)| (key, json.0)).collect(),
                next_cursor,
            ));
        }

        let rows: Vec<String> = self
            .connection
            .fetch_scalar(
                sqlx::query(BRIEF_SQL)
                    .bind(pattern)
                    .bind(pager.limit())
                    .bind(pager.offset()),
                database,
            )
            .await?;
        let (functions, next_cursor) = pager.paginate(rows);
        Ok(ListFunctionsResponse::brief(functions, next_cursor))
    }
}