micromegas_analytics/lakehouse/
retire_partitions_table_function.rs

1use anyhow::Context;
2use chrono::DateTime;
3use chrono::Utc;
4use datafusion::catalog::TableFunctionArgs;
5use datafusion::catalog::TableFunctionImpl;
6use datafusion::catalog::TableProvider;
7use datafusion::common::plan_err;
8use micromegas_ingestion::data_lake_connection::DataLakeConnection;
9use micromegas_tracing::prelude::*;
10use std::sync::Arc;
11
12use crate::dfext::expressions::exp_to_string;
13use crate::dfext::expressions::exp_to_timestamp;
14use crate::dfext::log_stream_table_provider::LogStreamTableProvider;
15use crate::dfext::task_log_exec_plan::TaskLogExecPlan;
16use crate::response_writer::LogSender;
17use crate::response_writer::Logger;
18
19use super::write_partition::{RetireMatch, retire_partitions};
20
21/// A DataFusion `TableFunctionImpl` for retiring lakehouse partitions.
22#[derive(Debug)]
23pub struct RetirePartitionsTableFunction {
24    lake: Arc<DataLakeConnection>,
25}
26
27impl RetirePartitionsTableFunction {
28    pub fn new(lake: Arc<DataLakeConnection>) -> Self {
29        Self { lake }
30    }
31}
32
33async fn retire_partitions_impl(
34    lake: Arc<DataLakeConnection>,
35    view_set_name: &str,
36    view_instance_id: &str,
37    begin_insert_time: DateTime<Utc>,
38    end_insert_time: DateTime<Utc>,
39    logger: Arc<dyn Logger>,
40) -> anyhow::Result<()> {
41    let mut tr = lake.db_pool.begin().await?;
42    // The public admin UDF: unchanged behavior, so RetireMatch::Containment (today's query) and no
43    // same-run exclusion (that mechanism only applies to RetireMatch::Overlap).
44    retire_partitions(
45        &mut tr,
46        view_set_name,
47        view_instance_id,
48        begin_insert_time,
49        end_insert_time,
50        RetireMatch::Containment,
51        &[],
52        logger,
53    )
54    .await?;
55    tr.commit().await.with_context(|| "commit")?;
56    Ok(())
57}
58
59impl TableFunctionImpl for RetirePartitionsTableFunction {
60    fn call_with_args(
61        &self,
62        args: TableFunctionArgs,
63    ) -> datafusion::error::Result<Arc<dyn TableProvider>> {
64        let args = args.exprs();
65        // an alternative would be to use coerce & create_physical_expr
66        let Some(view_set_name) = args.first().map(exp_to_string).transpose()? else {
67            return plan_err!("Missing first argument, expected view_set_name: String");
68        };
69        let Some(view_instance_id) = args.get(1).map(exp_to_string).transpose()? else {
70            return plan_err!("Missing 2nd argument, expected view_instance_id: String");
71        };
72        let Some(begin) = args.get(2).map(exp_to_timestamp).transpose()? else {
73            return plan_err!("Missing 3rd argument, expected a UTC nanoseconds timestamp");
74        };
75        let Some(end) = args.get(3).map(exp_to_timestamp).transpose()? else {
76            return plan_err!("Missing 4th argument, expected a UTC nanoseconds timestamp");
77        };
78
79        let lake = self.lake.clone();
80
81        let spawner = move || {
82            let (tx, rx) = tokio::sync::mpsc::channel(100);
83            let logger = Arc::new(LogSender::new(tx));
84            spawn_with_context(async move {
85                if let Err(e) = retire_partitions_impl(
86                    lake,
87                    &view_set_name,
88                    &view_instance_id,
89                    begin,
90                    end,
91                    logger,
92                )
93                .await
94                .with_context(|| "retire_partitions_impl")
95                {
96                    error!("{e:?}");
97                }
98            });
99            rx
100        };
101
102        Ok(Arc::new(LogStreamTableProvider {
103            log_stream: Arc::new(TaskLogExecPlan::new(Box::new(spawner))),
104        }))
105    }
106}
micromegas_analytics/lakehouse/retire_partitions_table_function.rs

micromegas_analytics/lakehouse/
retire_partitions_table_function.rs
