Skip to main content

micromegas_analytics/lakehouse/
async_events_view.rs

1use super::{
2    batch_update::PartitionCreationStrategy,
3    block_partition_spec::{BlockProcessor, BlockProcessorMap},
4    blocks_view::BlocksView,
5    dataframe_time_bounds::{DataFrameTimeBounds, NamedColumnsTimeBounds},
6    jit_partitions::{JitPartitionConfig, write_partition_from_blocks},
7    lakehouse_context::LakehouseContext,
8    partition_cache::PartitionCache,
9    view::{PartitionSpec, View, ViewMetadata},
10    view_factory::{ViewFactory, ViewMaker},
11};
12use crate::{
13    async_events_table::async_events_table_schema,
14    lakehouse::jit_partitions::{
15        BlockOrder, find_up_to_date_partitions, generate_process_jit_partitions,
16    },
17    metadata::find_process_with_latest_timing,
18    time::{TimeRange, datetime_to_scalar, make_time_converter_from_latest_timing},
19};
20use anyhow::{Context, Result};
21use async_trait::async_trait;
22use chrono::{DateTime, TimeDelta, Utc};
23use datafusion::{
24    arrow::datatypes::Schema,
25    logical_expr::{Between, Expr, col},
26};
27use micromegas_ingestion::web_ingestion_service::FORMAT_TRANSIT;
28use micromegas_tracing::prelude::*;
29use std::collections::HashMap;
30use std::sync::Arc;
31use uuid::Uuid;
32
33use super::async_events_block_processor::AsyncEventsBlockProcessor;
34
35const VIEW_SET_NAME: &str = "async_events";
36const SCHEMA_VERSION: u8 = 3; // Version 3: widened dictionary keys to Int32
37lazy_static::lazy_static! {
38    static ref TIME_COLUMN: Arc<String> = Arc::new(String::from("time"));
39}
40
41/// A `ViewMaker` for creating `AsyncEventsView` instances.
42#[derive(Debug)]
43pub struct AsyncEventsViewMaker {
44    view_factory: Arc<ViewFactory>,
45}
46
47impl AsyncEventsViewMaker {
48    pub fn new(view_factory: Arc<ViewFactory>) -> Self {
49        Self { view_factory }
50    }
51}
52
53impl ViewMaker for AsyncEventsViewMaker {
54    fn make_view(&self, view_instance_id: &str) -> Result<Arc<dyn View>> {
55        Ok(Arc::new(AsyncEventsView::new(
56            view_instance_id,
57            self.view_factory.clone(),
58        )?))
59    }
60
61    fn get_schema_hash(&self) -> Vec<u8> {
62        vec![SCHEMA_VERSION]
63    }
64
65    fn get_schema(&self) -> Arc<Schema> {
66        Arc::new(async_events_table_schema())
67    }
68}
69
70/// A view of async span events.
71#[derive(Debug)]
72pub struct AsyncEventsView {
73    view_set_name: Arc<String>,
74    view_instance_id: Arc<String>,
75    process_id: Option<sqlx::types::Uuid>,
76    view_factory: Arc<ViewFactory>,
77}
78
79impl AsyncEventsView {
80    pub fn new(view_instance_id: &str, view_factory: Arc<ViewFactory>) -> Result<Self> {
81        if view_instance_id == "global" {
82            anyhow::bail!("AsyncEventsView does not support global view access");
83        }
84
85        let process_id =
86            Some(Uuid::parse_str(view_instance_id).with_context(|| "Uuid::parse_str")?);
87
88        Ok(Self {
89            view_set_name: Arc::new(String::from(VIEW_SET_NAME)),
90            view_instance_id: Arc::new(view_instance_id.into()),
91            process_id,
92            view_factory,
93        })
94    }
95}
96
97#[async_trait]
98impl View for AsyncEventsView {
99    fn get_view_set_name(&self) -> Arc<String> {
100        self.view_set_name.clone()
101    }
102
103    fn get_view_instance_id(&self) -> Arc<String> {
104        self.view_instance_id.clone()
105    }
106
107    async fn make_batch_partition_spec(
108        &self,
109        _lakehouse: Arc<LakehouseContext>,
110        _existing_partitions: Arc<PartitionCache>,
111        _insert_range: TimeRange,
112    ) -> Result<Arc<dyn PartitionSpec>> {
113        anyhow::bail!("AsyncEventsView does not support batch partition specs");
114    }
115
116    fn get_file_schema_hash(&self) -> Vec<u8> {
117        vec![SCHEMA_VERSION]
118    }
119
120    fn get_file_schema(&self) -> Arc<Schema> {
121        Arc::new(async_events_table_schema())
122    }
123
124    #[span_fn]
125    async fn jit_update(
126        &self,
127        lakehouse: Arc<LakehouseContext>,
128        query_range: Option<TimeRange>,
129    ) -> Result<()> {
130        let (process, last_block_end_ticks, last_block_end_time) = find_process_with_latest_timing(
131            lakehouse.clone(),
132            self.view_factory.clone(),
133            &self
134                .process_id
135                .with_context(|| "getting a view's process_id")?,
136            query_range,
137        )
138        .await
139        .with_context(|| "find_process_with_latest_timing")?;
140
141        let process = Arc::new(process);
142        let query_range =
143            query_range.unwrap_or_else(|| TimeRange::new(process.start_time, last_block_end_time));
144
145        // Create a consistent ConvertTicks using the latest timing information
146        let convert_ticks = Arc::new(
147            make_time_converter_from_latest_timing(
148                &process,
149                last_block_end_ticks,
150                last_block_end_time,
151            )
152            .with_context(|| "make_time_converter_from_latest_timing")?,
153        );
154
155        // Use process-based partition generation to get all streams for this process
156        let blocks_view = BlocksView::new(lakehouse.default_audience())?;
157        let all_partitions = generate_process_jit_partitions(
158            &JitPartitionConfig::default(),
159            lakehouse.clone(),
160            &blocks_view,
161            &query_range,
162            process.clone(),
163            "cpu",
164        )
165        .await
166        .with_context(|| "generate_process_jit_partitions")?;
167        let view_meta = ViewMetadata {
168            view_set_name: self.get_view_set_name(),
169            view_instance_id: self.get_view_instance_id(),
170            file_schema_hash: self.get_file_schema_hash(),
171        };
172
173        // async_events only sources native transit streams (tagged "cpu");
174        // OTel spans go to the dedicated otel_spans view.
175        let mut block_processors: BlockProcessorMap = HashMap::new();
176        block_processors.insert(
177            FORMAT_TRANSIT,
178            Arc::new(AsyncEventsBlockProcessor::new(convert_ticks.clone()))
179                as Arc<dyn BlockProcessor>,
180        );
181        let block_processors = Arc::new(block_processors);
182
183        let up_to_date = find_up_to_date_partitions(
184            &lakehouse.lake().db_pool,
185            view_meta.clone(),
186            BlockOrder::InsertTime,
187            &all_partitions,
188        )
189        .await
190        .with_context(|| "find_up_to_date_partitions")?;
191        for (part, up_to_date) in all_partitions.into_iter().zip(up_to_date) {
192            if !up_to_date {
193                write_partition_from_blocks(
194                    lakehouse.lake().clone(),
195                    view_meta.clone(),
196                    self.get_file_schema(),
197                    part,
198                    block_processors.clone(),
199                )
200                .await
201                .with_context(|| "write_partition_from_blocks")?;
202            }
203        }
204        Ok(())
205    }
206
207    fn make_time_filter(&self, begin: DateTime<Utc>, end: DateTime<Utc>) -> Result<Vec<Expr>> {
208        Ok(vec![Expr::Between(Between::new(
209            col("time").into(),
210            false,
211            Expr::Literal(datetime_to_scalar(begin), None).into(),
212            Expr::Literal(datetime_to_scalar(end), None).into(),
213        ))])
214    }
215
216    fn get_time_bounds(&self) -> Arc<dyn DataFrameTimeBounds> {
217        Arc::new(NamedColumnsTimeBounds::new(
218            TIME_COLUMN.clone(),
219            TIME_COLUMN.clone(),
220        ))
221    }
222
223    fn get_update_group(&self) -> Option<i32> {
224        None // Process-specific views don't use update groups
225    }
226
227    fn get_max_partition_time_delta(&self, _strategy: &PartitionCreationStrategy) -> TimeDelta {
228        TimeDelta::hours(1)
229    }
230}