1use super::builtin_function_dispatch_name;
10use crate::{
11 ast::RelationPersistence,
12 catalog::domain::StoredDomain,
13 plan::{QueryPlan, UnifiedPlan},
14};
15use crate::{expr::EngineHook, routines::RoutineResolution, SQLError};
16use std::collections::BTreeSet;
17
18pub trait QueryEffectCatalog: RoutineResolution + EngineHook {
20 fn registered_runtime_function_may_mutate_engine(&self, name: &str) -> bool;
21 fn domain_by_oid(&self, oid: u32) -> Option<StoredDomain>;
22 fn sequence_persistence(&self, name: &str) -> Result<Option<RelationPersistence>, String>;
23 fn table_persistence(&self, name: &str) -> Result<Option<RelationPersistence>, String>;
24 fn view_plan(&self, name: &str) -> Result<Option<QueryPlan>, SQLError>;
25 fn lookup_prepared(&self, name: &str) -> Option<UnifiedPlan>;
26}
27
28mod domains;
29pub struct QueryEffectContext<'a> {
31 pub catalog: &'a dyn QueryEffectCatalog,
32 pub optimizer_effects: fn(&QueryPlan) -> bool,
33 pub graph_effects: fn(&str) -> Result<bool, SQLError>,
34}
35
36pub mod read_only;
37mod routines;
38
39use routines::plpgsql_function_may_mutate_engine;
40
41#[derive(Clone, Copy)]
42struct MutabilityClassification {
43 include_session_mutations: bool,
44 include_transaction_scopes: bool,
45}
46
47impl MutabilityClassification {
48 const DATABASE_WRITES: Self = Self {
49 include_session_mutations: false,
50 include_transaction_scopes: false,
51 };
52 const ENGINE_MUTATIONS: Self = Self {
53 include_session_mutations: true,
54 include_transaction_scopes: false,
55 };
56 const STATEMENT_TRANSACTION: Self = Self {
57 include_session_mutations: true,
58 include_transaction_scopes: true,
59 };
60}
61
62pub fn query_may_mutate_engine(
64 context: &QueryEffectContext<'_>,
65 query: &crate::plan::QueryPlan,
66) -> Result<bool, SQLError> {
67 query_may_mutate_engine_inner(
68 context,
69 query,
70 &mut BTreeSet::new(),
71 &mut BTreeSet::new(),
72 MutabilityClassification::ENGINE_MUTATIONS,
73 )
74}
75
76pub fn query_requires_statement_transaction(
78 context: &QueryEffectContext<'_>,
79 query: &crate::plan::QueryPlan,
80) -> Result<bool, SQLError> {
81 query_may_mutate_engine_inner(
82 context,
83 query,
84 &mut BTreeSet::new(),
85 &mut BTreeSet::new(),
86 MutabilityClassification::STATEMENT_TRANSACTION,
87 )
88}
89
90pub fn query_may_write_database(
92 context: &QueryEffectContext<'_>,
93 query: &crate::plan::QueryPlan,
94) -> Result<bool, SQLError> {
95 query_may_mutate_engine_inner(
96 context,
97 query,
98 &mut BTreeSet::new(),
99 &mut BTreeSet::new(),
100 MutabilityClassification::DATABASE_WRITES,
101 )
102}
103
104pub fn command_payload_may_write_database(
106 context: &QueryEffectContext<'_>,
107 command: &crate::plan::CommandPlan,
108) -> Result<bool, SQLError> {
109 let mut plan = crate::plan::UnifiedPlan::Command(Box::new(command.clone()));
110 let mut writes = false;
111 let mut classification_error = None;
112 plan.rewrite_scalar_expressions(&mut |expression| {
113 if classification_error.is_some() {
114 return;
115 }
116 match scalar_node_may_mutate_engine(
117 context,
118 expression,
119 &mut BTreeSet::new(),
120 &mut BTreeSet::new(),
121 MutabilityClassification::DATABASE_WRITES,
122 ) {
123 Ok(value) => writes |= value,
124 Err(error) => classification_error = Some(error),
125 }
126 });
127 if let Some(error) = classification_error {
128 return Err(error);
129 }
130 if writes {
131 return Ok(true);
132 }
133 match command {
134 crate::plan::CommandPlan::Insert(plan) => {
135 if ctes_may_write_database(context, &plan.ctes)? {
136 return Ok(true);
137 }
138 if let Some(source) = plan.source.as_deref() {
139 if query_may_write_database(context, source)? {
140 return Ok(true);
141 }
142 }
143 queries_may_write_database(context, &plan.subqueries)
144 }
145 crate::plan::CommandPlan::Update(plan) => {
146 if ctes_may_write_database(context, &plan.ctes)? {
147 return Ok(true);
148 }
149 if let Some(source) = plan.source.as_deref() {
150 if source_may_mutate_engine(
151 context,
152 source,
153 &mut BTreeSet::new(),
154 &mut BTreeSet::new(),
155 MutabilityClassification::DATABASE_WRITES,
156 )? {
157 return Ok(true);
158 }
159 }
160 queries_may_write_database(context, &plan.subqueries)
161 }
162 crate::plan::CommandPlan::Delete(plan) => {
163 if ctes_may_write_database(context, &plan.ctes)? {
164 return Ok(true);
165 }
166 if let Some(source) = plan.source.as_deref() {
167 if source_may_mutate_engine(
168 context,
169 source,
170 &mut BTreeSet::new(),
171 &mut BTreeSet::new(),
172 MutabilityClassification::DATABASE_WRITES,
173 )? {
174 return Ok(true);
175 }
176 }
177 queries_may_write_database(context, &plan.subqueries)
178 }
179 crate::plan::CommandPlan::Merge(plan) => {
180 if ctes_may_write_database(context, &plan.ctes)? {
181 return Ok(true);
182 }
183 if source_may_mutate_engine(
184 context,
185 &plan.source,
186 &mut BTreeSet::new(),
187 &mut BTreeSet::new(),
188 MutabilityClassification::DATABASE_WRITES,
189 )? {
190 return Ok(true);
191 }
192 queries_may_write_database(context, &plan.subqueries)
193 }
194 _ => Ok(false),
195 }
196}
197
198fn command_may_write_database(
199 context: &QueryEffectContext<'_>,
200 command: &crate::plan::CommandPlan,
201) -> Result<bool, SQLError> {
202 read_only::forbidden_command(
203 context,
204 &crate::plan::UnifiedPlan::Command(Box::new(command.clone())),
205 )
206 .map(|forbidden| forbidden.is_some())
207}
208
209fn ctes_may_write_database(
210 context: &QueryEffectContext<'_>,
211 ctes: &[crate::plan::CtePlan],
212) -> Result<bool, SQLError> {
213 for cte in ctes {
214 let writes = match &cte.body {
215 crate::plan::CtePlanBody::Query(query) => query_may_write_database(context, query)?,
216 crate::plan::CtePlanBody::Command(command) => {
217 command_may_write_database(context, command)?
218 }
219 };
220 if writes {
221 return Ok(true);
222 }
223 }
224 Ok(false)
225}
226
227fn queries_may_write_database<'a>(
228 context: &QueryEffectContext<'_>,
229 queries: impl IntoIterator<Item = &'a crate::plan::QueryPlan>,
230) -> Result<bool, SQLError> {
231 for query in queries {
232 if query_may_write_database(context, query)? {
233 return Ok(true);
234 }
235 }
236 Ok(false)
237}
238
239fn query_may_mutate_engine_inner(
240 context: &QueryEffectContext<'_>,
241 query: &crate::plan::QueryPlan,
242 visiting_views: &mut BTreeSet<String>,
243 visiting_routines: &mut BTreeSet<String>,
244 classification: MutabilityClassification,
245) -> Result<bool, SQLError> {
246 if query_source_may_mutate_engine(
247 context,
248 query,
249 visiting_views,
250 visiting_routines,
251 classification,
252 )? {
253 return Ok(true);
254 }
255 let mut plan = crate::plan::UnifiedPlan::Query(Box::new(query.clone()));
256 let mut mutates = (context.optimizer_effects)(query);
257 let mut classification_error = None;
258 plan.rewrite_scalar_expressions(&mut |expression| {
259 if classification_error.is_some() {
260 return;
261 }
262 match scalar_node_may_mutate_engine(
263 context,
264 expression,
265 visiting_views,
266 visiting_routines,
267 classification,
268 ) {
269 Ok(value) => mutates |= value,
270 Err(error) => classification_error = Some(error),
271 }
272 });
273 classification_error.map_or(Ok(mutates), Err)
274}
275
276fn scalar_node_may_mutate_engine(
277 context: &QueryEffectContext<'_>,
278 expression: &crate::ScalarExpr,
279 visiting_views: &mut BTreeSet<String>,
280 visiting_routines: &mut BTreeSet<String>,
281 classification: MutabilityClassification,
282) -> Result<bool, SQLError> {
283 match expression {
284 crate::ScalarExpr::Func {
285 name,
286 binding,
287 args,
288 ..
289 } => function_may_mutate_engine(
290 context,
291 name,
292 binding.as_ref(),
293 args,
294 visiting_views,
295 visiting_routines,
296 classification,
297 ),
298 crate::ScalarExpr::Cast { ty, .. } => {
299 crate::expr::EngineHook::resolve_type_name(context.catalog, ty)
300 .ok()
301 .flatten()
302 .map_or(Ok(false), |ty| {
303 domains::domain_cast_may_mutate(
304 context,
305 &ty,
306 visiting_views,
307 visiting_routines,
308 classification,
309 )
310 })
311 }
312 _ => Ok(false),
313 }
314}
315
316fn function_may_mutate_engine(
317 context: &QueryEffectContext<'_>,
318 name: &str,
319 binding: Option<&crate::ast::FunctionBinding>,
320 args: &[crate::ScalarExpr],
321 visiting_views: &mut BTreeSet<String>,
322 visiting_routines: &mut BTreeSet<String>,
323 classification: MutabilityClassification,
324) -> Result<bool, SQLError> {
325 let identity = name.to_ascii_lowercase();
326 let dispatch_name = builtin_function_dispatch_name(&identity);
327 let cypher_mutates = if dispatch_name == "cypher" {
328 match args.get(1) {
329 Some(crate::ScalarExpr::Literal(uqa_core::Value::Str(query))) => {
330 (context.graph_effects)(query)?
331 }
332 _ => classification.include_session_mutations,
333 }
334 } else {
335 false
336 };
337 let mutates_database_directly = cypher_mutates
338 || matches!(
339 dispatch_name.as_str(),
340 "create_analyzer"
341 | "drop_analyzer"
342 | "set_table_analyzer"
343 | "graph_create"
344 | "graph_drop"
345 | "create_graph"
346 | "drop_graph"
347 | "create_vlabel"
348 | "create_elabel"
349 | "drop_label"
350 | "alter_graph"
351 | "deep_learn"
352 | "bayesian_match"
353 | "bayesian_match_with_prior"
354 )
355 || context
356 .catalog
357 .registered_runtime_function_may_mutate_engine(&identity);
358 let temporary_sequence_call = matches!(dispatch_name.as_str(), "nextval" | "setval")
359 && sequence_function_targets_temporary(context, args)?;
360 let builtin_requires_mutating_execution =
361 matches!(
362 dispatch_name.as_str(),
363 "random" | "setseed" | "pg_notify" | "fts_match" | "multi_field_match"
364 ) || (matches!(dispatch_name.as_str(), "nextval" | "setval") && !temporary_sequence_call);
365 let sql_routine_mutates = classification.include_session_mutations
366 && sql_routine_may_mutate_engine(
367 context,
368 &identity,
369 binding,
370 visiting_views,
371 visiting_routines,
372 classification,
373 )?;
374 let sequence_lock_scope = classification.include_transaction_scopes
375 && matches!(
376 dispatch_name.as_str(),
377 "nextval" | "currval" | "lastval" | "setval"
378 );
379 Ok(mutates_database_directly
380 || sequence_lock_scope
381 || (classification.include_session_mutations
382 && (builtin_requires_mutating_execution || sql_routine_mutates)))
383}
384
385fn sequence_function_targets_temporary(
386 context: &QueryEffectContext<'_>,
387 args: &[crate::ScalarExpr],
388) -> Result<bool, SQLError> {
389 fn literal_reference(expression: &crate::ScalarExpr) -> Option<&str> {
390 match expression {
391 crate::ScalarExpr::Literal(uqa_core::Value::Str(reference)) => Some(reference),
392 crate::ScalarExpr::Cast { expr, .. } => literal_reference(expr),
393 _ => None,
394 }
395 }
396
397 let Some(reference) = args.first().and_then(literal_reference) else {
398 return Ok(false);
399 };
400 Ok(context
401 .catalog
402 .sequence_persistence(reference)
403 .map_err(|error| {
404 SQLError::Internal(format!(
405 "resolve sequence `{reference}` while classifying query mutability: {error}"
406 ))
407 })?
408 .is_some_and(|persistence| persistence == crate::ast::RelationPersistence::Temporary))
409}
410
411fn sql_routine_may_mutate_engine(
412 context: &QueryEffectContext<'_>,
413 name: &str,
414 binding: Option<&crate::ast::FunctionBinding>,
415 visiting_views: &mut BTreeSet<String>,
416 visiting_routines: &mut BTreeSet<String>,
417 classification: MutabilityClassification,
418) -> Result<bool, SQLError> {
419 if binding.is_some_and(|binding| binding.builtin) {
420 return Ok(false);
421 }
422 let overloads = match binding {
423 Some(binding) => context
424 .catalog
425 .lookup_bound_sql_functions_by_binding(binding),
426 None => context
427 .catalog
428 .lookup_visible_sql_functions_for_analysis(name)?,
429 };
430 let Some(overloads) = overloads else {
431 return Ok(false);
432 };
433 for function in overloads {
434 if function.def.is_procedure
435 || binding.is_some_and(|binding| {
436 crate::routines::routine_signature_types(&function.def) != binding.argument_types
437 })
438 {
439 continue;
440 }
441 let signature = crate::routines::routine_signature_types(&function.def);
442 let key = format!("{}({})", function.def.name, signature.join(","));
443 if !visiting_routines.insert(key.clone()) {
444 continue;
445 }
446 let result = (|| {
447 if domains::routine_coercions_may_mutate(
448 context,
449 &function.def,
450 visiting_views,
451 visiting_routines,
452 classification,
453 )? {
454 return Ok(true);
455 }
456 let Some(body) = crate::routines::analyzable_routine_body(context.catalog, &function)?
457 else {
458 return Ok(function.def.language == "plpgsql");
461 };
462 match &*body {
463 crate::routines::CompiledFunctionBody::SQL(plans) => (|| {
464 let mut mutates = false;
465 for plan in plans {
466 match plan {
467 crate::plan::UnifiedPlan::Query(query) => {
468 if query_may_mutate_engine_inner(
469 context,
470 query,
471 visiting_views,
472 visiting_routines,
473 classification,
474 )? {
475 mutates = true;
476 break;
477 }
478 }
479 crate::plan::UnifiedPlan::Command(_) => {
480 mutates = true;
481 break;
482 }
483 }
484 }
485 Ok(mutates)
486 })(),
487 crate::routines::CompiledFunctionBody::PLpgSQL(function) => {
488 plpgsql_function_may_mutate_engine(
489 context,
490 function,
491 visiting_views,
492 visiting_routines,
493 classification,
494 )
495 }
496 }
497 })();
498 visiting_routines.remove(&key);
499 if result? {
500 return Ok(true);
501 }
502 }
503 Ok(false)
504}
505
506fn query_source_may_mutate_engine(
507 context: &QueryEffectContext<'_>,
508 query: &crate::plan::QueryPlan,
509 visiting_views: &mut BTreeSet<String>,
510 visiting_routines: &mut BTreeSet<String>,
511 classification: MutabilityClassification,
512) -> Result<bool, SQLError> {
513 for cte in &query.ctes {
514 let mutates = match &cte.body {
515 crate::plan::CtePlanBody::Query(query) => query_source_may_mutate_engine(
516 context,
517 query,
518 visiting_views,
519 visiting_routines,
520 classification,
521 )?,
522 crate::plan::CtePlanBody::Command(command) => {
523 classification.include_session_mutations
524 || classification.include_transaction_scopes
525 || command_may_write_database(context, command)?
526 }
527 };
528 if mutates {
529 return Ok(true);
530 }
531 }
532 match &query.root {
533 crate::plan::RelationalPlan::QueryBlock(block) => {
534 if let Some(source) = block.from.as_ref() {
535 if source_may_mutate_engine(
536 context,
537 source,
538 visiting_views,
539 visiting_routines,
540 classification,
541 )? {
542 return Ok(true);
543 }
544 }
545 for subquery in &block.subqueries {
546 if query_source_may_mutate_engine(
547 context,
548 subquery,
549 visiting_views,
550 visiting_routines,
551 classification,
552 )? {
553 return Ok(true);
554 }
555 }
556 }
557 crate::plan::RelationalPlan::SetOp {
558 left,
559 right,
560 subqueries,
561 ..
562 } => {
563 if query_source_may_mutate_engine(
564 context,
565 left,
566 visiting_views,
567 visiting_routines,
568 classification,
569 )? || query_source_may_mutate_engine(
570 context,
571 right,
572 visiting_views,
573 visiting_routines,
574 classification,
575 )? {
576 return Ok(true);
577 }
578 for subquery in subqueries {
579 if query_source_may_mutate_engine(
580 context,
581 subquery,
582 visiting_views,
583 visiting_routines,
584 classification,
585 )? {
586 return Ok(true);
587 }
588 }
589 }
590 crate::plan::RelationalPlan::Values { subqueries, .. } => {
591 for subquery in subqueries {
592 if query_source_may_mutate_engine(
593 context,
594 subquery,
595 visiting_views,
596 visiting_routines,
597 classification,
598 )? {
599 return Ok(true);
600 }
601 }
602 }
603 }
604 Ok(false)
605}
606
607fn source_may_mutate_engine(
608 context: &QueryEffectContext<'_>,
609 source: &crate::plan::SourcePlan,
610 visiting_views: &mut BTreeSet<String>,
611 visiting_routines: &mut BTreeSet<String>,
612 classification: MutabilityClassification,
613) -> Result<bool, SQLError> {
614 match source {
615 crate::plan::SourcePlan::Function {
616 name,
617 binding,
618 args,
619 ..
620 } => function_may_mutate_engine(
621 context,
622 name,
623 binding.as_ref(),
624 args,
625 visiting_views,
626 visiting_routines,
627 classification,
628 ),
629 crate::plan::SourcePlan::FunctionGroup { functions, .. } => {
630 for function in functions {
631 if function_may_mutate_engine(
632 context,
633 &function.name,
634 function.binding.as_ref(),
635 &function.args,
636 visiting_views,
637 visiting_routines,
638 classification,
639 )? {
640 return Ok(true);
641 }
642 }
643 Ok(false)
644 }
645 crate::plan::SourcePlan::Join { left, right, .. } => Ok(source_may_mutate_engine(
646 context,
647 left,
648 visiting_views,
649 visiting_routines,
650 classification,
651 )? || source_may_mutate_engine(
652 context,
653 right,
654 visiting_views,
655 visiting_routines,
656 classification,
657 )?),
658 crate::plan::SourcePlan::Subquery { body, .. } => query_source_may_mutate_engine(
659 context,
660 body,
661 visiting_views,
662 visiting_routines,
663 classification,
664 ),
665 crate::plan::SourcePlan::Table { name, .. } => {
666 let key = name.to_ascii_lowercase();
667 if !visiting_views.insert(key.clone()) {
668 return Ok(false);
669 }
670 let result = match context.catalog.view_plan(name) {
671 Ok(Some(view)) => query_may_mutate_engine_inner(
672 context,
673 &view,
674 visiting_views,
675 visiting_routines,
676 classification,
677 ),
678 Ok(None) => Ok(false),
679 Err(error) => Err(error),
680 };
681 visiting_views.remove(&key);
682 result
683 }
684 crate::plan::SourcePlan::Values { .. } => Ok(false),
685 }
686}
687
688pub fn is_transaction_control(plan: &crate::plan::UnifiedPlan) -> bool {
689 matches!(
690 plan,
691 crate::plan::UnifiedPlan::Command(command)
692 if matches!(command.as_ref(), crate::plan::CommandPlan::Transaction(_))
693 )
694}
695
696pub mod transaction_blocks;