Skip to main content

micromegas_analytics/
metadata.rs

1use anyhow::{Context, Result};
2use chrono::{DateTime, Utc};
3use datafusion::arrow::array::{
4    BinaryArray, GenericListArray, Int32Array, Int64Array, RecordBatch, StringArray,
5    TimestampNanosecondArray,
6};
7use micromegas_telemetry::{
8    property::Property, stream_info::StreamInfo, types::block::BlockMetadata,
9};
10use micromegas_tracing::prelude::*;
11use micromegas_transit::{UserDefinedType, uuid_utils::parse_optional_uuid};
12use sqlx::Row;
13use std::sync::Arc;
14use uuid::Uuid;
15
16use crate::{
17    arrow_properties::serialize_properties_to_jsonb,
18    audience::coalesced_audience_column,
19    dfext::{string_column_accessor::string_column_by_name, typed_column::typed_column_by_name},
20    lakehouse::{
21        lakehouse_context::LakehouseContext, partition_cache::LivePartitionProvider,
22        query::make_session_context, read_scope::CallerContext,
23        session_configurator::NoOpSessionConfigurator, view_factory::ViewFactory,
24    },
25    properties::properties_column_accessor::properties_column_by_name,
26    time::TimeRange,
27};
28/// Type alias for shared, pre-serialized JSONB data.
29/// This represents JSONB properties that have been serialized once and can be reused.
30pub type SharedJsonbSerialized = Arc<Vec<u8>>;
31
32/// Analytics-optimized process metadata.
33///
34/// This struct is designed for analytics use cases where process properties need to be
35/// efficiently serialized to JSONB format multiple times. Unlike `ProcessInfo`, which
36/// uses `HashMap<String, String>` for properties, this struct stores pre-serialized
37/// JSONB data to eliminate redundant serialization overhead.
38#[derive(Debug, Clone)]
39pub struct ProcessMetadata {
40    // Core fields (same as ProcessInfo)
41    pub process_id: uuid::Uuid,
42    pub exe: String,
43    pub username: String,
44    pub realname: String,
45    pub computer: String,
46    pub distro: String,
47    pub cpu_brand: String,
48    pub tsc_frequency: i64,
49    pub start_time: chrono::DateTime<chrono::Utc>,
50    pub start_ticks: i64,
51    pub parent_process_id: Option<uuid::Uuid>,
52    pub properties: SharedJsonbSerialized,
53    /// The owning process's own `audience` column -- always present, never empty. A process
54    /// registered through the HTTP ingestion path always carries a real, non-NULL `audience`
55    /// column now: a credential with no bound audience is stamped with the resolved deployment
56    /// default explicitly. Only a legacy, pre-v8 row keeps a NULL column; every producer of
57    /// this struct resolves that case to `MICROMEGAS_DEFAULT_AUDIENCE` before it gets here, so
58    /// the `Option` stops at the database boundary either way. Appended last: a required field
59    /// with no default, so the compiler enumerates every construction site.
60    pub audience: Arc<str>,
61}
62
63/// Analytics-optimized stream metadata.
64///
65/// This struct is designed for analytics use cases where stream properties need to be
66/// efficiently serialized to JSONB format multiple times. Unlike `StreamInfo`, which
67/// uses `HashMap<String, String>` for properties, this struct stores pre-serialized
68/// JSONB data to eliminate redundant serialization overhead.
69#[derive(Debug, Clone)]
70pub struct StreamMetadata {
71    // Core fields (same as StreamInfo)
72    pub process_id: Uuid,
73    pub stream_id: Uuid,
74    pub dependencies_metadata: Vec<UserDefinedType>,
75    pub objects_metadata: Vec<UserDefinedType>,
76    pub tags: Vec<String>,
77    pub properties: SharedJsonbSerialized,
78}
79
80impl StreamMetadata {
81    /// Creates StreamMetadata from StreamInfo by converting properties to JSONB format.
82    pub fn from_stream_info(stream_info: &StreamInfo) -> Result<Self> {
83        let properties = serialize_properties_to_jsonb(&stream_info.properties)
84            .with_context(|| "serializing stream properties to JSONB")?;
85        Ok(Self {
86            process_id: stream_info.process_id,
87            stream_id: stream_info.stream_id,
88            dependencies_metadata: stream_info.dependencies_metadata.clone(),
89            objects_metadata: stream_info.objects_metadata.clone(),
90            tags: stream_info.tags.clone(),
91            properties: Arc::new(properties),
92        })
93    }
94}
95
96/// Returns the thread name associated with the stream, if available.
97/// This function is only meaningful for streams associated with CPU threads.
98/// Returns format: "thread-name-thread-id" (e.g., "main-12345") or just "thread-id" if no name.
99pub fn get_thread_name_from_stream_metadata(stream: &StreamMetadata) -> Result<String> {
100    use jsonb::RawJsonb;
101
102    const THREAD_NAME_KEY: &str = "thread-name";
103    const THREAD_ID_KEY: &str = "thread-id";
104
105    if stream.properties.is_empty() {
106        return Ok(format!("{}", stream.stream_id));
107    }
108
109    let jsonb = RawJsonb::new(&stream.properties);
110
111    let thread_name = if let Ok(Some(thread_name_value)) = jsonb.get_by_name(THREAD_NAME_KEY, false)
112        && let Ok(Some(name)) = thread_name_value.as_raw().as_str()
113    {
114        Some(name.to_string())
115    } else {
116        None
117    };
118
119    // thread_id falls back to stream_id
120    let thread_id = if let Ok(Some(thread_id_value)) = jsonb.get_by_name(THREAD_ID_KEY, false)
121        && let Ok(Some(id)) = thread_id_value.as_raw().as_str()
122    {
123        id.to_string()
124    } else {
125        stream.stream_id.to_string()
126    };
127
128    // Return "name-id" if name exists, otherwise just "id"
129    match thread_name {
130        Some(name) => Ok(format!("{}-{}", name, thread_id)),
131        None => Ok(thread_id),
132    }
133}
134
135/// Creates a `StreamMetadata` from a row of the `streams` view's (unprefixed) columns.
136#[span_fn]
137pub fn stream_metadata_from_batch_row(batch: &RecordBatch, row: usize) -> Result<StreamMetadata> {
138    let stream_id_column = string_column_by_name(batch, "stream_id")?;
139    let process_id_column = string_column_by_name(batch, "process_id")?;
140    let dependencies_metadata_column: &BinaryArray =
141        typed_column_by_name(batch, "dependencies_metadata")?;
142    let objects_metadata_column: &BinaryArray = typed_column_by_name(batch, "objects_metadata")?;
143    let tags_column: &GenericListArray<i32> = typed_column_by_name(batch, "tags")?;
144    let properties_accessor = properties_column_by_name(batch, "properties")
145        .with_context(|| "accessing properties column")?;
146
147    let stream_id =
148        Uuid::parse_str(stream_id_column.value(row)?).with_context(|| "parsing stream_id")?;
149    let process_id =
150        Uuid::parse_str(process_id_column.value(row)?).with_context(|| "parsing process_id")?;
151
152    let dependencies_metadata: Vec<UserDefinedType> =
153        ciborium::from_reader(dependencies_metadata_column.value(row))
154            .with_context(|| "decoding dependencies_metadata")?;
155    let objects_metadata: Vec<UserDefinedType> =
156        ciborium::from_reader(objects_metadata_column.value(row))
157            .with_context(|| "decoding objects_metadata")?;
158    let tags: Vec<String> = tags_column
159        .value(row)
160        .as_any()
161        .downcast_ref::<StringArray>()
162        .with_context(|| "casting tags")?
163        .iter()
164        .map(|item| String::from(item.unwrap_or_default()))
165        .collect();
166
167    let properties = properties_accessor
168        .jsonb_value(row)
169        .with_context(|| "extracting JSONB from properties column")?;
170
171    Ok(StreamMetadata {
172        stream_id,
173        process_id,
174        dependencies_metadata,
175        objects_metadata,
176        tags,
177        properties: Arc::new(properties),
178    })
179}
180
181/// Finds a stream and its metadata using DataFusion (reads the `streams` view).
182#[span_fn]
183pub async fn find_stream_from_view(
184    lakehouse: Arc<LakehouseContext>,
185    view_factory: Arc<ViewFactory>,
186    stream_id: &Uuid,
187    query_range: Option<TimeRange>,
188) -> Result<StreamMetadata> {
189    let partition_provider = Arc::new(LivePartitionProvider::new(lakehouse.lake().db_pool.clone()));
190
191    // `ReadScope::All` here is intentional, not a latent bypass: this function is
192    // only ever called from `thread_spans_view.rs`'s `jit_update` implementation, which
193    // `MaterializedView::scan` runs *before* the caller's own scan, and it returns stream
194    // *metadata* used to build a partition of the very view instance the caller named -- never
195    // rows returned to the caller directly. `thread_spans` is registered only via
196    // `add_view_set`, so this function is reachable only through a guarded, non-`'global'`
197    // `view_instance(...)` scan -- `OwnershipRewrite` (the row-level filter) injects no predicate there; the
198    // caller's own read of that instance is authorized instead by
199    // `AudienceGuard::authorize_view_instance` one frame up, the sole enforcement for this view
200    // set. `MaterializedView::scan` runs that check before calling `jit_update` at all, so a
201    // denied caller never reaches this function.
202    let ctx = make_session_context(
203        lakehouse,
204        partition_provider,
205        query_range,
206        view_factory,
207        Arc::new(NoOpSessionConfigurator),
208        CallerContext::internal(),
209    )
210    .await
211    .with_context(|| "creating DataFusion session context")?;
212
213    let sql = format!(
214        "SELECT stream_id, process_id, dependencies_metadata, objects_metadata, tags, properties
215         FROM streams
216         WHERE stream_id = '{stream_id}'"
217    );
218
219    let df = ctx
220        .sql(&sql)
221        .await
222        .with_context(|| "executing SQL query for stream")?;
223
224    let results = df
225        .collect()
226        .await
227        .with_context(|| "collecting results from DataFusion")?;
228
229    if results.is_empty() || results[0].num_rows() == 0 {
230        anyhow::bail!("Stream not found");
231    }
232
233    stream_metadata_from_batch_row(&results[0], 0)
234}
235
236/// Creates a `ProcessMetadata` from a database row with pre-serialized JSONB properties.
237#[span_fn]
238pub fn process_metadata_from_row(row: &sqlx::postgres::PgRow) -> Result<ProcessMetadata> {
239    let properties: Vec<Property> = row.try_get("process_properties")?;
240    let properties_map = micromegas_telemetry::property::into_hashmap(properties);
241    let serialized_properties = serialize_properties_to_jsonb(&properties_map)
242        .with_context(|| "serializing process properties to JSONB")?;
243    // A NULL here (`try_get` on a String column) is a `try_get` error, and that is correct: both
244    // producers of this row -- `find_process` and the `blocks` view's `data_sql` -- wrap the
245    // audience column in `COALESCE(..., <default>)`, so a NULL can only mean a bug in that
246    // expression or a query that bypassed it, never a legitimately never-stamped process.
247    let audience: String = row.try_get("audience")?;
248
249    Ok(ProcessMetadata {
250        process_id: row.try_get("process_id")?,
251        exe: row.try_get("exe")?,
252        username: row.try_get("username")?,
253        realname: row.try_get("realname")?,
254        computer: row.try_get("computer")?,
255        distro: row.try_get("distro")?,
256        cpu_brand: row.try_get("cpu_brand")?,
257        tsc_frequency: row.try_get("tsc_frequency")?,
258        start_time: row.try_get("start_time")?,
259        start_ticks: row.try_get("start_ticks")?,
260        parent_process_id: row.try_get("parent_process_id")?,
261        properties: Arc::new(serialized_properties),
262        audience: Arc::from(audience),
263    })
264}
265
266/// Finds a process by its ID and returns it as ProcessMetadata with pre-serialized JSONB properties.
267///
268/// One of the read sites that resolves a row's own `audience` column: `default_audience` is
269/// bound as `$2` and resolves a never-stamped (legacy, pre-v8) process's NULL column, so
270/// `view_instance('log_entries'|'measures'|'images'|'otel_spans', <process_id>)` works for such
271/// a process instead of failing in `jit_update` when `process_metadata_from_row` reads a NULL
272/// audience. Callers source it from `LakehouseContext::default_audience`.
273#[span_fn]
274pub async fn find_process(
275    pool: &sqlx::Pool<sqlx::Postgres>,
276    process_id: &sqlx::types::Uuid,
277    default_audience: &str,
278) -> Result<ProcessMetadata> {
279    let sql = format!(
280        "SELECT process_id,
281                exe,
282                username,
283                realname,
284                computer,
285                distro,
286                cpu_brand,
287                tsc_frequency,
288                start_time,
289                start_ticks,
290                parent_process_id,
291                properties as process_properties,
292                {audience_column} AS audience
293         FROM processes
294         WHERE process_id = $1;",
295        audience_column = coalesced_audience_column("processes", 2),
296    );
297    let row = instrument_named!(
298        sqlx::query(&sql)
299            .bind(process_id)
300            .bind(default_audience)
301            .fetch_one(pool),
302        "sql_select_process"
303    )
304    .await
305    .with_context(|| "select from processes")?;
306    process_metadata_from_row(&row)
307}
308
309/// Finds a process and its latest timing information using DataFusion (optimized version).
310/// Returns (ProcessMetadata, last_block_end_ticks, last_block_end_time)
311#[span_fn]
312pub async fn find_process_with_latest_timing(
313    lakehouse: Arc<LakehouseContext>,
314    view_factory: Arc<ViewFactory>,
315    process_id: &Uuid,
316    query_range: Option<TimeRange>,
317) -> Result<(ProcessMetadata, i64, DateTime<Utc>)> {
318    let partition_provider = Arc::new(LivePartitionProvider::new(lakehouse.lake().db_pool.clone()));
319
320    // `ReadScope::All` here is intentional, not a latent bypass: this function is called from
321    // three JIT callers -- `net_spans_view.rs`, `async_events_view.rs`, and
322    // `thread_spans_view.rs` -- and all three view sets are registered only via `add_view_set`,
323    // reachable only through a guarded, non-`'global'` `view_instance(...)` scan, so
324    // `AudienceGuard::authorize_view_instance` is the sole enforcement for all three, same
325    // reasoning as `find_stream_from_view`'s comment above though not identical to it. Each call
326    // happens before the caller's own scan, and returns process *metadata* used to build the
327    // caller's own partition rather than rows returned to the caller directly.
328    let ctx = make_session_context(
329        lakehouse,
330        partition_provider,
331        query_range,
332        view_factory,
333        Arc::new(NoOpSessionConfigurator),
334        CallerContext::internal(),
335    )
336    .await
337    .with_context(|| "creating DataFusion session context")?;
338
339    let sql = format!(
340        "SELECT process_id, exe, username, realname, computer, distro, cpu_brand,
341                tsc_frequency, start_time, start_ticks, parent_process_id, properties,
342                last_block_end_ticks, last_block_end_time, audience
343         FROM processes
344         WHERE process_id = '{}'",
345        process_id
346    );
347
348    let df = ctx
349        .sql(&sql)
350        .await
351        .with_context(|| "executing SQL query for process with timing")?;
352
353    let results = df
354        .collect()
355        .await
356        .with_context(|| "collecting results from DataFusion")?;
357
358    if results.is_empty() || results[0].num_rows() == 0 {
359        anyhow::bail!("Process not found");
360    }
361
362    let batch = &results[0];
363
364    // Extract all the required columns
365    let process_id_column = string_column_by_name(batch, "process_id")?;
366    let exe_column = string_column_by_name(batch, "exe")?;
367    let username_column = string_column_by_name(batch, "username")?;
368    let realname_column = string_column_by_name(batch, "realname")?;
369    let computer_column = string_column_by_name(batch, "computer")?;
370    let distro_column = string_column_by_name(batch, "distro")?;
371    let cpu_brand_column = string_column_by_name(batch, "cpu_brand")?;
372    let tsc_frequency_column: &Int64Array = typed_column_by_name(batch, "tsc_frequency")?;
373    let start_time_column: &TimestampNanosecondArray = typed_column_by_name(batch, "start_time")?;
374    let start_ticks_column: &Int64Array = typed_column_by_name(batch, "start_ticks")?;
375    let last_block_end_ticks_column: &Int64Array =
376        typed_column_by_name(batch, "last_block_end_ticks")?;
377    let last_block_end_time_column: &TimestampNanosecondArray =
378        typed_column_by_name(batch, "last_block_end_time")?;
379    let parent_process_id_column = string_column_by_name(batch, "parent_process_id")?;
380    let audience_column = string_column_by_name(batch, "audience")?;
381
382    let parent_process_id = if parent_process_id_column.is_null(0) {
383        None
384    } else {
385        parse_optional_uuid(parent_process_id_column.value(0)?)?
386    };
387
388    // Handle properties column using PropertiesColumnAccessor
389    let properties_accessor = properties_column_by_name(batch, "properties")
390        .with_context(|| "accessing properties column")?;
391
392    // Get JSONB bytes directly from the properties column
393    let properties_jsonb = Arc::new(
394        properties_accessor
395            .jsonb_value(0)
396            .with_context(|| "extracting JSONB from properties column")?,
397    );
398
399    let process_metadata = ProcessMetadata {
400        process_id: parse_optional_uuid(process_id_column.value(0)?)?
401            .ok_or_else(|| anyhow::anyhow!("process_id cannot be empty"))?,
402        exe: exe_column.value(0)?.to_string(),
403        username: username_column.value(0)?.to_string(),
404        realname: realname_column.value(0)?.to_string(),
405        computer: computer_column.value(0)?.to_string(),
406        distro: distro_column.value(0)?.to_string(),
407        cpu_brand: cpu_brand_column.value(0)?.to_string(),
408        tsc_frequency: tsc_frequency_column.value(0),
409        start_time: DateTime::from_timestamp_nanos(start_time_column.value(0)),
410        start_ticks: start_ticks_column.value(0),
411        parent_process_id,
412        properties: properties_jsonb,
413        audience: Arc::from(audience_column.value(0)?),
414    };
415
416    let last_block_end_ticks = last_block_end_ticks_column.value(0);
417    let last_block_end_time = DateTime::from_timestamp_nanos(last_block_end_time_column.value(0));
418
419    Ok((process_metadata, last_block_end_ticks, last_block_end_time))
420}
421/// Creates a `BlockMetadata` from a recordbatch row.
422#[span_fn]
423pub fn block_from_batch_row(rb: &RecordBatch, row: usize) -> Result<BlockMetadata> {
424    let block_id_column = string_column_by_name(rb, "block_id")?;
425    let stream_id_column = string_column_by_name(rb, "stream_id")?;
426    let process_id_column = string_column_by_name(rb, "process_id")?;
427    let begin_time_column: &TimestampNanosecondArray = typed_column_by_name(rb, "begin_time")?;
428    let begin_ticks_column: &Int64Array = typed_column_by_name(rb, "begin_ticks")?;
429    let end_time_column: &TimestampNanosecondArray = typed_column_by_name(rb, "end_time")?;
430    let end_ticks_column: &Int64Array = typed_column_by_name(rb, "end_ticks")?;
431    let nb_objects_column: &Int32Array = typed_column_by_name(rb, "nb_objects")?;
432    let object_offset_column: &Int64Array = typed_column_by_name(rb, "object_offset")?;
433    let payload_size_column: &Int64Array = typed_column_by_name(rb, "payload_size")?;
434    let insert_time_column: &TimestampNanosecondArray = typed_column_by_name(rb, "insert_time")?;
435    Ok(BlockMetadata {
436        block_id: Uuid::parse_str(block_id_column.value(row)?)?,
437        stream_id: Uuid::parse_str(stream_id_column.value(row)?)?,
438        process_id: Uuid::parse_str(process_id_column.value(row)?)?,
439        begin_time: DateTime::from_timestamp_nanos(begin_time_column.value(row)),
440        end_time: DateTime::from_timestamp_nanos(end_time_column.value(row)),
441        begin_ticks: begin_ticks_column.value(row),
442        end_ticks: end_ticks_column.value(row),
443        nb_objects: nb_objects_column.value(row),
444        object_offset: object_offset_column.value(row),
445        payload_size: payload_size_column.value(row),
446        insert_time: DateTime::from_timestamp_nanos(insert_time_column.value(row)),
447    })
448}