polyc-query 2026.9.6

//! AST statement-type allowlist for the SQL front end.
//!
//! This gate accepts only statements that parse as a single query — an AST
//! statement-type allowlist, not a prefix check. It must
//! reject:
//!
//! - DDL (`CREATE`, `DROP`, `ALTER`, ...)
//! - DML (`INSERT`, `UPDATE`, `DELETE`, `MERGE`)
//! - `EXPLAIN ANALYZE` (it executes the query and stays banned everywhere)
//! - `EXPLAIN` of anything other than a query (e.g. `EXPLAIN INSERT ...`)
//! - `EXPLAIN` at all, unless the caller opts in — plain `EXPLAIN` is
//!   side-effect-free plan output and is allowed only on maintainer
//!   surfaces that pass `allow_explain = true`
//! - `SHOW`, `COPY`, `SET`, and transaction-control statements
//! - Multi-statement batches (more than one statement in one submission)
//! - SQL that fails to parse at all
//! - `WITH RECURSIVE` in any position — top level, subquery, CTE body, or
//!   the inner query of an allowed `EXPLAIN` (POLY-374). A recursive CTE
//!   loops until the request timeout while it holds a concurrency permit,
//!   so the AST refuses it outright; the pinned
//!   `datafusion.execution.enable_recursive_ctes = false` on the composed
//!   session is the second line for the same refusal.
//!
//! `information_schema` is scoped away from non-admin surfaces and ad hoc
//! external-table references are disabled everywhere — both are
//! enforced elsewhere in the real implementation, not by this gate alone.
//! No composed session enables `information_schema` today: the ordering
//! guard below stays as defence in depth, not because a live caller can
//! reach it.
//! Negative tests across each rejected statement family are a stated
//! verification seam.
//!
//! This gate also rejects one narrower shape within an otherwise-allowed
//! `Statement::Query`: a query that both references `information_schema`
//! and orders its output anywhere in its tree (#1540). `DataFusion` 54
//! silently drops the alphabetically-first row from an ordered
//! `information_schema` scan once combined with `crate::core_execution`'s
//! automatic row-cap `LIMIT` push — a wrong-but-plausible result is worse
//! than an explicit refusal, so [`reject_ordered_information_schema`]
//! rejects the shape outright rather than merely documenting it.

use datafusion::sql::sqlparser::ast::{
    Expr, ObjectName, PipeOperator, Query, Select, SelectItem, Statement, TableFactor, Visit,
    Visitor,
};
use datafusion::sql::sqlparser::dialect::GenericDialect;
use datafusion::sql::sqlparser::parser::Parser;
use std::collections::BTreeMap;
use std::ops::ControlFlow;

/// Why a submitted SQL string was rejected before reaching the planner.
#[derive(Debug, thiserror::Error)]
pub(crate) enum StatementRejected {
    /// The statement parsed, but its kind is not on the allowlist — DDL,
    /// DML, `EXPLAIN` (disallowed or of a non-query), `EXPLAIN ANALYZE`,
    /// `SHOW`, `COPY`, `SET`, transaction control, a multi-statement batch,
    /// or an empty submission.
    #[error("statement kind not allowed: {0}")]
    DisallowedKind(String),
    /// The SQL failed to parse at all.
    #[error("could not parse SQL: {0}")]
    ParseError(String),
    /// The statement parsed as an allowed kind but exceeds a complexity
    /// bound (POLY-371). `bound` names the limit so the refusal metric and
    /// the warn log can count by reason without seeing the SQL text.
    #[error("statement exceeds the {bound} complexity bound")]
    Complexity {
        /// The metric-safe name of the bound that refuses the statement.
        bound: &'static str,
    },
    /// The statement calls a name the closed function surface does not
    /// register (POLY-371): a refused function by canonical name or alias,
    /// an unclassified name, `OVERLAY` syntax a planner would desugar to
    /// the refused `overlay`, or a `FROM` table-function call — the
    /// composed session registers no table functions. The refusal carries
    /// no name so nothing the caller wrote can reach a log label.
    #[error("statement calls a function outside the closed surface")]
    FunctionSurface,
}

impl StatementRejected {
    /// A bounded, SQL-free label for the refusal — safe for the warn log
    /// and the `reason` label of
    /// [`crate::metrics::record_statement_refusal`].
    pub(crate) const fn reason_key(&self) -> &'static str {
        match self {
            Self::DisallowedKind(_) => "statement_kind",
            Self::ParseError(_) => "statement_parse",
            Self::Complexity { bound } => bound,
            Self::FunctionSurface => "function_surface",
        }
    }
}

/// Which kind of statement [`check_statement_allowed`] accepted.
///
/// `crate::core_execution` needs this
/// to decide whether pushing the row-cap `Limit` into the logical plan makes
/// sense: it does for a query's own row stream ([`AllowedStatement::Query`]),
/// but not for `EXPLAIN`'s plan-text output
/// ([`AllowedStatement::Explain`]) — wrapping a `Limit` around an `Explain`
/// plan node doesn't bound the wrapped query's execution, and the plan text
/// itself is always a handful of rows regardless of cap.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum AllowedStatement {
    /// A plain query (`SELECT`/`WITH`/bare `VALUES`).
    Query,
    /// An opted-in plain `EXPLAIN` of a query.
    Explain,
}

/// Check that `sql` parses as a single allowed query statement.
///
/// Allows exactly one [`Statement::Query`] (covers `SELECT`/`WITH`/bare
/// `VALUES`). When `allow_explain` is `true`, also allows a single plain
/// `EXPLAIN` (`analyze == false`) whose inner statement is itself a
/// [`Statement::Query`] — maintainer surfaces opt in by passing `true`;
/// every other surface must pass `false`. `EXPLAIN ANALYZE` is rejected
/// unconditionally because it executes the query.
///
/// # Errors
///
/// Returns [`StatementRejected::ParseError`] when `sql` does not parse, and
/// [`StatementRejected::DisallowedKind`] when it parses to anything but a
/// single allowed query (or allowed `EXPLAIN`) statement — including empty
/// input and multi-statement batches.
pub(crate) fn check_statement_allowed(
    sql: &str,
    allow_explain: bool,
) -> Result<AllowedStatement, StatementRejected> {
    let statements = Parser::parse_sql(&GenericDialect {}, sql)
        .map_err(|err| StatementRejected::ParseError(err.to_string()))?;

    match statements.as_slice() {
        [Statement::Query(query)] => {
            reject_recursive_cte(query)?;
            reject_ordered_information_schema(query)?;
            reject_outside_surface(query)?;
            reject_over_complex(query)?;
            Ok(AllowedStatement::Query)
        }
        [Statement::Explain { analyze: true, .. }] => Err(StatementRejected::DisallowedKind(
            "EXPLAIN ANALYZE executes the query and is never allowed".to_string(),
        )),
        [
            Statement::Explain {
                analyze: false,
                statement,
                ..
            },
        ] => {
            if !allow_explain {
                return Err(StatementRejected::DisallowedKind(
                    "EXPLAIN is not enabled on this surface".to_string(),
                ));
            }
            match statement.as_ref() {
                Statement::Query(query) => {
                    reject_recursive_cte(query)?;
                    reject_outside_surface(query)?;
                    reject_over_complex(query)?;
                    Ok(AllowedStatement::Explain)
                }
                other => Err(StatementRejected::DisallowedKind(format!(
                    "EXPLAIN of a {} statement is not allowed",
                    statement_kind(other)
                ))),
            }
        }
        [other] => Err(StatementRejected::DisallowedKind(format!(
            "{} statements are not allowed",
            statement_kind(other)
        ))),
        [] => Err(StatementRejected::DisallowedKind(
            "no statement found".to_string(),
        )),
        multiple => Err(StatementRejected::DisallowedKind(format!(
            "expected exactly one statement, found {}",
            multiple.len()
        ))),
    }
}

/// Reject `query` if any `WITH` clause in its tree is recursive
/// (POLY-374).
///
/// One [`RecursiveCteScan`] pass — via `sqlparser`'s
/// [`Visit`]/[`Visitor`] machinery — covers every position a recursive
/// clause can occupy: `pre_visit_query` fires for the outer query, each
/// CTE body, each subquery, and each set-operation arm, and each of those
/// nodes carries its own `with` field. The composed session also pins
/// `datafusion.execution.enable_recursive_ctes = false`, so a statement
/// that bypassed this gate still fails at planning; the gate exists so
/// the refusal reaches the caller as the collapsed invalid-query error.
///
/// # Errors
///
/// Returns [`StatementRejected::DisallowedKind`] when any `WITH` clause
/// is recursive.
fn reject_recursive_cte(query: &Query) -> Result<(), StatementRejected> {
    let mut scan = RecursiveCteScan::default();
    // `Visitor::Break` is `()` here and the hook never returns `Break` —
    // the walk always completes, so the `ControlFlow` result carries
    // nothing this function needs.
    let _: ControlFlow<()> = query.visit(&mut scan);

    if scan.has_recursive_with {
        return Err(StatementRejected::DisallowedKind(
            "WITH RECURSIVE is not supported".to_string(),
        ));
    }
    Ok(())
}

