Skip to main content

systemprompt_database/admin/
admin_sql.rs

1//! Parser/validator for admin-supplied SQL strings.
2//!
3//! Two parse modes are exposed:
4//! - [`AdminSql::parse_readonly`] — only `SELECT`/`WITH`/`EXPLAIN`/`SHOW`/
5//!   `TABLE`/`VALUES` queries with no forbidden keywords.
6//! - [`AdminSql::parse_unrestricted`] — single statement, otherwise free-form.
7//!
8//! Copyright (c) systemprompt.io — Business Source License 1.1.
9//! See <https://systemprompt.io> for licensing details.
10
11use 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}