database_mcp_postgres/tools/
get_table_schema.rs

1//! MCP tool: `get_table_schema`.
2
3use std::borrow::Cow;
4use std::collections::HashMap;
5
6use database_mcp_server::AppError;
7use database_mcp_server::types::{GetTableSchemaRequest, TableSchemaResponse};
8use database_mcp_sql::identifier::validate_identifier;
9use database_mcp_sql::timeout::execute_with_timeout;
10use rmcp::handler::server::router::tool::{AsyncTool, ToolBase};
11use rmcp::model::{ErrorData, ToolAnnotations};
12use serde_json::{Value, json};
13use sqlx::Row;
14use sqlx::postgres::PgRow;
15
16use crate::PostgresHandler;
17
18/// Marker type for the `get_table_schema` MCP tool.
19pub(crate) struct GetTableSchemaTool;
20
21impl GetTableSchemaTool {
22    const NAME: &'static str = "get_table_schema";
23    const DESCRIPTION: &'static str = "Get column definitions (type, nullable, key, default) and foreign key\nrelationships for a table. Requires `database_name` and `table_name`.";
24}
25
26impl ToolBase for GetTableSchemaTool {
27    type Parameter = GetTableSchemaRequest;
28    type Output = TableSchemaResponse;
29    type Error = ErrorData;
30
31    fn name() -> Cow<'static, str> {
32        Self::NAME.into()
33    }
34
35    fn description() -> Option<Cow<'static, str>> {
36        Some(Self::DESCRIPTION.into())
37    }
38
39    fn annotations() -> Option<ToolAnnotations> {
40        Some(
41            ToolAnnotations::new()
42                .read_only(true)
43                .destructive(false)
44                .idempotent(true)
45                .open_world(false),
46        )
47    }
48}
49
50impl AsyncTool<PostgresHandler> for GetTableSchemaTool {
51    async fn invoke(handler: &PostgresHandler, params: Self::Parameter) -> Result<Self::Output, Self::Error> {
52        Ok(handler.get_table_schema(&params).await?)
53    }
54}
55
56impl PostgresHandler {
57    /// Returns column definitions with foreign key relationships.
58    ///
59    /// # Errors
60    ///
61    /// Returns [`AppError`] if validation fails or the query errors.
62    pub async fn get_table_schema(&self, request: &GetTableSchemaRequest) -> Result<TableSchemaResponse, AppError> {
63        let table = &request.table_name;
64        validate_identifier(table)?;
65        let db = if request.database_name.is_empty() {
66            None
67        } else {
68            Some(request.database_name.as_str())
69        };
70        let pool = self.get_pool(db).await?;
71
72        // 1. Get basic schema
73        let schema_sql = r"SELECT column_name, data_type, is_nullable, column_default,
74                      character_maximum_length
75               FROM information_schema.columns
76               WHERE table_schema = 'public' AND table_name = $1
77               ORDER BY ordinal_position";
78        let rows: Vec<PgRow> = execute_with_timeout(
79            self.config.query_timeout,
80            schema_sql,
81            sqlx::query(schema_sql).bind(table).fetch_all(&pool),
82        )
83        .await?;
84
85        if rows.is_empty() {
86            return Err(AppError::TableNotFound(table.clone()));
87        }
88
89        let mut columns: HashMap<String, Value> = HashMap::new();
90        for row in &rows {
91            let col_name: String = row.try_get("column_name").unwrap_or_default();
92            let data_type: String = row.try_get("data_type").unwrap_or_default();
93            let nullable: String = row.try_get("is_nullable").unwrap_or_default();
94            let default: Option<String> = row.try_get("column_default").ok();
95            columns.insert(
96                col_name,
97                json!({
98                    "type": data_type,
99                    "nullable": nullable.to_uppercase() == "YES",
100                    "key": Value::Null,
101                    "default": default,
102                    "extra": Value::Null,
103                    "foreign_key": Value::Null,
104                }),
105            );
106        }
107
108        // 2. Get FK relationships
109        let fk_sql = r"SELECT
110                kcu.column_name,
111                tc.constraint_name,
112                ccu.table_name AS referenced_table,
113                ccu.column_name AS referenced_column,
114                rc.update_rule AS on_update,
115                rc.delete_rule AS on_delete
116            FROM information_schema.table_constraints tc
117            JOIN information_schema.key_column_usage kcu
118                ON tc.constraint_name = kcu.constraint_name
119                AND tc.table_schema = kcu.table_schema
120            JOIN information_schema.constraint_column_usage ccu
121                ON ccu.constraint_name = tc.constraint_name
122                AND ccu.table_schema = tc.table_schema
123            JOIN information_schema.referential_constraints rc
124                ON rc.constraint_name = tc.constraint_name
125                AND rc.constraint_schema = tc.table_schema
126            WHERE tc.constraint_type = 'FOREIGN KEY'
127                AND tc.table_name = $1
128                AND tc.table_schema = 'public'";
129        let fk_rows: Vec<PgRow> = execute_with_timeout(
130            self.config.query_timeout,
131            fk_sql,
132            sqlx::query(fk_sql).bind(table).fetch_all(&pool),
133        )
134        .await?;
135
136        for fk_row in &fk_rows {
137            let col_name: String = fk_row.try_get("column_name").unwrap_or_default();
138            if let Some(col_info) = columns.get_mut(&col_name)
139                && let Some(obj) = col_info.as_object_mut()
140            {
141                obj.insert(
142                    "foreign_key".to_string(),
143                    json!({
144                        "constraint_name": fk_row.try_get::<String, _>("constraint_name").ok(),
145                        "referenced_table": fk_row.try_get::<String, _>("referenced_table").ok(),
146                        "referenced_column": fk_row.try_get::<String, _>("referenced_column").ok(),
147                        "on_update": fk_row.try_get::<String, _>("on_update").ok(),
148                        "on_delete": fk_row.try_get::<String, _>("on_delete").ok(),
149                    }),
150                );
151            }
152        }
153
154        Ok(TableSchemaResponse {
155            table_name: table.clone(),
156            columns: json!(columns),
157        })
158    }
159}
database_mcp_postgres/tools/get_table_schema.rs

database_mcp_postgres/tools/
get_table_schema.rs