/// Accumulates, over one [`Visit`] walk of a query's AST, whether any
/// `WITH` clause in the tree is recursive — the condition
/// [`reject_recursive_cte`] refuses. The hook only ever sets a flag to
/// `true` and continues; nothing here short-circuits the walk, since a
/// later [`Query`] node carries the flag its own `with` declares.
#[derive(Debug, Default)]
struct RecursiveCteScan {
    /// Set once any visited [`Query`] node carries `WITH RECURSIVE`.
    has_recursive_with: bool,
}

impl Visitor for RecursiveCteScan {
    type Break = ();

    /// Fires for every [`Query`] node in the tree — the outer query, each
    /// CTE body, each subquery, each set-operation arm — with that node's
    /// own `with` field directly in hand, so a `WITH RECURSIVE` nested in
    /// a CTE body or a `FROM`/`WHERE` subquery is found without a
    /// separate hook for either position.
    fn pre_visit_query(&mut self, query: &Query) -> ControlFlow<Self::Break> {
        if query.with.as_ref().is_some_and(|with| with.recursive) {
            self.has_recursive_with = true;
        }
        ControlFlow::Continue(())
    }
}

/// Reject `query` if it both references `information_schema` anywhere in
/// its tree and orders its output anywhere in its tree (#1540).
///
/// A single [`InformationSchemaOrderingScan`] pass over `query` — via
/// `sqlparser`'s [`Visit`]/[`Visitor`] machinery, not string matching — finds
/// both conditions together: the walk already descends into every subquery,
/// CTE body, set-operation arm, and pipe-syntax stage in the tree, so a
/// second, separate pass would only repeat that traversal. Deliberately
/// conservative about *where* the two conditions occur: they need not share
/// a subquery — an `information_schema` reference in one branch and an
/// `ORDER BY` in an unrelated branch of the same statement still rejects.
/// A false positive here only asks the caller to drop an `ORDER BY` a
/// catalog-introspection query rarely needs; a false negative would let
/// `DataFusion` 54's row-drop reach the caller silently, which is the worse
/// failure mode.
///
/// # Errors
///
/// Returns [`StatementRejected::DisallowedKind`] when both conditions hold.
fn reject_ordered_information_schema(query: &Query) -> Result<(), StatementRejected> {
    let mut scan = InformationSchemaOrderingScan::default();
    // `Visitor::Break` is `()` here and no hook ever returns `Break` — this
    // walk always completes, so the `ControlFlow` result carries nothing
    // this function needs.
    let _: ControlFlow<()> = query.visit(&mut scan);

    if scan.has_information_schema_reference && scan.has_ordering {
        return Err(StatementRejected::DisallowedKind(
            "ordering an information_schema query is not supported: it can silently drop rows \
             from the result — remove the ORDER BY (or SORT BY, or pipe ORDER BY) and query \
             information_schema without ordering it"
                .to_string(),
        ));
    }
    Ok(())
}

/// Accumulates, over one [`Visit`] walk of a query's AST, whether it
/// references `information_schema` anywhere and whether it orders its
/// output anywhere — the two conditions
/// [`reject_ordered_information_schema`] combines. Each `pre_visit_*` hook
/// only ever sets a flag to `true` and continues; nothing here short-
/// circuits the walk, since a later part of the tree may still supply the
/// other condition.
#[derive(Debug, Default)]
struct InformationSchemaOrderingScan {
    /// Set once any relation the query references resolves to
    /// `information_schema` — see [`references_information_schema`].
    has_information_schema_reference: bool,
    /// Set once any ordering operation appears anywhere in the tree: a
    /// query's (or a CTE's, or a subquery's) own `ORDER BY` — including
    /// `ORDER BY ALL` — Hive `SORT BY`, or pipe-syntax `|> ORDER BY`.
    has_ordering: bool,
}

impl Visitor for InformationSchemaOrderingScan {
    type Break = ();

    /// Fires for every table/view reference in the tree — direct `FROM`,
    /// joins, and subquery/CTE bodies alike — with the relation's own
    /// name, never an `AS` alias layered over it, so an alias cannot hide
    /// an underlying `information_schema` relation from this check.
    fn pre_visit_relation(&mut self, relation: &ObjectName) -> ControlFlow<Self::Break> {
        if references_information_schema(relation) {
            self.has_information_schema_reference = true;
        }
        ControlFlow::Continue(())
    }

    /// Fires for every [`Query`] node in the tree — the outer query, each
    /// CTE body, each subquery, each set-operation arm — with that node's
    /// own `order_by` and `pipe_operators` fields directly in hand, so both
    /// standard `ORDER BY` (including `ORDER BY ALL`, folded into
    /// `Query::order_by` regardless of kind) and pipe-syntax `|> ORDER BY`
    /// are checked here without a separate hook for either.
    fn pre_visit_query(&mut self, query: &Query) -> ControlFlow<Self::Break> {
        if query.order_by.is_some()
            || query
                .pipe_operators
                .iter()
                .any(|operator| matches!(operator, PipeOperator::OrderBy { .. }))
        {
            self.has_ordering = true;
        }
        ControlFlow::Continue(())
    }

    /// Fires for every [`Select`] node in the tree — catches Hive
    /// `SORT BY`, which orders a select's output but, unlike standard
    /// `ORDER BY`, lives on `Select` rather than the enclosing `Query`.
    fn pre_visit_select(&mut self, select: &Select) -> ControlFlow<Self::Break> {
        if !select.sort_by.is_empty() {
            self.has_ordering = true;
        }
        ControlFlow::Continue(())
    }
}

/// Reject `query` if it calls a name the closed function surface does not
/// register (POLY-371).
///
/// The composed session's registries hold only
/// [`crate::function_surface::ALLOWED`]'s implementations, so anything else
/// would fail at planning anyway — the registry is the enforcement and
/// this scan is what makes the refusal deterministic and countable. It
/// fires before planning so a refused or unclassified name surfaces as
/// [`StatementRejected::FunctionSurface`] — a bounded reason the warn log
/// and the refusal metric can count — instead of an engine lookup error.
///
/// Three AST positions can name a function:
///
/// - `Expr::Function` — ordinary calls, qualified or not, in any case.
///   The name's last component, lowercased, must be in
///   [`crate::function_surface::allowed_names`], which covers every
///   canonical name and alias the composed session registers.
/// - `Expr::Overlay` — `OVERLAY(x PLACING y FROM z)` syntax. The planner
///   desugars it to the refused `overlay` implementation directly,
///   bypassing the registry, so the gate refuses the node itself.
/// - `FROM`-position function calls — `TableFactor::Table` with arguments
///   (`FROM range(1, 10)`) and `TableFactor::Function`
///   (`FROM FLATTEN(...)`) resolve through the table-function registry,
///   which the composed session leaves empty. `TableFactor::UNNEST` is
///   not refused: it expands one input array per row — input-bounded —
///   and the walk still checks every expression inside it.
///
/// # Errors
///
/// Returns [`StatementRejected::FunctionSurface`] when any of those
/// positions names something outside the closed surface.
fn reject_outside_surface(query: &Query) -> Result<(), StatementRejected> {
    let mut scan = FunctionSurfaceScan::default();
    // `Visitor::Break` is `()` here and the hook never returns `Break` —
    // the walk always completes, so the `ControlFlow` result carries
    // nothing this function needs.
    let _: ControlFlow<()> = query.visit(&mut scan);
    if scan.outside {
        return Err(StatementRejected::FunctionSurface);
    }
    Ok(())
}

/// Accumulates, over one [`Visit`] walk of a query's AST, whether any
/// function call names something the closed surface does not register —
/// the condition [`reject_outside_surface`] refuses. The flag only ever
/// goes `true`; nothing here short-circuits the walk.
#[derive(Debug, Default)]
struct FunctionSurfaceScan {
    /// Set once an [`Expr::Function`] name, an [`Expr::Overlay`] node, or
    /// a `FROM`-position table function sits outside the closed surface.
    outside: bool,
}

impl Visitor for FunctionSurfaceScan {
    type Break = ();

    /// Fires for every [`Expr`] node — including the expressions inside
    /// `TableFactor::UNNEST` arguments, subqueries, and CTE bodies —
    /// with each function call's own name in hand. `Expr::Overlay` is
    /// refused regardless of name: it reaches the refused `overlay`
    /// implementation without a registry lookup.
    fn pre_visit_expr(&mut self, expr: &Expr) -> ControlFlow<Self::Break> {
        match expr {
            Expr::Function(function) => {
                let name = function
                    .name
                    .0
                    .last()
                    .and_then(|part| part.as_ident())
                    .map_or_else(String::new, |ident| ident.value.to_ascii_lowercase());
                if !crate::function_surface::allowed_names().contains(&name) {
                    self.outside = true;
                }
            }
            Expr::Overlay { .. } => self.outside = true,
            _ => {}
        }
        ControlFlow::Continue(())
    }

