Skip to main content

mcp_postgres/actions/
migration_helpers.rs

1use crate::errors::Result as MCPResult;
2use serde_json::{Value, json};
3use tokio_postgres::Client;
4
5const MAX_IDENTIFIER_LEN: usize = 255;
6
7pub async fn generate_create_table_ddl(
8    client: &Client,
9    params: &Option<&Value>,
10) -> MCPResult<Value> {
11    let table = params
12        .as_ref()
13        .and_then(|p| p.get("table").and_then(|v| v.as_str()))
14        .ok_or_else(|| {
15            crate::errors::MCPError::InvalidParams("Missing 'table' parameter".into())
16        })?;
17    let schema = params
18        .as_ref()
19        .and_then(|p| p.get("schema").and_then(|v| v.as_str()))
20        .unwrap_or("public");
21
22    if table.is_empty() || table.len() > MAX_IDENTIFIER_LEN {
23        return Err(crate::errors::MCPError::InvalidParams(format!(
24            "'table' must be 1-{MAX_IDENTIFIER_LEN} characters"
25        )));
26    }
27
28    let rows = client
29        .query(
30            "SELECT column_name, data_type, character_maximum_length,
31                    is_nullable, column_default, ordinal_position
32             FROM information_schema.columns
33             WHERE table_schema = $1 AND table_name = $2
34             ORDER BY ordinal_position",
35            &[&schema, &table],
36        )
37        .await?;
38
39    if rows.is_empty() {
40        return Err(crate::errors::MCPError::InvalidParams(format!(
41            "Table {}.{} not found",
42            schema, table
43        )));
44    }
45
46    let mut ddl = format!(
47        "CREATE TABLE {}.{} (\n",
48        crate::validation::quote_ident(schema),
49        crate::validation::quote_ident(table)
50    );
51    let mut cols = Vec::new();
52
53    for row in &rows {
54        let col_name: String = row.get(0);
55        let data_type: String = row.get(1);
56        let max_len: Option<i32> = row.get(2);
57        let nullable: String = row.get(3);
58        let default: Option<String> = row.get(4);
59
60        let q_col = crate::validation::quote_ident(&col_name);
61        let mut col = format!("    {q_col} {data_type}");
62        if let Some(len) = max_len {
63            col = format!("    {q_col} {data_type}({len})");
64        }
65        if let Some(d) = default {
66            col.push_str(&format!(" DEFAULT {}", d));
67        }
68        if nullable == "NO" {
69            col.push_str(" NOT NULL");
70        }
71        cols.push(col);
72    }
73
74    // Get primary key
75    let pk_rows = client
76        .query(
77            "SELECT kcu.column_name
78             FROM information_schema.table_constraints tc
79             JOIN information_schema.key_column_usage kcu ON kcu.constraint_name = tc.constraint_name
80             WHERE tc.table_schema = $1 AND tc.table_name = $2 AND tc.constraint_type = 'PRIMARY KEY'
81             ORDER BY kcu.ordinal_position",
82            &[&schema, &table],
83        )
84        .await?;
85
86    if !pk_rows.is_empty() {
87        let pk_cols: Vec<String> = pk_rows
88            .iter()
89            .map(|r| crate::validation::quote_ident(&r.get::<_, String>(0)))
90            .collect();
91        cols.push(format!("    PRIMARY KEY ({})", pk_cols.join(", ")));
92    }
93
94    ddl.push_str(&cols.join(",\n"));
95    ddl.push_str("\n);");
96
97    Ok(json!({ "ddl": ddl, "table": table, "schema": schema }))
98}
99
100pub async fn generate_create_index_ddl(
101    client: &Client,
102    params: &Option<&Value>,
103) -> MCPResult<Value> {
104    let table = params
105        .as_ref()
106        .and_then(|p| p.get("table").and_then(|v| v.as_str()))
107        .ok_or_else(|| {
108            crate::errors::MCPError::InvalidParams("Missing 'table' parameter".into())
109        })?;
110    let schema = params
111        .as_ref()
112        .and_then(|p| p.get("schema").and_then(|v| v.as_str()))
113        .unwrap_or("public");
114
115    let rows = client
116        .query(
117            "SELECT i.indexrelid::regclass::text AS index_name,
118                    pg_get_indexdef(i.indexrelid) AS index_def
119             FROM pg_index i
120             JOIN pg_class c ON c.oid = i.indrelid
121             JOIN pg_namespace n ON n.oid = c.relnamespace
122             WHERE n.nspname = $1 AND c.relname = $2
123             ORDER BY i.indexrelid::regclass::text",
124            &[&schema, &table],
125        )
126        .await?;
127
128    if rows.is_empty() {
129        return Ok(json!({ "indexes": [], "table": table, "message": "No indexes found" }));
130    }
131
132    let indexes: Vec<Value> = rows
133        .iter()
134        .map(|row| {
135            json!({
136                "name": row.get::<_, String>(0),
137                "ddl": row.get::<_, String>(1),
138            })
139        })
140        .collect();
141
142    Ok(json!({ "indexes": indexes, "table": table }))
143}
144
145pub async fn table_dependencies(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
146    let table = params
147        .as_ref()
148        .and_then(|p| p.get("table").and_then(|v| v.as_str()))
149        .ok_or_else(|| {
150            crate::errors::MCPError::InvalidParams("Missing 'table' parameter".into())
151        })?;
152    let schema = params
153        .as_ref()
154        .and_then(|p| p.get("schema").and_then(|v| v.as_str()))
155        .unwrap_or("public");
156
157    let rows = client
158        .query(
159            "SELECT
160                cl.relname AS dependent_object,
161                CASE
162                    WHEN cl.relkind = 'r' THEN 'table'
163                    WHEN cl.relkind = 'v' THEN 'view'
164                    WHEN cl.relkind = 'm' THEN 'materialized_view'
165                    WHEN cl.relkind = 'i' THEN 'index'
166                    WHEN cl.relkind = 'S' THEN 'sequence'
167                    ELSE 'other'
168                END AS object_type,
169                n.nspname AS object_schema,
170                d.deptype::text AS dependency_type
171             FROM pg_depend d
172             JOIN pg_class cl ON cl.oid = d.objid
173             JOIN pg_namespace n ON n.oid = cl.relnamespace
174             JOIN pg_class ref_cl ON ref_cl.oid = d.refobjid
175             WHERE d.refobjid = (SELECT c.oid FROM pg_class c
176                                 JOIN pg_namespace n2 ON n2.oid = c.relnamespace
177                                 WHERE n2.nspname = $1 AND c.relname = $2)
178               AND d.deptype IN ('n', 'a')
179             ORDER BY object_type, dependent_object",
180            &[&schema, &table],
181        )
182        .await?;
183
184    let dependencies: Vec<Value> = rows
185        .iter()
186        .map(|row| {
187            json!({
188                "object": row.get::<_, String>(0),
189                "type": row.get::<_, String>(1),
190                "schema": row.get::<_, String>(2),
191                "dependency": if row.get::<_, String>(3) == "n" { "normal" } else { "automatic" },
192            })
193        })
194        .collect();
195
196    // Also get what this table depends on
197    let dep_on_rows = client
198        .query(
199            "SELECT cl.relname AS referenced_object,
200                    CASE
201                        WHEN cl.relkind = 'r' THEN 'table'
202                        WHEN cl.relkind = 'v' THEN 'view'
203                        WHEN cl.relkind = 'S' THEN 'sequence'
204                        ELSE 'other'
205                    END AS object_type,
206                    n.nspname AS object_schema
207             FROM pg_depend d
208             JOIN pg_class cl ON cl.oid = d.refobjid
209             JOIN pg_namespace n ON n.oid = cl.relnamespace
210             WHERE d.objid = (SELECT c.oid FROM pg_class c
211                              JOIN pg_namespace n2 ON n2.oid = c.relnamespace
212                              WHERE n2.nspname = $1 AND c.relname = $2)
213               AND cl.relname != $2
214               AND d.deptype = 'n'
215             ORDER BY object_type, referenced_object",
216            &[&schema, &table],
217        )
218        .await?;
219
220    let depends_on: Vec<Value> = dep_on_rows
221        .iter()
222        .map(|row| {
223            json!({
224                "object": row.get::<_, String>(0),
225                "type": row.get::<_, String>(1),
226                "schema": row.get::<_, String>(2),
227            })
228        })
229        .collect();
230
231    Ok(json!({
232        "table": table,
233        "schema": schema,
234        "referenced_by": dependencies,
235        "depends_on": depends_on,
236    }))
237}