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

micromegas_analytics/lakehouse/
export_log_view.rs
