mcp_postgres/actions/
migration_helpers.rs

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

mcp_postgres/actions/
migration_helpers.rs
