micromegas_analytics/lakehouse/
perfetto_trace_table_function.rs

1use super::{
2    audience_guard::AudienceGuard, lakehouse_context::LakehouseContext,
3    partition_cache::QueryPartitionProvider, view_factory::ViewFactory,
4};
5use crate::{
6    dfext::expressions::{exp_to_string, exp_to_timestamp},
7    time::TimeRange,
8};
9use datafusion::{
10    arrow::datatypes::{DataType, Field, Schema},
11    catalog::{TableFunctionArgs, TableFunctionImpl, TableProvider},
12    common::plan_err,
13};
14use micromegas_tracing::prelude::*;
15use std::sync::Arc;
16use uuid::Uuid;
17
18/// `PerfettoTraceTableFunction` generates Perfetto trace chunks from process telemetry data.
19///
20/// SQL Interface:
21/// ```sql
22/// SELECT chunk_id, chunk_data
23/// FROM perfetto_trace_chunks(
24///     'process_id',                              -- Process UUID (required)
25///     'span_types',                              -- 'thread', 'async', or 'both' (required)
26///     TIMESTAMP '2024-01-01T00:00:00Z',          -- Start time as UTC timestamp (required)
27///     TIMESTAMP '2024-01-01T01:00:00Z'           -- End time as UTC timestamp (required)
28/// ) ORDER BY chunk_id
29/// ```
30///
31/// Returns a table with schema:
32/// - chunk_id: Int32 - Sequential chunk identifier
33/// - chunk_data: Binary - Binary protobuf TracePacket data
34///
35#[derive(Debug)]
36pub struct PerfettoTraceTableFunction {
37    lakehouse: Arc<LakehouseContext>,
38    view_factory: Arc<ViewFactory>,
39    part_provider: Arc<dyn QueryPartitionProvider>,
40    guard: Arc<AudienceGuard>,
41}
42
43impl PerfettoTraceTableFunction {
44    pub fn new(
45        lakehouse: Arc<LakehouseContext>,
46        view_factory: Arc<ViewFactory>,
47        part_provider: Arc<dyn QueryPartitionProvider>,
48        guard: Arc<AudienceGuard>,
49    ) -> Self {
50        Self {
51            lakehouse,
52            view_factory,
53            part_provider,
54            guard,
55        }
56    }
57
58    /// Create the output schema for the table function
59    pub fn output_schema() -> Arc<Schema> {
60        Arc::new(Schema::new(vec![
61            Field::new("chunk_id", DataType::Int32, false),
62            Field::new("chunk_data", DataType::Binary, false),
63        ]))
64    }
65}
66
67impl TableFunctionImpl for PerfettoTraceTableFunction {
68    #[span_fn]
69    fn call_with_args(
70        &self,
71        args: TableFunctionArgs,
72    ) -> datafusion::error::Result<Arc<dyn TableProvider>> {
73        let exprs = args.exprs();
74        // Parse process_id (arg 1)
75        let arg1 = exprs.first().map(exp_to_string);
76        let Some(Ok(process_id_arg)) = arg1 else {
77            return plan_err!(
78                "First argument to perfetto_trace_chunks must be a string (the process ID), given {:?}",
79                arg1
80            );
81        };
82        // Parsed as a `Uuid` at plan time -- see `process_spans_table_function.rs`'s identical
83        // comment for why (closes an injection vector into the inner session's SQL, matches
84        // `parse_block`'s existing block-id parsing). The canonical hyphenated rendering is kept
85        // for the inner literal comparisons `perfetto_trace_execution_plan.rs` builds.
86        let Ok(process_uuid) = Uuid::parse_str(&process_id_arg) else {
87            return plan_err!(
88                "First argument to perfetto_trace_chunks must be a valid UUID (the process ID), given '{process_id_arg}'"
89            );
90        };
91        let process_id = process_uuid.hyphenated().to_string();
92
93        // Parse span_types (arg 2)
94        let arg2 = exprs.get(1).map(exp_to_string);
95        let Some(Ok(span_types_str)) = arg2 else {
96            return plan_err!(
97                "Second argument to perfetto_trace_chunks must be a string ('thread', 'async', or 'both'), given {:?}",
98                arg2
99            );
100        };
101
102        let span_types = match span_types_str.as_str() {
103            "thread" => SpanTypes::Thread,
104            "async" => SpanTypes::Async,
105            "both" => SpanTypes::Both,
106            _ => {
107                return plan_err!(
108                    "span_types must be 'thread', 'async', or 'both', given: {}",
109                    span_types_str
110                );
111            }
112        };
113
114        // Parse start_time (arg 3) - expecting a timestamp expression
115        let arg3 = exprs.get(2).map(exp_to_timestamp);
116        let Some(Ok(start_time)) = arg3 else {
117            return plan_err!(
118                "Third argument to perfetto_trace_chunks must be a timestamp (start time), given {:?}",
119                arg3
120            );
121        };
122
123        // Parse end_time (arg 4) - expecting a timestamp expression
124        let arg4 = exprs.get(3).map(exp_to_timestamp);
125        let Some(Ok(end_time)) = arg4 else {
126            return plan_err!(
127                "Fourth argument to perfetto_trace_chunks must be a timestamp (end time), given {:?}",
128                arg4
129            );
130        };
131
132        // Create time range from parsed timestamps
133        let time_range = TimeRange {
134            begin: start_time,
135            end: end_time,
136        };
137
138        // Create the execution plan that will generate the trace chunks
139        let execution_plan = Arc::new(PerfettoTraceExecutionPlan::new(
140            Self::output_schema(),
141            process_id,
142            process_uuid,
143            span_types,
144            time_range,
145            self.lakehouse.clone(),
146            self.view_factory.clone(),
147            self.part_provider.clone(),
148            self.guard.clone(),
149        ));
150
151        // Wrap it in a TableProvider
152        Ok(Arc::new(PerfettoTraceTableProvider::new(execution_plan)))
153    }
154}
155
156// Import the execution plan
157use super::perfetto_trace_execution_plan::{
158    PerfettoTraceExecutionPlan, PerfettoTraceTableProvider, SpanTypes,
159};
micromegas_analytics/lakehouse/perfetto_trace_table_function.rs

micromegas_analytics/lakehouse/
perfetto_trace_table_function.rs
