pylon_core/query/mod.rs
1//
2// This source file is part of the Pylon open source project.
3//
4// Copyright (c) 2026 Jaldis B.V.
5//
6// Licensed under the MIT OR Apache-2.0 license (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10// https://opensource.org/licenses/MIT
11// https://www.apache.org/licenses/LICENSE-2.0
12//
13// Unless required by applicable law or agreed to in writing, software
14// distributed under the License is distributed on an "AS IS" BASIS,
15// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16// See the License for the specific language governing permissions and
17// limitations under the License.
18//
19
20use lru::LruCache;
21use std::collections::hash_map::DefaultHasher;
22use std::hash::{Hash, Hasher};
23use std::num::NonZeroUsize;
24use std::sync::{Arc, OnceLock, RwLock};
25
26use crate::analyze::ShapePathAlias;
27use crate::error::PyQLError;
28use crate::schema::SchemaDescriptor;
29use crate::{analyze, ir, parse, sql};
30
31const CACHE_CAPACITY: usize = 1024;
32
33/// Number of independently-locked cache shards.
34///
35/// `LruCache::get` has to reorder the recency list, so it needs a *write*
36/// lock even on a hit — meaning a single map serializes every compile in the
37/// process behind one lock. Sharding by key hash keeps LRU semantics while
38/// cutting that contention by this factor. Must be a power of two, since
39/// shard selection masks the low bits of the hash.
40const CACHE_SHARDS: usize = 16;
41
42/// One cached compilation, stored behind an `Arc` so a hit hands out a
43/// pointer rather than deep-cloning the SQL string, parameter names, tag
44/// list and entire shape tree (measured at 0.94 µs of a 1.05 µs cache hit).
45///
46/// The key is kept alongside the value and re-checked on every hit: lookup
47/// is by hash alone (so it costs no allocation), and a hash collision must
48/// read as a miss, never as "here is some other query's SQL."
49struct CacheEntry {
50 query: String,
51 config: ir::SessionConfig,
52 compiled: Arc<CompiledQuery>,
53}
54
55// Keyed by (query text, session config) — not query text alone. Compilation
56// success/shape/SQL can depend on the config (e.g. an INSERT assigning `id`
57// compiles under allow_user_specified_id=true and errors otherwise), so two
58// requests for the same query text under different configs must never share
59// a cache entry.
60type CacheShard = RwLock<LruCache<u64, CacheEntry>>;
61
62static QUERY_CACHE: OnceLock<Vec<CacheShard>> = OnceLock::new();
63
64fn query_cache() -> &'static [CacheShard] {
65 QUERY_CACHE.get_or_init(|| {
66 let per_shard = NonZeroUsize::new(CACHE_CAPACITY / CACHE_SHARDS).unwrap();
67 (0..CACHE_SHARDS)
68 .map(|_| RwLock::new(LruCache::new(per_shard)))
69 .collect()
70 })
71}
72
73fn cache_key_hash(query: &str, config: &ir::SessionConfig) -> u64 {
74 let mut hasher = DefaultHasher::new();
75 query.hash(&mut hasher);
76 config.hash(&mut hasher);
77 hasher.finish()
78}
79
80/// Discard all cached compiled queries. Call when the schema is reloaded.
81pub fn clear_query_cache() {
82 if let Some(shards) = QUERY_CACHE.get() {
83 for shard in shards {
84 shard.write().unwrap().clear();
85 }
86 }
87}
88
89#[derive(Debug, Clone, PartialEq, Eq)]
90pub enum Cardinality {
91 Required,
92 Optional,
93 Many,
94}
95
96/// Internal tree describing one position in the query output shape.
97/// Opaque to Python — only the Rust deserializer inspects it.
98#[derive(Debug, Clone)]
99pub enum ShapeNode {
100 /// Leaf value; native PG type.
101 Scalar { name: String, position: usize },
102 /// The `result` column IS the value — not wrapped in ROW(). Used for array literals
103 /// where an array is returned as its own top-level column.
104 RawScalar,
105 /// Like RawScalar but the value is a decoded JSON object (from a <json> cast).
106 /// REPL displays it as `Json("...")`.
107 JsonScalar,
108 /// Object shape.
109 /// `type_name = Some(s)` → named schema type decoded to a registered dataclass.
110 /// `type_name = None` → free type decoded to Pylon's generic Object dataclass.
111 /// `name` is the pointer name within the parent (empty string for the root).
112 Object {
113 name: String,
114 type_name: Option<String>,
115 position: usize,
116 cardinality: Cardinality,
117 pointers: Vec<ShapeNode>,
118 /// True when `pointers[0]` is an `id` the compiler added to a shape
119 /// that did not select one — see `IrScalarPointer::implicit_id`. Every
120 /// decoder hands it to the caller like any other property; only JSON
121 /// output drops it.
122 has_implicit_id: bool,
123 },
124 /// `record[]` column decoded to a Python list.
125 Array {
126 name: String,
127 position: usize,
128 element: Box<ShapeNode>,
129 },
130 /// Anonymous positional tuple decoded to a Python tuple. No type name — no registry lookup.
131 /// A composite tuple row, read by indexing rather than out of jsonb.
132 /// `names` is set for a named tuple that had to be emitted as a composite
133 /// because an element holds an object -- jsonb has no member kind for one
134 /// (see `JsonMemberKind`) -- and decides whether this hydrates to a plain
135 /// tuple or a `NamedTupleValue`.
136 Tuple {
137 position: usize,
138 elements: Vec<ShapeNode>,
139 names: Option<Vec<String>>,
140 },
141 /// Named tuple decoded from jsonb. When `type_name` is Some, hydrated to the registered class.
142 /// `members` carries the full per-member decode plan when statically known (a registered
143 /// NamedTupleDescriptor's members, a structural `pylon.Tuple[...]` property's tuple_members,
144 /// or a `<tuple<...>>`/nominal cast's own target type) — `None` falls back to decoding the
145 /// raw jsonb value generically (dict/list, no per-member typing).
146 NamedTuple {
147 name: String,
148 position: usize,
149 type_name: Option<String>,
150 members: Option<Vec<JsonMember>>,
151 /// True when this value came from `{ x := 1.0 }` (curly-brace free-
152 /// object syntax) rather than `(x := 1.0)` (paren tuple syntax) —
153 /// decoded identically either way (both are jsonb), but the
154 /// frontend's own value-shape-tag tree (pylon/query.py's
155 /// shape_value_tags) uses this to tell the Python API layer to
156 /// describe it as an "object" rather than a "namedTuple", so
157 /// JsonTree renders an expandable `Object {x: 1.0}` instead of a
158 /// non-expandable `(x := 1.0)` tuple literal.
159 is_free_object: bool,
160 },
161 /// Enum value arrived as text; hydrated to the Python enum class keyed by `enum_type`.
162 Enum {
163 name: String,
164 position: usize,
165 /// Pylon-qualified name, e.g. `default::Gender`.
166 enum_type: String,
167 },
168 /// Result of a `vector::search` statement.
169 /// The outer `result` tuple has three slots:
170 /// 0 → NULL (virtual type, no registry lookup)
171 /// `object_position` → the object sub-tuple (decoded as a Pylon object)
172 /// `distance_position` → the distance scalar (float64)
173 VectorSearch {
174 object_position: usize,
175 distance_position: usize,
176 object_node: Box<ShapeNode>,
177 },
178 /// Result of a `fts::search` statement.
179 /// Outer tuple layout mirrors `VectorSearch`: pos 0 = NULL, pos 1 = object, pos 2 = score.
180 FtsSearch {
181 object_position: usize,
182 rank_position: usize,
183 object_node: Box<ShapeNode>,
184 },
185 /// Result of a `group` statement: each row is a free object with key/grouping/elements.
186 Group {
187 /// One ShapeNode per grouping key (carries name, position, and type).
188 /// Positions are 1-based in the outer tuple (pos 0 is the NULL type slot).
189 key_nodes: Vec<ShapeNode>,
190 /// Position of the `ARRAY[key_names...]::text[]` in the outer tuple.
191 grouping_position: usize,
192 /// Position of the `array_agg(elements)` in the outer tuple.
193 elements_position: usize,
194 /// Shape node for each element in the elements array.
195 element: Box<ShapeNode>,
196 },
197}
198
199/// One member's decode plan within a jsonb-backed tuple value
200/// (`ShapeNode::NamedTuple.members`) — recursive so a member can itself be a
201/// nested tuple.
202#[derive(Debug, Clone)]
203pub struct JsonMember {
204 /// `None` for a positional/unnamed element of a structural tuple
205 /// (the value is a jsonb array); `Some` for a named member (a jsonb
206 /// object key) — either a nominal named-tuple field or a named
207 /// structural element.
208 pub key: Option<String>,
209 pub kind: JsonMemberKind,
210}
211
212#[derive(Debug, Clone)]
213pub enum JsonMemberKind {
214 /// Plain scalar — the jsonb value's own native JSON type is already
215 /// correct (number/string/bool), used as-is.
216 Scalar,
217 /// Value arrived as a jsonb string; hydrate to the Python enum class
218 /// keyed by `enum_type` (Pylon-qualified name, e.g. `default::Gender`).
219 Enum { enum_type: String },
220 /// Nested tuple member — recurse. `type_name` hydrates to a registered
221 /// dataclass when present (nominal); `None` decodes to a plain tuple
222 /// (all-positional members) or a dynamically-built dataclass (named,
223 /// unregistered structural).
224 Tuple {
225 type_name: Option<String>,
226 members: Vec<JsonMember>,
227 },
228}
229
230/// Opaque handle to the output shape of a compiled query.
231#[derive(Debug, Clone)]
232pub struct ShapeDescriptor {
233 pub root: ShapeNode,
234}
235
236/// Typed query parameter value produced during compilation.
237#[derive(Debug, Clone)]
238pub enum QueryParam {
239 Null,
240 Bool(bool),
241 Int(i64),
242 Float(f64),
243 Text(String),
244 Bytes(Vec<u8>),
245 Uuid([u8; 16]),
246}
247
248/// Pre-execution inference plan — set when the query requires an external model call
249/// before the SQL can be executed. Python switches on the variant.
250#[derive(Debug, Clone)]
251pub enum InferencePlan {
252 /// `fts::search` with a remote backend (OpenSearch or Meilisearch).
253 /// Python fetches (id, score) pairs from the backend, then injects them as
254 /// `__deferred_ids__` / `__deferred_scores__` params and runs `sql` against Postgres.
255 Search {
256 /// Remote backend identifier: `"opensearch"` or `"meilisearch"`.
257 backend: String,
258 /// Remote index name.
259 index_name: String,
260 /// Name of the user's query-text param; empty string when an inline literal.
261 query_param_name: String,
262 /// Inline literal query text.
263 query_literal: Option<String>,
264 /// Requested result size (limit), if known at compile time.
265 size: Option<usize>,
266 },
267 /// `vector::search(TypeName, query := $text)` text overload.
268 /// Python embeds the text via the configured model provider, then injects the
269 /// resulting vector as `__deferred_vec__` and runs `sql` against Postgres.
270 Embedding {
271 /// Embedding model identifier from the schema, e.g. `"mistral-embed"`.
272 model_name: String,
273 /// Qualified type name for provider lookup, e.g. `"default::Product"`.
274 type_name: String,
275 /// Vector index name for provider lookup (`None` = default index).
276 index_name: Option<String>,
277 /// Name of the user's `query :=` param; empty string when an inline literal.
278 query_param_name: String,
279 /// Inline literal query text.
280 query_literal: Option<String>,
281 },
282}
283
284/// The output of a successful PyQL compilation.
285/// Immutable and safe to cache and reuse across requests.
286#[derive(Debug, Clone)]
287pub struct CompiledQuery {
288 /// PostgreSQL SQL string ready for execution.
289 pub sql: String,
290 /// Ordered parameter names matching $1, $2, … in the SQL.
291 /// The client uses this to map kwargs to positional arguments.
292 pub param_names: Vec<String>,
293 /// Typed bound parameters — populated at execution time, empty after compilation.
294 pub params: Vec<QueryParam>,
295 /// Opaque shape handle — consumed by the Rust deserializer.
296 pub shape: ShapeDescriptor,
297 /// Non-fatal warnings produced during compilation.
298 pub warnings: Vec<String>,
299 /// Set when the query requires a pre-execution model call.
300 pub inference_plan: Option<InferencePlan>,
301 /// Every schema-qualified table (`"schema.table"`) this statement reads
302 /// from or writes to — see `ir::tags::collect_tags`. For a SELECT, the
303 /// set of tags to cache this result under; for an INSERT/UPDATE/DELETE,
304 /// the set of tags a cache layer must invalidate after the write commits.
305 pub tags: Vec<String>,
306 /// True when executing this statement writes to any of `tags`. Lets a
307 /// cache layer act on the invalidation duty described above without
308 /// re-parsing the SQL to guess whether a write happened.
309 pub mutates: bool,
310 /// `Some` only for `analyze <query>` — the shape-path↔SQL-alias map an
311 /// `analyze` execution needs to correlate Postgres's `EXPLAIN` plan
312 /// nodes back to the query's own shape (see `analyze` module). `None`
313 /// for every other query, which doesn't pay for this extra shape walk.
314 pub analyze_paths: Option<Vec<ShapePathAlias>>,
315 /// This query's stable shape id — see `crate::shape_id` and `shape_id()`.
316 /// Computed once at compile time; fill it with `derive_shape_id` if you
317 /// ever build a `CompiledQuery` outside `compile_uncached`.
318 pub shape_id: Arc<str>,
319}
320
321/// The shape id for a given SQL string and result shape — the derivation
322/// `compile_uncached` uses to populate `CompiledQuery::shape_id`, exposed so
323/// anything else constructing a `CompiledQuery` by hand can fill that field
324/// consistently rather than inventing its own value.
325pub fn derive_shape_id(sql: &str, shape: &ShapeDescriptor) -> Arc<str> {
326 crate::shape_id::query_shape_id(sql, &format!("{shape:?}")).into()
327}
328
329impl CompiledQuery {
330 /// This query's stable shape id — see `crate::shape_id`.
331 ///
332 /// Independent of the values bound to it, which is what makes it safe as
333 /// a metric label: a query run a million times with different parameters
334 /// reports one label value, not a million.
335 ///
336 /// Precomputed rather than derived per call. Every execution reports this
337 /// label, and deriving it meant `Debug`-formatting the entire shape tree
338 /// into a throwaway `String` and hashing it — measured at 3.65 µs per
339 /// execution on a 20-field shape, all of it repeated work, since a
340 /// `CompiledQuery` is immutable.
341 pub fn shape_id(&self) -> Arc<str> {
342 self.shape_id.clone()
343 }
344
345 /// A stable hash of this query's SQL, for callers that need to key a
346 /// cache entry on the statement without moving the SQL text itself
347 /// around. Same derivation as `shape_id`, minus the shape.
348 pub fn sql_id(&self) -> &str {
349 // The shape id already covers the SQL — a query's shape can't change
350 // without its SQL changing — so a second hash would be redundant.
351 &self.shape_id
352 }
353}
354
355/// A coarse bucket for a failed query, for use as a metric label.
356///
357/// Deliberately coarse: an unbounded label (a raw error message, a SQLSTATE,
358/// a constraint name) is exactly the thing that blows up a metrics backend's
359/// cardinality. Anything finer belongs on a span or in a log line.
360#[derive(Debug, Clone, Copy, PartialEq, Eq)]
361pub enum ErrorClass {
362 /// A constraint the write violated — unique, foreign key, check, not-null.
363 ConstraintViolation,
364 /// Serialization failure or deadlock: the caller can retry.
365 Contention,
366 /// Statement or lock timeout.
367 Timeout,
368 /// Couldn't reach or stay connected to the database.
369 Connection,
370 /// The query never reached the database — bad PyQL, unknown type/pointer.
371 Compile,
372 Other,
373}
374
375impl ErrorClass {
376 /// The label value. `&'static str` so it can't accidentally become
377 /// unbounded.
378 pub fn as_label(self) -> &'static str {
379 match self {
380 ErrorClass::ConstraintViolation => "constraint_violation",
381 ErrorClass::Contention => "contention",
382 ErrorClass::Timeout => "timeout",
383 ErrorClass::Connection => "connection",
384 ErrorClass::Compile => "compile",
385 ErrorClass::Other => "other",
386 }
387 }
388
389 /// Buckets a SQLSTATE by its two-character class, which is how the
390 /// standard already groups them — so a code this was never written
391 /// against still lands somewhere sensible instead of in `Other`.
392 pub fn from_sqlstate(code: &str) -> Self {
393 match code {
394 "40001" => ErrorClass::Contention,
395 "40P01" => ErrorClass::Contention,
396 "57014" => ErrorClass::Timeout,
397 "55P03" => ErrorClass::Timeout,
398 _ => match code.get(..2) {
399 // 23 = integrity constraint violation.
400 Some("23") => ErrorClass::ConstraintViolation,
401 // 08 = connection exception.
402 Some("08") => ErrorClass::Connection,
403 // 40 = transaction rollback.
404 Some("40") => ErrorClass::Contention,
405 _ => ErrorClass::Other,
406 },
407 }
408 }
409}
410
411/// Whether a query succeeded, and if not, how it failed.
412#[derive(Debug, Clone, Copy, PartialEq, Eq)]
413pub enum Outcome {
414 Ok,
415 Error(ErrorClass),
416}
417
418impl Outcome {
419 pub fn as_label(self) -> &'static str {
420 match self {
421 Outcome::Ok => "success",
422 Outcome::Error(_) => "error",
423 }
424 }
425}
426
427/// Timing and result facts about one query execution.
428///
429/// Returned alongside the result rather than reported from inside the
430/// execution path, so the caller decides what to do with it — record a
431/// metric, attach it to a span, log it, or ignore it. Keeping the decision
432/// out here is what lets the same execution path serve an instrumented
433/// server and an uninstrumented script.
434#[derive(Debug, Clone)]
435pub struct ExecutionMetadata {
436 /// PyQL → SQL compilation. Zero on a compile-cache hit, which is itself
437 /// the signal that the cache is working.
438 pub compile_duration: std::time::Duration,
439 /// Time in the database, from handing over the SQL to having the rows.
440 pub execute_duration: std::time::Duration,
441 /// See `CompiledQuery::shape_id`.
442 pub query_shape_id: String,
443 /// `None` for a statement that returns no rows, which is distinct from
444 /// `Some(0)` — a query that ran and matched nothing.
445 pub rows_returned: Option<u64>,
446 pub outcome: Outcome,
447}
448
449impl ExecutionMetadata {
450 /// Compile plus execute — what a caller timing "the query" means.
451 pub fn total_duration(&self) -> std::time::Duration {
452 self.compile_duration + self.execute_duration
453 }
454}
455
456/// Compile a PyQL expression string in the context of a named type to a bare SQL
457/// expression suitable for use in an UPDATE SET clause.
458///
459/// Pointer references (`.name`) are emitted without a table alias because UPDATE
460/// SET expressions reference the current row directly. Query parameters (`$name`)
461/// are rejected — fill expressions must be literal values or pointer references.
462pub fn compile_fill_expr(
463 type_name: &str,
464 expr_str: &str,
465 schema: &SchemaDescriptor,
466) -> Result<String, crate::error::PyQLError> {
467 let expr_ast = parse::parse_expr(expr_str)?;
468 let (ir_expr, params) = ir::compile_expr_unaliased(&expr_ast, type_name, schema)?;
469 if !params.is_empty() {
470 return Err(crate::error::PyQLError::Syntax(crate::error::PyQLSyntaxError {
471 message: "fill expressions may not contain query parameters".into(),
472 position: crate::error::Position { line: 0, col: 0 },
473 }));
474 }
475 Ok(sql::emit_expr(&ir_expr))
476}
477
478/// Compile a schema `Trigger`'s `handler` PyQL statement (e.g. `insert Note
479/// { note := __new__.name }`) to a full SQL statement, for embedding in the
480/// generated plpgsql trigger function body — see `ir::compile_trigger_handler`
481/// for the `__new__`/`__old__` row-context binding rules. `on_mask` is
482/// Pylon's `On` bitmask (1=Insert, 2=Update, 4=Delete), matching
483/// `TriggerDescriptor::on`. Query parameters (`$name`) are rejected, same
484/// rule as `compile_fill_expr` — a trigger handler has no caller to supply
485/// them.
486pub fn compile_trigger_handler(
487 handler: &str,
488 type_name: &str,
489 on_mask: u8,
490 schema: &SchemaDescriptor,
491) -> Result<String, crate::error::PyQLError> {
492 let ir_out = ir::compile_trigger_handler(handler, type_name, on_mask, schema)?;
493 if !ir_out.params.is_empty() {
494 return Err(crate::error::PyQLError::Syntax(crate::error::PyQLSyntaxError {
495 message: "trigger handlers may not contain query parameters".into(),
496 position: crate::error::Position { line: 0, col: 0 },
497 }));
498 }
499 Ok(sql::emit(&ir_out).sql)
500}
501
502/// Compile a PyQL query string to SQL against `schema`, using default
503/// session config (see `ir::SessionConfig`) — for schema-time/test callers
504/// with no live client-supplied config. `compile_with_config` is the real
505/// entry a query request uses.
506///
507/// Results are cached in a process-global LRU (capacity 1024). Call
508/// `clear_query_cache()` when the schema is reloaded to avoid stale entries.
509/// Synchronous — compilation is CPU-bound; async lives at the DB execution layer.
510/// Raises `PyQLError` on any grammar, type, or resolution failure.
511pub fn compile(query: &str, schema: &SchemaDescriptor) -> Result<Arc<CompiledQuery>, PyQLError> {
512 compile_with_config(query, schema, &ir::SessionConfig::default())
513}
514
515/// Every statement of a script, compiled in the order written.
516///
517/// A single statement comes back as a one-element list, and takes the cached
518/// path — so this is safe to call for any query, not only ones with semicolons
519/// in them.
520pub fn compile_script(
521 query: &str,
522 schema: &SchemaDescriptor,
523 config: &ir::SessionConfig,
524) -> Result<Vec<Arc<CompiledQuery>>, PyQLError> {
525 let statements = parse::parse_script(query)?;
526 if statements.len() == 1 {
527 return Ok(vec![compile_with_config(query, schema, config)?]);
528 }
529 statements
530 .iter()
531 .map(|ast| compile_ast(ast, schema, config).map(Arc::new))
532 .collect()
533}
534
535/// Like `compile`, but honors a caller-supplied `SessionConfig` for this query.
536///
537/// Returns an `Arc` — callers share one immutable compilation rather than
538/// each getting a deep copy of it.
539pub fn compile_with_config(
540 query: &str,
541 schema: &SchemaDescriptor,
542 config: &ir::SessionConfig,
543) -> Result<Arc<CompiledQuery>, PyQLError> {
544 let hash = cache_key_hash(query, config);
545 let shard = &query_cache()[(hash as usize) % CACHE_SHARDS];
546 {
547 let mut cache = shard.write().unwrap();
548 if let Some(entry) = cache.get(&hash) {
549 // Verify, don't assume: a 64-bit collision is vanishingly rare
550 // but handing back the wrong query's SQL would be silent and
551 // catastrophic, so a mismatch falls through to a real compile.
552 if entry.query == query && &entry.config == config {
553 return Ok(entry.compiled.clone());
554 }
555 }
556 }
557 let compiled = Arc::new(compile_uncached(query, schema, config)?);
558 shard.write().unwrap().put(
559 hash,
560 CacheEntry {
561 query: query.to_string(),
562 config: config.clone(),
563 compiled: compiled.clone(),
564 },
565 );
566 Ok(compiled)
567}
568
569fn compile_uncached(
570 query: &str,
571 schema: &SchemaDescriptor,
572 config: &ir::SessionConfig,
573) -> Result<CompiledQuery, PyQLError> {
574 let ast = parse::parse(query)?;
575 compile_ast(&ast, schema, config)
576}
577
578/// Compile one already-parsed statement. Shared by the single-statement path
579/// and by `compile_script`, which has several of them and no source text to
580/// hand back to the parser per statement.
581fn compile_ast(
582 ast: &parse::Stmt,
583 schema: &SchemaDescriptor,
584 config: &ir::SessionConfig,
585) -> Result<CompiledQuery, PyQLError> {
586 let is_analyze = matches!(ast, parse::Stmt::Analyze(_));
587 let ir_out = ir::compile_with_config(ast, schema, config)?;
588 // `analyze`'s own shape-path walk is skipped for every other query — no
589 // reason to pay for it when nothing will read `analyze_paths`.
590 let analyze_paths = is_analyze.then(|| {
591 let mut paths = analyze::collect_shape_path_aliases(&ir_out.stmt);
592 // The root marker has no IR pointer of its own to carry it (see
593 // `root_marker_offset`'s own doc comment) — filled in here from the
594 // original AST, still in scope at this point.
595 if let Some(root) = paths.iter_mut().find(|p| p.path == "root") {
596 root.marker_offset = analyze::root_marker_offset(ast);
597 }
598 paths
599 });
600 let tags = ir::tags::collect_tags(&ir_out);
601 let mutates = stmt_mutates(&ir_out.stmt) || ir_out.ctes.iter().any(|c| stmt_mutates(&c.stmt));
602 let sql_out = sql::emit(&ir_out);
603 let shape_id = derive_shape_id(&sql_out.sql, &sql_out.shape);
604 Ok(CompiledQuery {
605 sql: sql_out.sql,
606 param_names: ir_out.params,
607 params: Vec::new(),
608 shape: sql_out.shape,
609 warnings: ir_out.warnings,
610 inference_plan: sql_out.inference_plan,
611 tags,
612 mutates,
613 analyze_paths,
614 shape_id,
615 })
616}
617
618/// Whether executing this statement writes to any of its `tags`.
619///
620/// Not just the outermost node: a `FOR` body, a `WITH` binding
621/// (`with c := (insert Company {...}) select c`), and a select over a DML
622/// source (`select (insert Person {...}) { id }`) all write while presenting
623/// as something else.
624fn stmt_mutates(stmt: &ir::IrStmt) -> bool {
625 match stmt {
626 ir::IrStmt::Insert(_) | ir::IrStmt::Update(_) | ir::IrStmt::Delete(_) => true,
627 ir::IrStmt::For(f) => stmt_mutates(&f.body),
628 ir::IrStmt::Select(sel) => sel.dml_source.as_deref().is_some_and(stmt_mutates),
629 _ => false,
630 }
631}
632
633#[cfg(test)]
634mod tests {
635 use super::*;
636 use crate::schema::{PropertyDescriptor, TypeDescriptor};
637
638 fn make_schema() -> SchemaDescriptor {
639 SchemaDescriptor {
640 types: vec![TypeDescriptor {
641 name: "Person".into(),
642 module: "default".into(),
643 table: "person".into(),
644 abstract_: false,
645 materialized: false,
646 description: None,
647 parents: vec![],
648 interfaces: vec![],
649 bases: vec![],
650 properties: vec![PropertyDescriptor {
651 name: "id".into(),
652 pg_type: "uuid".into(),
653 nullable: false,
654 default_sql: Some("uuidv7()".into()),
655 default_pyql: None,
656 description: None,
657 check_constraints: vec![],
658 is_exclusive: true,
659 is_pk: true,
660 is_readonly: true,
661 rewrites: vec![],
662 tuple_members: None,
663 column_type: None,
664 }],
665 links: vec![],
666 multilinks: vec![],
667 computed: vec![],
668 constraints: vec![],
669 indexes: vec![],
670 partition: None,
671 vector_indexes: vec![],
672 search_indexes: vec![],
673 triggers: vec![],
674 junction: false,
675 signals: vec![],
676 }],
677 scalars: vec![],
678 enums: vec![],
679 named_tuples: vec![],
680 globals: vec![],
681 functions: vec![],
682 aliases: vec![],
683 channels: vec![],
684 ..Default::default()
685 }
686 }
687
688 #[test]
689 fn test_analyze_paths_is_none_for_a_plain_query() {
690 let schema = make_schema();
691 let compiled = compile("select Person { id }", &schema).unwrap();
692 assert!(compiled.analyze_paths.is_none());
693 }
694
695 #[test]
696 fn test_analyze_paths_is_populated_for_an_analyze_query() {
697 let schema = make_schema();
698 let query = "analyze select Person { id }";
699 let compiled = compile(query, &schema).unwrap();
700 let paths = compiled
701 .analyze_paths
702 .as_ref()
703 .expect("analyze query should populate analyze_paths");
704 assert_eq!(paths.len(), 1);
705 assert_eq!(paths[0].path, "root");
706 let offset = paths[0].marker_offset.expect("root path should carry a marker offset");
707 assert_eq!(&query[offset..offset + "Person".len()], "Person");
708 }
709}