micromegas_analytics/lakehouse/
migration.rs

1use crate::arrow_utils::parse_parquet_metadata;
2use anyhow::{Context, Result};
3use micromegas_ingestion::remote_data_lake::acquire_lock;
4use micromegas_tracing::prelude::*;
5use sqlx::Executor;
6use sqlx::Row;
7
8pub const LATEST_LAKEHOUSE_SCHEMA_VERSION: i32 = 7;
9
10async fn read_lakehouse_schema_version(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> i32 {
11    match sqlx::query(
12        "SELECT version
13         FROM lakehouse_migration;",
14    )
15    .fetch_one(&mut **tr)
16    .await
17    {
18        Ok(row) => row.get("version"),
19        Err(e) => {
20            info!(
21                "Error reading data lake schema version, assuming version 0: {}",
22                e
23            );
24            0
25        }
26    }
27}
28
29/// Migrates the lakehouse schema to the latest version.
30pub async fn migrate_lakehouse(pool: sqlx::Pool<sqlx::Postgres>) -> Result<()> {
31    let mut tr = pool.begin().await?;
32    let mut current_version = read_lakehouse_schema_version(&mut tr).await;
33    drop(tr);
34    info!("current lakehouse schema: {}", current_version);
35    if current_version != LATEST_LAKEHOUSE_SCHEMA_VERSION {
36        let mut tr = pool.begin().await?;
37        acquire_lock(&mut tr, 1).await?;
38        current_version = read_lakehouse_schema_version(&mut pool.begin().await?).await;
39        if LATEST_LAKEHOUSE_SCHEMA_VERSION == current_version {
40            return Ok(());
41        }
42        if let Err(e) = execute_lakehouse_migration(pool.clone()).await {
43            error!("Error migrating database: {}", e);
44            return Err(e);
45        }
46        current_version = read_lakehouse_schema_version(&mut tr).await;
47    }
48    assert_eq!(current_version, LATEST_LAKEHOUSE_SCHEMA_VERSION);
49    Ok(())
50}
51
52/// Executes the lakehouse migration steps.
53async fn execute_lakehouse_migration(pool: sqlx::Pool<sqlx::Postgres>) -> Result<()> {
54    let mut current_version = read_lakehouse_schema_version(&mut pool.begin().await?).await;
55    if 0 == current_version {
56        info!("creating v1 lakehouse_schema");
57        let mut tr = pool.begin().await?;
58        create_tables(&mut tr).await?;
59        current_version = read_lakehouse_schema_version(&mut tr).await;
60        tr.commit().await?;
61    }
62    if 1 == current_version {
63        info!("upgrade lakehouse schema to v2");
64        let mut tr = pool.begin().await?;
65        upgrade_v1_to_v2(&mut tr).await?;
66        current_version = read_lakehouse_schema_version(&mut tr).await;
67        tr.commit().await?;
68    }
69    if 2 == current_version {
70        info!("upgrade lakehouse schema to v3");
71        let mut tr = pool.begin().await?;
72        upgrade_v2_to_v3(&mut tr).await?;
73        current_version = read_lakehouse_schema_version(&mut tr).await;
74        tr.commit().await?;
75    }
76    if 3 == current_version {
77        info!("upgrade lakehouse schema to v4");
78        let mut tr = pool.begin().await?;
79        upgrade_v3_to_v4(&mut tr).await?;
80        current_version = read_lakehouse_schema_version(&mut tr).await;
81        tr.commit().await?;
82    }
83    if 4 == current_version {
84        info!("upgrade lakehouse schema to v5");
85        let mut tr = pool.begin().await?;
86        upgrade_v4_to_v5(&mut tr).await?;
87        current_version = read_lakehouse_schema_version(&mut tr).await;
88        tr.commit().await?;
89    }
90    if 5 == current_version {
91        info!("upgrade lakehouse schema to v6");
92        let mut tr = pool.begin().await?;
93        upgrade_v5_to_v6(&mut tr).await?;
94        current_version = read_lakehouse_schema_version(&mut tr).await;
95        tr.commit().await?;
96    }
97    if 6 == current_version {
98        info!("upgrade lakehouse schema to v7");
99        let mut tr = pool.begin().await?;
100        upgrade_v6_to_v7(&mut tr).await?;
101        current_version = read_lakehouse_schema_version(&mut tr).await;
102        tr.commit().await?;
103    }
104    assert_eq!(current_version, LATEST_LAKEHOUSE_SCHEMA_VERSION);
105    Ok(())
106}
107
108async fn create_partitions_table(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
109    // Every table in the set shares the same schema - that schema can change over time, in which case the partition has to be rebuilt.
110    // The instance id is a unique key to the table within the table set.
111    //    * Example 1: if there is a process_log table for each process, the instance id could be the process_id.
112    //      It would not clash with an instance of process_metrics table for the same process.
113    //    * Example 2:  if there is a table instance for each metric, the view_instance_id could be the name of the metric.
114
115    // source_data_hash can be used to detect that the partition is out of date (sha1 of the block_ids, for example)
116    tr.execute("
117         CREATE TABLE lakehouse_partitions(
118                  view_set_name VARCHAR(255),
119                  view_instance_id VARCHAR(255),
120                  begin_insert_time TIMESTAMPTZ,
121                  end_insert_time TIMESTAMPTZ,
122                  min_event_time TIMESTAMPTZ,
123                  max_event_time TIMESTAMPTZ,
124                  updated TIMESTAMPTZ,
125                  file_path VARCHAR(2047),
126                  file_size BIGINT,
127                  file_schema_hash bytea,
128                  source_data_hash bytea
129                  );
130         CREATE INDEX lh_part_begin_insert on lakehouse_partitions(view_set_name, view_instance_id, begin_insert_time);
131         CREATE INDEX lh_part_end_insert on lakehouse_partitions(view_set_name, view_instance_id, end_insert_time);
132         CREATE INDEX lh_part_min_time on lakehouse_partitions(view_set_name, view_instance_id, min_event_time);
133         CREATE INDEX lh_part_max_time on lakehouse_partitions(view_set_name, view_instance_id, max_event_time);
134")
135        .await
136        .with_context(|| "Creating table blocks and its indices")?;
137    Ok(())
138}
139
140async fn create_temp_files_table(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
141    // partitions that are out of date can still be referenced until they expire
142    tr.execute(
143        "
144         CREATE TABLE temporary_files(
145                  file_path VARCHAR(2047),
146                  file_size BIGINT,
147                  expiration TIMESTAMPTZ );
148         CREATE INDEX temporary_files_expiration on temporary_files(expiration);
149",
150    )
151    .await
152    .with_context(|| "Creating temporary_files table")?;
153    Ok(())
154}
155
156async fn create_migration_table(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
157    sqlx::query("CREATE table lakehouse_migration(version integer);")
158        .execute(&mut **tr)
159        .await
160        .with_context(|| "Creating table lakehouse_migration")?;
161    sqlx::query("INSERT INTO lakehouse_migration VALUES(1);")
162        .execute(&mut **tr)
163        .await
164        .with_context(|| "Recording the initial lakehouse schema version")?;
165    Ok(())
166}
167
168async fn create_tables(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
169    create_partitions_table(tr).await?;
170    create_temp_files_table(tr).await?;
171    create_migration_table(tr).await?;
172    Ok(())
173}
174
175async fn upgrade_v1_to_v2(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
176    // file_metadata is meant to be the serialized ParquetMetaData
177    // which can be found in the footer of the file
178    tr.execute("ALTER TABLE lakehouse_partitions ADD file_metadata bytea;")
179        .await
180        .with_context(|| "adding column file_metadata to lakehouse_partitions table")?;
181    tr.execute("UPDATE lakehouse_migration SET version=2;")
182        .await
183        .with_context(|| "Updating lakehouse schema version to 2")?;
184    Ok(())
185}
186
187async fn upgrade_v2_to_v3(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
188    // Add num_rows column to store row count separately from file_metadata
189    tr.execute("ALTER TABLE lakehouse_partitions ADD num_rows BIGINT;")
190        .await
191        .with_context(|| "adding column num_rows to lakehouse_partitions table")?;
192
193    // Add index on file_path for efficient on-demand metadata loading
194    tr.execute("CREATE INDEX lakehouse_partitions_file_path ON lakehouse_partitions(file_path);")
195        .await
196        .with_context(|| "creating index on file_path")?;
197
198    // Populate num_rows column for existing partitions
199    populate_num_rows_column(tr)
200        .await
201        .with_context(|| "populating num_rows column")?;
202
203    // Make num_rows NOT NULL after populating existing data
204    tr.execute("ALTER TABLE lakehouse_partitions ALTER COLUMN num_rows SET NOT NULL;")
205        .await
206        .with_context(|| "setting num_rows column to NOT NULL")?;
207
208    tr.execute("UPDATE lakehouse_migration SET version=3;")
209        .await
210        .with_context(|| "Updating lakehouse schema version to 3")?;
211    Ok(())
212}
213
214async fn populate_num_rows_column(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
215    info!("populating num_rows column for existing partitions");
216
217    let mut total_count = 0;
218    let batch_size = 1000;
219
220    loop {
221        // Fetch partitions in batches to avoid loading all metadata into memory
222        let rows = sqlx::query("SELECT file_path, file_metadata FROM lakehouse_partitions WHERE file_metadata IS NOT NULL AND num_rows IS NULL LIMIT $1")
223            .bind(batch_size)
224            .fetch_all(&mut **tr)
225            .await?;
226
227        if rows.is_empty() {
228            break;
229        }
230
231        let mut batch_count = 0;
232        for row in rows {
233            let file_path: String = row.try_get("file_path")?;
234            let file_metadata_buffer: Vec<u8> = row.try_get("file_metadata")?;
235
236            // Parse metadata only for this partition
237            match parse_parquet_metadata(&file_metadata_buffer.into()) {
238                Ok(file_metadata) => {
239                    let num_rows = file_metadata.file_metadata().num_rows();
240
241                    // Update just this partition
242                    if let Err(e) = sqlx::query(
243                        "UPDATE lakehouse_partitions SET num_rows = $1 WHERE file_path = $2",
244                    )
245                    .bind(num_rows)
246                    .bind(&file_path)
247                    .execute(&mut **tr)
248                    .await
249                    {
250                        error!(
251                            "failed to update num_rows for partition {}: {}",
252                            file_path, e
253                        );
254                        continue;
255                    }
256
257                    batch_count += 1;
258                }
259                Err(e) => {
260                    error!(
261                        "failed to parse metadata for partition {}: {}",
262                        file_path, e
263                    );
264                    // For partitions with unparseable metadata, set num_rows to 0 as a fallback
265                    if let Err(e2) = sqlx::query(
266                        "UPDATE lakehouse_partitions SET num_rows = 0 WHERE file_path = $1",
267                    )
268                    .bind(&file_path)
269                    .execute(&mut **tr)
270                    .await
271                    {
272                        error!(
273                            "failed to set fallback num_rows for partition {}: {}",
274                            file_path, e2
275                        );
276                    }
277                }
278            }
279        }
280
281        total_count += batch_count;
282        info!(
283            "populated num_rows for {} partitions (total: {})",
284            batch_count, total_count
285        );
286    }
287
288    info!("populated num_rows for {} total partitions", total_count);
289    Ok(())
290}
291
292async fn upgrade_v3_to_v4(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
293    // Create dedicated partition_metadata table
294    tr.execute(
295        "CREATE TABLE partition_metadata(
296            file_path VARCHAR(2047) PRIMARY KEY,
297            metadata bytea NOT NULL,
298            insert_time TIMESTAMPTZ NOT NULL
299        );",
300    )
301    .await
302    .with_context(|| "creating partition_metadata table")?;
303
304    // Migrate existing metadata from lakehouse_partitions to partition_metadata
305    migrate_metadata_to_new_table(tr)
306        .await
307        .with_context(|| "migrating metadata to partition_metadata table")?;
308
309    // Drop the file_metadata column from lakehouse_partitions after successful migration
310    tr.execute("ALTER TABLE lakehouse_partitions DROP COLUMN file_metadata;")
311        .await
312        .with_context(|| "dropping file_metadata column from lakehouse_partitions")?;
313
314    tr.execute("UPDATE lakehouse_migration SET version=4;")
315        .await
316        .with_context(|| "Updating lakehouse schema version to 4")?;
317    Ok(())
318}
319
320async fn migrate_metadata_to_new_table(
321    tr: &mut sqlx::Transaction<'_, sqlx::Postgres>,
322) -> Result<()> {
323    info!("migrating metadata to partition_metadata table");
324
325    // First, get all file paths that have metadata (small data)
326    let file_paths: Vec<String> = sqlx::query_scalar(
327        "SELECT file_path 
328         FROM lakehouse_partitions 
329         WHERE file_metadata IS NOT NULL
330         ORDER BY file_path",
331    )
332    .fetch_all(&mut **tr)
333    .await?;
334
335    let total_to_migrate = file_paths.len();
336    info!(
337        "found {} partitions with metadata to migrate",
338        total_to_migrate
339    );
340
341    let mut total_count = 0;
342    let batch_size = 10; // Small batch size since metadata can be large
343
344    // Process in batches to avoid loading too much metadata at once
345    for chunk in file_paths.chunks(batch_size) {
346        // Build a query to fetch just this batch
347        let placeholders: Vec<String> = (1..=chunk.len()).map(|i| format!("${}", i)).collect();
348        let query_str = format!(
349            "SELECT file_path, file_metadata, updated 
350             FROM lakehouse_partitions 
351             WHERE file_path IN ({})",
352            placeholders.join(", ")
353        );
354
355        let mut query = sqlx::query(&query_str);
356        for path in chunk {
357            query = query.bind(path);
358        }
359
360        let rows = query.fetch_all(&mut **tr).await?;
361
362        for row in rows {
363            let file_path: String = row.try_get("file_path")?;
364            let file_metadata: Vec<u8> = row.try_get("file_metadata")?;
365            let updated: chrono::DateTime<chrono::Utc> = row.try_get("updated")?;
366
367            // Insert into new partition_metadata table (with ON CONFLICT for migration safety)
368            if let Err(e) = sqlx::query(
369                "INSERT INTO partition_metadata (file_path, metadata, insert_time) 
370                 VALUES ($1, $2, $3)
371                 ON CONFLICT (file_path) DO NOTHING",
372            )
373            .bind(&file_path)
374            .bind(&file_metadata)
375            .bind(updated)
376            .execute(&mut **tr)
377            .await
378            {
379                error!(
380                    "failed to migrate metadata for partition {}: {}",
381                    file_path, e
382                );
383                continue;
384            }
385
386            total_count += 1;
387        }
388
389        if total_count % 100 == 0 || total_count == total_to_migrate {
390            info!(
391                "migrated {}/{} partition metadata entries",
392                total_count, total_to_migrate
393            );
394        }
395    }
396
397    info!("migrated metadata for {} total partitions", total_count);
398    Ok(())
399}
400
401async fn upgrade_v4_to_v5(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
402    // Add partition_format_version column to lakehouse_partitions
403    // Default to 1 for existing partitions (Arrow 56.0 format)
404    tr.execute(
405        "ALTER TABLE lakehouse_partitions
406         ADD COLUMN partition_format_version INTEGER NOT NULL DEFAULT 1;",
407    )
408    .await
409    .with_context(|| "adding partition_format_version to lakehouse_partitions")?;
410    // Add partition_format_version column to partition_metadata
411    // Default to 1 for existing metadata (Arrow 56.0 format)
412    tr.execute(
413        "ALTER TABLE partition_metadata
414         ADD COLUMN partition_format_version INTEGER NOT NULL DEFAULT 1;",
415    )
416    .await
417    .with_context(|| "adding partition_format_version to partition_metadata")?;
418    tr.execute("UPDATE lakehouse_migration SET version=5;")
419        .await
420        .with_context(|| "Updating lakehouse schema version to 5")?;
421    info!("added partition_format_version columns to both tables");
422    Ok(())
423}
424
425async fn upgrade_v5_to_v6(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
426    tr.execute("DROP TABLE partition_metadata;")
427        .await
428        .with_context(|| "dropping partition_metadata table")?;
429    tr.execute("UPDATE lakehouse_migration SET version=6;")
430        .await
431        .with_context(|| "Updating lakehouse schema version to 6")?;
432    Ok(())
433}
434
435async fn upgrade_v6_to_v7(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
436    // Nullable, no backfill needed: NULL means "no ordering guarantee", which is automatically
437    // correct for every partition written before this column existed (see
438    // tasks/blocks_view_ordered_merges_plan.md's Design §4).
439    tr.execute("ALTER TABLE lakehouse_partitions ADD COLUMN sort_order TEXT[];")
440        .await
441        .with_context(|| "adding sort_order column to lakehouse_partitions")?;
442    tr.execute("CREATE EXTENSION IF NOT EXISTS btree_gist;")
443        .await
444        .with_context(|| "creating btree_gist extension (required by the partition overlap exclusion constraint); on PostgreSQL <= 12, or if this role lacks CREATE on the database, a superuser must run CREATE EXTENSION btree_gist once")?;
445    // Detect-then-fail: the exclusion constraint below cannot be added NOT VALID, so if legacy
446    // overlapping partitions exist, surface exactly which rows conflict instead of failing with a
447    // raw constraint error. Zero-width ranges are excluded to match tstzrange semantics (an empty
448    // range never overlaps anything). Scoped to same-schema pairs, matching the constraint: a
449    // cross-schema overlap is the legal schema-rollout coexistence state, not a conflict.
450    let conflicts = sqlx::query(
451        "SELECT a.view_set_name, a.view_instance_id,
452                a.file_path AS file_path_a, b.file_path AS file_path_b,
453                a.begin_insert_time AS begin_a, a.end_insert_time AS end_a,
454                b.begin_insert_time AS begin_b, b.end_insert_time AS end_b
455         FROM lakehouse_partitions a
456         JOIN lakehouse_partitions b
457           ON a.view_set_name = b.view_set_name
458          AND a.view_instance_id = b.view_instance_id
459          AND a.file_schema_hash = b.file_schema_hash
460          AND a.ctid < b.ctid
461          AND a.begin_insert_time < b.end_insert_time
462          AND a.end_insert_time > b.begin_insert_time
463          AND a.begin_insert_time < a.end_insert_time
464          AND b.begin_insert_time < b.end_insert_time
465         LIMIT 20;",
466    )
467    .fetch_all(&mut **tr)
468    .await
469    .with_context(|| "detecting overlapping partitions before adding exclusion constraint")?;
470    if !conflicts.is_empty() {
471        let mut msg = String::from(
472            "lakehouse_partitions contains partitions with overlapping insert-time ranges; \
473             retire them (e.g. retire_partition_by_metadata) and restart to complete the \
474             migration. Conflicting pairs (first 20):",
475        );
476        for row in &conflicts {
477            let view_set_name: String = row.try_get("view_set_name")?;
478            let view_instance_id: String = row.try_get("view_instance_id")?;
479            let file_path_a: Option<String> = row.try_get("file_path_a")?;
480            let file_path_b: Option<String> = row.try_get("file_path_b")?;
481            let begin_a: chrono::DateTime<chrono::Utc> = row.try_get("begin_a")?;
482            let end_a: chrono::DateTime<chrono::Utc> = row.try_get("end_a")?;
483            let begin_b: chrono::DateTime<chrono::Utc> = row.try_get("begin_b")?;
484            let end_b: chrono::DateTime<chrono::Utc> = row.try_get("end_b")?;
485            msg.push_str(&format!(
486                "\n  {view_set_name}/{view_instance_id}: {file_path_a:?} [{begin_a}, {end_a}] overlaps {file_path_b:?} [{begin_b}, {end_b}]"
487            ));
488        }
489        anyhow::bail!(msg);
490    }
491    // Enforced by the index even across concurrent transactions (the second conflicting writer
492    // blocks, then errors) -- closes the race a SELECT-based check under READ COMMITTED cannot:
493    // two concurrent writers whose snapshots don't see each other's uncommitted partitions.
494    // tstzrange is '[)' so adjacent partitions sharing a boundary do not conflict, and the write
495    // path's retire+insert runs in one transaction so replacing a partition never self-conflicts.
496    // Scoped by file_schema_hash (btree_gist supports bytea equality): after a schema-hash bump,
497    // old-schema partitions legally coexist with overlapping new-schema writes until
498    // retire_incompatible_partitions cleans them up -- e.g. an old-schema daily merged partition
499    // vs a new hourly write that overlaps it without containing it, which the insert path's
500    // containment-based retire cannot remove. Queries filter partitions by schema hash, so
501    // cross-schema overlap never yields duplicate rows; only same-schema overlap is a bug.
502    tr.execute(
503        "ALTER TABLE lakehouse_partitions ADD CONSTRAINT lakehouse_partitions_no_overlap
504         EXCLUDE USING gist (
505             view_set_name WITH =,
506             view_instance_id WITH =,
507             file_schema_hash WITH =,
508             tstzrange(begin_insert_time, end_insert_time) WITH &&
509         );",
510    )
511    .await
512    .with_context(|| "adding partition overlap exclusion constraint")?;
513    tr.execute("UPDATE lakehouse_migration SET version=7;")
514        .await
515        .with_context(|| "Updating lakehouse schema version to 7")?;
516    Ok(())
517}
micromegas_analytics/lakehouse/migration.rs

micromegas_analytics/lakehouse/
migration.rs
