Skip to main content

micromegas_analytics/lakehouse/
view_instance_table_function.rs

1use super::{
2    audience_guard::AudienceGuard, lakehouse_context::LakehouseContext,
3    materialized_view::MaterializedView, partition_cache::QueryPartitionProvider,
4    view_factory::ViewFactory,
5};
6use crate::{dfext::expressions::exp_to_string, time::TimeRange};
7use datafusion::{
8    catalog::{TableFunctionArgs, TableFunctionImpl, TableProvider},
9    common::plan_err,
10    error::DataFusionError,
11};
12use micromegas_tracing::prelude::*;
13use std::sync::Arc;
14
15/// `ViewInstanceTableFunction` gives access to any view instance using a [ViewFactory].
16///
17/// ```python
18/// # Python code showing the usage of `view_instance(view_set_name, view_instance_id)`
19/// sql = """
20/// SELECT *
21/// FROM view_instance('thread_spans', '{stream_id}')
22/// ;""".format(stream_id=stream_id)
23/// df_spans = client.query(sql, begin_spans, end_spans)
24/// ```
25///
26/// The only site that supplies an [`AudienceGuard`] to the [`MaterializedView`]s it builds:
27/// every caller-named instance goes through this table function, so this is where the
28/// scan-time audience check gets wired in.
29#[derive(Debug)]
30pub struct ViewInstanceTableFunction {
31    lakehouse: Arc<LakehouseContext>,
32    view_factory: Arc<ViewFactory>,
33    part_provider: Arc<dyn QueryPartitionProvider>,
34    query_range: Option<TimeRange>,
35    guard: Arc<AudienceGuard>,
36}
37
38impl ViewInstanceTableFunction {
39    pub fn new(
40        lakehouse: Arc<LakehouseContext>,
41        view_factory: Arc<ViewFactory>,
42        part_provider: Arc<dyn QueryPartitionProvider>,
43        query_range: Option<TimeRange>,
44        guard: Arc<AudienceGuard>,
45    ) -> Self {
46        Self {
47            lakehouse,
48            view_factory,
49            part_provider,
50            query_range,
51            guard,
52        }
53    }
54}
55
56impl TableFunctionImpl for ViewInstanceTableFunction {
57    #[span_fn]
58    fn call_with_args(
59        &self,
60        args: TableFunctionArgs,
61    ) -> datafusion::error::Result<Arc<dyn TableProvider>> {
62        let exprs = args.exprs();
63        let arg1 = exprs.first().map(exp_to_string);
64        let Some(Ok(view_set_name)) = arg1 else {
65            return plan_err!(
66                "First argument to view_instance must be a string (the view set name), given {:?}",
67                arg1
68            );
69        };
70        let arg2 = exprs.get(1).map(exp_to_string);
71        let Some(Ok(view_instance_id)) = arg2 else {
72            return plan_err!(
73                "Second argument to view_instance must be a string (the view instance id), given {:?}",
74                arg2
75            );
76        };
77
78        let view = self
79            .view_factory
80            .make_view(&view_set_name, &view_instance_id)
81            .map_err(|e| DataFusionError::Plan(format!("error making view {e:?}")))?;
82
83        Ok(Arc::new(MaterializedView::new(
84            self.lakehouse.clone(),
85            self.lakehouse.reader_factory().clone(),
86            view,
87            self.part_provider.clone(),
88            self.query_range,
89            Some(self.guard.clone()),
90        )))
91    }
92}