micromegas_analytics/lakehouse/
query.rs

1use super::{
2    answer::Answer,
3    deny_queries_table_function::DenyQueriesTableFunction,
4    get_payload_function::GetPayload,
5    lakehouse_context::LakehouseContext,
6    list_audience_grants_table_function::{GrantVisibility, ListAudienceGrantsTableFunction},
7    list_partitions_table_function::ListPartitionsTableFunction,
8    list_query_denials_table_function::ListQueryDenialsTableFunction,
9    list_view_set_definitions_table_function::ListViewSetDefinitionsTableFunction,
10    list_view_sets_table_function::ListViewSetsTableFunction,
11    materialize_partitions_table_function::MaterializePartitionsTableFunction,
12    parse_block_table_function::ParseBlockTableFunction,
13    partition::Partition,
14    partition_cache::QueryPartitionProvider,
15    partitioned_table_provider::PartitionedTableProvider,
16    perfetto_trace_table_function::PerfettoTraceTableFunction,
17    process_spans_table_function::ProcessSpansTableFunction,
18    reader_factory::ReaderFactory,
19    regenerate_partitions_table_function::RegeneratePartitionsTableFunction,
20    remove_query_denial_udf::make_remove_query_denial_udf,
21    retire_partition_by_file_udf::make_retire_partition_by_file_udf,
22    retire_partition_by_metadata_udf::make_retire_partition_by_metadata_udf,
23    retire_partitions_table_function::RetirePartitionsTableFunction,
24    session_configurator::SessionConfigurator,
25    view::View,
26    view_factory::ViewFactory,
27};
28use crate::{
29    lakehouse::{
30        audience_guard::AudienceGuard,
31        materialized_view::MaterializedView,
32        ownership_rewrite::OwnershipRewrite,
33        read_scope::{CallerContext, ReadScope},
34        table_scan_rewrite::TableScanRewrite,
35        view_instance_table_function::ViewInstanceTableFunction,
36    },
37    properties::{
38        properties_to_dict_udf::PropertiesToDict, properties_to_jsonb_udf::PropertiesToJsonb,
39    },
40    time::TimeRange,
41};
42use anyhow::{Context, Result};
43use datafusion::{
44    arrow::{array::RecordBatch, datatypes::SchemaRef},
45    common::TableReference,
46    datasource::DefaultTableSource,
47    execution::{context::SessionContext, object_store::ObjectStoreUrl, runtime_env::RuntimeEnv},
48    logical_expr::{ScalarUDF, TableSource, async_udf::AsyncScalarUDF},
49    prelude::*,
50};
51use micromegas_tracing::prelude::*;
52use std::sync::Arc;
53
54#[span_fn]
55async fn register_table(
56    lakehouse: Arc<LakehouseContext>,
57    reader_factory: Arc<ReaderFactory>,
58    part_provider: Arc<dyn QueryPartitionProvider>,
59    query_range: Option<TimeRange>,
60    ctx: &SessionContext,
61    view: Arc<dyn View>,
62) -> Result<()> {
63    let table = MaterializedView::new(
64        lakehouse,
65        reader_factory,
66        view.clone(),
67        part_provider,
68        query_range,
69        None,
70    );
71    view.register_table(ctx, table).await
72}
73
74/// query_partitions_context returns a context to run queries using the partitions as the "source" table
75#[span_fn]
76pub async fn query_partitions_context(
77    runtime: Arc<RuntimeEnv>,
78    reader_factory: Arc<ReaderFactory>,
79    object_store: Arc<dyn object_store::ObjectStore>,
80    schema: SchemaRef,
81    partitions: Arc<Vec<Partition>>,
82) -> Result<SessionContext> {
83    let table = PartitionedTableProvider::new(schema, reader_factory, partitions);
84    let object_store_url = ObjectStoreUrl::parse("obj://lakehouse/").unwrap();
85    let ctx = SessionContext::new_with_config_rt(SessionConfig::default(), runtime);
86    ctx.register_object_store(object_store_url.as_ref(), object_store);
87    ctx.register_table(
88        TableReference::Bare {
89            table: "source".into(),
90        },
91        Arc::new(table),
92    )?;
93    register_extension_functions(&ctx);
94    Ok(ctx)
95}
96
97// query_partitions returns a dataframe, leaving the option of streaming the results
98#[span_fn]
99pub async fn query_partitions(
100    runtime: Arc<RuntimeEnv>,
101    reader_factory: Arc<ReaderFactory>,
102    object_store: Arc<dyn object_store::ObjectStore>,
103    schema: SchemaRef,
104    partitions: Arc<Vec<Partition>>,
105    sql: &str,
106) -> Result<DataFrame> {
107    let ctx =
108        query_partitions_context(runtime, reader_factory, object_store, schema, partitions).await?;
109    Ok(ctx.sql(sql).await?)
110}
111
112/// register functions that are part of the lakehouse architecture
113#[span_fn]
114pub fn register_lakehouse_functions(
115    ctx: &SessionContext,
116    lakehouse: Arc<LakehouseContext>,
117    part_provider: Arc<dyn QueryPartitionProvider>,
118    query_range: Option<TimeRange>,
119    view_factory: Arc<ViewFactory>,
120    caller: &CallerContext,
121) {
122    // The call-level guard: one guard, shared (via `Arc`) across every
123    // arg-addressed UDTF/UDF this call registers below -- `view_instance`, `process_spans`,
124    // `perfetto_trace_chunks`, `parse_block`, `get_payload`, and `list_partitions`' row filter.
125    // `ReadScope::All` makes every one of its checks a no-op, so this costs nothing for
126    // internal/maintenance callers or an auth-unset deployment.
127    let lakehouse_admin = caller.is_admin;
128    let audience_guard = Arc::new(AudienceGuard::new(
129        caller.read_scope.clone(),
130        lakehouse_admin,
131        caller.isolation_config.public_view_sets.clone(),
132        lakehouse.audience_index().clone(),
133    ));
134    ctx.register_udtf(
135        "view_instance",
136        Arc::new(ViewInstanceTableFunction::new(
137            lakehouse.clone(),
138            view_factory.clone(),
139            part_provider.clone(),
140            query_range,
141            audience_guard.clone(),
142        )),
143    );
144    ctx.register_udtf(
145        "list_partitions",
146        Arc::new(ListPartitionsTableFunction::new(
147            lakehouse.lake().clone(),
148            audience_guard.clone(),
149        )),
150    );
151    ctx.register_udtf(
152        "list_view_sets",
153        Arc::new(ListViewSetsTableFunction::new(view_factory.clone())),
154    );
155    // Caller-scoped, not admin-gated: registered for every caller, admin or not, with the
156    // visibility split decided here rather than inside the function -- an admin sees every
157    // row, a non-admin only the pairs their `grant_selectors` hold.
158    ctx.register_udtf(
159        "list_audience_grants",
160        Arc::new(ListAudienceGrantsTableFunction::new(
161            lakehouse.lake().db_pool.clone(),
162            if caller.is_admin {
163                GrantVisibility::All
164            } else {
165                GrantVisibility::Held(caller.grant_selectors.clone())
166            },
167        )),
168    );
169    ctx.register_udtf(
170        "perfetto_trace_chunks",
171        Arc::new(PerfettoTraceTableFunction::new(
172            lakehouse.clone(),
173            view_factory.clone(),
174            part_provider.clone(),
175            audience_guard.clone(),
176        )),
177    );
178    ctx.register_udtf(
179        "parse_block",
180        Arc::new(ParseBlockTableFunction::new(
181            lakehouse.clone(),
182            view_factory.clone(),
183            part_provider.clone(),
184            query_range,
185            audience_guard.clone(),
186        )),
187    );
188    ctx.register_udtf(
189        "process_spans",
190        Arc::new(ProcessSpansTableFunction::new(
191            lakehouse.clone(),
192            view_factory.clone(),
193            part_provider.clone(),
194            query_range,
195            audience_guard.clone(),
196        )),
197    );
198    ctx.register_udf(
199        AsyncScalarUDF::new(Arc::new(GetPayload::new(
200            lakehouse.lake().clone(),
201            audience_guard.clone(),
202        )))
203        .into_scalar_udf(),
204    );
205    if lakehouse_admin {
206        ctx.register_udtf(
207            "retire_partitions",
208            Arc::new(RetirePartitionsTableFunction::new(lakehouse.lake().clone())),
209        );
210        ctx.register_udtf(
211            "materialize_partitions",
212            Arc::new(MaterializePartitionsTableFunction::new(
213                lakehouse.clone(),
214                view_factory.clone(),
215            )),
216        );
217        ctx.register_udtf(
218            "regenerate_partitions",
219            Arc::new(RegeneratePartitionsTableFunction::new(
220                lakehouse.clone(),
221                view_factory.clone(),
222            )),
223        );
224        ctx.register_udf(
225            make_retire_partition_by_file_udf(lakehouse.lake().clone()).into_scalar_udf(),
226        );
227        ctx.register_udf(
228            make_retire_partition_by_metadata_udf(lakehouse.lake().clone()).into_scalar_udf(),
229        );
230        // Admin-managed query deny list: same admin gate as the five functions above.
231        ctx.register_udtf(
232            "list_query_denials",
233            Arc::new(ListQueryDenialsTableFunction::new(
234                lakehouse.query_denials().clone(),
235            )),
236        );
237        ctx.register_udtf(
238            "deny_queries",
239            Arc::new(DenyQueriesTableFunction::new(
240                lakehouse.query_denials().clone(),
241                caller.identity.clone(),
242            )),
243        );
244        ctx.register_udf(
245            make_remove_query_denial_udf(lakehouse.query_denials().clone()).into_scalar_udf(),
246        );
247        // DDL-defined view sets: lists every row of
248        // `lakehouse_view_set_definitions` straight from Postgres, so an admin can see a
249        // definition that exists on disk but failed to load.
250        ctx.register_udtf(
251            "list_view_set_definitions",
252            Arc::new(ListViewSetDefinitionsTableFunction::new(
253                lakehouse.lake().clone(),
254            )),
255        );
256    }
257}
258
259/// register functions that are not depended on the lakehouse architecture
260#[span_fn]
261pub fn register_extension_functions(ctx: &SessionContext) {
262    ctx.register_udf(ScalarUDF::from(PropertiesToDict::new()));
263    ctx.register_udf(ScalarUDF::from(PropertiesToJsonb::new()));
264    micromegas_datafusion_extensions::register_extension_udfs(ctx);
265}
266
267#[span_fn]
268pub fn register_functions(
269    ctx: &SessionContext,
270    lakehouse: Arc<LakehouseContext>,
271    part_provider: Arc<dyn QueryPartitionProvider>,
272    query_range: Option<TimeRange>,
273    view_factory: Arc<ViewFactory>,
274    caller: &CallerContext,
275) {
276    register_lakehouse_functions(
277        ctx,
278        lakehouse,
279        part_provider,
280        query_range,
281        view_factory,
282        caller,
283    );
284    register_extension_functions(ctx);
285}
286
287#[span_fn]
288pub async fn make_session_context(
289    lakehouse: Arc<LakehouseContext>,
290    part_provider: Arc<dyn QueryPartitionProvider>,
291    query_range: Option<TimeRange>,
292    view_factory: Arc<ViewFactory>,
293    configurator: Arc<dyn SessionConfigurator>,
294    caller: CallerContext,
295) -> Result<SessionContext> {
296    // Disable page index reading for backward compatibility with legacy Parquet files
297    // Legacy files may have incomplete ColumnIndex metadata (missing null_pages field)
298    // which causes errors in DataFusion 51+ with Arrow 57.0 when reading page indexes
299    let config = SessionConfig::default()
300        .set_bool("datafusion.execution.parquet.enable_page_index", false)
301        // Populate `Diagnostic`/`Span` on plan-time errors (unknown column, ambiguous
302        // reference, type mismatch, ...) so a caller-facing message can point at a
303        // line/column in their SQL text instead of just naming the problem. Off by
304        // default in DataFusion.
305        .set_bool("datafusion.sql_parser.collect_spans", true)
306        .with_information_schema(true);
307    let ctx = SessionContext::new_with_config_rt(config, lakehouse.runtime().clone());
308    if let Some(range) = &query_range {
309        ctx.add_analyzer_rule(Arc::new(TableScanRewrite::new(*range)));
310    }
311    let object_store_url = ObjectStoreUrl::parse("obj://lakehouse/").unwrap();
312    let object_store = lakehouse.lake().blob_storage.inner();
313    ctx.register_object_store(object_store_url.as_ref(), object_store);
314    let reader_factory = lakehouse.reader_factory().clone();
315    register_functions(
316        &ctx,
317        lakehouse.clone(),
318        part_provider.clone(),
319        query_range,
320        view_factory.clone(),
321        &caller,
322    );
323    for view in view_factory.get_global_views() {
324        register_table(
325            lakehouse.clone(),
326            reader_factory.clone(),
327            part_provider.clone(),
328            query_range,
329            &ctx,
330            view.clone(),
331        )
332        .await?;
333    }
334    if caller.read_scope != ReadScope::All {
335        // ReadScope::All is the internal/maintenance marker -- OwnershipRewrite would no-op for
336        // it anyway, so skip resolving `processes`/`streams` sources and registering the rule
337        // entirely rather than requiring every ReadScope::All caller's ViewFactory to carry them.
338        //
339        // Must be registered *after* the TableScanRewrite registration above (`query_range.is_some()`
340        // block): TableScanRewrite::analyze walks the whole plan with transform_up_with_subqueries
341        // and would time-bound the audience lookup's own processes/streams scans if it ran after
342        // OwnershipRewrite injected them (analyzer rules each run exactly once, in registration
343        // order). The audience lookup must stay time-unbounded (query_range: None below).
344        let processes_view = view_factory.get_global_view("processes").with_context(
345            || "OwnershipRewrite requires the `processes` global view to be registered",
346        )?;
347        let streams_view = view_factory.get_global_view("streams").with_context(
348            || "OwnershipRewrite requires the `streams` global view to be registered",
349        )?;
350        let processes_source: Arc<dyn TableSource> =
351            Arc::new(DefaultTableSource::new(Arc::new(MaterializedView::new(
352                lakehouse.clone(),
353                reader_factory.clone(),
354                processes_view,
355                part_provider.clone(),
356                None,
357                None,
358            ))));
359        let streams_source: Arc<dyn TableSource> =
360            Arc::new(DefaultTableSource::new(Arc::new(MaterializedView::new(
361                lakehouse.clone(),
362                reader_factory.clone(),
363                streams_view,
364                part_provider.clone(),
365                None,
366                None,
367            ))));
368        ctx.add_analyzer_rule(Arc::new(OwnershipRewrite::new(
369            caller.read_scope.clone(),
370            caller.isolation_config.public_view_sets.clone(),
371            processes_source,
372            streams_source,
373        )));
374    }
375    // Apply custom configuration
376    configurator.configure(&ctx).await?;
377    Ok(ctx)
378}
379
380#[span_fn]
381pub async fn query(
382    lakehouse: Arc<LakehouseContext>,
383    part_provider: Arc<dyn QueryPartitionProvider>,
384    query_range: Option<TimeRange>,
385    sql: &str,
386    view_factory: Arc<ViewFactory>,
387    configurator: Arc<dyn SessionConfigurator>,
388    caller: CallerContext,
389) -> Result<Answer> {
390    info!("query sql={sql}");
391    let ctx = make_session_context(
392        lakehouse,
393        part_provider,
394        query_range,
395        view_factory,
396        configurator,
397        caller,
398    )
399    .await
400    .with_context(|| "make_session_context")?;
401    let df = ctx.sql(sql).await?;
402    let schema = df.schema().inner().clone();
403    let batches: Vec<RecordBatch> = df.collect().await?;
404    Ok(Answer::new(schema, batches))
405}
micromegas_analytics/lakehouse/query.rs

micromegas_analytics/lakehouse/
query.rs
