database_mcp_postgres/
schema.rs

1//! `PostgreSQL` table schema introspection.
2
3use std::collections::HashMap;
4
5use database_mcp_backend::error::AppError;
6use database_mcp_backend::identifier::validate_identifier;
7use serde_json::{Value, json};
8use sqlx::Row;
9use sqlx::postgres::PgRow;
10
11use super::PostgresBackend;
12
13impl PostgresBackend {
14    /// Returns column definitions with foreign key relationships.
15    ///
16    /// # Errors
17    ///
18    /// Returns [`AppError`] if validation fails or the query errors.
19    pub async fn get_table_schema(&self, database: &str, table: &str) -> Result<Value, AppError> {
20        validate_identifier(table)?;
21        let db = if database.is_empty() { None } else { Some(database) };
22        let pool = self.get_pool(db).await?;
23
24        // 1. Get basic schema
25        let rows: Vec<PgRow> = sqlx::query(
26            r"SELECT column_name, data_type, is_nullable, column_default,
27                      character_maximum_length
28               FROM information_schema.columns
29               WHERE table_schema = 'public' AND table_name = $1
30               ORDER BY ordinal_position",
31        )
32        .bind(table)
33        .fetch_all(&pool)
34        .await
35        .map_err(|e| AppError::Query(e.to_string()))?;
36
37        if rows.is_empty() {
38            return Err(AppError::TableNotFound(table.to_string()));
39        }
40
41        let mut columns: HashMap<String, Value> = HashMap::new();
42        for row in &rows {
43            let col_name: String = row.try_get("column_name").unwrap_or_default();
44            let data_type: String = row.try_get("data_type").unwrap_or_default();
45            let nullable: String = row.try_get("is_nullable").unwrap_or_default();
46            let default: Option<String> = row.try_get("column_default").ok();
47            columns.insert(
48                col_name,
49                json!({
50                    "type": data_type,
51                    "nullable": nullable.to_uppercase() == "YES",
52                    "key": Value::Null,
53                    "default": default,
54                    "extra": Value::Null,
55                    "foreign_key": Value::Null,
56                }),
57            );
58        }
59
60        // 2. Get FK relationships
61        let fk_rows: Vec<PgRow> = sqlx::query(
62            r"SELECT
63                kcu.column_name,
64                tc.constraint_name,
65                ccu.table_name AS referenced_table,
66                ccu.column_name AS referenced_column,
67                rc.update_rule AS on_update,
68                rc.delete_rule AS on_delete
69            FROM information_schema.table_constraints tc
70            JOIN information_schema.key_column_usage kcu
71                ON tc.constraint_name = kcu.constraint_name
72                AND tc.table_schema = kcu.table_schema
73            JOIN information_schema.constraint_column_usage ccu
74                ON ccu.constraint_name = tc.constraint_name
75                AND ccu.table_schema = tc.table_schema
76            JOIN information_schema.referential_constraints rc
77                ON rc.constraint_name = tc.constraint_name
78                AND rc.constraint_schema = tc.table_schema
79            WHERE tc.constraint_type = 'FOREIGN KEY'
80                AND tc.table_name = $1
81                AND tc.table_schema = 'public'",
82        )
83        .bind(table)
84        .fetch_all(&pool)
85        .await
86        .map_err(|e| AppError::Query(e.to_string()))?;
87
88        for fk_row in &fk_rows {
89            let col_name: String = fk_row.try_get("column_name").unwrap_or_default();
90            if let Some(col_info) = columns.get_mut(&col_name)
91                && let Some(obj) = col_info.as_object_mut()
92            {
93                obj.insert(
94                    "foreign_key".to_string(),
95                    json!({
96                        "constraint_name": fk_row.try_get::<String, _>("constraint_name").ok(),
97                        "referenced_table": fk_row.try_get::<String, _>("referenced_table").ok(),
98                        "referenced_column": fk_row.try_get::<String, _>("referenced_column").ok(),
99                        "on_update": fk_row.try_get::<String, _>("on_update").ok(),
100                        "on_delete": fk_row.try_get::<String, _>("on_delete").ok(),
101                    }),
102                );
103            }
104        }
105
106        Ok(json!({
107            "table_name": table,
108            "columns": columns,
109        }))
110    }
111}
database_mcp_postgres/schema.rs

database_mcp_postgres/
schema.rs
