dbmcp-postgres 0.14.0

//! MCP tool: `listProcedures`.

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

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

const NAME: &str = "listProcedures";
const TITLE: &str = "List Procedures";
const DESCRIPTION_PINNED: &str = include_str!("../../assets/tools/list_procedures/pinned.md");
const DESCRIPTION_UNPINNED: &str = include_str!("../../assets/tools/list_procedures/unpinned.md");

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

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

impl ToolBase for PinnedListProceduresTool {
    type Parameter = ListProceduresRequest;
    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 PinnedListProceduresTool {
    async fn invoke(handler: &PostgresHandler, params: Self::Parameter) -> Result<Self::Output, Self::Error> {
        handler
            .list_procedures(None, params.cursor, params.search, params.detailed)
            .await
    }
}

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

impl ToolBase for UnpinnedListProceduresTool {
    type Parameter = ListProceduresRequest;
    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 UnpinnedListProceduresTool {
    async fn invoke(handler: &PostgresHandler, params: Self::Parameter) -> Result<Self::Output, Self::Error> {
        handler
            .list_procedures(params.database, params.cursor, params.search, params.detailed)
            .await
    }
}

/// Brief-mode SQL: `pg_proc` scan with `ILIKE` filter on procedure name.
///
/// `n.nspname = 'public'` plus `p.prokind = 'p'` keep functions / aggregates /
/// window functions 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 = 'p'
      AND ($1::text IS NULL OR p.proname ILIKE '%' || $1 || '%')
    ORDER BY p.proname, p.oid
    LIMIT $2 OFFSET $3";

/// Detailed-mode SQL: per-procedure `json_build_object` projection.
///
/// `n.nspname = 'public'` and `p.prokind = 'p'` filter to user-defined
/// procedures 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`) 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.
///
/// `prosecdef` is a boolean. Procedure-only fields (`returnType`, `volatility`,
/// `strict`, `parallelSafety`) are deliberately absent — see the tool description above.
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,
            'security',    CASE WHEN p.prosecdef THEN 'DEFINER' ELSE 'INVOKER' 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 = 'p'
      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 procedures, 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_procedures(
        &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 (procedures, next_cursor) = pager.paginate(rows);
        Ok(ListEntriesResponse::brief(procedures, next_cursor))
    }
}