micromegas_analytics/lakehouse/
lakehouse_context.rs

1use super::metadata_cache::MetadataCache;
2use super::migration::migrate_lakehouse;
3use super::reader_factory::ReaderFactory;
4use super::runtime::make_runtime_env;
5use anyhow::Context;
6use anyhow::Result;
7use datafusion::execution::runtime_env::RuntimeEnv;
8use micromegas_ingestion::data_lake_config::DataLakeConfig;
9use micromegas_ingestion::data_lake_connection::{DataLakeConnection, connect_to_data_lake};
10use micromegas_tracing::prelude::*;
11use std::sync::Arc;
12
13/// Default metadata cache size in MB
14const DEFAULT_METADATA_CACHE_SIZE_MB: u64 = 50;
15
16/// Bundles all runtime resources needed for lakehouse query execution.
17///
18/// This struct holds the data lake connection, metadata cache, and DataFusion runtime,
19/// providing a single context object that can be passed through the query path. Parquet
20/// byte-range caching is handled by the in-process L1 cache wrapped around the reader
21/// factory's object store (see `object_cache::l1_wrap`), not by this struct.
22#[derive(Clone)]
23pub struct LakehouseContext {
24    lake: Arc<DataLakeConnection>,
25    metadata_cache: Arc<MetadataCache>,
26    runtime: Arc<RuntimeEnv>,
27    reader_factory: Arc<ReaderFactory>,
28}
29
30impl LakehouseContext {
31    /// Builds a lakehouse context from an already-connected `DataLakeConnection`.
32    ///
33    /// Runs `migrate_lakehouse` on the supplied connection (idempotent) and
34    /// creates the DataFusion runtime.  The caller is responsible for running
35    /// `migrate_db` (ingestion schema) before this if both migrations are needed
36    /// — the monolith does this via `connect_to_remote_data_lake`.
37    pub async fn from_connection(lake: Arc<DataLakeConnection>) -> Result<Arc<Self>> {
38        migrate_lakehouse(lake.db_pool.clone())
39            .await
40            .with_context(|| "migrate_lakehouse")?;
41        let runtime = Arc::new(make_runtime_env()?);
42        Ok(Arc::new(Self::new(lake, runtime)))
43    }
44
45    /// Reads MICROMEGAS_SQL_CONNECTION_STRING and MICROMEGAS_OBJECT_STORE_URI,
46    /// connects to the data lake, runs lakehouse migrations, and creates the
47    /// runtime environment.
48    pub async fn from_env() -> Result<Arc<Self>> {
49        let cfg = DataLakeConfig::from_env()?;
50        let data_lake = Arc::new(
51            connect_to_data_lake(&cfg.sql_connection_string, &cfg.object_store_uri).await?,
52        );
53        migrate_lakehouse(data_lake.db_pool.clone())
54            .await
55            .with_context(|| "migrate_lakehouse")?;
56        let runtime = Arc::new(make_runtime_env()?);
57        Ok(Arc::new(Self::new(data_lake, runtime)))
58    }
59
60    /// Creates a new lakehouse context with a default-sized metadata cache.
61    pub fn new(lake: Arc<DataLakeConnection>, runtime: Arc<RuntimeEnv>) -> Self {
62        let metadata_cache_mb = match std::env::var("MICROMEGAS_METADATA_CACHE_MB") {
63            Ok(s) => s.parse::<u64>().unwrap_or_else(|_| {
64                warn!(
65                    "Invalid MICROMEGAS_METADATA_CACHE_MB value '{s}', using default {DEFAULT_METADATA_CACHE_SIZE_MB} MB"
66                );
67                DEFAULT_METADATA_CACHE_SIZE_MB
68            }),
69            Err(_) => DEFAULT_METADATA_CACHE_SIZE_MB,
70        };
71
72        let metadata_cache = Arc::new(MetadataCache::new(metadata_cache_mb * 1024 * 1024));
73
74        let reader_factory = Arc::new(ReaderFactory::new(
75            micromegas_object_cache::l1_wrap(lake.blob_storage.inner(), "lakehouse"),
76            metadata_cache.clone(),
77        ));
78        Self {
79            lake,
80            metadata_cache,
81            runtime,
82            reader_factory,
83        }
84    }
85
86    /// Creates a new lakehouse context with a custom metadata cache.
87    pub fn with_caches(
88        lake: Arc<DataLakeConnection>,
89        runtime: Arc<RuntimeEnv>,
90        metadata_cache: Arc<MetadataCache>,
91    ) -> Self {
92        let reader_factory = Arc::new(ReaderFactory::new(
93            micromegas_object_cache::l1_wrap(lake.blob_storage.inner(), "lakehouse"),
94            metadata_cache.clone(),
95        ));
96        Self {
97            lake,
98            metadata_cache,
99            runtime,
100            reader_factory,
101        }
102    }
103
104    /// Returns the data lake connection.
105    pub fn lake(&self) -> &Arc<DataLakeConnection> {
106        &self.lake
107    }
108
109    /// Returns the metadata cache.
110    pub fn metadata_cache(&self) -> &Arc<MetadataCache> {
111        &self.metadata_cache
112    }
113
114    /// Returns the DataFusion runtime environment.
115    pub fn runtime(&self) -> &Arc<RuntimeEnv> {
116        &self.runtime
117    }
118
119    /// Returns the shared `ReaderFactory`.
120    pub fn reader_factory(&self) -> &Arc<ReaderFactory> {
121        &self.reader_factory
122    }
123}
124
125impl std::fmt::Debug for LakehouseContext {
126    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
127        f.debug_struct("LakehouseContext")
128            .field("metadata_cache", &self.metadata_cache)
129            .finish()
130    }
131}
micromegas_analytics/lakehouse/lakehouse_context.rs

micromegas_analytics/lakehouse/
lakehouse_context.rs