    /// Fires for every `FROM` table factor. A `Table` with `args` is a
    /// table-valued call (`FROM range(1, 10)`); `Function` is the named
    /// table-function form (`FROM FLATTEN(...)`) — both resolve through
    /// the empty table-function registry, so the gate refuses them here.
    /// `UNNEST` is the one admitted table factor: input-bounded rows, and
    /// its argument expressions are visited and checked like any other.
    fn pre_visit_table_factor(&mut self, factor: &TableFactor) -> ControlFlow<Self::Break> {
        match factor {
            TableFactor::Table { args: Some(_), .. } | TableFactor::Function { .. } => {
                self.outside = true;
            }
            _ => {}
        }
        ControlFlow::Continue(())
    }
}

/// Largest number of [`Expr`] nodes one statement may carry (POLY-371).
///
/// The bounds in this module exist because the session's other limits —
/// the row cap, the decoded-source byte ceiling, the bounded memory pool,
/// the deadline — all meter work that `DataFusion` reports or that
/// yields. A scalar expression whose argument chooses its output size
/// allocates inside one `poll_next`, outside the pool, before any ceiling
/// is checked. The closed function surface refuses the functions whose
/// arguments do that outright; these bounds refuse the amplification a
/// statement can still build out of *allowed* functions — `c || c || …`
/// multiplies input width with no refused name in sight.
///
/// Each number is drawn from a census of the fixed statements under
/// `crates/query-model/src/statements` and every `public_view_sql` —
/// `fixed_statements_fit_the_surface_and_bounds` in this module's tests
/// and the view sweep in `core_resolution`'s tests hold every one under
/// the bounds — with headroom, not
/// from what a caller might plausibly want. Observed maxima at
/// introduction:
/// 543 nodes / depth 5 / 13 calls (`composite_trace`), 169 items
/// (`composite_trace`), 7 queries (`routine_overview`), 2 output
/// references (`dashboard_spend`), fanout 4 and statement fanout 6
/// (`composite_trace`).
pub(crate) const MAX_EXPRESSION_NODES: usize = 1024;
/// Deepest one expression may nest. Expressions nest inside subqueries
/// too, so this counts through query boundaries.
pub(crate) const MAX_EXPRESSION_DEPTH: usize = 32;
/// Largest number of function calls (`Expr::Function`) in one statement.
pub(crate) const MAX_FUNCTION_CALLS: usize = 128;
/// Largest select-list width of any one `Select` node.
pub(crate) const MAX_SELECT_ITEMS: usize = 512;
/// Largest number of [`Query`] nodes: the outer query plus every CTE body,
/// subquery, and set-operation arm.
pub(crate) const MAX_QUERY_NODES: usize = 64;
/// Largest number of emit-position references to the same column inside
/// one expression node — see [`output_references`].
///
/// This is the width bound the module's arithmetic hangs on. Only a
/// reference whose bytes can reach the node's own output amplifies, and
/// only a repeat amplifies: `c || c` emits `c`'s bytes twice, while
/// `a || b` emits each column once — bounded by the decoded input itself
/// — and `a.x = b.x` emits a boolean and no input bytes at all. A node's
/// output is at most the largest same-name count times the decoded
/// input, so `MAX_EXPRESSION_OUTPUT_REFERENCES` bounds one expression's
/// output at `4 * 32 MiB = 128 MiB` against the 32 MiB decoded-source
/// ceiling.
pub(crate) const MAX_EXPRESSION_OUTPUT_REFERENCES: usize = 4;
/// Largest number of emit-position references to the same unqualified
/// column name across one `Select`'s projection items.
///
/// One expression's bound does not cover `SELECT c || c, c || c, …`: each
/// item is its own expression, but their outputs share one output batch.
/// That batch's size is at most the sum over each referenced column of
/// (references to it) times (its decoded bytes), which is at most
/// `MAX_SELECT_COLUMN_FANOUT` times the decoded input — 128 MiB. The key
/// is the last identifier component, lowercased, so `s.c`, `t.c`, and
/// bare `c` share one budget: an alias cannot spend it twice.
///
/// Together the two width bounds hold one query's un-metered peak at
/// `MAX_EXPRESSION_OUTPUT_REFERENCES + MAX_SELECT_COLUMN_FANOUT` times
/// the decoded input — `8 * 32 MiB = 256 MiB`: the projection's output
/// batch plus the largest live expression. Four concurrent executions
/// therefore hold at most ~1 GiB un-metered beside the 768 MiB memory
/// pool, inside the pod's 2 GiB limit with ~256 MiB of process and
/// encode overhead to spare.
pub(crate) const MAX_SELECT_COLUMN_FANOUT: usize = 4;
/// Largest number of emit-position references to the same unqualified
/// column name across every select-item expression in the statement.
///
/// The per-select bound covers one output batch; this bound covers
/// amplified columns that pass through a CTE or derived table into a
/// later projection under the same name. It does not cover a column
/// renamed with `AS` between stages — that residual is documented in
/// `docs/decisions/0028-closed-function-surface.md`.
pub(crate) const MAX_STATEMENT_COLUMN_REFERENCES: usize = 8;

/// Reject `query` when any bound above is exceeded (POLY-371).
///
/// One [`ComplexityScan`] pass counts every bound over the whole tree —
/// CTE bodies, subqueries, set-operation arms, and pipe stages included,
/// since `sqlparser`'s visitor already descends through all of them.
///
/// # Errors
///
/// Returns [`StatementRejected::Complexity`] naming the first bound the
/// statement exceeds.
fn reject_over_complex(query: &Query) -> Result<(), StatementRejected> {
    let mut scan = ComplexityScan::default();
    // `Visitor::Break` is `()` here — the walk always completes so the
    // census reports the true maximum rather than the first overage.
    let _: ControlFlow<()> = query.visit(&mut scan);
    scan.check()
}

/// Accumulates the complexity counts over one [`Visit`] walk.
///
/// Expression *roots* — the topmost [`Expr`] of each select item, `WHERE`,
/// `HAVING`, join constraint, and so on — are the nodes visited while
/// `expression_depth == 0`. Each root launches one [`output_references`]
/// pass over its own subtree, which records the largest emit-position
/// reference count of any single node — a `concat` nested under a
/// comparison still allocates its amplified output.
#[derive(Debug, Default)]
struct ComplexityScan {
    /// [`Expr`] nodes visited so far.
    expression_nodes: usize,
    /// Current expression nesting depth — `pre_visit_expr` pushes,
    /// `post_visit_expr` pops.
    expression_depth: usize,
    /// Deepest `expression_depth` reached.
    max_expression_depth: usize,
    /// `Expr::Function` nodes visited so far.
    function_calls: usize,
    /// Widest `select.projection` seen.
    max_select_items: usize,
    /// [`Query`] nodes visited so far.
    query_nodes: usize,
    /// Largest emit-reference count inside one expression node.
    max_expression_output_references: usize,
    /// Largest same-name emit-reference count across one select's items.
    max_column_fanout: usize,
    /// Emit references per unqualified name across every select item in
    /// the statement, merged as the walk visits each `Select`.
    statement_column_references: BTreeMap<String, usize>,
}

impl ComplexityScan {
    /// Refuse when any bound is exceeded; otherwise allow.
    fn check(&self) -> Result<(), StatementRejected> {
        let statement_fanout = self
            .statement_column_references
            .values()
            .max()
            .copied()
            .unwrap_or(0);
        let checks = [
            (
                self.expression_nodes,
                MAX_EXPRESSION_NODES,
                "expression_nodes",
            ),
            (
                self.max_expression_depth,
                MAX_EXPRESSION_DEPTH,
                "expression_depth",
            ),
            (self.function_calls, MAX_FUNCTION_CALLS, "function_calls"),
            (self.max_select_items, MAX_SELECT_ITEMS, "select_items"),
            (self.query_nodes, MAX_QUERY_NODES, "query_nodes"),
            (
                self.max_expression_output_references,
                MAX_EXPRESSION_OUTPUT_REFERENCES,
                "expression_output_references",
            ),
            (
                self.max_column_fanout,
                MAX_SELECT_COLUMN_FANOUT,
                "column_fanout",
            ),
            (
                statement_fanout,
                MAX_STATEMENT_COLUMN_REFERENCES,
                "statement_column_references",
            ),
        ];
        for &(observed, limit, bound) in &checks {
            if observed > limit {
                return Err(StatementRejected::Complexity { bound });
            }
        }
        Ok(())
    }

    /// Merges one item's emit map into the select's fanout and the
    /// statement-wide totals.
    fn merge_emit_map(&mut self, fanout: &mut BTreeMap<String, usize>, expr: &Expr) {
        let mut emit = EmitScan::default();
        for (name, count) in output_references(expr, &mut emit) {
            *fanout.entry(name.clone()).or_default() += count;
            *self.statement_column_references.entry(name).or_default() += count;
        }
        self.max_expression_output_references =
            self.max_expression_output_references.max(emit.max_node);
    }
}

impl Visitor for ComplexityScan {
    type Break = ();

