fraiseql_functions/host/live/sql_classifier/
mod.rs1use fraiseql_error::{FraiseQLError, Result};
9
10#[derive(Debug, Clone, PartialEq, Eq)]
12pub enum SqlClassification {
13 ReadOnly,
15 Rejected(RejectionReason),
17}
18
19#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum RejectionReason {
22 WriteStatement(String),
24 DdlStatement(String),
26 WritableCte,
28 PrivilegeEscalation,
30 ProceduralBlock,
32 ProcedureCall,
34 CopyStatement,
36 ExplainAnalyze,
38 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
60pub 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 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
99fn classify_statement(stmt: &sqlparser::ast::Statement) -> Result<SqlClassification> {
101 use sqlparser::ast::Statement;
102
103 match stmt {
104 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 Statement::Explain { analyze, .. } => {
118 if *analyze {
119 Ok(SqlClassification::Rejected(RejectionReason::ExplainAnalyze))
120 } else {
121 Ok(SqlClassification::ReadOnly)
122 }
123 },
124
125 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 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 Statement::Set(_) => Ok(SqlClassification::Rejected(RejectionReason::PrivilegeEscalation)),
190
191 Statement::Call(_) => Ok(SqlClassification::Rejected(RejectionReason::ProcedureCall)),
193
194 Statement::Copy { .. } | Statement::CopyIntoSnowflake { .. } => {
196 Ok(SqlClassification::Rejected(RejectionReason::CopyStatement))
197 },
198
199 Statement::Analyze { .. } => Ok(SqlClassification::Rejected(RejectionReason::Unknown(
201 "ANALYZE statement not allowed".to_string(),
202 ))),
203
204 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 _ => Ok(SqlClassification::Rejected(RejectionReason::Unknown(
247 format!("{:?}", stmt).chars().take(50).collect(),
248 ))),
249 }
250}
251
252fn query_is_data_modifying(query: &sqlparser::ast::Query) -> bool {
266 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 set_expr_is_data_modifying(&query.body)
275}
276
277fn set_expr_is_data_modifying(set_expr: &sqlparser::ast::SetExpr) -> bool {
280 use sqlparser::ast::SetExpr;
281
282 match set_expr {
283 SetExpr::Insert(_) | SetExpr::Update(_) | SetExpr::Delete(_) | SetExpr::Merge(_) => true,
285
286 SetExpr::Query(nested) => query_is_data_modifying(nested),
289
290 SetExpr::SetOperation { left, right, .. } => {
292 set_expr_is_data_modifying(left) || set_expr_is_data_modifying(right)
293 },
294
295 SetExpr::Select(select) => select.from.iter().any(table_with_joins_is_data_modifying),
298
299 SetExpr::Values(_) | SetExpr::Table(_) => false,
301 }
302}
303
304fn 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
311fn table_factor_is_data_modifying(factor: &sqlparser::ast::TableFactor) -> bool {
314 use sqlparser::ast::TableFactor;
315
316 match factor {
317 TableFactor::Derived { subquery, .. } => query_is_data_modifying(subquery),
319
320 TableFactor::NestedJoin {
322 table_with_joins, ..
323 } => table_with_joins_is_data_modifying(table_with_joins),
324
325 _ => false,
327 }
328}
329
330#[cfg(test)]
331#[allow(clippy::unwrap_used)] mod tests;