mcp_postgres/actions/
migration_helpers.rs1use 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 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 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}