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};
28pub type SharedJsonbSerialized = Arc<Vec<u8>>;
31
32#[derive(Debug, Clone)]
39pub struct ProcessMetadata {
40 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 pub audience: Arc<str>,
61}
62
63#[derive(Debug, Clone)]
70pub struct StreamMetadata {
71 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 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
96pub 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 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 match thread_name {
130 Some(name) => Ok(format!("{}-{}", name, thread_id)),
131 None => Ok(thread_id),
132 }
133}
134
135#[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#[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 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#[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 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#[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#[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 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 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 let properties_accessor = properties_column_by_name(batch, "properties")
390 .with_context(|| "accessing properties column")?;
391
392 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#[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}