dbmcp-postgres 0.14.0

//! MCP tool: `listFunctions`.

use dbmcp_server::pagination::{Cursor, Pager};

use super::prelude::*;
use crate::types::{ListEntriesResponse, ListFunctionsRequest};

const NAME: &str = "listFunctions";
const TITLE: &str = "List Functions";
const DESCRIPTION_PINNED: &str = include_str!("../../assets/tools/list_functions/pinned.md");
const DESCRIPTION_UNPINNED: &str = include_str!("../../assets/tools/list_functions/unpinned.md");

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

/// Marker type for the `listFunctions` MCP tool (pinned variant — no `database` field).
pub(crate) struct PinnedListFunctionsTool;

impl ToolBase for PinnedListFunctionsTool {
    type Parameter = ListFunctionsRequest;
    type Output = ListEntriesResponse;
    type Error = ErrorData;

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

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

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

    fn annotations() -> Option<ToolAnnotations> {
        Some(annotations())
    }

    fn input_schema() -> Option<Arc<JsonObject>> {
        Some(input_schema::<Self::Parameter>(true))
    }

    fn output_schema() -> Option<Arc<JsonObject>> {
        Some(output_schema::<Self::Output>())
    }
}

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

/// Marker type for the `listFunctions` MCP tool (unpinned variant — carries `database`).
pub(crate) struct UnpinnedListFunctionsTool;

impl ToolBase for UnpinnedListFunctionsTool {
    type Parameter = ListFunctionsRequest;
    type Output = ListEntriesResponse;
    type Error = ErrorData;

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

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

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

    fn annotations() -> Option<ToolAnnotations> {
        Some(annotations())
    }

    fn input_schema() -> Option<Arc<JsonObject>> {
        Some(input_schema::<Self::Parameter>(false))
    }

    fn output_schema() -> Option<Arc<JsonObject>> {
        Some(output_schema::<Self::Output>())
    }
}

impl AsyncTool<PostgresHandler> for UnpinnedListFunctionsTool {
    async fn invoke(handler: &PostgresHandler, params: Self::Parameter) -> Result<Self::Output, Self::Error> {
        handler
            .list_functions(params.database, params.cursor, params.search, params.detailed)
            .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,
        database: Option<String>,
        cursor: Option<Cursor>,
        search: Option<String>,
        detailed: bool,
    ) -> Result<ListEntriesResponse, 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(ListEntriesResponse::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(ListEntriesResponse::brief(functions, next_cursor))
    }
}