micromegas_analytics/lakehouse/
parse_block_table_function.rs

1use super::{
2    audience_guard::{AudienceGuard, IdKind},
3    block_object_decoder::{BlockObjectDecoderMap, ObjectVisitor, default_block_object_decoders},
4    lakehouse_context::LakehouseContext,
5    partition_cache::QueryPartitionProvider,
6    read_scope::CallerContext,
7    session_configurator::NoOpSessionConfigurator,
8    view_factory::ViewFactory,
9};
10use crate::{
11    dfext::{string_column_accessor::string_column_by_name, typed_column::typed_column_by_name},
12    metadata::StreamMetadata,
13    payload::fetch_block_payload,
14    time::TimeRange,
15};
16use anyhow::Context;
17use async_trait::async_trait;
18use datafusion::{
19    arrow::{
20        array::{BinaryBuilder, Int64Array, Int64Builder, RecordBatch, StringBuilder},
21        datatypes::{DataType, Field, Schema, SchemaRef},
22    },
23    catalog::{Session, TableFunctionArgs, TableFunctionImpl, TableProvider},
24    common::plan_err,
25    datasource::{
26        TableType,
27        memory::{DataSourceExec, MemorySourceConfig},
28    },
29    error::DataFusionError,
30    physical_plan::ExecutionPlan,
31    prelude::Expr,
32};
33use jsonb::Value as JsonbValue;
34use micromegas_transit::{UserDefinedType, value::Value as TransitValue};
35use std::{borrow::Cow, collections::BTreeMap, sync::Arc};
36use uuid::Uuid;
37
38use crate::dfext::expressions::exp_to_string;
39
40fn output_schema() -> SchemaRef {
41    Arc::new(Schema::new(vec![
42        Field::new("object_index", DataType::Int64, false),
43        Field::new("type_name", DataType::Utf8, false),
44        Field::new("value", DataType::Binary, false),
45    ]))
46}
47
48/// Converts a `transit::Value` to a `jsonb::Value`.
49pub fn transit_value_to_jsonb(value: TransitValue<'_>) -> JsonbValue<'_> {
50    match value {
51        TransitValue::String(s) => JsonbValue::String(Cow::Borrowed(s)),
52        TransitValue::Object(obj) => {
53            let mut map = BTreeMap::new();
54            map.insert(
55                "__type".to_string(),
56                JsonbValue::String(Cow::Borrowed(obj.type_name)),
57            );
58            for &(name, val) in obj.members {
59                map.insert(name.to_string(), transit_value_to_jsonb(val));
60            }
61            JsonbValue::Object(map)
62        }
63        TransitValue::U8(v) => JsonbValue::Number(jsonb::Number::UInt64(u64::from(v))),
64        TransitValue::U32(v) => JsonbValue::Number(jsonb::Number::UInt64(u64::from(v))),
65        TransitValue::U64(v) => JsonbValue::Number(jsonb::Number::UInt64(v)),
66        TransitValue::I64(v) => JsonbValue::Number(jsonb::Number::Int64(v)),
67        TransitValue::F64(v) => JsonbValue::Number(jsonb::Number::Float64(v)),
68        TransitValue::None => JsonbValue::Null,
69        TransitValue::Bytes(b) => {
70            JsonbValue::String(Cow::Owned(format!("<binary {} bytes>", b.len())))
71        }
72    }
73}
74
75/// Queries the global blocks view for a block's metadata and constructs a `StreamMetadata`.
76/// Returns `None` if the block is not found in `blocks` for the queried range.
77async fn fetch_block_metadata(
78    lakehouse: Arc<LakehouseContext>,
79    part_provider: Arc<dyn QueryPartitionProvider>,
80    query_range: Option<TimeRange>,
81    view_factory: Arc<ViewFactory>,
82    block_id: Uuid,
83    caller: CallerContext,
84) -> anyhow::Result<Option<(i64, String, StreamMetadata)>> {
85    // Runs under the witness's internal caller (`caller`, threaded in from `scan` after
86    // `AudienceGuard::authorize` succeeds), not the caller's own scope: the query below is
87    // server-constructed and confined to the single, already-authorized `block_id` -- if that
88    // block's process is readable, everything this statement can reach is readable too. A
89    // deliberate deviation from naive scope inheritance.
90    let ctx = super::query::make_session_context(
91        lakehouse,
92        part_provider,
93        query_range,
94        view_factory,
95        Arc::new(NoOpSessionConfigurator),
96        caller,
97    )
98    .await?;
99
100    // Interpolate the canonical hyphenated rendering (not the caller's original
101    // string) so a valid-but-non-canonical form (braced, URN, bare-hex) still
102    // matches the `blocks` view's canonical rendering.
103    let sql = format!(
104        "SELECT stream_id, process_id, object_offset,
105                \"streams.dependencies_metadata\", \"streams.objects_metadata\", \"streams.format\"
106         FROM blocks
107         WHERE block_id = '{block_id}'"
108    );
109    let df = ctx.sql(&sql).await?;
110    let batches = df.collect().await?;
111
112    if batches.is_empty() || batches[0].num_rows() == 0 {
113        return Ok(None);
114    }
115
116    let batch = &batches[0];
117
118    let stream_id_col = string_column_by_name(batch, "stream_id")?;
119    let process_id_col = string_column_by_name(batch, "process_id")?;
120    let object_offset_col: &Int64Array = typed_column_by_name(batch, "object_offset")?;
121    let format_col = string_column_by_name(batch, "streams.format")?;
122    let format = format_col.value(0)?.to_string();
123
124    let stream_id = Uuid::parse_str(stream_id_col.value(0)?)?;
125    let process_id = Uuid::parse_str(process_id_col.value(0)?)?;
126    let object_offset = object_offset_col.value(0);
127
128    let deps_col = batch
129        .column_by_name("streams.dependencies_metadata")
130        .context("streams.dependencies_metadata column not found")?;
131    let deps_binary: &datafusion::arrow::array::BinaryArray = deps_col
132        .as_any()
133        .downcast_ref()
134        .context("failed to cast dependencies_metadata to BinaryArray")?;
135    let deps_bytes = deps_binary.value(0);
136    let dependencies_metadata: Vec<UserDefinedType> =
137        ciborium::from_reader(deps_bytes).context("decoding dependencies_metadata")?;
138
139    let objs_col = batch
140        .column_by_name("streams.objects_metadata")
141        .context("streams.objects_metadata column not found")?;
142    let objs_binary: &datafusion::arrow::array::BinaryArray = objs_col
143        .as_any()
144        .downcast_ref()
145        .context("failed to cast objects_metadata to BinaryArray")?;
146    let objs_bytes = objs_binary.value(0);
147    let objects_metadata: Vec<UserDefinedType> =
148        ciborium::from_reader(objs_bytes).context("decoding objects_metadata")?;
149
150    let stream_metadata = StreamMetadata {
151        process_id,
152        stream_id,
153        dependencies_metadata,
154        objects_metadata,
155        tags: vec![],
156        properties: Arc::new(vec![]),
157    };
158
159    Ok(Some((object_offset, format, stream_metadata)))
160}
161
162/// `ObjectVisitor` that owns the Arrow builders, the running `object_index`,
163/// and the early-limit check for `parse_block` — row construction stays in one
164/// place regardless of which `BlockObjectDecoder` drives it.
165struct ParseBlockRowBuilder {
166    index_builder: Int64Builder,
167    name_builder: StringBuilder,
168    value_builder: BinaryBuilder,
169    object_offset: i64,
170    local_index: i64,
171    nb_objects: usize,
172    early_limit: Option<usize>,
173}
174
175impl ParseBlockRowBuilder {
176    fn new(object_offset: i64, early_limit: Option<usize>) -> Self {
177        Self {
178            index_builder: Int64Builder::new(),
179            name_builder: StringBuilder::new(),
180            value_builder: BinaryBuilder::new(),
181            object_offset,
182            local_index: 0,
183            nb_objects: 0,
184            early_limit,
185        }
186    }
187
188    /// Whether iteration should continue, based on rows emitted so far
189    /// (`nb_objects`) — mirrors the pre-registry `parse_block_objects` check,
190    /// which applied to both emitted and skipped entries alike.
191    fn continue_iterating(&self) -> bool {
192        match self.early_limit {
193            Some(lim) => self.nb_objects < lim,
194            None => true,
195        }
196    }
197
198    fn finish(mut self) -> anyhow::Result<RecordBatch> {
199        Ok(RecordBatch::try_new(
200            output_schema(),
201            vec![
202                Arc::new(self.index_builder.finish()),
203                Arc::new(self.name_builder.finish()),
204                Arc::new(self.value_builder.finish()),
205            ],
206        )?)
207    }
208}
209
210impl ObjectVisitor for ParseBlockRowBuilder {
211    fn visit(&mut self, type_name: &str, value: &[u8]) -> anyhow::Result<bool> {
212        self.index_builder
213            .append_value(self.object_offset + self.local_index);
214        self.name_builder.append_value(type_name);
215        self.value_builder.append_value(value);
216        self.nb_objects += 1;
217        self.local_index += 1;
218        Ok(self.continue_iterating())
219    }
220
221    fn skip(&mut self) -> anyhow::Result<bool> {
222        self.local_index += 1;
223        Ok(self.continue_iterating())
224    }
225}
226
227/// A DataFusion `TableFunctionImpl` that parses a block's payload — through the
228/// `BlockObjectDecoder` registered for its `streams.format` — and returns each
229/// decoded object as a row with its type name and full content as JSONB.
230#[derive(Debug)]
231pub struct ParseBlockTableFunction {
232    lakehouse: Arc<LakehouseContext>,
233    view_factory: Arc<ViewFactory>,
234    part_provider: Arc<dyn QueryPartitionProvider>,
235    query_range: Option<TimeRange>,
236    decoders: Arc<BlockObjectDecoderMap>,
237    guard: Arc<AudienceGuard>,
238}
239
240impl ParseBlockTableFunction {
241    pub fn new(
242        lakehouse: Arc<LakehouseContext>,
243        view_factory: Arc<ViewFactory>,
244        part_provider: Arc<dyn QueryPartitionProvider>,
245        query_range: Option<TimeRange>,
246        guard: Arc<AudienceGuard>,
247    ) -> Self {
248        Self {
249            lakehouse,
250            view_factory,
251            part_provider,
252            query_range,
253            decoders: default_block_object_decoders(),
254            guard,
255        }
256    }
257}
258
259impl TableFunctionImpl for ParseBlockTableFunction {
260    fn call_with_args(
261        &self,
262        args: TableFunctionArgs,
263    ) -> datafusion::error::Result<Arc<dyn TableProvider>> {
264        let exprs = args.exprs();
265        let arg = exprs.first().map(exp_to_string);
266        let Some(Ok(block_id)) = arg else {
267            return plan_err!(
268                "First argument to parse_block must be a string (the block ID), given {:?}",
269                arg
270            );
271        };
272        Ok(Arc::new(ParseBlockProvider {
273            block_id,
274            lakehouse: self.lakehouse.clone(),
275            view_factory: self.view_factory.clone(),
276            part_provider: self.part_provider.clone(),
277            query_range: self.query_range,
278            decoders: self.decoders.clone(),
279            guard: self.guard.clone(),
280        }))
281    }
282}
283
284#[derive(Debug)]
285struct ParseBlockProvider {
286    block_id: String,
287    lakehouse: Arc<LakehouseContext>,
288    view_factory: Arc<ViewFactory>,
289    part_provider: Arc<dyn QueryPartitionProvider>,
290    query_range: Option<TimeRange>,
291    decoders: Arc<BlockObjectDecoderMap>,
292    guard: Arc<AudienceGuard>,
293}
294
295#[async_trait]
296impl TableProvider for ParseBlockProvider {
297    fn schema(&self) -> SchemaRef {
298        output_schema()
299    }
300
301    fn table_type(&self) -> TableType {
302        TableType::Temporary
303    }
304
305    async fn scan(
306        &self,
307        _state: &dyn Session,
308        projection: Option<&Vec<usize>>,
309        filters: &[Expr],
310        limit: Option<usize>,
311    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
312        let block_id_str = &self.block_id;
313
314        let Ok(block_id) = Uuid::parse_str(block_id_str) else {
315            return plan_err!("parse_block: '{block_id_str}' is not a valid block id");
316        };
317
318        let authorized = self
319            .guard
320            .authorize(block_id, IdKind::Block, "parse_block")
321            .await?;
322
323        let Some((object_offset, format, stream_metadata)) = fetch_block_metadata(
324            self.lakehouse.clone(),
325            self.part_provider.clone(),
326            self.query_range,
327            self.view_factory.clone(),
328            block_id,
329            authorized.internal_caller(),
330        )
331        .await
332        .map_err(|e| DataFusionError::External(e.into()))?
333        else {
334            let message = match self.query_range {
335                Some(range) => format!(
336                    "parse_block: block '{block_id_str}' not found in `blocks` for the queried \
337                     range [{}, {}]. The block may be outside the query's time range — widen \
338                     the range to include it.",
339                    range.begin.to_rfc3339(),
340                    range.end.to_rfc3339()
341                ),
342                None => format!("parse_block: block '{block_id_str}' not found in `blocks`."),
343            };
344            return plan_err!("{message}");
345        };
346
347        let Some(decoder) = self.decoders.get(format.as_str()) else {
348            let mut known: Vec<&str> = self.decoders.keys().copied().collect();
349            known.sort_unstable();
350            return plan_err!(
351                "parse_block: no decoder for streams.format='{format}' (known formats: {})",
352                known.join(", ")
353            );
354        };
355
356        // Fetch and decode the block payload
357        let blob_storage = self.lakehouse.lake().blob_storage.clone();
358        let payload = fetch_block_payload(
359            blob_storage,
360            sqlx::types::Uuid::from_bytes(*stream_metadata.process_id.as_bytes()),
361            sqlx::types::Uuid::from_bytes(*stream_metadata.stream_id.as_bytes()),
362            sqlx::types::Uuid::from_bytes(*block_id.as_bytes()),
363        )
364        .await
365        .map_err(|e| DataFusionError::External(e.into()))?;
366
367        let early_limit = if filters.is_empty() { limit } else { None };
368        let mut builder = ParseBlockRowBuilder::new(object_offset, early_limit);
369        decoder
370            .decode(&stream_metadata, &payload, &mut builder)
371            .with_context(|| format!("parsing block {block_id_str}"))
372            .map_err(|e| DataFusionError::External(e.into()))?;
373        let rb = builder
374            .finish()
375            .with_context(|| format!("building record batch for block {block_id_str}"))
376            .map_err(|e| DataFusionError::External(e.into()))?;
377
378        let source = MemorySourceConfig::try_new(
379            &[vec![rb]],
380            self.schema(),
381            projection.map(|v| v.to_owned()),
382        )?;
383        Ok(DataSourceExec::from_data_source(source))
384    }
385}
micromegas_analytics/lakehouse/parse_block_table_function.rs

micromegas_analytics/lakehouse/
parse_block_table_function.rs
