tap_mcp/tools/
database_tools.rs

1//! Database query tools for read-only SQL access
2
3use crate::error::{Error, Result};
4use crate::mcp::protocol::{CallToolResult, Tool};
5use crate::tap_integration::TapIntegration;
6use crate::tools::{error_text_response, success_text_response, ToolHandler};
7use async_trait::async_trait;
8use serde::{Deserialize, Serialize};
9use serde_json::{json, Value};
10use sqlx::{Column, Connection, Row, SqliteConnection};
11use std::sync::Arc;
12use tracing::{debug, error};
13
14/// Tool for executing read-only SQL queries
15pub struct QueryDatabaseTool {
16    tap_integration: Arc<TapIntegration>,
17}
18
19/// Parameters for querying the database
20#[derive(Debug, Deserialize)]
21struct QueryDatabaseParams {
22    agent_did: String,
23    query: String,
24}
25
26/// Response for database query
27#[derive(Debug, Serialize)]
28struct QueryDatabaseResponse {
29    columns: Vec<String>,
30    rows: Vec<Vec<Value>>,
31    row_count: usize,
32    query: String,
33}
34
35impl QueryDatabaseTool {
36    pub fn new(tap_integration: Arc<TapIntegration>) -> Self {
37        Self { tap_integration }
38    }
39
40    /// Validate that a table name contains only safe identifier characters
41    fn validate_table_name(name: &str) -> bool {
42        !name.is_empty()
43            && name.len() <= 128
44            && name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
45            && name
46                .chars()
47                .next()
48                .is_some_and(|c| c.is_ascii_alphabetic() || c == '_')
49    }
50
51    /// Check if a query is read-only
52    fn is_read_only_query(query: &str) -> bool {
53        let query_upper = query.trim().to_uppercase();
54        let forbidden_keywords = [
55            "INSERT",
56            "UPDATE",
57            "DELETE",
58            "DROP",
59            "CREATE",
60            "ALTER",
61            "TRUNCATE",
62            "REPLACE",
63            "MERGE",
64            "CALL",
65            "EXECUTE",
66            "EXEC",
67            "BEGIN",
68            "COMMIT",
69            "ROLLBACK",
70            "SAVEPOINT",
71            "GRANT",
72            "REVOKE",
73            "DENY",
74            "ATTACH",
75            "DETACH",
76        ];
77
78        // Check if query starts with SELECT, WITH, or PRAGMA (for schema queries)
79        let allowed_starts = ["SELECT", "WITH", "PRAGMA", "EXPLAIN"];
80        let starts_with_allowed = allowed_starts
81            .iter()
82            .any(|&start| query_upper.starts_with(start));
83
84        // Check for forbidden keywords
85        let contains_forbidden = forbidden_keywords.iter().any(|&keyword| {
86            // Check for whole word matches to avoid false positives
87            query_upper.split_whitespace().any(|word| word == keyword)
88        });
89
90        starts_with_allowed && !contains_forbidden
91    }
92}
93
94#[async_trait]
95impl ToolHandler for QueryDatabaseTool {
96    async fn handle(&self, arguments: Option<Value>) -> Result<CallToolResult> {
97        let params: QueryDatabaseParams = match arguments {
98            Some(args) => serde_json::from_value(args)
99                .map_err(|e| Error::invalid_parameter(format!("Invalid parameters: {}", e)))?,
100            None => {
101                return Ok(error_text_response(
102                    "Missing required parameters".to_string(),
103                ))
104            }
105        };
106
107        debug!(
108            "Executing query for agent {}: {}",
109            params.agent_did, params.query
110        );
111
112        // Validate query is read-only
113        if !Self::is_read_only_query(&params.query) {
114            return Ok(error_text_response(
115                "Only read-only queries are allowed. Query must start with SELECT, WITH, PRAGMA, or EXPLAIN and cannot contain modification keywords.".to_string(),
116            ));
117        }
118
119        // Get agent storage
120        let storage = match self
121            .tap_integration
122            .storage_for_agent(&params.agent_did)
123            .await
124        {
125            Ok(storage) => storage,
126            Err(e) => {
127                error!(
128                    "Failed to get storage for agent {}: {}",
129                    params.agent_did, e
130                );
131                return Ok(error_text_response(format!(
132                    "Failed to get storage for agent {}: {}",
133                    params.agent_did, e
134                )));
135            }
136        };
137
138        // Get database path from storage
139        let db_path = storage.db_path();
140        let db_url = format!("sqlite://{}?mode=ro", db_path.display());
141
142        // Connect to database in read-only mode
143        let mut conn = match SqliteConnection::connect(&db_url).await {
144            Ok(conn) => conn,
145            Err(e) => {
146                error!("Failed to connect to database: {}", e);
147                return Ok(error_text_response(format!(
148                    "Failed to connect to database: {}",
149                    e
150                )));
151            }
152        };
153
154        // Enforce read-only mode at the database level
155        if let Err(e) = sqlx::query("PRAGMA query_only = ON")
156            .execute(&mut conn)
157            .await
158        {
159            error!("Failed to set query_only pragma: {}", e);
160            return Ok(error_text_response(format!(
161                "Failed to set read-only mode: {}",
162                e
163            )));
164        }
165
166        // Execute query
167        match sqlx::query(&params.query).fetch_all(&mut conn).await {
168            Ok(rows) => {
169                let mut columns = Vec::new();
170                let mut result_rows = Vec::new();
171
172                if !rows.is_empty() {
173                    // Get column names from the first row
174                    let first_row = &rows[0];
175                    for column in first_row.columns().iter() {
176                        columns.push(column.name().to_string());
177                    }
178
179                    // Process all rows
180                    for row in &rows {
181                        let mut row_values = Vec::new();
182                        for i in 0..columns.len() {
183                            // Try different types in order of likelihood
184                            let value = if let Ok(v) = row.try_get::<Option<i64>, _>(i) {
185                                v.map(Value::from).unwrap_or(Value::Null)
186                            } else if let Ok(v) = row.try_get::<Option<f64>, _>(i) {
187                                v.map(Value::from).unwrap_or(Value::Null)
188                            } else if let Ok(v) = row.try_get::<Option<String>, _>(i) {
189                                v.map(Value::from).unwrap_or(Value::Null)
190                            } else if let Ok(v) = row.try_get::<Option<bool>, _>(i) {
191                                v.map(Value::from).unwrap_or(Value::Null)
192                            } else if let Ok(v) = row.try_get::<Option<Vec<u8>>, _>(i) {
193                                v.map(|bytes| {
194                                    // Try to convert bytes to string if possible
195                                    if let Ok(s) = String::from_utf8(bytes.clone()) {
196                                        Value::String(s)
197                                    } else {
198                                        // Return as base64 encoded string
199                                        use base64::Engine;
200                                        Value::String(
201                                            base64::engine::general_purpose::STANDARD.encode(bytes),
202                                        )
203                                    }
204                                })
205                                .unwrap_or(Value::Null)
206                            } else {
207                                Value::Null
208                            };
209                            row_values.push(value);
210                        }
211                        result_rows.push(row_values);
212                    }
213                }
214
215                let response = QueryDatabaseResponse {
216                    columns,
217                    row_count: result_rows.len(),
218                    rows: result_rows,
219                    query: params.query,
220                };
221
222                let response_json = serde_json::to_string_pretty(&response).map_err(|e| {
223                    Error::tool_execution(format!("Failed to serialize response: {}", e))
224                })?;
225
226                Ok(success_text_response(response_json))
227            }
228            Err(e) => {
229                error!("Failed to execute query: {}", e);
230                Ok(error_text_response(format!(
231                    "Failed to execute query: {}",
232                    e
233                )))
234            }
235        }
236    }
237
238    fn get_definition(&self) -> Tool {
239        Tool {
240            name: "tap_query_database".to_string(),
241            description: "Executes read-only SQL queries on an agent's database. Only SELECT, WITH, PRAGMA, and EXPLAIN queries are allowed.".to_string(),
242            input_schema: json!({
243                "type": "object",
244                "properties": {
245                    "agent_did": {
246                        "type": "string",
247                        "description": "The DID of the agent whose database to query"
248                    },
249                    "query": {
250                        "type": "string",
251                        "description": "The read-only SQL query to execute"
252                    }
253                },
254                "required": ["agent_did", "query"],
255                "additionalProperties": false
256            }),
257        }
258    }
259}
260
261/// Tool for getting database schema
262pub struct GetDatabaseSchemaTool {
263    tap_integration: Arc<TapIntegration>,
264}
265
266/// Parameters for getting database schema
267#[derive(Debug, Deserialize)]
268struct GetDatabaseSchemaParams {
269    agent_did: String,
270    #[serde(default)]
271    table_name: Option<String>,
272}
273
274/// Table information
275#[derive(Debug, Serialize)]
276struct TableInfo {
277    name: String,
278    columns: Vec<ColumnInfo>,
279    indexes: Vec<IndexInfo>,
280    row_count: i64,
281}
282
283/// Column information
284#[derive(Debug, Serialize)]
285struct ColumnInfo {
286    cid: i32,
287    name: String,
288    #[serde(rename = "type")]
289    column_type: String,
290    notnull: bool,
291    dflt_value: Option<String>,
292    pk: bool,
293}
294
295/// Index information
296#[derive(Debug, Serialize)]
297struct IndexInfo {
298    name: String,
299    unique: bool,
300    origin: String,
301    partial: bool,
302}
303
304/// Response for database schema
305#[derive(Debug, Serialize)]
306struct GetDatabaseSchemaResponse {
307    database_path: String,
308    tables: Vec<TableInfo>,
309}
310
311impl GetDatabaseSchemaTool {
312    pub fn new(tap_integration: Arc<TapIntegration>) -> Self {
313        Self { tap_integration }
314    }
315}
316
317#[async_trait]
318impl ToolHandler for GetDatabaseSchemaTool {
319    async fn handle(&self, arguments: Option<Value>) -> Result<CallToolResult> {
320        let params: GetDatabaseSchemaParams = match arguments {
321            Some(args) => serde_json::from_value(args)
322                .map_err(|e| Error::invalid_parameter(format!("Invalid parameters: {}", e)))?,
323            None => {
324                return Ok(error_text_response(
325                    "Missing required parameters".to_string(),
326                ))
327            }
328        };
329
330        debug!("Getting database schema for agent {}", params.agent_did);
331
332        // Get agent storage
333        let storage = match self
334            .tap_integration
335            .storage_for_agent(&params.agent_did)
336            .await
337        {
338            Ok(storage) => storage,
339            Err(e) => {
340                error!(
341                    "Failed to get storage for agent {}: {}",
342                    params.agent_did, e
343                );
344                return Ok(error_text_response(format!(
345                    "Failed to get storage for agent {}: {}",
346                    params.agent_did, e
347                )));
348            }
349        };
350
351        // Get database path from storage
352        let db_path = storage.db_path();
353        let db_url = format!("sqlite://{}?mode=ro", db_path.display());
354
355        // Connect to database in read-only mode
356        let mut conn = match SqliteConnection::connect(&db_url).await {
357            Ok(conn) => conn,
358            Err(e) => {
359                error!("Failed to connect to database: {}", e);
360                return Ok(error_text_response(format!(
361                    "Failed to connect to database: {}",
362                    e
363                )));
364            }
365        };
366
367        let mut tables = Vec::new();
368
369        // Validate table name if provided
370        if let Some(ref table_name) = params.table_name {
371            if !QueryDatabaseTool::validate_table_name(table_name) {
372                return Ok(error_text_response(
373                    "Invalid table name. Table names must contain only alphanumeric characters and underscores.".to_string(),
374                ));
375            }
376        }
377
378        // Get list of tables using parameterized query
379        let table_rows = if let Some(ref table_name) = params.table_name {
380            match sqlx::query(
381                "SELECT name FROM sqlite_master WHERE type='table' AND name=?1 ORDER BY name",
382            )
383            .bind(table_name)
384            .fetch_all(&mut conn)
385            .await
386            {
387                Ok(rows) => rows,
388                Err(e) => {
389                    error!("Failed to get tables: {}", e);
390                    return Ok(error_text_response(format!("Failed to get tables: {}", e)));
391                }
392            }
393        } else {
394            match sqlx::query("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
395                .fetch_all(&mut conn)
396                .await
397            {
398                Ok(rows) => rows,
399                Err(e) => {
400                    error!("Failed to get tables: {}", e);
401                    return Ok(error_text_response(format!("Failed to get tables: {}", e)));
402                }
403            }
404        };
405
406        for table_row in table_rows {
407            let table_name: String = table_row.try_get("name").unwrap_or_default();
408
409            // Validate table name from database as defense-in-depth
410            if !QueryDatabaseTool::validate_table_name(&table_name) {
411                continue;
412            }
413
414            // Get columns for this table (safe: table_name validated above)
415            let column_query = format!("PRAGMA table_info('{}')", table_name);
416            let column_rows = match sqlx::query(&column_query).fetch_all(&mut conn).await {
417                Ok(rows) => rows,
418                Err(e) => {
419                    error!("Failed to get columns for table {}: {}", table_name, e);
420                    continue;
421                }
422            };
423
424            let mut columns = Vec::new();
425            for col_row in column_rows {
426                columns.push(ColumnInfo {
427                    cid: col_row.try_get("cid").unwrap_or(0),
428                    name: col_row.try_get("name").unwrap_or_default(),
429                    column_type: col_row.try_get("type").unwrap_or_default(),
430                    notnull: col_row.try_get::<i32, _>("notnull").unwrap_or(0) != 0,
431                    dflt_value: col_row.try_get("dflt_value").ok(),
432                    pk: col_row.try_get::<i32, _>("pk").unwrap_or(0) != 0,
433                });
434            }
435
436            // Get indexes for this table
437            let index_query = format!("PRAGMA index_list('{}')", table_name);
438            let index_rows = match sqlx::query(&index_query).fetch_all(&mut conn).await {
439                Ok(rows) => rows,
440                Err(e) => {
441                    error!("Failed to get indexes for table {}: {}", table_name, e);
442                    vec![]
443                }
444            };
445
446            let mut indexes = Vec::new();
447            for idx_row in index_rows {
448                indexes.push(IndexInfo {
449                    name: idx_row.try_get("name").unwrap_or_default(),
450                    unique: idx_row.try_get::<i32, _>("unique").unwrap_or(0) != 0,
451                    origin: idx_row.try_get("origin").unwrap_or_default(),
452                    partial: idx_row.try_get::<i32, _>("partial").unwrap_or(0) != 0,
453                });
454            }
455
456            // Get row count
457            let count_query = format!("SELECT COUNT(*) as count FROM '{}'", table_name);
458            let row_count = match sqlx::query(&count_query).fetch_one(&mut conn).await {
459                Ok(row) => row.try_get::<i64, _>("count").unwrap_or(0),
460                Err(e) => {
461                    error!("Failed to get row count for table {}: {}", table_name, e);
462                    0
463                }
464            };
465
466            tables.push(TableInfo {
467                name: table_name,
468                columns,
469                indexes,
470                row_count,
471            });
472        }
473
474        let response = GetDatabaseSchemaResponse {
475            database_path: db_path.display().to_string(),
476            tables,
477        };
478
479        let response_json = serde_json::to_string_pretty(&response)
480            .map_err(|e| Error::tool_execution(format!("Failed to serialize response: {}", e)))?;
481
482        Ok(success_text_response(response_json))
483    }
484
485    fn get_definition(&self) -> Tool {
486        Tool {
487            name: "tap_get_database_schema".to_string(),
488            description: "Gets the schema of an agent's database, including all tables, columns, indexes, and row counts. Optionally filter by table name.".to_string(),
489            input_schema: json!({
490                "type": "object",
491                "properties": {
492                    "agent_did": {
493                        "type": "string",
494                        "description": "The DID of the agent whose database schema to retrieve"
495                    },
496                    "table_name": {
497                        "type": "string",
498                        "description": "Optional specific table name to get schema for"
499                    }
500                },
501                "required": ["agent_did"],
502                "additionalProperties": false
503            }),
504        }
505    }
506}
tap_mcp/tools/database_tools.rs

tap_mcp/tools/
database_tools.rs