    /// Fires for every [`Expr`] node, top-down. A node visited at
    /// `expression_depth == 0` is an expression root: count the emit
    /// references of every node in its subtree against the per-node
    /// bound.
    fn pre_visit_expr(&mut self, expr: &Expr) -> ControlFlow<Self::Break> {
        if self.expression_depth == 0 {
            let mut emit = EmitScan::default();
            let _ = output_references(expr, &mut emit);
            self.max_expression_output_references =
                self.max_expression_output_references.max(emit.max_node);
        }
        if matches!(expr, Expr::Function(_)) {
            self.function_calls += 1;
        }
        self.expression_nodes += 1;
        self.expression_depth += 1;
        self.max_expression_depth = self.max_expression_depth.max(self.expression_depth);
        ControlFlow::Continue(())
    }

    fn post_visit_expr(&mut self, _expr: &Expr) -> ControlFlow<Self::Break> {
        self.expression_depth -= 1;
        ControlFlow::Continue(())
    }

    /// Fires for every [`Select`] node — the outer select, each CTE body,
    /// each subquery, each set-operation arm — with that node's projection
    /// in hand, so both the width bound and the same-name fanout bound see
    /// every projection in the tree.
    fn pre_visit_select(&mut self, select: &Select) -> ControlFlow<Self::Break> {
        self.max_select_items = self.max_select_items.max(select.projection.len());
        let mut fanout: BTreeMap<String, usize> = BTreeMap::new();
        let mut wildcard = false;
        for item in &select.projection {
            match item {
                SelectItem::UnnamedExpr(expr)
                | SelectItem::ExprWithAlias { expr, .. }
                | SelectItem::ExprWithAliases { expr, .. }
                | SelectItem::QualifiedWildcard(
                    datafusion::sql::sqlparser::ast::SelectItemQualifiedWildcardKind::Expr(expr),
                    _,
                ) => self.merge_emit_map(&mut fanout, expr),
                SelectItem::QualifiedWildcard(
                    datafusion::sql::sqlparser::ast::SelectItemQualifiedWildcardKind::ObjectName(_),
                    _,
                )
                | SelectItem::Wildcard(_) => wildcard = true,
            }
        }
        // `*` contributes each covered column once, which can repeat a name
        // an explicit item already referenced — the `+ 1` keeps the bound
        // on the safe side without needing the schema.
        let select_fanout = fanout.values().max().copied().unwrap_or(0) + usize::from(wildcard);
        self.max_column_fanout = self.max_column_fanout.max(select_fanout);
        ControlFlow::Continue(())
    }

    /// Fires once per [`Query`] node — the outer query, each CTE body,
    /// each subquery, each set-operation arm — so the count covers CTEs
    /// and subqueries together.
    fn pre_visit_query(&mut self, _query: &Query) -> ControlFlow<Self::Break> {
        self.query_nodes += 1;
        ControlFlow::Continue(())
    }
}

/// Shared state for [`output_references`].
#[derive(Debug, Default)]
struct EmitScan {
    /// Largest same-name emit-reference count seen at any single
    /// expression node.
    max_node: usize,
}

/// Counts the column references whose bytes can reach `expr`'s own
/// output, keyed by unqualified name, and records the largest same-name
/// count at any node in the subtree in `emit.max_node`.
///
/// A reference *emits* when every operator on the path to `expr`'s root
/// passes input bytes through: identifiers, function arguments (an
/// allowed function's output is at most proportional to its inputs),
/// `CASE` result branches, `||`, the JSON access operators, and literal
/// constructors (`ARRAY`, `STRUCT`, `MAP`, tuples). Operators that yield
/// a fixed-size value — comparisons, `AND`/`OR`, arithmetic, `LIKE`,
/// `IN`, `BETWEEN`, `IS …`, `EXISTS` — emit nothing into their own
/// output, but their children are still visited so a `concat` nested
/// under a comparison keeps its own bound.
///
/// The match is exhaustive on purpose: a `sqlparser` upgrade that adds
/// an [`Expr`] variant fails to compile until the new form is classified.
/// Its length is the variant count — splitting the arms into helper
/// functions would scatter the one table a reviewer must read whole.
#[allow(clippy::too_many_lines)]
fn output_references(expr: &Expr, emit: &mut EmitScan) -> BTreeMap<String, usize> {
    /// Probes `child` for its own node bound without emitting into the
    /// parent's output.
    macro_rules! probe {
        ($child:expr) => {{
            let _ = output_references($child, emit);
        }};
    }
    /// Folds `child`'s emit map into `into`.
    fn merge_into(into: &mut BTreeMap<String, usize>, child: BTreeMap<String, usize>) {
        for (name, count) in child {
            *into.entry(name).or_default() += count;
        }
    }
    /// One emit reference to `name`'s last component.
    fn one(name: &str) -> BTreeMap<String, usize> {
        let mut map = BTreeMap::new();
        map.insert(name.to_ascii_lowercase(), 1);
        map
    }
    /// Merges `output_references` of each child into one map.
    fn emit_children<'a>(
        children: impl IntoIterator<Item = &'a Expr>,
        emit: &mut EmitScan,
    ) -> BTreeMap<String, usize> {
        let mut map = BTreeMap::new();
        for child in children {
            merge_into(&mut map, output_references(child, emit));
        }
        map
    }
    /// Emits `arguments`' expression arguments into `map`, probing the
    /// rest.
    fn emit_function_args(
        map: &mut BTreeMap<String, usize>,
        emit: &mut EmitScan,
        arguments: &datafusion::sql::sqlparser::ast::FunctionArguments,
    ) {
        use datafusion::sql::sqlparser::ast::{FunctionArg, FunctionArgExpr, FunctionArguments};
        if let FunctionArguments::List(list) = arguments {
            for arg in &list.args {
                match arg {
                    FunctionArg::Named { arg, .. } | FunctionArg::Unnamed(arg) => {
                        if let FunctionArgExpr::Expr(arg) = arg {
                            merge_into(map, output_references(arg, emit));
                        }
                    }
                    FunctionArg::ExprNamed { name, arg, .. } => {
                        let _ = output_references(name, emit);
                        if let FunctionArgExpr::Expr(arg) = arg {
                            merge_into(map, output_references(arg, emit));
                        }
                    }
                }
            }
        }
    }
    let local = match expr {
        Expr::Identifier(ident) => one(&ident.value),
        Expr::CompoundIdentifier(idents) => idents
            .last()
            .map_or_else(BTreeMap::new, |last| one(&last.value)),
        // Value extraction passes input bytes through.
        Expr::CompoundFieldAccess { root, access_chain } => {
            let map = output_references(root, emit);
            for access in access_chain {
                if let datafusion::sql::sqlparser::ast::AccessExpr::Subscript(subscript) = access {
                    use datafusion::sql::sqlparser::ast::Subscript;
                    match subscript {
                        Subscript::Index { index } => probe!(index),
                        Subscript::Slice {
                            lower_bound,
                            upper_bound,
                            stride,
                        } => {
                            for bound in [lower_bound, upper_bound, stride].into_iter().flatten() {
                                probe!(bound);
                            }
                        }
                    }
                }
            }
            map
        }
        Expr::JsonAccess { value, .. } => output_references(value, emit),
        Expr::Function(function) => {
            let mut map = BTreeMap::new();
            emit_function_args(&mut map, emit, &function.parameters);
            emit_function_args(&mut map, emit, &function.args);
            // `FILTER`, `OVER`, and `WITHIN GROUP` shape the aggregate,
            // not the emitted bytes — but their expressions still carry
            // their own per-node bound.
            if let Some(filter) = &function.filter {
                probe!(filter);
            }
            if let Some(datafusion::sql::sqlparser::ast::WindowType::WindowSpec(spec)) =
                &function.over
            {
                for key in &spec.partition_by {
                    probe!(key);
                }
                for order in &spec.order_by {
                    probe!(&order.expr);
                }
            }
            for order in &function.within_group {
                probe!(&order.expr);
            }
            map
        }
        Expr::Case {
            operand,
            conditions,
            else_result,
            ..
        } => {
            // Conditions compare; results emit.
            if let Some(operand) = operand {
                probe!(operand);
            }
            let mut map = BTreeMap::new();
            for when in conditions {
                probe!(&when.condition);
                merge_into(&mut map, output_references(&when.result, emit));
            }
            if let Some(else_result) = else_result {
                merge_into(&mut map, output_references(else_result, emit));
            }
            map
        }
        Expr::BinaryOp { left, op, right } => {
            use datafusion::sql::sqlparser::ast::BinaryOperator as B;
            match op {
                // `||` and the JSON access/deletion operators pass input
                // bytes to the output.
                B::StringConcat
                | B::Arrow
                | B::LongArrow
                | B::HashArrow
                | B::HashLongArrow
                | B::HashMinus
                | B::DoubleHash
                | B::Custom(_)
                | B::PGCustomBinaryOperator(_) => {
                    emit_children([left.as_ref(), right.as_ref()], emit)
                }
                // Comparisons, boolean logic, arithmetic, and bit
                // operators all produce a fixed-size value.
                _ => {
                    probe!(left);
                    probe!(right);
                    BTreeMap::new()
                }
            }
        }
        // Fixed-size output; children keep their own bounds.
        Expr::UnaryOp { expr: inner, .. }
        | Expr::Ceil { expr: inner, .. }
        | Expr::Floor { expr: inner, .. }
        | Expr::Extract { expr: inner, .. }
        | Expr::OuterJoin(inner)
        // Boolean-producing and fixed-size parents: children keep their
        // own bounds, nothing emits upward.
        | Expr::IsFalse(inner)
        | Expr::IsNotFalse(inner)
        | Expr::IsTrue(inner)
        | Expr::IsNotTrue(inner)
        | Expr::IsNull(inner)
        | Expr::IsNotNull(inner)
        | Expr::IsUnknown(inner)
        | Expr::IsNotUnknown(inner)
        | Expr::IsNormalized { expr: inner, .. }
        | Expr::InSubquery { expr: inner, .. }
        | Expr::InUnnest { expr: inner, .. }
        | Expr::Position { expr: inner, .. } => {
            probe!(inner);
            BTreeMap::new()
        }
        // Passthrough forms carry their operand's bytes unchanged.
        Expr::Nested(inner)
        | Expr::Collate { expr: inner, .. }
        | Expr::Named { expr: inner, .. }
        | Expr::Prior(inner)
        | Expr::Cast { expr: inner, .. }
        | Expr::Prefixed { value: inner, .. } => output_references(inner, emit),
        // Substring-family output is at most its source's bytes; the
        // position arguments are fixed-size.
        Expr::Substring {
            expr: inner,
            substring_from,
            substring_for,
            ..
        } => {
            for position in [substring_from, substring_for].into_iter().flatten() {
                probe!(position);
            }
            output_references(inner, emit)
        }
        Expr::Trim {
            expr: inner,
            trim_what,
            trim_characters,
            ..
        } => {
            if let Some(what) = trim_what {
                probe!(what);
            }
            for character in trim_characters.iter().flatten() {
                probe!(character);
            }
            output_references(inner, emit)
        }
        Expr::Overlay {
            expr: inner,
            overlay_what,
            overlay_from,
            overlay_for,
        } => {
            probe!(overlay_from);
            if let Some(for_expr) = overlay_for {
                probe!(for_expr);
            }
            emit_children([inner.as_ref(), overlay_what.as_ref()], emit)
        }
        Expr::Convert {
            expr: inner,
            styles,
            ..
        } => {
            for style in styles {
                probe!(style);
            }
            output_references(inner, emit)
        }
        // Literal constructors emit their element bytes.
        Expr::Tuple(children)
        | Expr::Array(datafusion::sql::sqlparser::ast::Array { elem: children, .. })
        | Expr::Struct {
            values: children, ..
        } => emit_children(children.iter(), emit),
        Expr::Dictionary(fields) => {
            emit_children(fields.iter().map(|field| field.value.as_ref()), emit)
        }
        Expr::Map(map) => emit_children(
            map.entries
                .iter()
                .flat_map(|entry| [entry.key.as_ref(), entry.value.as_ref()]),
            emit,
        ),
        // A scalar subquery's cell is bounded by its own inner emit
        // bound, so it weighs as much as one node may emit.
        Expr::Subquery(_) => {
            let mut map = BTreeMap::new();
            map.insert(
                SCALAR_SUBQUERY_KEY.to_owned(),
                MAX_EXPRESSION_OUTPUT_REFERENCES,
            );
            map
        }
        Expr::IsDistinctFrom(left, right)
        | Expr::IsNotDistinctFrom(left, right)
        | Expr::AnyOp { left, right, .. }
        | Expr::AllOp { left, right, .. } => {
            probe!(left);
            probe!(right);
            BTreeMap::new()
        }
        Expr::InList {
            expr: inner, list, ..
        } => {
            probe!(inner);
            for item in list {
                probe!(item);
            }
            BTreeMap::new()
        }
        Expr::Between {
            expr: inner,
            low,
            high,
            ..
        } => {
            probe!(inner);
            probe!(low);
            probe!(high);
            BTreeMap::new()
        }
        Expr::Like {
            expr: inner,
            pattern,
            ..
        }
        | Expr::ILike {
            expr: inner,
            pattern,
            ..
        }
        | Expr::SimilarTo {
            expr: inner,
            pattern,
            ..
        }
        | Expr::RLike {
            expr: inner,
            pattern,
            ..
        } => {
            probe!(inner);
            probe!(pattern);
            BTreeMap::new()
        }
        Expr::AtTimeZone {
            timestamp,
            time_zone,
        } => {
            probe!(timestamp);
            probe!(time_zone);
            BTreeMap::new()
        }
        // Grouping constructs appear only under `GROUP BY`, which does
        // not emit to the output batch.
        Expr::GroupingSets(sets) | Expr::Cube(sets) | Expr::Rollup(sets) => {
            for set in sets {
                for member in set {
                    probe!(member);
                }
            }
            BTreeMap::new()
        }
        // Literals and wildcards carry no column bytes; `EXISTS`,
        // `MATCH AGAINST`, lambdas, and interval values are fixed-size.
        Expr::Exists { .. }
        | Expr::MatchAgainst { .. }
        | Expr::Value(_)
        | Expr::TypedString(_)
        | Expr::Wildcard(_)
        | Expr::QualifiedWildcard(..)
        | Expr::Lambda(_) => BTreeMap::new(),
        Expr::Interval(interval) => {
            probe!(&interval.value);
            BTreeMap::new()
        }
        Expr::MemberOf(member_of) => {
            probe!(&member_of.value);
            probe!(&member_of.array);
            BTreeMap::new()
        }
    };
    emit.max_node = emit
        .max_node
        .max(local.values().max().copied().unwrap_or(0));
    local
}

