micromegas_analytics/lakehouse/
list_view_set_definitions_table_function.rs

1//! `list_view_set_definitions()` -- admin UDTF listing every row of `lakehouse_view_set_definitions`,
2//! straight from Postgres rather than from `ViewRegistry::current()` -- so an admin can see a
3//! definition that exists on disk but failed to load (present here, absent from
4//! `list_view_sets()`). Registered inside the same admin-gated block as `list_query_denials()`'s
5//! siblings.
6
7use async_trait::async_trait;
8use datafusion::arrow::array::{Int32Array, RecordBatch, StringArray, TimestampNanosecondArray};
9use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
10use datafusion::catalog::Session;
11use datafusion::catalog::TableFunctionArgs;
12use datafusion::catalog::TableFunctionImpl;
13use datafusion::catalog::TableProvider;
14use datafusion::datasource::TableType;
15use datafusion::datasource::memory::{DataSourceExec, MemorySourceConfig};
16use datafusion::error::DataFusionError;
17use datafusion::physical_plan::ExecutionPlan;
18use datafusion::prelude::Expr;
19use micromegas_ingestion::data_lake_connection::DataLakeConnection;
20use sqlx::Row;
21use std::sync::Arc;
22
23fn schema() -> SchemaRef {
24    Arc::new(Schema::new(vec![
25        Field::new("view_set_name", DataType::Utf8, false),
26        Field::new("definition_sql", DataType::Utf8, false),
27        Field::new("update_group", DataType::Int32, false),
28        Field::new(
29            "updated_at",
30            DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
31            false,
32        ),
33        Field::new("updated_by", DataType::Utf8, true),
34    ]))
35}
36
37/// A DataFusion `TableFunctionImpl` for `list_view_set_definitions()`.
38#[derive(Debug)]
39pub struct ListViewSetDefinitionsTableFunction {
40    lake: Arc<DataLakeConnection>,
41}
42
43impl ListViewSetDefinitionsTableFunction {
44    pub fn new(lake: Arc<DataLakeConnection>) -> Self {
45        Self { lake }
46    }
47}
48
49impl TableFunctionImpl for ListViewSetDefinitionsTableFunction {
50    fn call_with_args(
51        &self,
52        _args: TableFunctionArgs,
53    ) -> datafusion::error::Result<Arc<dyn TableProvider>> {
54        Ok(Arc::new(ListViewSetDefinitionsTableProvider {
55            lake: self.lake.clone(),
56        }))
57    }
58}
59
60#[derive(Debug)]
61struct ListViewSetDefinitionsTableProvider {
62    lake: Arc<DataLakeConnection>,
63}
64
65#[async_trait]
66impl TableProvider for ListViewSetDefinitionsTableProvider {
67    fn schema(&self) -> SchemaRef {
68        schema()
69    }
70
71    fn table_type(&self) -> TableType {
72        TableType::Temporary
73    }
74
75    async fn scan(
76        &self,
77        _state: &dyn Session,
78        projection: Option<&Vec<usize>>,
79        _filters: &[Expr],
80        limit: Option<usize>,
81    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
82        let rows = sqlx::query(
83            "SELECT view_set_name, definition_sql, update_group, updated_at, updated_by \
84             FROM lakehouse_view_set_definitions ORDER BY view_set_name",
85        )
86        .fetch_all(&self.lake.db_pool)
87        .await
88        .map_err(|e| DataFusionError::External(e.into()))?;
89        let mut view_set_name = Vec::with_capacity(rows.len());
90        let mut definition_sql = Vec::with_capacity(rows.len());
91        let mut update_group = Vec::with_capacity(rows.len());
92        let mut updated_at = Vec::with_capacity(rows.len());
93        let mut updated_by = Vec::with_capacity(rows.len());
94        for row in &rows {
95            view_set_name.push(
96                row.try_get::<String, _>("view_set_name")
97                    .map_err(|e| DataFusionError::External(e.into()))?,
98            );
99            definition_sql.push(
100                row.try_get::<String, _>("definition_sql")
101                    .map_err(|e| DataFusionError::External(e.into()))?,
102            );
103            update_group.push(
104                row.try_get::<i32, _>("update_group")
105                    .map_err(|e| DataFusionError::External(e.into()))?,
106            );
107            let ts: chrono::DateTime<chrono::Utc> = row
108                .try_get("updated_at")
109                .map_err(|e| DataFusionError::External(e.into()))?;
110            updated_at.push(ts.timestamp_nanos_opt().unwrap_or_default());
111            updated_by.push(
112                row.try_get::<Option<String>, _>("updated_by")
113                    .map_err(|e| DataFusionError::External(e.into()))?,
114            );
115        }
116        if let Some(n) = limit {
117            view_set_name.truncate(n);
118            definition_sql.truncate(n);
119            update_group.truncate(n);
120            updated_at.truncate(n);
121            updated_by.truncate(n);
122        }
123        let rb = RecordBatch::try_new(
124            schema(),
125            vec![
126                Arc::new(StringArray::from(view_set_name)),
127                Arc::new(StringArray::from(definition_sql)),
128                Arc::new(Int32Array::from(update_group)),
129                Arc::new(
130                    TimestampNanosecondArray::from(updated_at).with_timezone("+00:00".to_string()),
131                ),
132                Arc::new(StringArray::from(updated_by)),
133            ],
134        )?;
135        let source =
136            MemorySourceConfig::try_new(&[vec![rb]], schema(), projection.map(|v| v.to_owned()))?;
137        Ok(DataSourceExec::from_data_source(source))
138    }
139}
micromegas_analytics/lakehouse/list_view_set_definitions_table_function.rs

micromegas_analytics/lakehouse/
list_view_set_definitions_table_function.rs
