micromegas_analytics/lakehouse/
remove_query_denial_udf.rs

1//! `remove_query_denial(rule_id)` -- admin scalar UDF that deletes a query-deny-list rule. The
2//! audit log is the durable record of what was denied and what it rejected, so the row itself
3//! does not need to survive its removal (hard delete, no `removed_at`/`removed_by` trail).
4
5use async_trait::async_trait;
6use datafusion::{
7    arrow::{
8        array::{Array, StringArray, StringBuilder},
9        datatypes::DataType,
10    },
11    common::internal_err,
12    error::DataFusionError,
13    logical_expr::{
14        ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
15        async_udf::AsyncScalarUDFImpl,
16    },
17};
18use std::sync::Arc;
19use uuid::Uuid;
20
21use super::query_deny_list::QueryDenyList;
22
23/// A scalar UDF that removes a single query-deny-list rule by id.
24#[derive(Debug)]
25pub struct RemoveQueryDenial {
26    signature: Signature,
27    query_denials: Arc<QueryDenyList>,
28}
29
30impl PartialEq for RemoveQueryDenial {
31    fn eq(&self, other: &Self) -> bool {
32        self.signature == other.signature
33    }
34}
35
36impl Eq for RemoveQueryDenial {}
37
38impl std::hash::Hash for RemoveQueryDenial {
39    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
40        self.signature.hash(state);
41    }
42}
43
44impl RemoveQueryDenial {
45    pub fn new(query_denials: Arc<QueryDenyList>) -> Self {
46        Self {
47            signature: Signature::exact(vec![DataType::Utf8], Volatility::Volatile),
48            query_denials,
49        }
50    }
51}
52
53impl ScalarUDFImpl for RemoveQueryDenial {
54    fn name(&self) -> &str {
55        "remove_query_denial"
56    }
57
58    fn signature(&self) -> &Signature {
59        &self.signature
60    }
61
62    fn return_type(&self, _arg_types: &[DataType]) -> datafusion::error::Result<DataType> {
63        Ok(DataType::Utf8)
64    }
65
66    fn invoke_with_args(
67        &self,
68        _args: ScalarFunctionArgs,
69    ) -> datafusion::error::Result<ColumnarValue> {
70        Err(DataFusionError::NotImplemented(
71            "remove_query_denial can only be called from async contexts".into(),
72        ))
73    }
74}
75
76#[async_trait]
77impl AsyncScalarUDFImpl for RemoveQueryDenial {
78    async fn invoke_async_with_args(
79        &self,
80        args: ScalarFunctionArgs,
81    ) -> datafusion::error::Result<ColumnarValue> {
82        let args = ColumnarValue::values_to_arrays(&args.args)?;
83        if args.len() != 1 {
84            return internal_err!("remove_query_denial expects exactly 1 argument: rule_id");
85        }
86        let rule_ids: &StringArray = args[0].as_any().downcast_ref::<_>().ok_or_else(|| {
87            DataFusionError::Internal("error casting rule_id argument as StringArray".into())
88        })?;
89
90        let mut builder = StringBuilder::with_capacity(rule_ids.len(), 64);
91        for index in 0..rule_ids.len() {
92            if rule_ids.is_null(index) {
93                builder.append_value("ERROR: rule_id cannot be null");
94                continue;
95            }
96            let raw = rule_ids.value(index);
97            let message = match Uuid::parse_str(raw) {
98                Err(_) => format!("ERROR: '{raw}' is not a valid rule id"),
99                Ok(rule_id) => match self.query_denials.delete(rule_id).await {
100                    Ok(true) => format!("SUCCESS: removed rule {rule_id}"),
101                    Ok(false) => format!("ERROR: no such rule: {rule_id}"),
102                    Err(e) => format!("ERROR: {e:?}"),
103                },
104            };
105            builder.append_value(message);
106        }
107
108        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
109    }
110}
111
112/// Creates a user-defined function to remove a single query-deny-list rule by id.
113///
114/// # Usage
115/// ```sql
116/// SELECT remove_query_denial('9f2c41ab-73de-0155-...') as result;
117/// ```
118pub fn make_remove_query_denial_udf(
119    query_denials: Arc<QueryDenyList>,
120) -> datafusion::logical_expr::async_udf::AsyncScalarUDF {
121    datafusion::logical_expr::async_udf::AsyncScalarUDF::new(Arc::new(RemoveQueryDenial::new(
122        query_denials,
123    )))
124}
micromegas_analytics/lakehouse/remove_query_denial_udf.rs

micromegas_analytics/lakehouse/
remove_query_denial_udf.rs
