micromegas_analytics/lakehouse/
deny_queries_table_function.rs

1//! `deny_queries(match_expr, reason)` -- admin UDTF that validates and inserts a new query-deny-
2//! list rule. `call_with_args` runs every validation
3//! synchronously and fail-loud with `plan_err!`: expression compilation (`compile_match_expr`,
4//! against the bare `SessionContext` `QueryDenyList` holds for exactly this, not the caller's own
5//! session -- rule compilation must produce the same result on every replica regardless of whose
6//! session triggered it), the empty-`reason` check, the caller-identity check
7//! (`CallerContext::identity` must be `Some`), and the rule-count check against the current
8//! snapshot. Only the DB write and the local snapshot refresh happen in the async body, behind
9//! `LogStreamTableProvider`/`TaskLogExecPlan` -- `QueryDenyList::insert` takes the already-compiled
10//! expression and cannot itself fail on a bad expression.
11
12use super::query_deny_list::{QueryDenyList, compile_match_expr, max_rules};
13use crate::dfext::expressions::exp_to_string;
14use crate::dfext::log_stream_table_provider::LogStreamTableProvider;
15use crate::dfext::task_log_exec_plan::TaskLogExecPlan;
16use datafusion::catalog::TableFunctionArgs;
17use datafusion::catalog::TableFunctionImpl;
18use datafusion::catalog::TableProvider;
19use datafusion::common::plan_err;
20use micromegas_tracing::prelude::*;
21use std::sync::Arc;
22
23/// A DataFusion `TableFunctionImpl` for `deny_queries(match_expr, reason)`.
24#[derive(Debug)]
25pub struct DenyQueriesTableFunction {
26    query_denials: Arc<QueryDenyList>,
27    /// `CallerContext::identity`, captured at registration time (once per request, in
28    /// `register_lakehouse_functions`) -- `deny_queries` requires `Some`.
29    identity: Option<String>,
30}
31
32impl DenyQueriesTableFunction {
33    pub fn new(query_denials: Arc<QueryDenyList>, identity: Option<String>) -> Self {
34        Self {
35            query_denials,
36            identity,
37        }
38    }
39}
40
41impl TableFunctionImpl for DenyQueriesTableFunction {
42    fn call_with_args(
43        &self,
44        args: TableFunctionArgs,
45    ) -> datafusion::error::Result<Arc<dyn TableProvider>> {
46        let exprs = args.exprs();
47        let Some(match_expr) = exprs.first().map(exp_to_string).transpose()? else {
48            return plan_err!("Missing first argument, expected match_expr: String");
49        };
50        let Some(reason) = exprs.get(1).map(exp_to_string).transpose()? else {
51            return plan_err!("Missing 2nd argument, expected reason: String");
52        };
53        if reason.trim().is_empty() {
54            return plan_err!("deny_queries: reason must not be empty");
55        }
56        let Some(created_by) = self.identity.clone() else {
57            return plan_err!(
58                "deny_queries requires an authenticated caller identity, which this session \
59                 does not carry"
60            );
61        };
62        // Checked against the local snapshot only, with no re-check under `insert`'s write lock
63        // and no DB-side constraint -- so concurrent inserts (or another replica's insert within
64        // one refresh tick) can overshoot the cap by a rule or two. Advisory by design: the cap
65        // bounds per-query evaluation cost (~0.45 us/rule), where an off-by-a-few is harmless.
66        let rule_count = self.query_denials.rule_count();
67        let cap = max_rules();
68        if rule_count >= cap {
69            return plan_err!(
70                "deny_queries: at the cap of {cap} rules (MICROMEGAS_QUERY_DENY_MAX_RULES); \
71                 remove an existing rule with remove_query_denial(rule_id) before adding another"
72            );
73        }
74        // Compiled here, synchronously, against the bare `SessionContext` the list holds for
75        // exactly this -- not `args.session()` -- so compilation is independent of which
76        // replica's session triggered it.
77        let compiled = compile_match_expr(self.query_denials.compile_ctx(), &match_expr)?;
78
79        let query_denials = self.query_denials.clone();
80        let spawner = move || {
81            let (tx, rx) = tokio::sync::mpsc::channel(1);
82            spawn_with_context(async move {
83                match query_denials
84                    .insert(&match_expr, compiled, &reason, &created_by)
85                    .await
86                {
87                    Ok(row) => {
88                        info!("deny_queries: inserted rule {}", row.rule_id);
89                        let _ = tx
90                            .send(Ok((chrono::Utc::now(), row.rule_id.to_string())))
91                            .await;
92                    }
93                    Err(e) => {
94                        error!("deny_queries: insert failed: {e:?}");
95                        let _ = tx.send(Err(format!("{e:?}"))).await;
96                    }
97                }
98            });
99            rx
100        };
101
102        Ok(Arc::new(LogStreamTableProvider {
103            log_stream: Arc::new(TaskLogExecPlan::new(Box::new(spawner))),
104        }))
105    }
106}
micromegas_analytics/lakehouse/deny_queries_table_function.rs

micromegas_analytics/lakehouse/
deny_queries_table_function.rs
