micromegas_analytics/lakehouse/
migration.rs1use crate::arrow_utils::parse_parquet_metadata;
2use crate::lakehouse::log_stats_view::{log_stats_ddl_text, log_stats_view_definition};
3use anyhow::{Context, Result};
4use micromegas_ingestion::remote_data_lake::acquire_lock;
5use micromegas_tracing::prelude::*;
6use sqlx::Executor;
7use sqlx::Row;
8
9pub const LATEST_LAKEHOUSE_SCHEMA_VERSION: i32 = 10;
10
11async fn read_lakehouse_schema_version(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> i32 {
12 match sqlx::query(
13 "SELECT version
14 FROM lakehouse_migration;",
15 )
16 .fetch_one(&mut **tr)
17 .await
18 {
19 Ok(row) => row.get("version"),
20 Err(e) => {
21 info!(
22 "Error reading data lake schema version, assuming version 0: {}",
23 e
24 );
25 0
26 }
27 }
28}
29
30pub async fn migrate_lakehouse(pool: sqlx::Pool<sqlx::Postgres>) -> Result<()> {
32 let mut tr = pool.begin().await?;
33 let mut current_version = read_lakehouse_schema_version(&mut tr).await;
34 drop(tr);
35 info!("current lakehouse schema: {}", current_version);
36 if current_version != LATEST_LAKEHOUSE_SCHEMA_VERSION {
37 let mut tr = pool.begin().await?;
38 acquire_lock(&mut tr, 1).await?;
39 current_version = read_lakehouse_schema_version(&mut pool.begin().await?).await;
40 if LATEST_LAKEHOUSE_SCHEMA_VERSION == current_version {
41 return Ok(());
42 }
43 if let Err(e) = execute_lakehouse_migration(pool.clone()).await {
44 error!("Error migrating database: {}", e);
45 return Err(e);
46 }
47 current_version = read_lakehouse_schema_version(&mut tr).await;
48 }
49 assert_eq!(current_version, LATEST_LAKEHOUSE_SCHEMA_VERSION);
50 Ok(())
51}
52
53async fn execute_lakehouse_migration(pool: sqlx::Pool<sqlx::Postgres>) -> Result<()> {
55 let mut current_version = read_lakehouse_schema_version(&mut pool.begin().await?).await;
56 if 0 == current_version {
57 info!("creating v1 lakehouse_schema");
58 let mut tr = pool.begin().await?;
59 create_tables(&mut tr).await?;
60 current_version = read_lakehouse_schema_version(&mut tr).await;
61 tr.commit().await?;
62 }
63 if 1 == current_version {
64 info!("upgrade lakehouse schema to v2");
65 let mut tr = pool.begin().await?;
66 upgrade_v1_to_v2(&mut tr).await?;
67 current_version = read_lakehouse_schema_version(&mut tr).await;
68 tr.commit().await?;
69 }
70 if 2 == current_version {
71 info!("upgrade lakehouse schema to v3");
72 let mut tr = pool.begin().await?;
73 upgrade_v2_to_v3(&mut tr).await?;
74 current_version = read_lakehouse_schema_version(&mut tr).await;
75 tr.commit().await?;
76 }
77 if 3 == current_version {
78 info!("upgrade lakehouse schema to v4");
79 let mut tr = pool.begin().await?;
80 upgrade_v3_to_v4(&mut tr).await?;
81 current_version = read_lakehouse_schema_version(&mut tr).await;
82 tr.commit().await?;
83 }
84 if 4 == current_version {
85 info!("upgrade lakehouse schema to v5");
86 let mut tr = pool.begin().await?;
87 upgrade_v4_to_v5(&mut tr).await?;
88 current_version = read_lakehouse_schema_version(&mut tr).await;
89 tr.commit().await?;
90 }
91 if 5 == current_version {
92 info!("upgrade lakehouse schema to v6");
93 let mut tr = pool.begin().await?;
94 upgrade_v5_to_v6(&mut tr).await?;
95 current_version = read_lakehouse_schema_version(&mut tr).await;
96 tr.commit().await?;
97 }
98 if 6 == current_version {
99 info!("upgrade lakehouse schema to v7");
100 let mut tr = pool.begin().await?;
101 upgrade_v6_to_v7(&mut tr).await?;
102 current_version = read_lakehouse_schema_version(&mut tr).await;
103 tr.commit().await?;
104 }
105 if 7 == current_version {
106 info!("upgrade lakehouse schema to v8");
107 let mut tr = pool.begin().await?;
108 upgrade_v7_to_v8(&mut tr).await?;
109 current_version = read_lakehouse_schema_version(&mut tr).await;
110 tr.commit().await?;
111 }
112 if 8 == current_version {
113 info!("upgrade lakehouse schema to v9");
114 let mut tr = pool.begin().await?;
115 upgrade_v8_to_v9(&mut tr).await?;
116 current_version = read_lakehouse_schema_version(&mut tr).await;
117 tr.commit().await?;
118 }
119 if 9 == current_version {
120 info!("upgrade lakehouse schema to v10");
121 let mut tr = pool.begin().await?;
122 upgrade_v9_to_v10(&mut tr).await?;
123 current_version = read_lakehouse_schema_version(&mut tr).await;
124 tr.commit().await?;
125 }
126 assert_eq!(current_version, LATEST_LAKEHOUSE_SCHEMA_VERSION);
127 Ok(())
128}
129
130async fn create_partitions_table(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
131 tr.execute("
139 CREATE TABLE lakehouse_partitions(
140 view_set_name VARCHAR(255),
141 view_instance_id VARCHAR(255),
142 begin_insert_time TIMESTAMPTZ,
143 end_insert_time TIMESTAMPTZ,
144 min_event_time TIMESTAMPTZ,
145 max_event_time TIMESTAMPTZ,
146 updated TIMESTAMPTZ,
147 file_path VARCHAR(2047),
148 file_size BIGINT,
149 file_schema_hash bytea,
150 source_data_hash bytea
151 );
152 CREATE INDEX lh_part_begin_insert on lakehouse_partitions(view_set_name, view_instance_id, begin_insert_time);
153 CREATE INDEX lh_part_end_insert on lakehouse_partitions(view_set_name, view_instance_id, end_insert_time);
154 CREATE INDEX lh_part_min_time on lakehouse_partitions(view_set_name, view_instance_id, min_event_time);
155 CREATE INDEX lh_part_max_time on lakehouse_partitions(view_set_name, view_instance_id, max_event_time);
156")
157 .await
158 .with_context(|| "Creating table blocks and its indices")?;
159 Ok(())
160}
161
162async fn create_temp_files_table(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
163 tr.execute(
165 "
166 CREATE TABLE temporary_files(
167 file_path VARCHAR(2047),
168 file_size BIGINT,
169 expiration TIMESTAMPTZ );
170 CREATE INDEX temporary_files_expiration on temporary_files(expiration);
171",
172 )
173 .await
174 .with_context(|| "Creating temporary_files table")?;
175 Ok(())
176}
177
178async fn create_migration_table(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
179 sqlx::query("CREATE table lakehouse_migration(version integer);")
180 .execute(&mut **tr)
181 .await
182 .with_context(|| "Creating table lakehouse_migration")?;
183 sqlx::query("INSERT INTO lakehouse_migration VALUES(1);")
184 .execute(&mut **tr)
185 .await
186 .with_context(|| "Recording the initial lakehouse schema version")?;
187 Ok(())
188}
189
190async fn create_tables(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
191 create_partitions_table(tr).await?;
192 create_temp_files_table(tr).await?;
193 create_migration_table(tr).await?;
194 Ok(())
195}
196
197async fn upgrade_v1_to_v2(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
198 tr.execute("ALTER TABLE lakehouse_partitions ADD file_metadata bytea;")
201 .await
202 .with_context(|| "adding column file_metadata to lakehouse_partitions table")?;
203 tr.execute("UPDATE lakehouse_migration SET version=2;")
204 .await
205 .with_context(|| "Updating lakehouse schema version to 2")?;
206 Ok(())
207}
208
209async fn upgrade_v2_to_v3(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
210 tr.execute("ALTER TABLE lakehouse_partitions ADD num_rows BIGINT;")
212 .await
213 .with_context(|| "adding column num_rows to lakehouse_partitions table")?;
214
215 tr.execute("CREATE INDEX lakehouse_partitions_file_path ON lakehouse_partitions(file_path);")
217 .await
218 .with_context(|| "creating index on file_path")?;
219
220 populate_num_rows_column(tr)
222 .await
223 .with_context(|| "populating num_rows column")?;
224
225 tr.execute("ALTER TABLE lakehouse_partitions ALTER COLUMN num_rows SET NOT NULL;")
227 .await
228 .with_context(|| "setting num_rows column to NOT NULL")?;
229
230 tr.execute("UPDATE lakehouse_migration SET version=3;")
231 .await
232 .with_context(|| "Updating lakehouse schema version to 3")?;
233 Ok(())
234}
235
236async fn populate_num_rows_column(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
237 info!("populating num_rows column for existing partitions");
238
239 let mut total_count = 0;
240 let batch_size = 1000;
241
242 loop {
243 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")
245 .bind(batch_size)
246 .fetch_all(&mut **tr)
247 .await?;
248
249 if rows.is_empty() {
250 break;
251 }
252
253 let mut batch_count = 0;
254 for row in rows {
255 let file_path: String = row.try_get("file_path")?;
256 let file_metadata_buffer: Vec<u8> = row.try_get("file_metadata")?;
257
258 match parse_parquet_metadata(&file_metadata_buffer.into()) {
260 Ok(file_metadata) => {
261 let num_rows = file_metadata.file_metadata().num_rows();
262
263 if let Err(e) = sqlx::query(
265 "UPDATE lakehouse_partitions SET num_rows = $1 WHERE file_path = $2",
266 )
267 .bind(num_rows)
268 .bind(&file_path)
269 .execute(&mut **tr)
270 .await
271 {
272 error!(
273 "failed to update num_rows for partition {}: {}",
274 file_path, e
275 );
276 continue;
277 }
278
279 batch_count += 1;
280 }
281 Err(e) => {
282 error!(
283 "failed to parse metadata for partition {}: {}",
284 file_path, e
285 );
286 if let Err(e2) = sqlx::query(
288 "UPDATE lakehouse_partitions SET num_rows = 0 WHERE file_path = $1",
289 )
290 .bind(&file_path)
291 .execute(&mut **tr)
292 .await
293 {
294 error!(
295 "failed to set fallback num_rows for partition {}: {}",
296 file_path, e2
297 );
298 }
299 }
300 }
301 }
302
303 total_count += batch_count;
304 info!(
305 "populated num_rows for {} partitions (total: {})",
306 batch_count, total_count
307 );
308 }
309
310 info!("populated num_rows for {} total partitions", total_count);
311 Ok(())
312}
313
314async fn upgrade_v3_to_v4(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
315 tr.execute(
317 "CREATE TABLE partition_metadata(
318 file_path VARCHAR(2047) PRIMARY KEY,
319 metadata bytea NOT NULL,
320 insert_time TIMESTAMPTZ NOT NULL
321 );",
322 )
323 .await
324 .with_context(|| "creating partition_metadata table")?;
325
326 migrate_metadata_to_new_table(tr)
328 .await
329 .with_context(|| "migrating metadata to partition_metadata table")?;
330
331 tr.execute("ALTER TABLE lakehouse_partitions DROP COLUMN file_metadata;")
333 .await
334 .with_context(|| "dropping file_metadata column from lakehouse_partitions")?;
335
336 tr.execute("UPDATE lakehouse_migration SET version=4;")
337 .await
338 .with_context(|| "Updating lakehouse schema version to 4")?;
339 Ok(())
340}
341
342async fn migrate_metadata_to_new_table(
343 tr: &mut sqlx::Transaction<'_, sqlx::Postgres>,
344) -> Result<()> {
345 info!("migrating metadata to partition_metadata table");
346
347 let file_paths: Vec<String> = sqlx::query_scalar(
349 "SELECT file_path
350 FROM lakehouse_partitions
351 WHERE file_metadata IS NOT NULL
352 ORDER BY file_path",
353 )
354 .fetch_all(&mut **tr)
355 .await?;
356
357 let total_to_migrate = file_paths.len();
358 info!(
359 "found {} partitions with metadata to migrate",
360 total_to_migrate
361 );
362
363 let mut total_count = 0;
364 let batch_size = 10; for chunk in file_paths.chunks(batch_size) {
368 let placeholders: Vec<String> = (1..=chunk.len()).map(|i| format!("${}", i)).collect();
370 let query_str = format!(
371 "SELECT file_path, file_metadata, updated
372 FROM lakehouse_partitions
373 WHERE file_path IN ({})",
374 placeholders.join(", ")
375 );
376
377 let mut query = sqlx::query(&query_str);
378 for path in chunk {
379 query = query.bind(path);
380 }
381
382 let rows = query.fetch_all(&mut **tr).await?;
383
384 for row in rows {
385 let file_path: String = row.try_get("file_path")?;
386 let file_metadata: Vec<u8> = row.try_get("file_metadata")?;
387 let updated: chrono::DateTime<chrono::Utc> = row.try_get("updated")?;
388
389 if let Err(e) = sqlx::query(
391 "INSERT INTO partition_metadata (file_path, metadata, insert_time)
392 VALUES ($1, $2, $3)
393 ON CONFLICT (file_path) DO NOTHING",
394 )
395 .bind(&file_path)
396 .bind(&file_metadata)
397 .bind(updated)
398 .execute(&mut **tr)
399 .await
400 {
401 error!(
402 "failed to migrate metadata for partition {}: {}",
403 file_path, e
404 );
405 continue;
406 }
407
408 total_count += 1;
409 }
410
411 if total_count % 100 == 0 || total_count == total_to_migrate {
412 info!(
413 "migrated {}/{} partition metadata entries",
414 total_count, total_to_migrate
415 );
416 }
417 }
418
419 info!("migrated metadata for {} total partitions", total_count);
420 Ok(())
421}
422
423async fn upgrade_v4_to_v5(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
424 tr.execute(
427 "ALTER TABLE lakehouse_partitions
428 ADD COLUMN partition_format_version INTEGER NOT NULL DEFAULT 1;",
429 )
430 .await
431 .with_context(|| "adding partition_format_version to lakehouse_partitions")?;
432 tr.execute(
435 "ALTER TABLE partition_metadata
436 ADD COLUMN partition_format_version INTEGER NOT NULL DEFAULT 1;",
437 )
438 .await
439 .with_context(|| "adding partition_format_version to partition_metadata")?;
440 tr.execute("UPDATE lakehouse_migration SET version=5;")
441 .await
442 .with_context(|| "Updating lakehouse schema version to 5")?;
443 info!("added partition_format_version columns to both tables");
444 Ok(())
445}
446
447async fn upgrade_v5_to_v6(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
448 tr.execute("DROP TABLE partition_metadata;")
449 .await
450 .with_context(|| "dropping partition_metadata table")?;
451 tr.execute("UPDATE lakehouse_migration SET version=6;")
452 .await
453 .with_context(|| "Updating lakehouse schema version to 6")?;
454 Ok(())
455}
456
457async fn upgrade_v6_to_v7(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
458 tr.execute("ALTER TABLE lakehouse_partitions ADD COLUMN sort_order TEXT[];")
461 .await
462 .with_context(|| "adding sort_order column to lakehouse_partitions")?;
463 tr.execute("CREATE EXTENSION IF NOT EXISTS btree_gist;")
464 .await
465 .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")?;
466 let conflicts = sqlx::query(
472 "SELECT a.view_set_name, a.view_instance_id,
473 a.file_path AS file_path_a, b.file_path AS file_path_b,
474 a.begin_insert_time AS begin_a, a.end_insert_time AS end_a,
475 b.begin_insert_time AS begin_b, b.end_insert_time AS end_b
476 FROM lakehouse_partitions a
477 JOIN lakehouse_partitions b
478 ON a.view_set_name = b.view_set_name
479 AND a.view_instance_id = b.view_instance_id
480 AND a.file_schema_hash = b.file_schema_hash
481 AND a.ctid < b.ctid
482 AND a.begin_insert_time < b.end_insert_time
483 AND a.end_insert_time > b.begin_insert_time
484 AND a.begin_insert_time < a.end_insert_time
485 AND b.begin_insert_time < b.end_insert_time
486 LIMIT 20;",
487 )
488 .fetch_all(&mut **tr)
489 .await
490 .with_context(|| "detecting overlapping partitions before adding exclusion constraint")?;
491 if !conflicts.is_empty() {
492 let mut msg = String::from(
493 "lakehouse_partitions contains partitions with overlapping insert-time ranges; \
494 retire them (e.g. retire_partition_by_metadata) and restart to complete the \
495 migration. Conflicting pairs (first 20):",
496 );
497 for row in &conflicts {
498 let view_set_name: String = row.try_get("view_set_name")?;
499 let view_instance_id: String = row.try_get("view_instance_id")?;
500 let file_path_a: Option<String> = row.try_get("file_path_a")?;
501 let file_path_b: Option<String> = row.try_get("file_path_b")?;
502 let begin_a: chrono::DateTime<chrono::Utc> = row.try_get("begin_a")?;
503 let end_a: chrono::DateTime<chrono::Utc> = row.try_get("end_a")?;
504 let begin_b: chrono::DateTime<chrono::Utc> = row.try_get("begin_b")?;
505 let end_b: chrono::DateTime<chrono::Utc> = row.try_get("end_b")?;
506 msg.push_str(&format!(
507 "\n {view_set_name}/{view_instance_id}: {file_path_a:?} [{begin_a}, {end_a}] overlaps {file_path_b:?} [{begin_b}, {end_b}]"
508 ));
509 }
510 anyhow::bail!(msg);
511 }
512 tr.execute(
524 "ALTER TABLE lakehouse_partitions ADD CONSTRAINT lakehouse_partitions_no_overlap
525 EXCLUDE USING gist (
526 view_set_name WITH =,
527 view_instance_id WITH =,
528 file_schema_hash WITH =,
529 tstzrange(begin_insert_time, end_insert_time) WITH &&
530 );",
531 )
532 .await
533 .with_context(|| "adding partition overlap exclusion constraint")?;
534 tr.execute("UPDATE lakehouse_migration SET version=7;")
535 .await
536 .with_context(|| "Updating lakehouse schema version to 7")?;
537 Ok(())
538}
539
540async fn upgrade_v7_to_v8(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
541 tr.execute("ALTER TABLE lakehouse_partitions ADD COLUMN max_sort_key_time TIMESTAMPTZ;")
546 .await
547 .with_context(|| "adding max_sort_key_time column to lakehouse_partitions")?;
548 tr.execute("UPDATE lakehouse_migration SET version=8;")
549 .await
550 .with_context(|| "Updating lakehouse schema version to 8")?;
551 Ok(())
552}
553
554async fn upgrade_v8_to_v9(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
555 tr.execute(
560 "CREATE TABLE query_deny_list(
561 rule_id UUID PRIMARY KEY,
562 created_at TIMESTAMPTZ NOT NULL,
563 created_by VARCHAR(255) NOT NULL,
564 reason TEXT NOT NULL,
565 match_expr TEXT NOT NULL,
566 last_hit_at TIMESTAMPTZ
567 );",
568 )
569 .await
570 .with_context(|| "creating query_deny_list table")?;
571 tr.execute("UPDATE lakehouse_migration SET version=9;")
572 .await
573 .with_context(|| "Updating lakehouse schema version to 9")?;
574 Ok(())
575}
576
577async fn upgrade_v9_to_v10(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Result<()> {
578 tr.execute(
584 "CREATE TABLE lakehouse_view_set_definitions(
585 view_set_name VARCHAR(255) PRIMARY KEY,
586 definition_sql TEXT NOT NULL,
587 extract_query TEXT NOT NULL,
588 count_src_query TEXT NOT NULL,
589 merge_partitions_query TEXT NOT NULL,
590 update_group INTEGER NOT NULL,
591 view_options TEXT NOT NULL DEFAULT '{}',
592 created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
593 updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
594 updated_by TEXT
595 );",
596 )
597 .await
598 .with_context(|| "creating lakehouse_view_set_definitions table")?;
599
600 let def = log_stats_view_definition();
608 let view_options = serde_json::to_string(&def.options)
609 .with_context(|| "serializing log_stats's seeded view_options")?;
610 sqlx::query(
611 "INSERT INTO lakehouse_view_set_definitions
612 (view_set_name, definition_sql, extract_query, count_src_query,
613 merge_partitions_query, update_group, view_options, updated_by)
614 VALUES ($1, $2, $3, $4, $5, $6, $7, 'migration')
615 ON CONFLICT (view_set_name) DO NOTHING;",
616 )
617 .bind(&def.view_set_name)
618 .bind(log_stats_ddl_text())
619 .bind(&def.extract_query)
620 .bind(&def.count_src_query)
621 .bind(&def.merge_partitions_query)
622 .bind(def.update_group)
623 .bind(view_options)
624 .execute(&mut **tr)
625 .await
626 .with_context(|| "seeding log_stats into lakehouse_view_set_definitions")?;
627
628 tr.execute("UPDATE lakehouse_migration SET version=10;")
629 .await
630 .with_context(|| "Updating lakehouse schema version to 10")?;
631 Ok(())
632}