/// Emit-map key for a scalar subquery's output cell. A scalar subquery
/// is bounded by its own inner emit bound, so it weighs as much as one
/// node's emit maximum — counted under a key no column can share.
const SCALAR_SUBQUERY_KEY: &str = "?subquery";
/// (`"Information_Schema"."TABLES"`). SQL identifier comparison is
/// case-insensitive by default and stays semantically the same schema even
/// when quoted, so this compares case-insensitively regardless of how the
/// identifier was written; `Ident::value` already holds the identifier with
/// quotes stripped, so no unquoting step is needed here.
fn references_information_schema(relation: &ObjectName) -> bool {
    relation.0.iter().any(|part| {
        part.as_ident()
            .is_some_and(|ident| ident.value.eq_ignore_ascii_case("information_schema"))
    })
}

/// A short, user-readable label for a rejected statement's kind, derived
/// from its leading keyword — good enough for an error message, not a
/// full statement-kind taxonomy.
fn statement_kind(statement: &Statement) -> String {
    statement
        .to_string()
        .split_whitespace()
        .next()
        .unwrap_or("unknown")
        .trim_end_matches(|c: char| !c.is_ascii_alphanumeric())
        .to_ascii_uppercase()
}

#[cfg(test)]
mod tests {
    use super::*;

    /// A plain `SELECT` is the baseline allowed statement — and must be
    /// classified as [`AllowedStatement::Query`] specifically, not merely
    /// accepted, so a misclassification (e.g. as `Explain`) would fail this
    /// test rather than slip through.
    #[test]
    fn select_is_allowed() {
        assert_eq!(
            check_statement_allowed("SELECT 1", false).expect("must be allowed"),
            AllowedStatement::Query
        );
    }

    /// A `WITH ... SELECT` common-table-expression query is still a
    /// `Statement::Query` and must be allowed and classified as
    /// [`AllowedStatement::Query`].
    #[test]
    fn with_select_is_allowed() {
        assert_eq!(
            check_statement_allowed("WITH t AS (SELECT 1) SELECT * FROM t", false)
                .expect("must be allowed"),
            AllowedStatement::Query
        );
    }

    /// Maintainer surfaces that opt in get plain `EXPLAIN` of a query,
    /// classified as [`AllowedStatement::Explain`] specifically — not
    /// `Query`, which would wrongly let `crate::core_execution`
    /// push a row-cap `Limit` around `EXPLAIN`'s plan-text output.
    #[test]
    fn explain_select_allowed_when_opted_in() {
        assert_eq!(
            check_statement_allowed("EXPLAIN SELECT 1", true).expect("must be allowed"),
            AllowedStatement::Explain
        );
    }

