micromegas_analytics/lakehouse/
export_log_view.rs

1use super::{
2    batch_update::PartitionCreationStrategy,
3    dataframe_time_bounds::{DataFrameTimeBounds, NamedColumnsTimeBounds},
4    lakehouse_context::LakehouseContext,
5    partition_cache::{NullPartitionProvider, PartitionCache},
6    query::make_session_context,
7    session_configurator::SessionConfigurator,
8    view::{PartitionSpec, View},
9    view_factory::ViewFactory,
10};
11use crate::{
12    lakehouse::{sql_partition_spec::fetch_sql_partition_spec, view::ViewMetadata},
13    record_batch_transformer::RecordBatchTransformer,
14    time::{TimeRange, datetime_to_scalar},
15};
16use anyhow::{Context, Result};
17use async_trait::async_trait;
18use chrono::{DateTime, TimeDelta, Utc};
19use datafusion::{
20    arrow::{
21        array::{PrimitiveBuilder, RecordBatch, StringBuilder},
22        datatypes::{DataType, Field, Int32Type, Schema, TimeUnit, TimestampNanosecondType},
23    },
24    execution::runtime_env::RuntimeEnv,
25    prelude::*,
26};
27use micromegas_ingestion::data_lake_connection::DataLakeConnection;
28use micromegas_tracing::levels::Level;
29use std::hash::Hash;
30use std::hash::Hasher;
31use std::{hash::DefaultHasher, sync::Arc};
32
33/// A builder for creating log entries for export.
34pub struct ExportLogBuilder {
35    times: PrimitiveBuilder<TimestampNanosecondType>,
36    levels: PrimitiveBuilder<Int32Type>,
37    msgs: StringBuilder,
38}
39
40impl ExportLogBuilder {
41    #[expect(clippy::new_without_default)]
42    pub fn new() -> Self {
43        Self {
44            times: PrimitiveBuilder::new(),
45            levels: PrimitiveBuilder::new(),
46            msgs: StringBuilder::new(),
47        }
48    }
49
50    pub fn append(&mut self, level: Level, msg: &str) {
51        let now = Utc::now();
52        self.times
53            .append_value(now.timestamp_nanos_opt().unwrap_or_default());
54        self.levels.append_value(level as i32);
55        self.msgs.append_value(msg);
56    }
57
58    pub fn finish(mut self) -> Result<RecordBatch> {
59        RecordBatch::try_new(
60            make_export_log_schema(),
61            vec![
62                Arc::new(self.times.finish().with_timezone_utc()),
63                Arc::new(self.levels.finish()),
64                Arc::new(self.msgs.finish()),
65            ],
66        )
67        .with_context(|| "building record batch")
68    }
69}
70
71/// A view for exporting log data.
72#[derive(Debug)]
73pub struct ExportLogView {
74    view_set_name: Arc<String>,
75    view_instance_id: Arc<String>,
76    time_column_name: Arc<String>,
77    count_src_query: Arc<String>,
78    extract_query: Arc<String>,
79    exporter: Arc<dyn RecordBatchTransformer>,
80    log_schema: Arc<Schema>,
81    view_factory: Arc<ViewFactory>,
82    session_configurator: Arc<dyn SessionConfigurator>,
83    update_group: Option<i32>,
84    max_partition_delta_from_source: TimeDelta,
85    max_partition_delta_from_merge: TimeDelta,
86}
87
88/// Creates the Arrow schema for the export log.
89pub fn make_export_log_schema() -> Arc<Schema> {
90    Arc::new(Schema::new(vec![
91        Field::new(
92            "time",
93            DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
94            false,
95        ),
96        Field::new("level", DataType::Int32, false),
97        Field::new("msg", DataType::Utf8, false),
98    ]))
99}
100
101impl ExportLogView {
102    #[expect(clippy::too_many_arguments)]
103    pub async fn new(
104        runtime: Arc<RuntimeEnv>,
105        view_set_name: Arc<String>,
106        count_src_query: Arc<String>,
107        extract_query: Arc<String>,
108        exporter: Arc<dyn RecordBatchTransformer>,
109        lake: Arc<DataLakeConnection>,
110        view_factory: Arc<ViewFactory>,
111        session_configurator: Arc<dyn SessionConfigurator>,
112        update_group: Option<i32>,
113        max_partition_delta_from_source: TimeDelta,
114        max_partition_delta_from_merge: TimeDelta,
115    ) -> Result<Self> {
116        let null_part_provider = Arc::new(NullPartitionProvider {});
117        let lakehouse = Arc::new(LakehouseContext::new(lake.clone(), runtime.clone()));
118        let ctx = make_session_context(
119            lakehouse,
120            null_part_provider,
121            None,
122            view_factory.clone(),
123            session_configurator.clone(),
124            true,
125        )
126        .await
127        .with_context(|| "make_session_context")?;
128        let now_str = Utc::now().to_rfc3339();
129        let sql = extract_query
130            .replace("{begin}", &now_str)
131            .replace("{end}", &now_str);
132        let _extracted_df = ctx.sql(&sql).await?;
133        Ok(Self {
134            view_set_name,
135            view_instance_id: Arc::new(String::from("global")),
136            time_column_name: Arc::new(String::from("time")),
137            count_src_query,
138            extract_query,
139            exporter,
140            log_schema: make_export_log_schema(),
141            view_factory,
142            session_configurator,
143            update_group,
144            max_partition_delta_from_source,
145            max_partition_delta_from_merge,
146        })
147    }
148}
149
150#[async_trait]
151impl View for ExportLogView {
152    fn get_view_set_name(&self) -> Arc<String> {
153        self.view_set_name.clone()
154    }
155
156    fn get_view_instance_id(&self) -> Arc<String> {
157        self.view_instance_id.clone()
158    }
159
160    async fn make_batch_partition_spec(
161        &self,
162        lakehouse: Arc<LakehouseContext>,
163        existing_partitions: Arc<PartitionCache>,
164        insert_range: TimeRange,
165    ) -> Result<Arc<dyn PartitionSpec>> {
166        let view_meta = ViewMetadata {
167            view_set_name: self.get_view_set_name(),
168            view_instance_id: self.get_view_instance_id(),
169            file_schema_hash: self.get_file_schema_hash(),
170        };
171        let partitions_in_range = Arc::new(existing_partitions.filter_insert_range(insert_range));
172        let ctx = make_session_context(
173            lakehouse,
174            partitions_in_range.clone(),
175            None,
176            self.view_factory.clone(),
177            self.session_configurator.clone(),
178            true,
179        )
180        .await
181        .with_context(|| "make_session_context")?;
182        let count_src_sql = self
183            .count_src_query
184            .replace("{begin}", &insert_range.begin.to_rfc3339())
185            .replace("{end}", &insert_range.end.to_rfc3339());
186        let extract_sql = self
187            .extract_query
188            .replace("{begin}", &insert_range.begin.to_rfc3339())
189            .replace("{end}", &insert_range.end.to_rfc3339());
190        Ok(Arc::new(
191            fetch_sql_partition_spec(
192                ctx,
193                self.exporter.clone(),
194                self.get_time_bounds(),
195                self.log_schema.clone(),
196                count_src_sql,
197                extract_sql,
198                view_meta,
199                insert_range,
200            )
201            .await
202            .with_context(|| "fetch_sql_partition_spec")?,
203        ))
204    }
205
206    fn get_file_schema_hash(&self) -> Vec<u8> {
207        let mut hasher = DefaultHasher::new();
208        self.log_schema.hash(&mut hasher);
209        hasher.finish().to_le_bytes().to_vec()
210    }
211
212    fn get_file_schema(&self) -> Arc<Schema> {
213        self.log_schema.clone()
214    }
215
216    async fn jit_update(
217        &self,
218        _lakehouse: Arc<LakehouseContext>,
219        _query_range: Option<TimeRange>,
220    ) -> Result<()> {
221        Ok(())
222    }
223
224    fn make_time_filter(&self, begin: DateTime<Utc>, end: DateTime<Utc>) -> Result<Vec<Expr>> {
225        Ok(vec![
226            col(&**self.time_column_name).lt_eq(lit(datetime_to_scalar(end))),
227            col(&**self.time_column_name).gt_eq(lit(datetime_to_scalar(begin))),
228        ])
229    }
230
231    fn get_time_bounds(&self) -> Arc<dyn DataFrameTimeBounds> {
232        Arc::new(NamedColumnsTimeBounds::new(
233            self.time_column_name.clone(),
234            self.time_column_name.clone(),
235        ))
236    }
237
238    fn get_update_group(&self) -> Option<i32> {
239        self.update_group
240    }
241
242    fn get_max_partition_time_delta(&self, strategy: &PartitionCreationStrategy) -> TimeDelta {
243        match strategy {
244            PartitionCreationStrategy::Abort | PartitionCreationStrategy::CreateFromSource => {
245                self.max_partition_delta_from_source
246            }
247            PartitionCreationStrategy::MergeExisting(_partitions) => {
248                self.max_partition_delta_from_merge
249            }
250        }
251    }
252}
micromegas_analytics/lakehouse/export_log_view.rs

micromegas_analytics/lakehouse/
export_log_view.rs
