micromegas_ingestion/
data_lake_connection.rs

1use anyhow::{Context, Result};
2use micromegas_object_cache::CacheClientStore;
3use micromegas_object_cache::prefetch::{ObjectPrefetch, PrefetchItem, PrefixPrefetch};
4use micromegas_telemetry::blob_storage::BlobStorage;
5use micromegas_tracing::prelude::*;
6use object_store::ObjectStore;
7use sqlx::PgPool;
8use std::sync::Arc;
9use tokio::task::JoinHandle;
10
11/// A connection to the data lake, including a database pool and a blob storage client.
12#[derive(Debug, Clone)]
13pub struct DataLakeConnection {
14    pub db_pool: PgPool,
15    pub blob_storage: Arc<BlobStorage>,
16    /// `Some` when the object cache is configured for this connection; used to
17    /// fire-and-forget warm freshly-written objects (`warm_object`).
18    /// `None` when the cache is not configured.
19    prefetch: Option<Arc<dyn ObjectPrefetch>>,
20}
21
22impl DataLakeConnection {
23    pub fn new(db_pool: PgPool, blob_storage: Arc<BlobStorage>) -> Self {
24        Self {
25            db_pool,
26            blob_storage,
27            prefetch: None,
28        }
29    }
30
31    /// Like `new`, but also wires the object cache's prefetch face for
32    /// write-time warming (see `warm_object`).
33    pub fn new_with_prefetch(
34        db_pool: PgPool,
35        blob_storage: Arc<BlobStorage>,
36        prefetch: Option<Arc<dyn ObjectPrefetch>>,
37    ) -> Self {
38        Self {
39            db_pool,
40            blob_storage,
41            prefetch,
42        }
43    }
44
45    /// Warm a freshly-written object in the object cache by key. Fire-and-forget
46    /// at prefetch priority: spawns a detached task and returns immediately, so the
47    /// caller's write path is never delayed or failed by a warm. No-op when the
48    /// cache is not configured or `size <= 0`. Returns the spawned task handle (or
49    /// None) purely so tests can await completion deterministically; production
50    /// callers ignore it.
51    ///
52    /// This is a general "warm any object" primitive — the write-partition path is
53    /// its first caller, but nothing here is partition-specific (e.g. the ingestion
54    /// service could warm raw payloads the same way). `key` is the lake-root-relative
55    /// object key; the configured prefetch handle applies the lake root prefix so the
56    /// warmed key matches the key demand reads produce.
57    pub fn warm_object(&self, key: &str, size: i64) -> Option<JoinHandle<()>> {
58        let prefetch = self.prefetch.as_ref()?.clone();
59        if size <= 0 {
60            return None; // nothing to warm
61        }
62        let key = key.to_string(); // owned copy: the spawned future must be 'static
63        let item = PrefetchItem {
64            key: key.clone(),
65            size: size as u64,
66            ranges: None,
67        };
68        imetric!("object_warm_requested", "count", 1_u64);
69        Some(spawn_with_context(async move {
70            match prefetch.prefetch(vec![item]).await {
71                Ok(resp) => debug!(
72                    "write-time warm enqueued accepted={} rejected={} dropped={}",
73                    resp.accepted, resp.rejected, resp.dropped
74                ),
75                // CacheClientStore::prefetch already bumps range_cache_client_prefetch_error;
76                // keep this at debug — a failed warm just means the first read is a cold miss.
77                Err(e) => debug!("write-time warm failed for {key}: {e}"),
78            }
79        }))
80    }
81}
82
83/// Wrap `direct` with the object cache when configured, returning the store
84/// layer and — when enabled — the same client's `ObjectPrefetch` face for
85/// write-time warming.
86pub(crate) fn make_cache(
87    direct: Arc<dyn ObjectStore>,
88) -> (Arc<dyn ObjectStore>, Option<Arc<dyn ObjectPrefetch>>) {
89    let cache_url = std::env::var("MICROMEGAS_OBJECT_CACHE_URL").ok();
90    let api_key = std::env::var("MICROMEGAS_OBJECT_CACHE_API_KEY").ok();
91    match cache_url {
92        Some(url) if api_key.is_some() => {
93            let client = Arc::new(CacheClientStore::new(url, api_key, direct));
94            (
95                client.clone() as Arc<dyn ObjectStore>,
96                Some(client as Arc<dyn ObjectPrefetch>),
97            )
98        }
99        Some(url) => {
100            // URL without key: disabled, warn (preserve current behavior)
101            warn!(
102                "MICROMEGAS_OBJECT_CACHE_URL is set ({url}) but MICROMEGAS_OBJECT_CACHE_API_KEY is missing: the object cache is disabled and requests will go directly to the store"
103            );
104            (direct, None)
105        }
106        None => (direct, None),
107    }
108}
109
110/// Connects to the data lake.
111pub async fn connect_to_data_lake(
112    db_uri: &str,
113    object_store_url: &str,
114) -> Result<DataLakeConnection> {
115    info!("connecting to blob storage");
116    let (raw_store, root) = BlobStorage::parse_url_opts(object_store_url)
117        .with_context(|| "connecting to blob storage")?;
118    let (layered, prefetch_client) = make_cache(raw_store);
119    let blob_storage = Arc::new(BlobStorage::new(layered, root.clone()));
120    let prefetch =
121        prefetch_client.map(|p| Arc::new(PrefixPrefetch::new(p, root)) as Arc<dyn ObjectPrefetch>);
122    let pool = sqlx::postgres::PgPoolOptions::new()
123        .connect(db_uri)
124        .await
125        .with_context(|| String::from("Connecting to telemetry database"))?;
126    Ok(DataLakeConnection::new_with_prefetch(
127        pool,
128        blob_storage,
129        prefetch,
130    ))
131}
micromegas_ingestion/data_lake_connection.rs

micromegas_ingestion/
data_lake_connection.rs
