micromegas_analytics/lakehouse/
list_query_denials_table_function.rs

1//! `list_query_denials()` -- admin UDTF listing every query-deny-list rule currently in force.
2//! Registered inside the same admin-gated block as `list_partitions()`'s mutating siblings.
3
4use super::query_deny_list::QueryDenyList;
5use async_trait::async_trait;
6use datafusion::arrow::array::{RecordBatch, StringArray, TimestampNanosecondArray};
7use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
8use datafusion::catalog::Session;
9use datafusion::catalog::TableFunctionArgs;
10use datafusion::catalog::TableFunctionImpl;
11use datafusion::catalog::TableProvider;
12use datafusion::datasource::TableType;
13use datafusion::datasource::memory::{DataSourceExec, MemorySourceConfig};
14use datafusion::error::DataFusionError;
15use datafusion::physical_plan::ExecutionPlan;
16use datafusion::prelude::Expr;
17use std::sync::Arc;
18
19fn schema() -> SchemaRef {
20    Arc::new(Schema::new(vec![
21        Field::new("rule_id", DataType::Utf8, false),
22        Field::new(
23            "created_at",
24            DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
25            false,
26        ),
27        Field::new("created_by", DataType::Utf8, false),
28        Field::new("reason", DataType::Utf8, false),
29        Field::new("match_expr", DataType::Utf8, false),
30        Field::new(
31            "last_hit_at",
32            DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
33            true,
34        ),
35    ]))
36}
37
38/// A DataFusion `TableFunctionImpl` for `list_query_denials()`.
39#[derive(Debug)]
40pub struct ListQueryDenialsTableFunction {
41    query_denials: Arc<QueryDenyList>,
42}
43
44impl ListQueryDenialsTableFunction {
45    pub fn new(query_denials: Arc<QueryDenyList>) -> Self {
46        Self { query_denials }
47    }
48}
49
50impl TableFunctionImpl for ListQueryDenialsTableFunction {
51    fn call_with_args(
52        &self,
53        _args: TableFunctionArgs,
54    ) -> datafusion::error::Result<Arc<dyn TableProvider>> {
55        Ok(Arc::new(ListQueryDenialsTableProvider {
56            query_denials: self.query_denials.clone(),
57        }))
58    }
59}
60
61#[derive(Debug)]
62struct ListQueryDenialsTableProvider {
63    query_denials: Arc<QueryDenyList>,
64}
65
66#[async_trait]
67impl TableProvider for ListQueryDenialsTableProvider {
68    fn schema(&self) -> SchemaRef {
69        schema()
70    }
71
72    fn table_type(&self) -> TableType {
73        TableType::Temporary
74    }
75
76    async fn scan(
77        &self,
78        _state: &dyn Session,
79        projection: Option<&Vec<usize>>,
80        _filters: &[Expr],
81        limit: Option<usize>,
82    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
83        let mut rows = self
84            .query_denials
85            .list()
86            .await
87            .map_err(|e| DataFusionError::External(e.into()))?;
88        if let Some(n) = limit {
89            rows.truncate(n);
90        }
91        let rule_id: StringArray = rows.iter().map(|r| Some(r.rule_id.to_string())).collect();
92        let created_at: TimestampNanosecondArray = rows
93            .iter()
94            .map(|r| r.created_at.timestamp_nanos_opt())
95            .collect::<TimestampNanosecondArray>()
96            .with_timezone("+00:00".to_string());
97        let created_by: StringArray = rows.iter().map(|r| Some(r.created_by.clone())).collect();
98        let reason: StringArray = rows.iter().map(|r| Some(r.reason.clone())).collect();
99        let match_expr: StringArray = rows.iter().map(|r| Some(r.match_expr.clone())).collect();
100        let last_hit_at: TimestampNanosecondArray = rows
101            .iter()
102            .map(|r| r.last_hit_at.and_then(|t| t.timestamp_nanos_opt()))
103            .collect::<TimestampNanosecondArray>()
104            .with_timezone("+00:00".to_string());
105        let rb = RecordBatch::try_new(
106            schema(),
107            vec![
108                Arc::new(rule_id),
109                Arc::new(created_at),
110                Arc::new(created_by),
111                Arc::new(reason),
112                Arc::new(match_expr),
113                Arc::new(last_hit_at),
114            ],
115        )?;
116        let source =
117            MemorySourceConfig::try_new(&[vec![rb]], schema(), projection.map(|v| v.to_owned()))?;
118        Ok(DataSourceExec::from_data_source(source))
119    }
120}
micromegas_analytics/lakehouse/list_query_denials_table_function.rs

micromegas_analytics/lakehouse/
list_query_denials_table_function.rs
