Skip to main content

micromegas_analytics/lakehouse/
list_audience_grants_table_function.rs

1//! `list_audience_grants()` -- caller-scoped UDTF listing rows of the `audience_grants` table.
2//! Modeled on `list_query_denials_table_function.rs`'s shape (a
3//! `TableFunctionImpl` whose sync `call_with_args` returns a lazy `TableProvider`, an async
4//! `scan` that runs the DB query and builds arrays), but registered **outside** the admin gate
5//! in `register_lakehouse_functions`, next to `list_view_sets()` -- every authenticated caller
6//! (and, per `is_admin(metadata)`'s documented `--disable-auth` convention, every caller when
7//! auth is disabled) can call it, scoped by [`crate::lakehouse::list_audience_grants_table_function::GrantVisibility`].
8//!
9//! **Visibility.** Admin: every row. Non-admin: every grant on each `(audience, axis)` pair the
10//! caller holds a matching grant on -- deliberately wider than "rows whose selector matches me":
11//! if a caller may read an audience, they may see who else may, which is exactly the "who can
12//! see this audience" question the Audience Access page (and this function) exists to answer.
13//! An empty selector list (`CallerContext::internal()`, or an identity-less caller once the
14//! always-present `"*"` is stripped -- see `scan`) yields zero rows. A maintenance caller and a
15//! request with no `AuthContext` at all (`--disable-auth`) are both `is_admin: true` and take
16//! the `All` branch below instead, never `Held`.
17//!
18//! **What this is not.** This function reads the DB table directly on every call, never
19//! `DbAudienceGrantsSource`'s TTL snapshot, so a write is visible immediately. It also does not
20//! (and structurally cannot) apply `analytics-web-srv`'s `self_service_mint_enabled` knob --
21//! see the REST `GET .../audience-grants/visible` route for the knob-aware narrowing that backs
22//! the Audience Access page's own list; this function is the always-on SQL auditing surface,
23//! like `list_query_denials()`.
24
25use async_trait::async_trait;
26use datafusion::arrow::array::{RecordBatch, StringArray, TimestampNanosecondArray};
27use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
28use datafusion::catalog::Session;
29use datafusion::catalog::TableFunctionArgs;
30use datafusion::catalog::TableFunctionImpl;
31use datafusion::catalog::TableProvider;
32use datafusion::datasource::TableType;
33use datafusion::datasource::memory::{DataSourceExec, MemorySourceConfig};
34use datafusion::error::DataFusionError;
35use datafusion::physical_plan::ExecutionPlan;
36use datafusion::prelude::Expr;
37use sqlx::Row;
38use std::sync::Arc;
39
40/// How much of the `audience_grants` table `scan` returns,
41/// decided once at registration time in `register_lakehouse_functions` from
42/// `CallerContext::is_admin`/`grant_selectors`.
43#[derive(Debug, Clone)]
44pub enum GrantVisibility {
45    /// Every row -- an admin caller.
46    All,
47    /// Every grant on each `(audience, axis)` pair covered by one of these selectors -- a
48    /// non-admin caller's `grant_selectors`, `"*"` included. Empty for `CallerContext::internal()`;
49    /// a maintenance caller and a request with no `AuthContext` at all are `is_admin: true` and
50    /// take [`Self::All`] instead, never reaching this variant. `scan` strips the leading `"*"`
51    /// before binding it into the held-pair query -- see the note there for why.
52    Held(Arc<[String]>),
53}
54
55fn schema() -> SchemaRef {
56    Arc::new(Schema::new(vec![
57        Field::new("audience", DataType::Utf8, false),
58        Field::new("axis", DataType::Utf8, false),
59        Field::new("selector", DataType::Utf8, false),
60        Field::new(
61            "created_at",
62            DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
63            false,
64        ),
65        Field::new("created_by", DataType::Utf8, false),
66    ]))
67}
68
69/// One `audience_grants` row, as returned by either of `scan`'s two queries.
70struct GrantRow {
71    audience: String,
72    axis: String,
73    selector: String,
74    created_at: chrono::DateTime<chrono::Utc>,
75    created_by: String,
76}
77
78/// Every row, admin-visibility query -- `ORDER BY audience, axis, selector` matches the schema's
79/// stable column order and gives deterministic output.
80const SELECT_ALL_SQL: &str = "
81    SELECT audience, axis, selector, created_at, created_by
82    FROM audience_grants
83    ORDER BY audience, axis, selector";
84
85/// Held-pair, non-admin-visibility query: every row on a `(audience, axis)` pair
86/// the caller holds a matching grant on, not just the caller's own rows -- deliberately wider,
87/// since "who else can see this" is the question this function answers, and it is the same set
88/// a non-admin may modify (`analytics-web-srv`'s write policy).
89const SELECT_HELD_SQL: &str = "
90    SELECT g.audience, g.axis, g.selector, g.created_at, g.created_by
91    FROM audience_grants g
92    WHERE EXISTS (
93      SELECT 1 FROM audience_grants h
94      WHERE h.audience = g.audience AND h.axis = g.axis AND h.selector = ANY($1)
95    )
96    ORDER BY g.audience, g.axis, g.selector";
97
98/// A DataFusion `TableFunctionImpl` for `list_audience_grants()`.
99#[derive(Debug)]
100pub struct ListAudienceGrantsTableFunction {
101    pool: sqlx::PgPool,
102    visibility: GrantVisibility,
103}
104
105impl ListAudienceGrantsTableFunction {
106    pub fn new(pool: sqlx::PgPool, visibility: GrantVisibility) -> Self {
107        Self { pool, visibility }
108    }
109}
110
111impl TableFunctionImpl for ListAudienceGrantsTableFunction {
112    fn call_with_args(
113        &self,
114        _args: TableFunctionArgs,
115    ) -> datafusion::error::Result<Arc<dyn TableProvider>> {
116        Ok(Arc::new(ListAudienceGrantsTableProvider {
117            pool: self.pool.clone(),
118            visibility: self.visibility.clone(),
119        }))
120    }
121}
122
123#[derive(Debug)]
124struct ListAudienceGrantsTableProvider {
125    pool: sqlx::PgPool,
126    visibility: GrantVisibility,
127}
128
129#[async_trait]
130impl TableProvider for ListAudienceGrantsTableProvider {
131    fn schema(&self) -> SchemaRef {
132        schema()
133    }
134
135    fn table_type(&self) -> TableType {
136        TableType::Temporary
137    }
138
139    async fn scan(
140        &self,
141        _state: &dyn Session,
142        projection: Option<&Vec<usize>>,
143        _filters: &[Expr],
144        limit: Option<usize>,
145    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
146        let mut rows = match &self.visibility {
147            GrantVisibility::All => sqlx::query(SELECT_ALL_SQL)
148                .fetch_all(&self.pool)
149                .await
150                .map_err(|e| DataFusionError::External(e.into()))?,
151            GrantVisibility::Held(selectors) => {
152                // `caller_selectors` always leads with `"*"` -- even for an identity-less caller
153                // (no email, no groups) -- so binding it unfiltered would match every pair that
154                // carries a `*` grant row and leak every sibling row on it. Strip `"*"` first,
155                // the same way `caller_holds_pair`'s write-side hold check does: a `*` grant is
156                // not something a caller individually *holds*, only something that lets them
157                // *read* via `AudienceReadPolicy`. An empty list after filtering (or an already-
158                // empty list, for internal/maintenance callers) still runs the query -- `=
159                // ANY($1)` over an empty array matches no row, giving zero rows rather than a
160                // separate short-circuit branch to keep in sync with the SQL above.
161                let identity_selectors: Vec<String> = selectors
162                    .iter()
163                    .filter(|s| s.as_str() != "*")
164                    .cloned()
165                    .collect();
166                sqlx::query(SELECT_HELD_SQL)
167                    .bind(&identity_selectors)
168                    .fetch_all(&self.pool)
169                    .await
170                    .map_err(|e| DataFusionError::External(e.into()))?
171            }
172        }
173        .into_iter()
174        .map(|row| {
175            Ok::<_, sqlx::Error>(GrantRow {
176                audience: row.try_get("audience")?,
177                axis: row.try_get("axis")?,
178                selector: row.try_get("selector")?,
179                created_at: row.try_get("created_at")?,
180                created_by: row.try_get("created_by")?,
181            })
182        })
183        .collect::<Result<Vec<_>, _>>()
184        .map_err(|e| DataFusionError::External(e.into()))?;
185        if let Some(n) = limit {
186            rows.truncate(n);
187        }
188        let audience: StringArray = rows.iter().map(|r| Some(r.audience.clone())).collect();
189        let axis: StringArray = rows.iter().map(|r| Some(r.axis.clone())).collect();
190        let selector: StringArray = rows.iter().map(|r| Some(r.selector.clone())).collect();
191        let created_at: TimestampNanosecondArray = rows
192            .iter()
193            .map(|r| r.created_at.timestamp_nanos_opt())
194            .collect::<TimestampNanosecondArray>()
195            .with_timezone("+00:00".to_string());
196        let created_by: StringArray = rows.iter().map(|r| Some(r.created_by.clone())).collect();
197        let rb = RecordBatch::try_new(
198            schema(),
199            vec![
200                Arc::new(audience),
201                Arc::new(axis),
202                Arc::new(selector),
203                Arc::new(created_at),
204                Arc::new(created_by),
205            ],
206        )?;
207        let source =
208            MemorySourceConfig::try_new(&[vec![rb]], schema(), projection.map(|v| v.to_owned()))?;
209        Ok(DataSourceExec::from_data_source(source))
210    }
211}