micromegas_ingestion/
remote_data_lake.rs

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

micromegas_ingestion/
remote_data_lake.rs
