systemprompt_database/admin/
admin_sql.rs1use thiserror::Error;
12
13pub const DEFAULT_READONLY_ROW_LIMIT: usize = 1000;
14
15const READONLY_PREFIXES: &[&str] = &["select", "with", "explain", "show", "table", "values"];
16
17const FORBIDDEN_KEYWORDS: &[&str] = &[
18 " drop ",
19 " delete ",
20 " insert ",
21 " update ",
22 " alter ",
23 " create ",
24 " truncate ",
25 " grant ",
26 " revoke ",
27 " copy ",
28 " vacuum ",
29 " call ",
30 " lock ",
31 " set ",
32 " reset ",
33 " rename ",
34];
35
36#[derive(Debug, Clone, Copy, Error)]
37pub enum AdminSqlError {
38 #[error("SQL query is empty")]
39 Empty,
40 #[error("SQL query contains multiple statements; only one is allowed")]
41 MultipleStatements,
42 #[error("SQL query must begin with SELECT, WITH, EXPLAIN, SHOW, TABLE, or VALUES")]
43 NotReadOnly,
44 #[error("SQL query contains forbidden keyword for read-only mode")]
45 ForbiddenKeyword,
46}
47
48#[derive(Debug, Clone, PartialEq, Eq)]
49pub struct AdminSql(String);
50
51impl AdminSql {
52 pub fn parse_readonly(raw: &str) -> Result<Self, AdminSqlError> {
53 let stripped = strip_comments(raw);
54 let trimmed = stripped.trim();
55 if trimmed.is_empty() {
56 return Err(AdminSqlError::Empty);
57 }
58
59 let without_terminator = trimmed.strip_suffix(';').unwrap_or(trimmed).trim_end();
60 if without_terminator.contains(';') {
61 return Err(AdminSqlError::MultipleStatements);
62 }
63
64 let lower = without_terminator.to_lowercase();
65 if !READONLY_PREFIXES
66 .iter()
67 .any(|p| starts_with_word(&lower, p))
68 {
69 return Err(AdminSqlError::NotReadOnly);
70 }
71
72 let padded = format!(" {lower} ");
73 if FORBIDDEN_KEYWORDS.iter().any(|kw| padded.contains(kw)) {
74 return Err(AdminSqlError::ForbiddenKeyword);
75 }
76
77 Ok(Self(without_terminator.to_owned()))
78 }
79
80 pub fn parse_unrestricted(raw: &str) -> Result<Self, AdminSqlError> {
81 let stripped = strip_comments(raw);
82 let trimmed = stripped.trim();
83 if trimmed.is_empty() {
84 return Err(AdminSqlError::Empty);
85 }
86
87 let without_terminator = trimmed.strip_suffix(';').unwrap_or(trimmed).trim_end();
88 if without_terminator.contains(';') {
89 return Err(AdminSqlError::MultipleStatements);
90 }
91
92 Ok(Self(without_terminator.to_owned()))
93 }
94
95 pub fn as_str(&self) -> &str {
96 &self.0
97 }
98}
99
100fn strip_comments(raw: &str) -> String {
101 let mut out = String::with_capacity(raw.len());
102 let mut chars = raw.chars().peekable();
103 while let Some(c) = chars.next() {
104 if c == '-' && chars.peek() == Some(&'-') {
105 for nc in chars.by_ref() {
106 if nc == '\n' {
107 out.push('\n');
108 break;
109 }
110 }
111 continue;
112 }
113 if c == '/' && chars.peek() == Some(&'*') {
114 chars.next();
115 let mut prev = '\0';
116 for nc in chars.by_ref() {
117 if prev == '*' && nc == '/' {
118 break;
119 }
120 prev = nc;
121 }
122 continue;
123 }
124 out.push(c);
125 }
126 out
127}
128
129fn starts_with_word(haystack: &str, needle: &str) -> bool {
130 if !haystack.starts_with(needle) {
131 return false;
132 }
133 haystack[needle.len()..]
134 .chars()
135 .next()
136 .is_none_or(|c| c.is_whitespace() || c == '(' || c == ';')
137}