micromegas_analytics/
replication.rs

1use anyhow::{Context, Result};
2use arrow_flight::decode::FlightRecordBatchStream;
3use chrono::DateTime;
4use datafusion::arrow::array::{
5    BinaryArray, GenericListArray, Int32Array, Int64Array, StringArray, TimestampNanosecondArray,
6};
7use futures::StreamExt;
8use micromegas_ingestion::data_lake_connection::DataLakeConnection;
9use micromegas_telemetry::blob_storage::PutIfAbsent;
10use micromegas_tracing::prelude::*;
11use std::sync::Arc;
12use uuid::Uuid;
13
14use crate::{
15    dfext::{string_column_accessor::string_column_by_name, typed_column::typed_column_by_name},
16    properties::{
17        properties_column_accessor::properties_column_by_name,
18        utils::extract_properties_from_properties_column,
19    },
20};
21async fn ingest_streams(
22    lake: Arc<DataLakeConnection>,
23    mut rb_stream: FlightRecordBatchStream,
24) -> Result<i64> {
25    let mut tr = lake.db_pool.begin().await?;
26    let mut nb_rows: i64 = 0;
27    while let Some(res) = rb_stream.next().await {
28        let b = res?;
29        nb_rows += b.num_rows() as i64;
30        let stream_id_column = string_column_by_name(&b, "stream_id")?;
31        let process_id_column = string_column_by_name(&b, "process_id")?;
32        let dependencies_metadata_column: &BinaryArray =
33            typed_column_by_name(&b, "dependencies_metadata")?;
34        let objects_metadata_column: &BinaryArray = typed_column_by_name(&b, "objects_metadata")?;
35        let tags_column: &GenericListArray<i32> = typed_column_by_name(&b, "tags")?;
36        let properties_accessor = properties_column_by_name(&b, "properties")?;
37        let insert_time_column: &TimestampNanosecondArray =
38            typed_column_by_name(&b, "insert_time")?;
39        // Hard failure (rather than a silent default) on missing `format` so a
40        // v3 source replicating into a v4 target surfaces the schema mismatch loudly.
41        let format_column = string_column_by_name(&b, "format")?;
42        // Same rationale, for `audience`: every source built from `main` exposes
43        // `streams.audience`, so a missing column means the source predates the v8 column and
44        // must fail loudly rather than silently write NULL into every replicated row.
45        let audience_column = string_column_by_name(&b, "audience")?;
46
47        for row in 0..b.num_rows() {
48            let stream_id = Uuid::parse_str(stream_id_column.value(row)?)?;
49            let process_id = Uuid::parse_str(process_id_column.value(row)?)?;
50            let tags: Vec<String> = tags_column
51                .value(row)
52                .as_any()
53                .downcast_ref::<StringArray>()
54                .with_context(|| "casting tags")?
55                .iter()
56                .map(|item| String::from(item.unwrap_or_default()))
57                .collect();
58            let properties_map =
59                extract_properties_from_properties_column(properties_accessor.as_ref(), row)?;
60            let properties = micromegas_telemetry::property::make_properties(&properties_map);
61
62            instrument_named!(
63                sqlx::query(
64                    "INSERT INTO streams (stream_id, process_id, dependencies_metadata, objects_metadata, tags, properties, insert_time, format, audience)
65                 VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9)
66                 ON CONFLICT (stream_id) DO NOTHING;",
67                )
68                .bind(stream_id)
69                .bind(process_id)
70                .bind(dependencies_metadata_column.value(row))
71                .bind(objects_metadata_column.value(row))
72                .bind(tags)
73                .bind(properties)
74                .bind(DateTime::from_timestamp_nanos(
75                    insert_time_column.value(row),
76                ))
77                .bind(format_column.value(row)?)
78                .bind(audience_column.value(row)?)
79                .execute(&mut *tr),
80                "sql_insert_stream"
81            )
82            .await
83            .with_context(|| "inserting into streams")?;
84        }
85    }
86    tr.commit().await?;
87    info!("ingested {nb_rows} streams");
88    Ok(nb_rows)
89}
90
91async fn ingest_processes(
92    lake: Arc<DataLakeConnection>,
93    mut rb_stream: FlightRecordBatchStream,
94) -> Result<i64> {
95    let mut tr = lake.db_pool.begin().await?;
96    let mut nb_rows: i64 = 0;
97    while let Some(res) = rb_stream.next().await {
98        let b = res?;
99        nb_rows += b.num_rows() as i64;
100        let process_id_column = string_column_by_name(&b, "process_id")?;
101        let exe_column = string_column_by_name(&b, "exe")?;
102        let username_column = string_column_by_name(&b, "username")?;
103        let realname_column = string_column_by_name(&b, "realname")?;
104        let computer_column = string_column_by_name(&b, "computer")?;
105        let distro_column = string_column_by_name(&b, "distro")?;
106        let cpu_brand_column = string_column_by_name(&b, "cpu_brand")?;
107        let process_tsc_freq_column: &Int64Array = typed_column_by_name(&b, "tsc_frequency")?;
108        let start_time_column: &TimestampNanosecondArray = typed_column_by_name(&b, "start_time")?;
109        let start_ticks_column: &Int64Array = typed_column_by_name(&b, "start_ticks")?;
110        let insert_time_column: &TimestampNanosecondArray =
111            typed_column_by_name(&b, "insert_time")?;
112        let parent_process_id_column = string_column_by_name(&b, "parent_process_id")?;
113        let properties_accessor = properties_column_by_name(&b, "properties")?;
114        // Hard failure (rather than a silent default) on missing `audience`, matching the
115        // precedent `ingest_streams` set for `format` in v4: every source built from `main`
116        // exposes `processes.audience`, so a missing column means the source predates the v8
117        // column.
118        let audience_column = string_column_by_name(&b, "audience")?;
119        for row in 0..b.num_rows() {
120            let process_id = Uuid::parse_str(process_id_column.value(row)?)?;
121            let parent_process_str = parent_process_id_column.value(row)?;
122            let parent_process_id = if parent_process_str.is_empty() {
123                None
124            } else {
125                Some(Uuid::parse_str(parent_process_str)?)
126            };
127            let properties_map =
128                extract_properties_from_properties_column(properties_accessor.as_ref(), row)?;
129            let properties = micromegas_telemetry::property::make_properties(&properties_map);
130            instrument_named!(
131                sqlx::query(
132                    "INSERT INTO processes
133                     (process_id, exe, username, realname, computer, distro, cpu_brand,
134                      tsc_frequency, start_time, start_ticks, insert_time, parent_process_id,
135                      properties, audience)
136                     VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14)
137                     ON CONFLICT (process_id) DO NOTHING;",
138                )
139                .bind(process_id)
140                .bind(exe_column.value(row)?)
141                .bind(username_column.value(row)?)
142                .bind(realname_column.value(row)?)
143                .bind(computer_column.value(row)?)
144                .bind(distro_column.value(row)?)
145                .bind(cpu_brand_column.value(row)?)
146                .bind(process_tsc_freq_column.value(row))
147                .bind(DateTime::from_timestamp_nanos(start_time_column.value(row)))
148                .bind(start_ticks_column.value(row))
149                .bind(DateTime::from_timestamp_nanos(
150                    insert_time_column.value(row),
151                ))
152                .bind(parent_process_id)
153                .bind(properties)
154                .bind(audience_column.value(row)?)
155                .execute(&mut *tr),
156                "sql_insert_process"
157            )
158            .await
159            .with_context(|| "executing sql insert into processes")?;
160        }
161    }
162    tr.commit().await?;
163    info!("ingested {nb_rows} processes");
164    Ok(nb_rows)
165}
166
167async fn ingest_payloads(
168    lake: Arc<DataLakeConnection>,
169    mut rb_stream: FlightRecordBatchStream,
170) -> Result<i64> {
171    let mut nb_rows: i64 = 0;
172    while let Some(res) = rb_stream.next().await {
173        let b = res?;
174        nb_rows += b.num_rows() as i64;
175        let process_id_column = string_column_by_name(&b, "process_id")?;
176        let stream_id_column = string_column_by_name(&b, "stream_id")?;
177        let block_id_column = string_column_by_name(&b, "block_id")?;
178        let payload_column: &BinaryArray = typed_column_by_name(&b, "payload")?;
179        for row in 0..b.num_rows() {
180            let process_id = process_id_column.value(row)?;
181            let stream_id = stream_id_column.value(row)?;
182            let block_id = block_id_column.value(row)?;
183            let obj_path = format!("blobs/{process_id}/{stream_id}/{block_id}");
184            let payload = bytes::Bytes::copy_from_slice(payload_column.value(row));
185            // Create-only: the block-object namespace is write-once and content-addressed
186            // (see `micromegas_telemetry::blob_storage::put_if_absent`). A colliding write
187            // here just means the object was already replicated; `ingest_blocks` handles
188            // row-level dedup independently via its own `ON CONFLICT DO NOTHING` INSERT.
189            let put_outcome = lake
190                .blob_storage
191                .put_if_absent(&obj_path, payload)
192                .await
193                .with_context(|| "Error writing block to blob storage")?;
194            if put_outcome == PutIfAbsent::AlreadyExists {
195                debug!("block payload already exists at {obj_path}, skipping");
196            }
197        }
198    }
199    info!("ingested {nb_rows} payloads");
200    Ok(nb_rows)
201}
202
203async fn ingest_blocks(
204    lake: Arc<DataLakeConnection>,
205    mut rb_stream: FlightRecordBatchStream,
206) -> Result<i64> {
207    let mut tr = lake.db_pool.begin().await?;
208    let mut nb_rows: i64 = 0;
209    while let Some(res) = rb_stream.next().await {
210        let b = res?;
211        nb_rows += b.num_rows() as i64;
212        let block_id_column = string_column_by_name(&b, "block_id")?;
213        let stream_id_column = string_column_by_name(&b, "stream_id")?;
214        let process_id_column = string_column_by_name(&b, "process_id")?;
215        let begin_time_column: &TimestampNanosecondArray = typed_column_by_name(&b, "begin_time")?;
216        let begin_ticks_column: &Int64Array = typed_column_by_name(&b, "begin_ticks")?;
217        let end_time_column: &TimestampNanosecondArray = typed_column_by_name(&b, "end_time")?;
218        let end_ticks_column: &Int64Array = typed_column_by_name(&b, "end_ticks")?;
219        let nb_objects_column: &Int32Array = typed_column_by_name(&b, "nb_objects")?;
220        let object_offset_column: &Int64Array = typed_column_by_name(&b, "object_offset")?;
221        let payload_size_column: &Int64Array = typed_column_by_name(&b, "payload_size")?;
222        let insert_time_column: &TimestampNanosecondArray =
223            typed_column_by_name(&b, "insert_time")?;
224        // Hard failure (rather than a silent default) on missing `audience`, matching the
225        // precedent `ingest_streams` set for `format` in v4: every source built from `main`
226        // exposes `blocks.audience`, so a missing column means the source predates the v8
227        // column.
228        let audience_column = string_column_by_name(&b, "audience")?;
229        for row in 0..b.num_rows() {
230            let block_id = Uuid::parse_str(block_id_column.value(row)?)?;
231            let stream_id = Uuid::parse_str(stream_id_column.value(row)?)?;
232            let process_id = Uuid::parse_str(process_id_column.value(row)?)?;
233            instrument_named!(
234                sqlx::query(
235                    "INSERT INTO blocks (block_id, stream_id, process_id, begin_time, begin_ticks,
236                      end_time, end_ticks, nb_objects, object_offset, payload_size, insert_time,
237                      audience)
238                     VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12)
239                     ON CONFLICT (block_id) DO NOTHING;",
240                )
241                .bind(block_id)
242                .bind(stream_id)
243                .bind(process_id)
244                .bind(DateTime::from_timestamp_nanos(begin_time_column.value(row)))
245                .bind(begin_ticks_column.value(row))
246                .bind(DateTime::from_timestamp_nanos(end_time_column.value(row)))
247                .bind(end_ticks_column.value(row))
248                .bind(nb_objects_column.value(row))
249                .bind(object_offset_column.value(row))
250                .bind(payload_size_column.value(row))
251                .bind(DateTime::from_timestamp_nanos(
252                    insert_time_column.value(row),
253                ))
254                .bind(audience_column.value(row)?)
255                .execute(&mut *tr),
256                "sql_insert_block"
257            )
258            .await
259            .with_context(|| "executing sql insert into blocks")?;
260        }
261    }
262    tr.commit().await?;
263    info!("ingested {nb_rows} blocks");
264    Ok(nb_rows)
265}
266
267/// Ingests data from a FlightRecordBatchStream into the specified table.
268pub async fn bulk_ingest(
269    lake: Arc<DataLakeConnection>,
270    table_name: &str,
271    rb_stream: FlightRecordBatchStream,
272) -> Result<i64> {
273    match table_name {
274        "processes" => ingest_processes(lake, rb_stream).await,
275        "streams" => ingest_streams(lake, rb_stream).await,
276        "blocks" => ingest_blocks(lake, rb_stream).await,
277        "payloads" => ingest_payloads(lake, rb_stream).await,
278        other => anyhow::bail!("bulk ingest for table {other} not supported"),
279    }
280}
micromegas_analytics/replication.rs

micromegas_analytics/
replication.rs
