micromegas_ingestion/
remote_data_lake.rs

1use crate::data_lake_connection::DataLakeConnection;
2use crate::data_lake_connection::make_cache;
3use crate::data_lake_connection::read_write_pool_options;
4use crate::sql_migration::LATEST_DATA_LAKE_SCHEMA_VERSION;
5use crate::sql_migration::execute_migration;
6use crate::sql_migration::read_data_lake_schema_version;
7use anyhow::{Context, Result};
8use micromegas_object_cache::prefetch::{ObjectPrefetch, PrefixPrefetch};
9use micromegas_telemetry::blob_storage::BlobStorage;
10use micromegas_tracing::prelude::*;
11use std::sync::Arc;
12
13/// Acquires a lock on the database to prevent concurrent migrations.
14pub async fn acquire_lock(tr: &mut sqlx::Transaction<'_, sqlx::Postgres>, key: i64) -> Result<()> {
15    sqlx::query("SELECT pg_advisory_xact_lock($1)")
16        .bind(key)
17        .execute(&mut **tr)
18        .await?;
19    Ok(())
20}
21
22/// Migrates the database to the latest schema version.
23pub async fn migrate_db(pool: sqlx::Pool<sqlx::Postgres>) -> Result<()> {
24    let mut tr = pool.begin().await?;
25    let mut current_version = read_data_lake_schema_version(&mut tr).await;
26    drop(tr);
27    info!("current data lake schema: {}", current_version);
28    if current_version != LATEST_DATA_LAKE_SCHEMA_VERSION {
29        let mut tr = pool.begin().await?;
30        acquire_lock(&mut tr, 0).await?;
31        current_version = read_data_lake_schema_version(&mut pool.begin().await?).await;
32        if LATEST_DATA_LAKE_SCHEMA_VERSION == current_version {
33            return Ok(());
34        }
35        if let Err(e) = execute_migration(pool.clone()).await {
36            error!("Error migrating database: {}", e);
37            return Err(e);
38        }
39        current_version = read_data_lake_schema_version(&mut tr).await;
40    }
41    assert_eq!(current_version, LATEST_DATA_LAKE_SCHEMA_VERSION);
42    Ok(())
43}
44
45/// Connects to the remote data lake, migrating the database if necessary.
46pub async fn connect_to_remote_data_lake(
47    db_uri: &str,
48    object_store_url: &str,
49) -> Result<DataLakeConnection> {
50    info!("connecting to blob storage");
51    let (raw_store, root) = BlobStorage::parse_url_opts(object_store_url)
52        .with_context(|| "connecting to blob storage")?;
53    let (layered, prefetch_client) = make_cache(raw_store);
54    let blob_storage = Arc::new(BlobStorage::new(layered, root.clone()));
55    let prefetch =
56        prefetch_client.map(|p| Arc::new(PrefixPrefetch::new(p, root)) as Arc<dyn ObjectPrefetch>);
57    let pool = read_write_pool_options()
58        .connect(db_uri)
59        .await
60        .with_context(|| String::from("Connecting to telemetry database"))?;
61    migrate_db(pool.clone()).await?;
62    Ok(DataLakeConnection::new_with_prefetch(
63        pool,
64        blob_storage,
65        prefetch,
66    ))
67}
micromegas_ingestion/remote_data_lake.rs

micromegas_ingestion/
remote_data_lake.rs
