micromegas_analytics/lakehouse/
get_payload_function.rs

1use super::audience_guard::{AudienceGuard, IdKind};
2use async_trait::async_trait;
3use datafusion::{
4    arrow::{
5        array::{Array, BinaryBuilder, StringArray},
6        datatypes::DataType,
7    },
8    common::{internal_err, not_impl_err, plan_err},
9    error::DataFusionError,
10    logical_expr::{
11        ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
12        async_udf::AsyncScalarUDFImpl,
13    },
14};
15use futures::stream::StreamExt;
16use micromegas_ingestion::data_lake_connection::DataLakeConnection;
17use micromegas_tracing::prelude::*;
18use std::collections::HashSet;
19use std::sync::Arc;
20use uuid::Uuid;
21
22/// A scalar UDF that retrieves the payload of a block from the data lake.
23#[derive(Debug)]
24pub struct GetPayload {
25    signature: Signature,
26    lake: Arc<DataLakeConnection>,
27    guard: Arc<AudienceGuard>,
28}
29
30impl PartialEq for GetPayload {
31    fn eq(&self, other: &Self) -> bool {
32        self.signature == other.signature
33    }
34}
35
36impl Eq for GetPayload {}
37
38impl std::hash::Hash for GetPayload {
39    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
40        self.signature.hash(state);
41    }
42}
43
44impl GetPayload {
45    pub fn new(lake: Arc<DataLakeConnection>, guard: Arc<AudienceGuard>) -> Self {
46        Self {
47            signature: Signature::exact(
48                vec![DataType::Utf8, DataType::Utf8, DataType::Utf8],
49                Volatility::Immutable,
50            ),
51            lake,
52            guard,
53        }
54    }
55}
56
57impl ScalarUDFImpl for GetPayload {
58    fn name(&self) -> &str {
59        "get_payload"
60    }
61
62    fn signature(&self) -> &Signature {
63        &self.signature
64    }
65
66    fn return_type(&self, _arg_types: &[DataType]) -> datafusion::error::Result<DataType> {
67        Ok(DataType::Binary)
68    }
69
70    fn invoke_with_args(
71        &self,
72        _args: datafusion::logical_expr::ScalarFunctionArgs,
73    ) -> datafusion::error::Result<ColumnarValue> {
74        not_impl_err!("GetPayload can only be called from async contexts")
75    }
76}
77
78#[async_trait]
79impl AsyncScalarUDFImpl for GetPayload {
80    async fn invoke_async_with_args(
81        &self,
82        args: ScalarFunctionArgs,
83    ) -> datafusion::error::Result<ColumnarValue> {
84        let args = ColumnarValue::values_to_arrays(&args.args)?;
85        if args.len() != 3 {
86            return internal_err!("wrong number of arguments to get_payload()");
87        }
88        let process_ids = args[0]
89            .as_any()
90            .downcast_ref::<StringArray>()
91            .ok_or_else(|| {
92                DataFusionError::Internal("downcasting process_ids in GetPayload".into())
93            })?
94            .clone();
95        let stream_ids = args[1]
96            .as_any()
97            .downcast_ref::<StringArray>()
98            .ok_or_else(|| {
99                DataFusionError::Internal("downcasting stream_ids in GetPayload".into())
100            })?
101            .clone();
102        let block_ids = args[2]
103            .as_any()
104            .downcast_ref::<StringArray>()
105            .ok_or_else(|| DataFusionError::Internal("downcasting block_ids in GetPayload".into()))?
106            .clone();
107
108        // The call-level guard: `get_payload` reads `blobs/{process_id}/{stream_id}
109        // /{block_id}` directly out of object storage, bypassing the lakehouse entirely -- arg 1
110        // (`process_id`) is therefore the whole check, and a complete one: a caller who names a
111        // readable process cannot reach another process's blob, because the foreign block simply
112        // isn't under that prefix. Distinct ids only (one resolution per unique process, not per
113        // row); a value that doesn't even parse as a UUID is denied the same way an unreadable
114        // one is, since it's the one input that could otherwise escape the `blobs/{process_id}`
115        // prefix the completeness argument above relies on. All-or-nothing over the batch, never
116        // a per-row `NULL`: a partially-filtered binary column would be indistinguishable from a
117        // missing payload.
118        let mut distinct_process_ids: Vec<Uuid> = Vec::new();
119        let mut seen: HashSet<Uuid> = HashSet::new();
120        for i in 0..process_ids.len() {
121            let raw = process_ids.value(i);
122            let Ok(id) = Uuid::parse_str(raw) else {
123                return plan_err!("get_payload: '{raw}' is not a valid process id");
124            };
125            if seen.insert(id) {
126                distinct_process_ids.push(id);
127            }
128        }
129        let readable = self
130            .guard
131            .readable_ids(&distinct_process_ids, IdKind::Process)
132            .await?;
133        for id in &distinct_process_ids {
134            if !readable.contains(id) {
135                return plan_err!("get_payload: process '{id}' not found or not accessible");
136            }
137        }
138
139        let lake = self.lake.clone();
140        let mut stream = futures::stream::iter(0..process_ids.len())
141            .map(|i| {
142                let process_id = process_ids.value(i);
143                let stream_id = stream_ids.value(i);
144                let block_id = block_ids.value(i);
145                let obj_path = format!("blobs/{process_id}/{stream_id}/{block_id}");
146                let lake = lake.clone();
147                spawn_with_context(async move { lake.blob_storage.read_blob(&obj_path).await })
148            })
149            .buffered(10);
150        let mut result_builder = BinaryBuilder::with_capacity(block_ids.len(), 1024 * 1024);
151        while let Some(res) = stream.next().await {
152            result_builder.append_value(
153                res.map_err(|e| {
154                    DataFusionError::Internal(format!("error downloading payload: {e:?}"))
155                })?
156                .map_err(|e| {
157                    DataFusionError::Internal(format!("error downloading payload: {e:?}"))
158                })?,
159            );
160        }
161        Ok(ColumnarValue::Array(Arc::new(result_builder.finish())))
162    }
163}
micromegas_analytics/lakehouse/get_payload_function.rs

micromegas_analytics/lakehouse/
get_payload_function.rs