    /// Non-maintainer surfaces must not get `EXPLAIN`, even of a query.
    #[test]
    fn explain_select_rejected_without_opt_in() {
        let err = check_statement_allowed("EXPLAIN SELECT 1", false).unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// `EXPLAIN ANALYZE` executes the query, so it is rejected even when
    /// the caller opted into plain `EXPLAIN`.
    #[test]
    fn explain_analyze_rejected_with_opt_in() {
        let err = check_statement_allowed("EXPLAIN ANALYZE SELECT 1", true).unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// `EXPLAIN ANALYZE` is also rejected without opt-in.
    #[test]
    fn explain_analyze_rejected_without_opt_in() {
        let err = check_statement_allowed("EXPLAIN ANALYZE SELECT 1", false).unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// `EXPLAIN` of a non-query statement is rejected even with opt-in.
    #[test]
    fn explain_of_insert_rejected_even_with_opt_in() {
        let err = check_statement_allowed("EXPLAIN INSERT INTO t VALUES (1)", true).unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// DDL family: `CREATE TABLE` is rejected.
    #[test]
    fn create_table_rejected() {
        let err = check_statement_allowed("CREATE TABLE t (a INT)", false).unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// DDL family: `DROP TABLE` is rejected.
    #[test]
    fn drop_table_rejected() {
        let err = check_statement_allowed("DROP TABLE t", false).unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// DML family: `INSERT` is rejected.
    #[test]
    fn insert_rejected() {
        let err = check_statement_allowed("INSERT INTO t VALUES (1)", false).unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// DML family: `UPDATE` is rejected.
    #[test]
    fn update_rejected() {
        let err = check_statement_allowed("UPDATE t SET a = 1", false).unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// DML family: `DELETE` is rejected.
    #[test]
    fn delete_rejected() {
        let err = check_statement_allowed("DELETE FROM t", false).unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// `SHOW` is rejected.
    #[test]
    fn show_rejected() {
        let err = check_statement_allowed("SHOW TABLES", false).unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// `COPY` is rejected.
    #[test]
    fn copy_rejected() {
        let err = check_statement_allowed("COPY t TO 'out.csv'", false).unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// `SET` is rejected.
    #[test]
    fn set_rejected() {
        let err = check_statement_allowed("SET timezone = 'UTC'", false).unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// A multi-statement batch is rejected even though every statement in
    /// it is individually an allowed query.
    #[test]
    fn multi_statement_batch_rejected() {
        let err = check_statement_allowed("SELECT 1; SELECT 2", false).unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// An empty submission has no statement to allow.
    #[test]
    fn empty_string_rejected() {
        let err = check_statement_allowed("", false).unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// Unparseable SQL surfaces as a parse error, not a panic.
    #[test]
    fn garbage_rejected_as_parse_error() {
        let err = check_statement_allowed("not even close to sql (((", false).unwrap_err();
        assert!(matches!(err, StatementRejected::ParseError(_)));
    }

    // --- #1540: information_schema + ordering anywhere in the tree -------

    /// Direct `FROM information_schema.tables ORDER BY ...` — the exact
    /// shape that triggers `DataFusion` 54's silent row drop — is rejected.
    #[test]
    fn information_schema_direct_order_by_rejected() {
        let err = check_statement_allowed(
            "SELECT table_name FROM information_schema.tables ORDER BY table_name",
            false,
        )
        .unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// A catalog-qualified relation (`<catalog>.information_schema.<table>`)
    /// still counts as an `information_schema` reference.
    #[test]
    fn catalog_qualified_information_schema_order_by_rejected() {
        let err = check_statement_allowed(
            "SELECT table_name FROM datafusion.information_schema.tables ORDER BY table_name",
            false,
        )
        .unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// Quoted, mixed-case identifiers still resolve to `information_schema`
    /// — SQL identifier comparison here is case-insensitive.
    #[test]
    fn quoted_mixed_case_information_schema_order_by_rejected() {
        let err = check_statement_allowed(
            r#"SELECT "TABLES"."TABLE_NAME" FROM "Information_Schema"."TABLES" ORDER BY "TABLES"."TABLE_NAME""#,
            false,
        )
        .unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// A table alias over `information_schema` does not hide the
    /// underlying relation from the check.
    #[test]
    fn aliased_information_schema_order_by_rejected() {
        let err = check_statement_allowed(
            "SELECT t.table_name FROM information_schema.tables t ORDER BY t.table_name",
            false,
        )
        .unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// A join between `information_schema` and an ordinary table, ordered,
    /// is rejected — the reference and the ordering need not be on the same
    /// side of the join.
    #[test]
    fn join_with_information_schema_order_by_rejected() {
        let err = check_statement_allowed(
            "SELECT t.table_name FROM information_schema.tables t \
             JOIN information_schema.columns c ON t.table_name = c.table_name \
             ORDER BY t.table_name",
            false,
        )
        .unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// Ordering inside a subquery that references `information_schema` is
    /// rejected even though the outer query has no `ORDER BY` of its own.
    #[test]
    fn ordering_in_subquery_over_information_schema_rejected() {
        let err = check_statement_allowed(
            "SELECT * FROM (SELECT table_name FROM information_schema.tables \
             ORDER BY table_name) sub",
            false,
        )
        .unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// An outer `ORDER BY` over a CTE whose body reads `information_schema`
    /// is rejected — the reference and the ordering are in different parts
    /// of the tree, and the rule is deliberately conservative about that.
    #[test]
    fn outer_order_by_over_cte_reading_information_schema_rejected() {
        let err = check_statement_allowed(
            "WITH t AS (SELECT table_name FROM information_schema.tables) \
             SELECT * FROM t ORDER BY table_name",
            false,
        )
        .unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// Pipe-style `|> ORDER BY` over `information_schema` is rejected too —
    /// the rule does not only look at the standard `ORDER BY` clause.
    #[test]
    fn pipe_style_order_by_over_information_schema_rejected() {
        let err = check_statement_allowed(
            "SELECT table_name FROM information_schema.tables |> ORDER BY table_name",
            false,
        )
        .unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// A set operation with `information_schema` in one arm and an outer
    /// `ORDER BY` over the whole union is rejected — the reference and the
    /// ordering live in different `SetExpr` branches, and the walk visits both.
    #[test]
    fn union_with_information_schema_arm_and_outer_order_by_rejected() {
        let err = check_statement_allowed(
            "SELECT table_name FROM information_schema.tables \
             UNION SELECT name FROM usage ORDER BY table_name",
            false,
        )
        .unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// An outer `ORDER BY` where `information_schema` appears only inside a
    /// `FROM (...)` derived-table subquery — the reference and the ordering
    /// straddle the subquery boundary, and both are still visited.
    #[test]
    fn outer_order_by_over_derived_subquery_reading_information_schema_rejected() {
        let err = check_statement_allowed(
            "SELECT table_name FROM (SELECT table_name FROM information_schema.tables) sub \
             ORDER BY table_name",
            false,
        )
        .unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// `information_schema` without any ordering is still allowed — the
    /// rule only fires when ordering is also present.
    #[test]
    fn information_schema_without_order_by_allowed() {
        assert_eq!(
            check_statement_allowed("SELECT table_name FROM information_schema.tables", false)
                .expect("must be allowed"),
            AllowedStatement::Query
        );
    }

    /// `ORDER BY` over an ordinary typed table is unaffected — the defect
    /// is specific to `information_schema`'s own physical scan.
    #[test]
    fn order_by_over_ordinary_table_allowed() {
        assert_eq!(
            check_statement_allowed("SELECT * FROM usage ORDER BY position", false)
                .expect("must be allowed"),
            AllowedStatement::Query
        );
    }

    /// Plain `EXPLAIN` of an ordered `information_schema` query is still
    /// allowed: `EXPLAIN` never receives the automatic row-cap `LIMIT`
    /// (`QueryEngine::execute` skips the `Limit` push for
    /// `AllowedStatement::Explain`), so it cannot hit the row-drop this rule
    /// guards against.
    #[test]
    fn explain_over_ordered_information_schema_allowed() {
        assert_eq!(
            check_statement_allowed(
                "EXPLAIN SELECT table_name FROM information_schema.tables ORDER BY table_name",
                true,
            )
            .expect("must be allowed"),
            AllowedStatement::Explain
        );
    }

    // --- POLY-374: WITH RECURSIVE in any position -----------------------

    /// A top-level `WITH RECURSIVE` is rejected. The clause lets one
    /// statement loop until the request timeout while it holds a
    /// concurrency permit.
    #[test]
    fn top_level_with_recursive_rejected() {
        let err = check_statement_allowed(
            "WITH RECURSIVE n AS (SELECT 1 AS v UNION ALL SELECT v + 1 FROM n) \
             SELECT v FROM n",
            false,
        )
        .unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// `WITH RECURSIVE` inside a `FROM` derived-table subquery is rejected
    /// — nesting the clause below the outer `SELECT` does not hide it.
    #[test]
    fn with_recursive_in_derived_subquery_rejected() {
        let err = check_statement_allowed(
            "SELECT v FROM \
             (WITH RECURSIVE n AS (SELECT 1 AS v UNION ALL SELECT v + 1 FROM n) \
              SELECT v FROM n) sub",
            false,
        )
        .unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// `WITH RECURSIVE` inside an `EXISTS` predicate subquery is rejected.
    #[test]
    fn with_recursive_in_exists_subquery_rejected() {
        let err = check_statement_allowed(
            "SELECT v FROM usage WHERE EXISTS \
             (WITH RECURSIVE n AS (SELECT 1 AS v UNION ALL SELECT v + 1 FROM n) \
              SELECT v FROM n)",
            false,
        )
        .unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// `WITH RECURSIVE` inside an ordinary CTE's body is rejected — a
    /// non-recursive outer `WITH` does not launder a recursive inner one.
    #[test]
    fn with_recursive_nested_in_cte_body_rejected() {
        let err = check_statement_allowed(
            "WITH outer_cte AS \
             (WITH RECURSIVE n AS (SELECT 1 AS v UNION ALL SELECT v + 1 FROM n) \
              SELECT v FROM n) \
             SELECT v FROM outer_cte",
            false,
        )
        .unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// `WITH RECURSIVE` on a query behind `IN` is rejected.
    #[test]
    fn with_recursive_in_in_subquery_rejected() {
        let err = check_statement_allowed(
            "SELECT v FROM usage WHERE position IN \
             (WITH RECURSIVE n AS (SELECT 1 AS v UNION ALL SELECT v + 1 FROM n) \
              SELECT v FROM n)",
            false,
        )
        .unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// `EXPLAIN` of a recursive query is rejected even with opt-in — the
    /// refusal is on the query itself, not on executing it.
    #[test]
    fn explain_of_with_recursive_rejected() {
        let err = check_statement_allowed(
            "EXPLAIN WITH RECURSIVE n AS (SELECT 1 AS v) SELECT v FROM n",
            true,
        )
        .unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    /// A plain non-recursive `WITH` stays allowed — the check refuses the
    /// `RECURSIVE` keyword, not the CTE shape.
    #[test]
    fn non_recursive_with_still_allowed() {
        assert_eq!(
            check_statement_allowed(
                "WITH t AS (SELECT position FROM usage) SELECT position FROM t",
                false,
            )
            .expect("must be allowed"),
            AllowedStatement::Query
        );
    }

    /// `WITH RECURSIVE` inside a `UNION` arm's subquery is rejected — the
    /// walk descends through set-operation branches.
    #[test]
    fn with_recursive_in_union_arm_rejected() {
        let err = check_statement_allowed(
            "SELECT position FROM usage UNION ALL SELECT v FROM \
             (WITH RECURSIVE n AS (SELECT 1 AS v UNION ALL SELECT v + 1 FROM n) \
              SELECT v FROM n) sub",
            false,
        )
        .unwrap_err();
        assert!(matches!(err, StatementRejected::DisallowedKind(_)));
    }

    // --- POLY-371: closed function surface --------------------------------

    /// A refused canonical name fails at the gate, whatever its case.
    #[test]
    fn refused_function_by_name_and_case_rejected() {
        for sql in [
            "SELECT repeat('x', 2)",
            "SELECT REPEAT('x', 2)",
            "SELECT Repeat('x', 2)",
        ] {
            let err = check_statement_allowed(sql, false).unwrap_err();
            assert!(
                matches!(err, StatementRejected::FunctionSurface),
                "{sql} must refuse as FunctionSurface, got {err:?}"
            );
        }
    }

    /// Every refused spelling — canonical and alias — fails at the gate.
    #[test]
    fn every_refused_name_and_alias_rejected() {
        for name in crate::function_surface::refused_names() {
            let sql = format!("SELECT {name}('x', 1)");
            let err = check_statement_allowed(&sql, false).unwrap_err();
            assert!(
                matches!(err, StatementRejected::FunctionSurface),
                "{sql} must refuse as FunctionSurface, got {err:?}"
            );
        }
    }

    /// A schema-qualified refused name still resolves to the refusal —
    /// a qualifier does not launder `repeat` into something else.
    #[test]
    fn qualified_refused_function_rejected() {
        let err = check_statement_allowed("SELECT pg_catalog.repeat('x', 2)", false).unwrap_err();
        assert!(matches!(err, StatementRejected::FunctionSurface));
    }

    /// `OVERLAY(... PLACING ... FROM ...)` reaches the refused `overlay`
    /// implementation through a planner desugar, not the registry — the
    /// gate refuses the AST node itself.
    #[test]
    fn overlay_syntax_rejected() {
        let err =
            check_statement_allowed("SELECT OVERLAY('abc' PLACING 'x' FROM 2)", false).unwrap_err();
        assert!(matches!(err, StatementRejected::FunctionSurface));
    }

    /// `FROM`-position table-function calls are refused: the composed
    /// session registers no table functions.
    #[test]
    fn table_function_in_from_rejected() {
        for sql in [
            "SELECT * FROM range(1, 10)",
            "SELECT * FROM generate_series(1, 10)",
        ] {
            let err = check_statement_allowed(sql, false).unwrap_err();
            assert!(
                matches!(err, StatementRejected::FunctionSurface),
                "{sql} must refuse as FunctionSurface, got {err:?}"
            );
        }
    }

    /// A name the surface has never classified fails closed — the gate
    /// admits by membership, not by blocklist.
    #[test]
    fn unclassified_function_rejected() {
        let err = check_statement_allowed("SELECT a_function_nobody_listed(1)", false).unwrap_err();
        assert!(matches!(err, StatementRejected::FunctionSurface));
    }

    /// `UNNEST` as a `FROM` factor stays allowed — it expands one input
    /// array per row and its argument expressions are name-checked like
    /// any other.
    #[test]
    fn unnest_table_factor_allowed() {
        assert_eq!(
            check_statement_allowed("SELECT * FROM UNNEST(make_array(1, 2))", false)
                .expect("must be allowed"),
            AllowedStatement::Query
        );
    }

    /// An allowed call keeps passing the gate — the surface check admits
    /// by membership, so the fixed statements' `count`/`sum`/`substr`
    /// shapes are unaffected.
    #[test]
    fn allowed_function_call_allowed() {
        assert_eq!(
            check_statement_allowed(
                "SELECT substr(name, 1, 2), count(*) FROM usage GROUP BY name",
                false,
            )
            .expect("must be allowed"),
            AllowedStatement::Query
        );
    }

    // --- POLY-371: statement complexity bounds ----------------------------

    /// Builds a `SELECT` with `n` comma-separated copies of `item` — one
    /// select list of the requested width.
    fn wide_select(item: &str, n: usize) -> String {
        let items: Vec<String> = (0..n).map(|_| item.to_owned()).collect();
        format!("SELECT {}", items.join(", "))
    }

    /// `n` parenthesized wrappers around `x` — `(((x)))` nests one
    /// [`Expr::Nested`] per pair.
    fn nested_parens(n: usize) -> String {
        format!("SELECT {}x{}", "(".repeat(n), ")".repeat(n))
    }

    /// The full-bound SQL builders below pair one statement exactly at a
    /// bound with the same shape one unit over it.
    #[test]
    fn expression_nodes_at_and_over_bound() {
        // `x IN (v, ...)` is one `InList` node plus one identifier and one
        // `Value` per element: `2 + n` nodes total.
        let at = format!(
            "SELECT x IN ({})",
            (0..MAX_EXPRESSION_NODES - 2)
                .map(|v| v.to_string())
                .collect::<Vec<_>>()
                .join(", ")
        );
        let over = format!(
            "SELECT x IN ({})",
            (0..MAX_EXPRESSION_NODES - 1)
                .map(|v| v.to_string())
                .collect::<Vec<_>>()
                .join(", ")
        );
        assert!(check_statement_allowed(&at, false).is_ok());
        let err = check_statement_allowed(&over, false).unwrap_err();
        assert!(
            matches!(
                err,
                StatementRejected::Complexity {
                    bound: "expression_nodes"
                }
            ),
            "got {err:?}"
        );
    }

    /// Depth counts every [`Expr`] level; `n` parens under the select
    /// item make the innermost identifier sit at depth `n + 1`.
    #[test]
    fn expression_depth_at_and_over_bound() {
        let at = nested_parens(MAX_EXPRESSION_DEPTH - 1);
        let over = nested_parens(MAX_EXPRESSION_DEPTH);
        assert!(check_statement_allowed(&at, false).is_ok());
        let err = check_statement_allowed(&over, false).unwrap_err();
        assert!(
            matches!(
                err,
                StatementRejected::Complexity {
                    bound: "expression_depth"
                }
            ),
            "got {err:?}"
        );
    }

    /// `abs(1)` is one `Expr::Function` with no column reference — the
    /// fanout bounds stay silent and only the call count grows.
    #[test]
    fn function_calls_at_and_over_bound() {
        let at = wide_select("abs(1)", MAX_FUNCTION_CALLS);
        let over = wide_select("abs(1)", MAX_FUNCTION_CALLS + 1);
        assert!(check_statement_allowed(&at, false).is_ok());
        let err = check_statement_allowed(&over, false).unwrap_err();
        assert!(
            matches!(
                err,
                StatementRejected::Complexity {
                    bound: "function_calls"
                }
            ),
            "got {err:?}"
        );
    }

    /// Literal-only items keep every other bound quiet while the
    /// projection width alone crosses the limit.
    #[test]
    fn select_items_at_and_over_bound() {
        let at = wide_select("1", MAX_SELECT_ITEMS);
        let over = wide_select("1", MAX_SELECT_ITEMS + 1);
        assert!(check_statement_allowed(&at, false).is_ok());
        let err = check_statement_allowed(&over, false).unwrap_err();
        assert!(
            matches!(
                err,
                StatementRejected::Complexity {
                    bound: "select_items"
                }
            ),
            "got {err:?}"
        );
    }

    /// `n` CTE bodies plus the outer query make `n + 1` [`Query`] nodes.
    #[test]
    fn query_nodes_at_and_over_bound() {
        let ctes = |n: usize| {
            let names: Vec<String> = (0..n).map(|i| format!("c{i} AS (SELECT 1)")).collect();
            format!("WITH {} SELECT 1", names.join(", "))
        };
        assert!(check_statement_allowed(&ctes(MAX_QUERY_NODES - 1), false).is_ok());
        let err = check_statement_allowed(&ctes(MAX_QUERY_NODES), false).unwrap_err();
        assert!(
            matches!(
                err,
                StatementRejected::Complexity {
                    bound: "query_nodes"
                }
            ),
            "got {err:?}"
        );
    }

    /// `||` concatenation emits every operand's bytes into the node:
    /// `n` references to the SAME column is `n` emit references to it.
    /// Distinct columns — `a || b || c` — each emit once and stay bounded
    /// by the decoded input itself, so only repeats trip this bound.
    #[test]
    fn expression_output_references_at_and_over_bound() {
        let at = "SELECT c || c || c || c FROM t";
        let over = "SELECT c || c || c || c || c FROM t";
        assert!(check_statement_allowed(at, false).is_ok());
        let err = check_statement_allowed(over, false).unwrap_err();
        assert!(
            matches!(
                err,
                StatementRejected::Complexity {
                    bound: "expression_output_references"
                }
            ),
            "got {err:?}"
        );
    }

    /// Repeating one column across one select's items spends the
    /// per-select fanout budget; a boolean predicate spends none.
    #[test]
    fn column_fanout_at_and_over_bound() {
        let at = wide_select("c", MAX_SELECT_COLUMN_FANOUT);
        let over = wide_select("c", MAX_SELECT_COLUMN_FANOUT + 1);
        assert!(check_statement_allowed(&at, false).is_ok());
        let err = check_statement_allowed(&over, false).unwrap_err();
        assert!(
            matches!(
                err,
                StatementRejected::Complexity {
                    bound: "column_fanout"
                }
            ),
            "got {err:?}"
        );
    }

    /// The same column emitted once per `UNION ALL` arm spends the
    /// statement-wide budget even though no select repeats it.
    #[test]
    fn statement_column_references_at_and_over_bound() {
        let union = |n: usize| {
            (0..n)
                .map(|_| "SELECT c FROM t".to_owned())
                .collect::<Vec<_>>()
                .join(" UNION ALL ")
        };
        assert!(check_statement_allowed(&union(MAX_STATEMENT_COLUMN_REFERENCES), false).is_ok());
        let err = check_statement_allowed(&union(MAX_STATEMENT_COLUMN_REFERENCES + 1), false)
            .unwrap_err();
        assert!(
            matches!(
                err,
                StatementRejected::Complexity {
                    bound: "statement_column_references"
                }
            ),
            "got {err:?}"
        );
    }

    /// A four-key join predicate emits a boolean, not eight copies of the
    /// key columns — join constraints of `composite_trace`'s shape must
    /// not count as amplification.
    #[test]
    fn join_predicate_does_not_consume_emit_budget() {
        assert!(
            check_statement_allowed(
                "SELECT s.position FROM l JOIN s ON \
                 l.a = s.a AND l.b = s.b AND l.c = s.c AND l.d = s.d AND \
                 l.e = s.e AND l.f = s.f AND l.g = s.g AND l.h = s.h",
                false,
            )
            .is_ok()
        );
    }

    /// A refused bound fires on the inner query of an opted-in `EXPLAIN`
    /// too — the explain path is not a bypass.
    #[test]
    fn explain_of_over_bound_query_rejected() {
        let sql = format!("EXPLAIN {}", wide_select("1", MAX_SELECT_ITEMS + 1));
        let err = check_statement_allowed(&sql, true).unwrap_err();
        assert!(
            matches!(
                err,
                StatementRejected::Complexity {
                    bound: "select_items"
                }
            ),
            "got {err:?}"
        );
    }

    /// A refused function inside an opted-in `EXPLAIN` is refused — the
    /// explain path applies the surface check to the inner query.
    #[test]
    fn explain_of_refused_function_rejected() {
        let err = check_statement_allowed("EXPLAIN SELECT repeat('x', 2)", true).unwrap_err();
        assert!(matches!(err, StatementRejected::FunctionSurface));
    }

    // --- POLY-371: every trusted caller's statement fits the bounds -------

    /// Every `reason_key` the gate can produce is one of the metric's
    /// pre-registered labels, and every registered label is reachable —
    /// the closed `reason` label set `record_statement_refusal` promises.
    /// A new `StatementRejected` variant or a new bound fails here until
    /// `REFUSAL_REASONS` names it.
    #[test]
    fn every_reason_key_is_a_registered_metric_label() {
        use std::collections::BTreeSet;
        let mut produced = BTreeSet::new();
        for rejected in [
            StatementRejected::DisallowedKind("x".to_owned()),
            StatementRejected::ParseError("x".to_owned()),
            StatementRejected::FunctionSurface,
        ] {
            produced.insert(rejected.reason_key());
        }
        // One statement one unit over each complexity bound.
        let union = |n: usize| {
            (0..n)
                .map(|_| "SELECT c FROM t".to_owned())
                .collect::<Vec<_>>()
                .join(" UNION ALL ")
        };
        let ctes = |n: usize| {
            let names: Vec<String> = (0..n).map(|i| format!("c{i} AS (SELECT 1)")).collect();
            format!("WITH {} SELECT 1", names.join(", "))
        };
        let over = [
            format!(
                "SELECT x IN ({})",
                (0..MAX_EXPRESSION_NODES - 1)
                    .map(|v| v.to_string())
                    .collect::<Vec<_>>()
                    .join(", ")
            ),
            nested_parens(MAX_EXPRESSION_DEPTH),
            wide_select("abs(1)", MAX_FUNCTION_CALLS + 1),
            wide_select("1", MAX_SELECT_ITEMS + 1),
            ctes(MAX_QUERY_NODES),
            "SELECT c || c || c || c || c FROM t".to_owned(),
            wide_select("c", MAX_SELECT_COLUMN_FANOUT + 1),
            union(MAX_STATEMENT_COLUMN_REFERENCES + 1),
        ];
        for sql in over {
            let err = check_statement_allowed(&sql, false).unwrap_err();
            produced.insert(err.reason_key());
        }
        let registered: BTreeSet<&str> = crate::metrics::REFUSAL_REASONS.iter().copied().collect();
        assert_eq!(produced, registered);
    }

    /// Every fixed statement passes the full gate — the surface check and
    /// the complexity bounds alike. Trusted callers get no bypass: a fixed
    /// statement that ever needs a refused function or a larger shape
    /// fails here, loudly. The redaction views' sweep lives beside
    /// `CoreTable` in `core_resolution`'s tests — this module must not
    /// reach the projected catalog directly.
    #[test]
    fn fixed_statements_fit_the_surface_and_bounds() {
        use polyc_query_model::statements::*;
        let statements: &[(&str, &str)] = &[
            ("composite_trace", COMPOSITE_TRACE_SQL),
            ("composite_trace_memory", COMPOSITE_TRACE_MEMORY_SQL),
            ("composite_trace_routines", COMPOSITE_TRACE_ROUTINES_SQL),
            ("dashboard_attribution", DASHBOARD_ATTRIBUTION_SQL),
            ("dashboard_context", DASHBOARD_CONTEXT_SQL),
            ("dashboard_conversations", DASHBOARD_CONVERSATIONS_SQL),
            ("dashboard_spend", DASHBOARD_SPEND_SQL),
            ("fleet_usage", FLEET_USAGE_SQL),
            ("persona_directory_detail", PERSONA_DIRECTORY_DETAIL_SQL),
            ("persona_directory_index", PERSONA_DIRECTORY_INDEX_SQL),
            ("routine_fire", ROUTINE_FIRE_SQL),
            ("routine_fire_count", ROUTINE_FIRE_COUNT_SQL),
            ("routine_fire_last", ROUTINE_FIRE_LAST_SQL),
            ("routine_fire_outcome", ROUTINE_FIRE_OUTCOME_SQL),
            ("routine_fires", ROUTINE_FIRES_SQL),
            ("routine_fires_index_admin", ROUTINE_FIRES_INDEX_ADMIN_SQL),
            ("routine_fires_index_owner", ROUTINE_FIRES_INDEX_OWNER_SQL),
            ("routine_lifecycle", ROUTINE_LIFECYCLE_SQL),
            ("routine_overview", ROUTINE_OVERVIEW_SQL),
            (
                "routine_owner_active_grants",
                ROUTINE_OWNER_ACTIVE_GRANTS_SQL,
            ),
            (
                "routine_owner_approval_aggregates",
                ROUTINE_OWNER_APPROVAL_AGGREGATES_SQL,
            ),
            (
                "routine_owner_fire_dispatch",
                ROUTINE_OWNER_FIRE_DISPATCH_SQL,
            ),
            ("routine_owner_refusals", ROUTINE_OWNER_REFUSALS_SQL),
            ("routine_owner_stopped_tool", ROUTINE_OWNER_STOPPED_TOOL_SQL),
        ];
        for &(name, sql) in statements {
            assert_eq!(
                check_statement_allowed(sql, false)
                    .unwrap_or_else(|err| panic!("{name} must be allowed, got {err:?}")),
                AllowedStatement::Query
            );
        }
    }
}