ahri_tre_lake/

restricted_query.rs

//! Dependency expansion runs before optimization. Accepted SQL names only
//! adapter-owned input tables and must execute in the restricted worker.
use crate::LakeError;
use sqlparser::{ast::*, dialect::DuckDbDialect, parser::Parser};
use std::{
    collections::{BTreeMap, BTreeSet},
    ops::ControlFlow,
};
use uuid::Uuid;

#[derive(Debug, Clone)]
pub enum QueryBinding {
    Dataset(Uuid),
    View(String),
}

pub struct RestrictedQueryPlan {
    sql: String,
    inputs: BTreeSet<Uuid>,
}
impl RestrictedQueryPlan {
    pub fn input_versions(&self) -> &BTreeSet<Uuid> {
        &self.inputs
    }
    /// Regenerated from the accepted AST, never the original untrusted text.
    pub fn restricted_sql(&self) -> &str {
        &self.sql
    }
}

pub fn analyze_disclosure_query(
    sql: &str,
    bindings: &BTreeMap<String, QueryBinding>,
) -> Result<RestrictedQueryPlan, LakeError> {
    if bindings.len() > 128
        || bindings
            .keys()
            .any(|name| name.is_empty() || name.len() > 768 || *name != name.to_ascii_lowercase())
    {
        return Err(LakeError::QueryNotAdmissible);
    }
    let mut inputs = BTreeSet::new();
    let mut remaining_bytes = AnalysisBudget {
        bytes: 1_048_576,
        tokens: 256,
    };
    let query = expand(
        sql,
        bindings,
        &mut Vec::new(),
        &mut inputs,
        &mut remaining_bytes,
    )?;
    let sql = query.to_string();
    if sql.len() > 1_048_576 {
        return Err(LakeError::QueryNotAdmissible);
    }
    Ok(RestrictedQueryPlan { sql, inputs })
}

struct AnalysisBudget {
    bytes: usize,
    tokens: usize,
}

fn expand(
    sql: &str,
    bindings: &BTreeMap<String, QueryBinding>,
    views: &mut Vec<String>,
    inputs: &mut BTreeSet<Uuid>,
    remaining_bytes: &mut AnalysisBudget,
) -> Result<Box<Query>, LakeError> {
    if sql.len() > 65536 || views.len() > 32 {
        return Err(LakeError::QueryNotAdmissible);
    }
    remaining_bytes.bytes = remaining_bytes
        .bytes
        .checked_sub(sql.len().max(1))
        .ok_or(LakeError::QueryNotAdmissible)?;
    // A left-associative expression may have shallow parser recursion and a
    // deep AST. Bound tokens before allocating any recursive syntax tree.
    let tokens = sqlparser::tokenizer::Tokenizer::new(&DuckDbDialect {}, sql)
        .tokenize()
        .map_err(|_| LakeError::QueryNotAdmissible)?;
    remaining_bytes.tokens = remaining_bytes
        .tokens
        .checked_sub(
            tokens
                .iter()
                .filter(|token| !matches!(token, sqlparser::tokenizer::Token::Whitespace(_)))
                .count(),
        )
        .ok_or(LakeError::QueryNotAdmissible)?;
    let mut statements = Parser::new(&DuckDbDialect {})
        .with_recursion_limit(64)
        .try_with_sql(sql)
        .and_then(|mut parser| parser.parse_statements())
        .map_err(|_| LakeError::QueryNotAdmissible)?;
    if statements.len() != 1 {
        return Err(LakeError::QueryNotAdmissible);
    }
    let Statement::Query(mut query) = statements.remove(0) else {
        return Err(LakeError::QueryNotAdmissible);
    };
    let mut analyzer = Analyzer {
        bindings,
        views,
        inputs,
        remaining_bytes,
        scopes: Vec::new(),
        nodes: 0,
    };
    if VisitMut::visit(&mut query, &mut analyzer).is_break() {
        return Err(LakeError::QueryNotAdmissible);
    }
    Ok(query)
}

