micromegas_analytics/lakehouse/
lakehouse_context.rs

1use super::audience_guard::{
2    AudienceIndex, DEFAULT_AUDIENCE_CACHE_ENTRIES, DEFAULT_AUDIENCE_CACHE_TTL,
3};
4use super::metadata_cache::MetadataCache;
5use super::migration::migrate_lakehouse;
6use super::query_deny_list::QueryDenyList;
7use super::reader_factory::ReaderFactory;
8use super::runtime::make_runtime_env;
9use crate::audience::default_audience_from_env;
10use anyhow::Context;
11use anyhow::Result;
12use datafusion::execution::runtime_env::RuntimeEnv;
13use micromegas_ingestion::data_lake_config::DataLakeConfig;
14use micromegas_ingestion::data_lake_connection::{
15    DataLakeConnection, WritablePolicy, connect_to_data_lake,
16};
17use micromegas_tracing::prelude::*;
18use std::sync::Arc;
19
20/// Default metadata cache size in MB
21const DEFAULT_METADATA_CACHE_SIZE_MB: u64 = 50;
22
23/// Bundles all runtime resources needed for lakehouse query execution.
24///
25/// This struct holds the data lake connection, metadata cache, and DataFusion runtime,
26/// providing a single context object that can be passed through the query path. Parquet
27/// byte-range caching is handled by the in-process L1 cache wrapped around the reader
28/// factory's object store (see `object_cache::l1_wrap`), not by this struct.
29#[derive(Clone)]
30pub struct LakehouseContext {
31    lake: Arc<DataLakeConnection>,
32    metadata_cache: Arc<MetadataCache>,
33    runtime: Arc<RuntimeEnv>,
34    reader_factory: Arc<ReaderFactory>,
35    /// The call-level guard -- resolves *any telemetry id -> its
36    /// owning process's audience* from Postgres, size- and TTL-bounded. Fixed shape, no
37    /// operational knob (see [`DEFAULT_AUDIENCE_CACHE_ENTRIES`]'s doc comment for why).
38    audience_index: Arc<AudienceIndex>,
39    /// Admin-managed query deny list. The refresh task that
40    /// keeps its snapshot warm is spawned only by the FlightSQL server builder -- every other
41    /// holder of a `LakehouseContext` (maintenance daemon, tests) keeps an empty snapshot and
42    /// `check` never denies anything.
43    query_denials: Arc<QueryDenyList>,
44    /// The audience a never-stamped process is read as (`MICROMEGAS_DEFAULT_AUDIENCE`).
45    /// Resolved once here and handed to all three sites that read an audience out of Postgres --
46    /// `BlocksView`'s `data_sql`, `metadata::find_process`, and [`AudienceIndex`]'s
47    /// `owner_query_sql` -- so one process can never resolve two different defaults depending on
48    /// which path reached it. It lives here rather than on `IsolationConfig` because the
49    /// maintenance daemon needs it and never builds one of those.
50    default_audience: Arc<str>,
51}
52
53impl LakehouseContext {
54    /// Builds a lakehouse context from an already-connected `DataLakeConnection`.
55    ///
56    /// Runs `migrate_lakehouse` on the supplied connection (idempotent) and
57    /// creates the DataFusion runtime.  The caller is responsible for running
58    /// `migrate_db` (ingestion schema) before this if both migrations are needed
59    /// — the monolith does this via `connect_to_remote_data_lake`.
60    pub async fn from_connection(lake: Arc<DataLakeConnection>) -> Result<Arc<Self>> {
61        migrate_lakehouse(lake.db_pool.clone())
62            .await
63            .with_context(|| "migrate_lakehouse")?;
64        let runtime = Arc::new(make_runtime_env()?);
65        Ok(Arc::new(Self::new(lake, runtime)?))
66    }
67
68    /// Reads MICROMEGAS_SQL_CONNECTION_STRING and MICROMEGAS_OBJECT_STORE_URI,
69    /// connects to the data lake under `policy`, runs lakehouse migrations, and creates the
70    /// runtime environment.
71    pub async fn from_env(policy: WritablePolicy) -> Result<Arc<Self>> {
72        let cfg = DataLakeConfig::from_env()?;
73        let data_lake = Arc::new(
74            connect_to_data_lake(policy, &cfg.sql_connection_string, &cfg.object_store_uri).await?,
75        );
76        migrate_lakehouse(data_lake.db_pool.clone())
77            .await
78            .with_context(|| "migrate_lakehouse")?;
79        let runtime = Arc::new(make_runtime_env()?);
80        Ok(Arc::new(Self::new(data_lake, runtime)?))
81    }
82
83    /// Creates a new lakehouse context with a default-sized metadata cache.
84    ///
85    /// Fails if `MICROMEGAS_DEFAULT_AUDIENCE` is set to a malformed value -- every caller is a
86    /// startup path, so a typo in that knob stops the role rather than silently relabelling
87    /// legacy data (see [`default_audience_from_env`]).
88    pub fn new(lake: Arc<DataLakeConnection>, runtime: Arc<RuntimeEnv>) -> Result<Self> {
89        let metadata_cache_mb = match std::env::var("MICROMEGAS_METADATA_CACHE_MB") {
90            Ok(s) => s.parse::<u64>().unwrap_or_else(|_| {
91                warn!(
92                    "Invalid MICROMEGAS_METADATA_CACHE_MB value '{s}', using default {DEFAULT_METADATA_CACHE_SIZE_MB} MB"
93                );
94                DEFAULT_METADATA_CACHE_SIZE_MB
95            }),
96            Err(_) => DEFAULT_METADATA_CACHE_SIZE_MB,
97        };
98
99        let metadata_cache = Arc::new(MetadataCache::new(metadata_cache_mb * 1024 * 1024));
100        let default_audience: Arc<str> = Arc::from(default_audience_from_env()?.as_str());
101
102        let reader_factory = Arc::new(ReaderFactory::new(
103            micromegas_object_cache::l1_wrap(lake.blob_storage.inner(), "lakehouse"),
104            metadata_cache.clone(),
105        ));
106        let audience_index = Arc::new(AudienceIndex::new(
107            lake.db_pool.clone(),
108            DEFAULT_AUDIENCE_CACHE_ENTRIES,
109            DEFAULT_AUDIENCE_CACHE_TTL,
110            default_audience.clone(),
111        ));
112        let query_denials = Arc::new(QueryDenyList::new(lake.db_pool.clone()));
113        Ok(Self {
114            lake,
115            metadata_cache,
116            runtime,
117            reader_factory,
118            audience_index,
119            query_denials,
120            default_audience,
121        })
122    }
123
124    /// Creates a new lakehouse context with a custom metadata cache.
125    ///
126    /// Takes `default_audience` rather than re-reading the environment: a context built this way
127    /// must resolve the same default as the one it is derived from (see the field's doc comment).
128    pub fn with_caches(
129        lake: Arc<DataLakeConnection>,
130        runtime: Arc<RuntimeEnv>,
131        metadata_cache: Arc<MetadataCache>,
132        default_audience: Arc<str>,
133    ) -> Self {
134        let reader_factory = Arc::new(ReaderFactory::new(
135            micromegas_object_cache::l1_wrap(lake.blob_storage.inner(), "lakehouse"),
136            metadata_cache.clone(),
137        ));
138        let audience_index = Arc::new(AudienceIndex::new(
139            lake.db_pool.clone(),
140            DEFAULT_AUDIENCE_CACHE_ENTRIES,
141            DEFAULT_AUDIENCE_CACHE_TTL,
142            default_audience.clone(),
143        ));
144        let query_denials = Arc::new(QueryDenyList::new(lake.db_pool.clone()));
145        Self {
146            lake,
147            metadata_cache,
148            runtime,
149            reader_factory,
150            audience_index,
151            query_denials,
152            default_audience,
153        }
154    }
155
156    /// Returns the audience a never-stamped process is read as
157    /// (`MICROMEGAS_DEFAULT_AUDIENCE`).
158    pub fn default_audience(&self) -> Arc<str> {
159        self.default_audience.clone()
160    }
161
162    /// Returns the data lake connection.
163    pub fn lake(&self) -> &Arc<DataLakeConnection> {
164        &self.lake
165    }
166
167    /// Returns the metadata cache.
168    pub fn metadata_cache(&self) -> &Arc<MetadataCache> {
169        &self.metadata_cache
170    }
171
172    /// Returns the DataFusion runtime environment.
173    pub fn runtime(&self) -> &Arc<RuntimeEnv> {
174        &self.runtime
175    }
176
177    /// Clones this context with `runtime` swapped, sharing the metadata cache and reader
178    /// factory. Deliberately a struct-update clone rather than a call to `new()`/
179    /// `with_caches()` — those rebuild the `MetadataCache` and `ReaderFactory`, which would
180    /// throw away the shared metadata cache per query.
181    pub fn with_runtime(&self, runtime: Arc<RuntimeEnv>) -> Arc<Self> {
182        Arc::new(Self {
183            runtime,
184            ..self.clone()
185        })
186    }
187
188    /// Returns the shared `ReaderFactory`.
189    pub fn reader_factory(&self) -> &Arc<ReaderFactory> {
190        &self.reader_factory
191    }
192
193    /// Returns the shared call-level-guard audience index.
194    pub fn audience_index(&self) -> &Arc<AudienceIndex> {
195        &self.audience_index
196    }
197
198    /// Returns the shared query deny list.
199    pub fn query_denials(&self) -> &Arc<QueryDenyList> {
200        &self.query_denials
201    }
202}
203
204impl std::fmt::Debug for LakehouseContext {
205    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
206        f.debug_struct("LakehouseContext")
207            .field("metadata_cache", &self.metadata_cache)
208            .finish()
209    }
210}
micromegas_analytics/lakehouse/lakehouse_context.rs

micromegas_analytics/lakehouse/
lakehouse_context.rs
