polyc-query 2026.9.0

//! 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
//!
//! `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.
//! 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
//! [`QueryEngine::execute`](crate::engine::QueryEngine::execute)'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::{
    ObjectName, PipeOperator, Query, Select, Statement, Visit, Visitor,
};
use datafusion::sql::sqlparser::dialect::GenericDialect;
use datafusion::sql::sqlparser::parser::Parser;
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),
}

/// Which kind of statement [`check_statement_allowed`] accepted.
///
/// [`QueryEngine::execute`](crate::engine::QueryEngine::execute) 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_ordered_information_schema(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(_) => 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 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(())
    }
}

/// `true` iff any part of `relation` is the identifier `information_schema`
/// — covers a bare `information_schema.tables`, a catalog-qualified
/// `<catalog>.information_schema.<table>`, and quoted/mixed-case spellings
/// (`"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::engine::QueryEngine::execute`]
    /// 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
        );
    }
}