struct Analyzer<'a> {
    bindings: &'a BTreeMap<String, QueryBinding>,
    views: &'a mut Vec<String>,
    inputs: &'a mut BTreeSet<Uuid>,
    remaining_bytes: &'a mut AnalysisBudget,
    scopes: Vec<(BTreeSet<String>, Option<With>)>,
    nodes: usize,
}
impl VisitorMut for Analyzer<'_> {
    type Break = ();
    fn pre_visit_query(&mut self, query: &mut Query) -> ControlFlow<()> {
        if !query.locks.is_empty()
            || query.for_clause.is_some()
            || query.settings.is_some()
            || query.format_clause.is_some()
            || !query.pipe_operators.is_empty()
            || !safe_set(&query.body)
        {
            return ControlFlow::Break(());
        }
        self.scopes.push((BTreeSet::new(), None));
        let mut retained_with = query.with.take();
        if let Some(with) = &mut retained_with {
            if with.recursive {
                return ControlFlow::Break(());
            }
            for cte in &mut with.cte_tables {
                let name = cte.alias.name.value.to_ascii_lowercase();
                if self.bindings.contains_key(&name)
                    || self.scopes.iter().any(|(names, _)| names.contains(&name))
                {
                    return ControlFlow::Break(());
                }
                // A nonrecursive CTE sees only earlier CTEs and outer scopes.
                // Visit here, then hide WITH from the generated traversal to
                // avoid visiting and rewriting its input relations twice.
                if VisitMut::visit(&mut cte.query, self).is_break() {
                    return ControlFlow::Break(());
                }
                self.scopes.last_mut().unwrap().0.insert(name);
            }
        }
        self.scopes.last_mut().unwrap().1 = retained_with;
        ControlFlow::Continue(())
    }
    fn post_visit_query(&mut self, query: &mut Query) -> ControlFlow<()> {
        query.with = self.scopes.pop().unwrap().1;
        ControlFlow::Continue(())
    }
    fn pre_visit_select(&mut self, select: &mut Select) -> ControlFlow<()> {
        if select.into.is_some()
            || !select.optimizer_hints.is_empty()
            || select.select_modifiers.is_some()
            || select.top.is_some()
            || !select.lateral_views.is_empty()
            || select.prewhere.is_some()
            || !select.connect_by.is_empty()
            || !select.cluster_by.is_empty()
            || !select.distribute_by.is_empty()
            || !select.sort_by.is_empty()
            || select.value_table_mode.is_some()
        {
            return ControlFlow::Break(());
        }
        ControlFlow::Continue(())
    }
    fn pre_visit_statement(&mut self, statement: &mut Statement) -> ControlFlow<()> {
        if matches!(statement, Statement::Query(_)) {
            ControlFlow::Continue(())
        } else {
            ControlFlow::Break(())
        }
    }
    fn pre_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow<()> {
        self.nodes += 1;
        if self.nodes > 8192 {
            return ControlFlow::Break(());
        }
        if let Expr::CompoundIdentifier(parts) = expr
            && parts.len() > 2
        {
            let prefix = parts[..parts.len() - 1]
                .iter()
                .map(|p| p.value.to_ascii_lowercase())
                .collect::<Vec<_>>()
                .join("\0");
            if self.bindings.contains_key(&prefix) {
                *parts = parts[parts.len() - 2..].to_vec();
            }
        }
        let allowed = match expr {
            Expr::Identifier(_)
            | Expr::CompoundIdentifier(_)
            | Expr::Value(_)
            | Expr::Nested(_)
            | Expr::IsNull(_)
            | Expr::IsNotNull(_)
            | Expr::IsTrue(_)
            | Expr::IsFalse(_)
            | Expr::IsNotTrue(_)
            | Expr::IsNotFalse(_)
            | Expr::IsUnknown(_)
            | Expr::IsNotUnknown(_)
            | Expr::IsDistinctFrom(..)
            | Expr::IsNotDistinctFrom(..)
            | Expr::InList { .. }
            | Expr::InSubquery { .. }
            | Expr::Between { .. }
            | Expr::Like { .. }
            | Expr::ILike { .. }
            | Expr::Cast { .. }
            | Expr::Extract { .. }
            | Expr::Ceil { .. }
            | Expr::Floor { .. }
            | Expr::Position { .. }
            | Expr::Substring { .. }
            | Expr::Trim { .. }
            | Expr::TypedString(_)
            | Expr::Case { .. }
            | Expr::Exists { .. }
            | Expr::Subquery(_)
            | Expr::Tuple(_)
            | Expr::Interval(_)
            | Expr::Wildcard(_)
            | Expr::QualifiedWildcard(..) => true,
            Expr::BinaryOp { op, .. } => matches!(
                op.to_string().as_str(),
                "+" | "-"
                    | "*"
                    | "/"
                    | "%"
                    | "="
                    | "<>"
                    | "!="
                    | "<"
                    | ">"
                    | "<="
                    | ">="
                    | "AND"
                    | "OR"
                    | "||"
            ),
            Expr::UnaryOp { op, .. } => matches!(op.to_string().as_str(), "+" | "-" | "NOT"),
            Expr::Function(function) => {
                single_name(&function.name).is_some_and(|name| {
                    matches!(
                        name.as_str(),
                        "abs"
                            | "ceil"
                            | "floor"
                            | "round"
                            | "coalesce"
                            | "nullif"
                            | "length"
                            | "lower"
                            | "upper"
                            | "trim"
                            | "substring"
                            | "count"
                            | "sum"
                            | "avg"
                            | "min"
                            | "max"
                            | "stddev"
                            | "stddev_samp"
                            | "stddev_pop"
                            | "variance"
                            | "row_number"
                            | "rank"
                            | "dense_rank"
                            | "lag"
                            | "lead"
                            | "first_value"
                            | "last_value"
                    )
                }) && !function.uses_odbc_syntax
                    && matches!(function.parameters, FunctionArguments::None)
            }
            _ => false,
        };
        if allowed {
            ControlFlow::Continue(())
        } else {
            ControlFlow::Break(())
        }
    }
    fn post_visit_table_factor(&mut self, table: &mut TableFactor) -> ControlFlow<()> {
        let TableFactor::Table {
            name,
            alias,
            args,
            with_hints,
            version,
            with_ordinality,
            partitions,
            json_path,
            sample,
            index_hints,
        } = table
        else {
            return if matches!(
                table,
                TableFactor::Derived { sample: None, .. } | TableFactor::NestedJoin { .. }
            ) {
                ControlFlow::Continue(())
            } else {
                ControlFlow::Break(())
            };
        };
        if args.is_some()
            || !with_hints.is_empty()
            || version.is_some()
            || *with_ordinality
            || !partitions.is_empty()
            || json_path.is_some()
            || sample.is_some()
            || !index_hints.is_empty()
        {
            return ControlFlow::Break(());
        }
        let Some(key) = relation_name(name) else {
            return ControlFlow::Break(());
        };
        if self
            .scopes
            .iter()
            .rev()
            .any(|(names, _)| names.contains(&key))
        {
            return ControlFlow::Continue(());
        }
        let Some(binding) = self.bindings.get(&key) else {
            return ControlFlow::Break(());
        };
        let retained_alias = alias.clone().unwrap_or_else(|| TableAlias {
            explicit: true,
            name: Ident::with_quote('"', key.rsplit('\0').next().unwrap().to_string()),
            columns: Vec::new(),
            at: None,
        });
        match binding {
            QueryBinding::Dataset(version) => {
                self.inputs.insert(*version);
                *name = ObjectName::from(vec![Ident::with_quote(
                    '"',
                    format!("input_{}", version.simple()),
                )]);
                *alias = Some(retained_alias);
            }
            QueryBinding::View(sql) => {
                if self.views.contains(&key) {
                    return ControlFlow::Break(());
                }
                self.views.push(key);
                let expanded = expand(
                    sql,
                    self.bindings,
                    self.views,
                    self.inputs,
                    self.remaining_bytes,
                );
                self.views.pop();
                let Ok(subquery) = expanded else {
                    return ControlFlow::Break(());
                };
                *table = TableFactor::Derived {
                    lateral: false,
                    subquery,
                    alias: Some(retained_alias),
                    sample: None,
                };
            }
        }
        ControlFlow::Continue(())
    }
}
fn relation_name(name: &ObjectName) -> Option<String> {
    if name.0.is_empty() || name.0.len() > 3 {
        return None;
    }
    name.0
        .iter()
        .map(|p| match p {
            ObjectNamePart::Identifier(id) => Some(id.value.to_ascii_lowercase()),
            _ => None,
        })
        .collect::<Option<Vec<_>>>()
        .map(|parts| parts.join("\0"))
}

