database_mcp_mysql/tools/
get_table_schema.rs1use std::borrow::Cow;
4use std::collections::HashMap;
5
6use database_mcp_server::AppError;
7use database_mcp_server::types::{GetTableSchemaRequest, TableSchemaResponse};
8use database_mcp_sql::Connection as _;
9use database_mcp_sql::identifier::validate_identifier;
10use rmcp::handler::server::router::tool::{AsyncTool, ToolBase};
11use rmcp::model::{ErrorData, ToolAnnotations};
12use serde_json::{Value, json};
13
14use crate::MysqlHandler;
15
16pub(crate) struct GetTableSchemaTool;
18
19impl GetTableSchemaTool {
20 const NAME: &'static str = "get_table_schema";
21 const DESCRIPTION: &'static str = r#"Get column definitions and foreign key relationships for a table. Requires `database_name` and `table_name` — call `list_databases` and `list_tables` first.
22
23<usecase>
24ALWAYS call this before writing queries to understand:
25- Column names and data types
26- Which columns are nullable, primary keys, or have defaults
27- Foreign key relationships for writing JOINs
28</usecase>
29
30<examples>
31✓ "What columns does the orders table have?" → get_table_schema(database_name="mydb", table_name="orders")
32✓ Before writing a SELECT → get_table_schema first to confirm column names
33✓ "How are users and orders related?" → check foreign keys in both tables
34</examples>
35
36<what_it_returns>
37A JSON object with table_name and columns keyed by column name, each containing type, nullable, key, default, and foreign_key info.
38</what_it_returns>"#;
39}
40
41impl ToolBase for GetTableSchemaTool {
42 type Parameter = GetTableSchemaRequest;
43 type Output = TableSchemaResponse;
44 type Error = ErrorData;
45
46 fn name() -> Cow<'static, str> {
47 Self::NAME.into()
48 }
49
50 fn description() -> Option<Cow<'static, str>> {
51 Some(Self::DESCRIPTION.into())
52 }
53
54 fn annotations() -> Option<ToolAnnotations> {
55 Some(
56 ToolAnnotations::new()
57 .read_only(true)
58 .destructive(false)
59 .idempotent(true)
60 .open_world(false),
61 )
62 }
63}
64
65impl AsyncTool<MysqlHandler> for GetTableSchemaTool {
66 async fn invoke(handler: &MysqlHandler, params: Self::Parameter) -> Result<Self::Output, Self::Error> {
67 Ok(handler.get_table_schema(¶ms).await?)
68 }
69}
70
71impl MysqlHandler {
72 pub async fn get_table_schema(&self, request: &GetTableSchemaRequest) -> Result<TableSchemaResponse, AppError> {
78 let database = &request.database_name;
79 let table = &request.table_name;
80 validate_identifier(database)?;
81 validate_identifier(table)?;
82
83 let describe_sql = format!(
85 "DESCRIBE {}.{}",
86 self.connection.quote_identifier(database),
87 self.connection.quote_identifier(table)
88 );
89 let schema_rows = self.connection.fetch(describe_sql.as_str(), None).await?;
90
91 if schema_rows.is_empty() {
92 return Err(AppError::TableNotFound(format!("{database}.{table}")));
93 }
94
95 let mut columns: HashMap<String, Value> = HashMap::new();
96 for row in &schema_rows {
97 if let Some(col_name) = row.get("Field").and_then(|v| v.as_str()) {
98 columns.insert(
99 col_name.to_string(),
100 json!({
101 "type": row.get("Type").unwrap_or(&Value::Null),
102 "nullable": row.get("Null").and_then(|v| v.as_str()).is_some_and(|s| s.to_uppercase() == "YES"),
103 "key": row.get("Key").unwrap_or(&Value::Null),
104 "default": row.get("Default").unwrap_or(&Value::Null),
105 "extra": row.get("Extra").unwrap_or(&Value::Null),
106 "foreign_key": Value::Null,
107 }),
108 );
109 }
110 }
111
112 let fk_sql = format!(
114 "SELECT
115 kcu.COLUMN_NAME as column_name,
116 kcu.CONSTRAINT_NAME as constraint_name,
117 kcu.REFERENCED_TABLE_NAME as referenced_table,
118 kcu.REFERENCED_COLUMN_NAME as referenced_column,
119 rc.UPDATE_RULE as on_update,
120 rc.DELETE_RULE as on_delete
121 FROM information_schema.KEY_COLUMN_USAGE kcu
122 INNER JOIN information_schema.REFERENTIAL_CONSTRAINTS rc
123 ON kcu.CONSTRAINT_NAME = rc.CONSTRAINT_NAME
124 AND kcu.CONSTRAINT_SCHEMA = rc.CONSTRAINT_SCHEMA
125 WHERE kcu.TABLE_SCHEMA = {}
126 AND kcu.TABLE_NAME = {}
127 AND kcu.REFERENCED_TABLE_NAME IS NOT NULL
128 ORDER BY kcu.CONSTRAINT_NAME, kcu.ORDINAL_POSITION",
129 self.connection.quote_string(database),
130 self.connection.quote_string(table),
131 );
132
133 let fk_rows = self.connection.fetch(fk_sql.as_str(), None).await?;
134
135 for fk_row in &fk_rows {
136 if let Some(col_name) = fk_row.get("column_name").and_then(|v| v.as_str())
137 && let Some(col_info) = columns.get_mut(col_name)
138 && let Some(obj) = col_info.as_object_mut()
139 {
140 obj.insert(
141 "foreign_key".to_string(),
142 json!({
143 "constraint_name": fk_row.get("constraint_name"),
144 "referenced_table": fk_row.get("referenced_table"),
145 "referenced_column": fk_row.get("referenced_column"),
146 "on_update": fk_row.get("on_update"),
147 "on_delete": fk_row.get("on_delete"),
148 }),
149 );
150 }
151 }
152
153 Ok(TableSchemaResponse {
154 table_name: table.clone(),
155 columns: json!(columns),
156 })
157 }
158}