Skip to main content

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}