/// Bind the existing Dataset SQL-transform source using its canonical relation
/// and logical alias. Any additional or external relation must fail resolution.
pub fn analyze_dataset_transform(
    sql: &str,
    study: ahri_tre_types::StudyId,
    name: &ahri_tre_types::NcName,
    version: ahri_tre_types::VersionId,
    major: i32,
    minor: i32,
    patch: i32,
) -> Result<RestrictedQueryPlan, LakeError> {
    let relation =
        crate::dataset_loading::dataset_table_relation(study, name.as_str(), major, minor, patch);
    let bindings = BTreeMap::from([
        (
            name.as_str().to_ascii_lowercase(),
            QueryBinding::Dataset(version.0),
        ),
        (
            relation.table_name.to_ascii_lowercase(),
            QueryBinding::Dataset(version.0),
        ),
        (
            format!("{}\0{}", relation.schema_name, relation.table_name).to_ascii_lowercase(),
            QueryBinding::Dataset(version.0),
        ),
        (
            format!(
                "{}\0{}\0{}",
                crate::LAKE_ALIAS,
                relation.schema_name,
                relation.table_name
            )
            .to_ascii_lowercase(),
            QueryBinding::Dataset(version.0),
        ),
    ]);
    analyze_disclosure_query(sql, &bindings)
}

fn single_name(name: &ObjectName) -> Option<String> {
    match name.0.as_slice() {
        [ObjectNamePart::Identifier(id)] => Some(id.value.to_ascii_lowercase()),
        _ => None,
    }
}
fn safe_set(set: &SetExpr) -> bool {
    match set {
        SetExpr::Select(_) | SetExpr::Query(_) | SetExpr::Values(_) => true,
        SetExpr::SetOperation { left, right, .. } => safe_set(left) && safe_set(right),
        _ => false,
    }
}