Skip to main content

fraiseql_functions/host/live/sql_classifier/
mod.rs

1//! SQL query classification for read-only enforcement.
2//!
3//! This module uses `sqlparser-rs` to parse SQL and classify it as either
4//! read-only or explicitly rejected. The approach is whitelist-only:
5//! only `SELECT` statements and `EXPLAIN` (without `ANALYZE`) are allowed.
6//! Everything else is rejected by default.
7
8use fraiseql_error::{FraiseQLError, Result};
9
10/// Classification result for a SQL statement.
11#[derive(Debug, Clone, PartialEq, Eq)]
12pub enum SqlClassification {
13    /// Statement is safe to execute (read-only).
14    ReadOnly,
15    /// Statement is not allowed.
16    Rejected(RejectionReason),
17}
18
19/// Reasons why a SQL statement was rejected.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum RejectionReason {
22    /// Write statement (INSERT, UPDATE, DELETE, MERGE).
23    WriteStatement(String),
24    /// DDL statement (CREATE, DROP, ALTER, TRUNCATE).
25    DdlStatement(String),
26    /// CTE containing writable statements.
27    WritableCte,
28    /// Privilege escalation (SET ROLE, SET SESSION AUTHORIZATION).
29    PrivilegeEscalation,
30    /// Procedural block (DO $$ ... $$).
31    ProceduralBlock,
32    /// Procedure call (CALL `procedure()`).
33    ProcedureCall,
34    /// COPY statement (can write).
35    CopyStatement,
36    /// EXPLAIN ANALYZE (actually executes the statement).
37    ExplainAnalyze,
38    /// Unknown statement type (not explicitly whitelisted).
39    Unknown(String),
40}
41
42impl std::fmt::Display for RejectionReason {
43    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44        match self {
45            Self::WriteStatement(stmt) => write!(f, "write statement not allowed: {}", stmt),
46            Self::DdlStatement(stmt) => write!(f, "DDL statement not allowed: {}", stmt),
47            Self::WritableCte => write!(f, "CTE with writable statement not allowed"),
48            Self::PrivilegeEscalation => write!(f, "privilege escalation not allowed"),
49            Self::ProceduralBlock => write!(f, "procedural block not allowed"),
50            Self::ProcedureCall => write!(f, "procedure call not allowed"),
51            Self::CopyStatement => write!(f, "COPY statement not allowed"),
52            Self::ExplainAnalyze => {
53                write!(f, "EXPLAIN ANALYZE not allowed (executes the statement)")
54            },
55            Self::Unknown(stmt) => write!(f, "unknown or disallowed statement: {}", stmt),
56        }
57    }
58}
59
60/// Classify a SQL statement as read-only or rejected.
61///
62/// Uses a whitelist-only approach: only `SELECT` queries and `EXPLAIN` (without `ANALYZE`)
63/// are allowed. All other statements are rejected.
64///
65/// # Arguments
66///
67/// * `sql` - The SQL statement to classify
68///
69/// # Returns
70///
71/// - `Ok(SqlClassification::ReadOnly)` if the statement is safe
72/// - `Ok(SqlClassification::Rejected(reason))` if the statement is not allowed
73/// - `Err` if parsing fails
74///
75/// # Errors
76///
77/// Returns a validation error if the SQL cannot be parsed.
78pub fn classify_sql(sql: &str) -> Result<SqlClassification> {
79    use sqlparser::{dialect::PostgreSqlDialect, parser::Parser};
80
81    let dialect = PostgreSqlDialect {};
82    let statements = Parser::parse_sql(&dialect, sql).map_err(|e| FraiseQLError::Validation {
83        message: format!("invalid SQL: {}", e),
84        path:    None,
85    })?;
86
87    // Check each statement in the batch
88    for stmt in statements {
89        let classification = classify_statement(&stmt)?;
90        match classification {
91            SqlClassification::ReadOnly => {},
92            SqlClassification::Rejected(reason) => return Ok(SqlClassification::Rejected(reason)),
93        }
94    }
95
96    Ok(SqlClassification::ReadOnly)
97}
98
99/// Classify a single parsed statement.
100fn classify_statement(stmt: &sqlparser::ast::Statement) -> Result<SqlClassification> {
101    use sqlparser::ast::Statement;
102
103    match stmt {
104        // SELECT queries are allowed *only* if no part of the query tree contains a
105        // data-modifying statement. PostgreSQL permits data-modifying CTEs
106        // (`WITH t AS (DELETE FROM x RETURNING *) SELECT * FROM t`); such a query
107        // parses as `Statement::Query` but must NOT be treated as read-only.
108        Statement::Query(q) => {
109            if query_is_data_modifying(q) {
110                Ok(SqlClassification::Rejected(RejectionReason::WritableCte))
111            } else {
112                Ok(SqlClassification::ReadOnly)
113            }
114        },
115
116        // EXPLAIN is allowed, but not EXPLAIN ANALYZE (which executes the statement)
117        Statement::Explain { analyze, .. } => {
118            if *analyze {
119                Ok(SqlClassification::Rejected(RejectionReason::ExplainAnalyze))
120            } else {
121                Ok(SqlClassification::ReadOnly)
122            }
123        },
124
125        // Explicitly reject write statements
126        Statement::Insert { .. } => Ok(SqlClassification::Rejected(
127            RejectionReason::WriteStatement("INSERT".to_string()),
128        )),
129        Statement::Update { .. } => Ok(SqlClassification::Rejected(
130            RejectionReason::WriteStatement("UPDATE".to_string()),
131        )),
132        Statement::Delete { .. } => Ok(SqlClassification::Rejected(
133            RejectionReason::WriteStatement("DELETE".to_string()),
134        )),
135
136        // Reject DDL
137        Statement::CreateTable { .. } => Ok(SqlClassification::Rejected(
138            RejectionReason::DdlStatement("CREATE TABLE".to_string()),
139        )),
140        Statement::CreateView { .. } => Ok(SqlClassification::Rejected(
141            RejectionReason::DdlStatement("CREATE VIEW".to_string()),
142        )),
143        Statement::CreateIndex { .. } => Ok(SqlClassification::Rejected(
144            RejectionReason::DdlStatement("CREATE INDEX".to_string()),
145        )),
146        Statement::CreateSchema { .. } => Ok(SqlClassification::Rejected(
147            RejectionReason::DdlStatement("CREATE SCHEMA".to_string()),
148        )),
149        Statement::CreateRole { .. } => Ok(SqlClassification::Rejected(
150            RejectionReason::DdlStatement("CREATE ROLE".to_string()),
151        )),
152        Statement::CreateExtension { .. } => Ok(SqlClassification::Rejected(
153            RejectionReason::DdlStatement("CREATE EXTENSION".to_string()),
154        )),
155        Statement::CreateSecret { .. } => Ok(SqlClassification::Rejected(
156            RejectionReason::DdlStatement("CREATE SECRET".to_string()),
157        )),
158        Statement::CreateVirtualTable { .. } => Ok(SqlClassification::Rejected(
159            RejectionReason::DdlStatement("CREATE VIRTUAL TABLE".to_string()),
160        )),
161        Statement::Drop { .. } => {
162            Ok(SqlClassification::Rejected(RejectionReason::DdlStatement("DROP".to_string())))
163        },
164        Statement::DropFunction { .. } => Ok(SqlClassification::Rejected(
165            RejectionReason::DdlStatement("DROP FUNCTION".to_string()),
166        )),
167        Statement::DropSecret { .. } => Ok(SqlClassification::Rejected(
168            RejectionReason::DdlStatement("DROP SECRET".to_string()),
169        )),
170        Statement::AlterTable { .. } => Ok(SqlClassification::Rejected(
171            RejectionReason::DdlStatement("ALTER TABLE".to_string()),
172        )),
173        Statement::AlterIndex { .. } => Ok(SqlClassification::Rejected(
174            RejectionReason::DdlStatement("ALTER INDEX".to_string()),
175        )),
176        Statement::AlterView { .. } => Ok(SqlClassification::Rejected(
177            RejectionReason::DdlStatement("ALTER VIEW".to_string()),
178        )),
179        Statement::AlterRole { .. } => Ok(SqlClassification::Rejected(
180            RejectionReason::DdlStatement("ALTER ROLE".to_string()),
181        )),
182        Statement::Truncate { .. } => Ok(SqlClassification::Rejected(
183            RejectionReason::DdlStatement("TRUNCATE".to_string()),
184        )),
185
186        // Reject privilege escalation. sqlparser 0.62 unified all SET-like statements
187        // (SetVariable, SetRole, SetTimeZone, SetNames, SetNamesDefault, plus session
188        // params and transaction settings) under `Statement::Set(Set)`.
189        Statement::Set(_) => Ok(SqlClassification::Rejected(RejectionReason::PrivilegeEscalation)),
190
191        // Reject procedure calls
192        Statement::Call(_) => Ok(SqlClassification::Rejected(RejectionReason::ProcedureCall)),
193
194        // Reject COPY
195        Statement::Copy { .. } | Statement::CopyIntoSnowflake { .. } => {
196            Ok(SqlClassification::Rejected(RejectionReason::CopyStatement))
197        },
198
199        // Reject ANALYZE (Hive statement that may cause writes)
200        Statement::Analyze { .. } => Ok(SqlClassification::Rejected(RejectionReason::Unknown(
201            "ANALYZE statement not allowed".to_string(),
202        ))),
203
204        // Reject other potentially dangerous statements
205        Statement::Install { .. } => Ok(SqlClassification::Rejected(RejectionReason::Unknown(
206            "INSTALL not allowed".to_string(),
207        ))),
208        Statement::Load { .. } => Ok(SqlClassification::Rejected(RejectionReason::Unknown(
209            "LOAD not allowed".to_string(),
210        ))),
211        Statement::Directory { .. } => Ok(SqlClassification::Rejected(RejectionReason::Unknown(
212            "DIRECTORY not allowed".to_string(),
213        ))),
214        Statement::AttachDatabase { .. } => Ok(SqlClassification::Rejected(
215            RejectionReason::Unknown("ATTACH DATABASE not allowed".to_string()),
216        )),
217        Statement::AttachDuckDBDatabase { .. } => Ok(SqlClassification::Rejected(
218            RejectionReason::Unknown("ATTACH DUCKDB DATABASE not allowed".to_string()),
219        )),
220        Statement::DetachDuckDBDatabase { .. } => Ok(SqlClassification::Rejected(
221            RejectionReason::Unknown("DETACH DUCKDB DATABASE not allowed".to_string()),
222        )),
223        Statement::Declare { .. } => Ok(SqlClassification::Rejected(RejectionReason::Unknown(
224            "DECLARE not allowed".to_string(),
225        ))),
226        Statement::Close { .. } => Ok(SqlClassification::Rejected(RejectionReason::Unknown(
227            "CLOSE not allowed".to_string(),
228        ))),
229        Statement::Fetch { .. } => Ok(SqlClassification::Rejected(RejectionReason::Unknown(
230            "FETCH not allowed".to_string(),
231        ))),
232        Statement::Flush { .. } => Ok(SqlClassification::Rejected(RejectionReason::Unknown(
233            "FLUSH not allowed".to_string(),
234        ))),
235        Statement::Discard { .. } => Ok(SqlClassification::Rejected(RejectionReason::Unknown(
236            "DISCARD not allowed".to_string(),
237        ))),
238        Statement::StartTransaction { .. } => Ok(SqlClassification::Rejected(
239            RejectionReason::Unknown("START TRANSACTION not allowed".to_string()),
240        )),
241        Statement::Msck { .. } => Ok(SqlClassification::Rejected(RejectionReason::Unknown(
242            "MSCK not allowed".to_string(),
243        ))),
244
245        // Reject everything else by default (whitelist-only approach)
246        _ => Ok(SqlClassification::Rejected(RejectionReason::Unknown(
247            format!("{:?}", stmt).chars().take(50).collect(),
248        ))),
249    }
250}
251
252/// Returns `true` if any part of the query tree contains a data-modifying
253/// statement (INSERT / UPDATE / DELETE / MERGE), reachable via:
254///
255/// - the query's `WITH` clause (CTE bodies), including data-modifying CTEs such as `WITH t AS
256///   (DELETE FROM x RETURNING *) SELECT * FROM t`
257/// - the query body itself and nested set operations / parenthesised subqueries
258/// - derived tables (parenthesised subqueries in `FROM`) and nested joins, which may themselves
259///   contain data-modifying CTEs
260///
261/// In sqlparser's AST a data-modifying statement embedded in a query body is
262/// represented as `SetExpr::Insert`, `SetExpr::Update`, `SetExpr::Delete`, or
263/// `SetExpr::Merge`. `DELETE … RETURNING` inside a CTE is modelled as
264/// `SetExpr::Delete`, not as a separate statement.
265fn query_is_data_modifying(query: &sqlparser::ast::Query) -> bool {
266    // Inspect every CTE body in the WITH clause.
267    if let Some(with) = &query.with {
268        if with.cte_tables.iter().any(|cte| query_is_data_modifying(&cte.query)) {
269            return true;
270        }
271    }
272
273    // Inspect the query body (and anything nested within it).
274    set_expr_is_data_modifying(&query.body)
275}
276
277/// Returns `true` if a `SetExpr` is — or transitively contains — a data-modifying
278/// statement.
279fn set_expr_is_data_modifying(set_expr: &sqlparser::ast::SetExpr) -> bool {
280    use sqlparser::ast::SetExpr;
281
282    match set_expr {
283        // A data-modifying statement appearing as a query body (e.g. inside a CTE).
284        SetExpr::Insert(_) | SetExpr::Update(_) | SetExpr::Delete(_) | SetExpr::Merge(_) => true,
285
286        // A parenthesised subquery — recurse into the nested query (which may carry
287        // its own WITH clause / data-modifying body).
288        SetExpr::Query(nested) => query_is_data_modifying(nested),
289
290        // UNION / EXCEPT / INTERSECT — both operands must be inspected.
291        SetExpr::SetOperation { left, right, .. } => {
292            set_expr_is_data_modifying(left) || set_expr_is_data_modifying(right)
293        },
294
295        // A SELECT may contain derived-table subqueries in its FROM clause that
296        // carry data-modifying CTEs of their own.
297        SetExpr::Select(select) => select.from.iter().any(table_with_joins_is_data_modifying),
298
299        // VALUES and TABLE cannot embed a data-modifying statement.
300        SetExpr::Values(_) | SetExpr::Table(_) => false,
301    }
302}
303
304/// Returns `true` if a `TableWithJoins` (a `FROM` entry and its joins) contains a
305/// derived subquery with a data-modifying statement.
306fn table_with_joins_is_data_modifying(twj: &sqlparser::ast::TableWithJoins) -> bool {
307    table_factor_is_data_modifying(&twj.relation)
308        || twj.joins.iter().any(|join| table_factor_is_data_modifying(&join.relation))
309}
310
311/// Returns `true` if a `TableFactor` is — or contains — a derived subquery with a
312/// data-modifying statement.
313fn table_factor_is_data_modifying(factor: &sqlparser::ast::TableFactor) -> bool {
314    use sqlparser::ast::TableFactor;
315
316    match factor {
317        // A parenthesised subquery in FROM: `... FROM (WITH t AS (DELETE …) …) sub`.
318        TableFactor::Derived { subquery, .. } => query_is_data_modifying(subquery),
319
320        // A parenthesised join expression — recurse into the wrapped relation.
321        TableFactor::NestedJoin {
322            table_with_joins, ..
323        } => table_with_joins_is_data_modifying(table_with_joins),
324
325        // No other table factor can embed a data-modifying query body.
326        _ => false,
327    }
328}
329
330#[cfg(test)]
331#[allow(clippy::unwrap_used)] // Reason: tests use unwrap for concise assertions
332mod tests;