micromegas/servers/
flight_sql_server.rs

1use anyhow::Result;
2use micromegas_analytics::lakehouse::lakehouse_context::LakehouseContext;
3use micromegas_analytics::lakehouse::partition_cache::LivePartitionProvider;
4use micromegas_analytics::lakehouse::session_configurator::SessionConfigurator;
5use micromegas_analytics::lakehouse::static_tables_configurator::StaticTablesConfigurator;
6use micromegas_analytics::lakehouse::view_factory::{ViewFactory, default_view_factory};
7use micromegas_auth::tower::AuthService;
8use micromegas_auth::types::AuthProvider;
9use micromegas_ingestion::data_lake_connection::DataLakeConnection;
10use micromegas_tracing::prelude::*;
11use std::future::Future;
12use std::net::SocketAddr;
13use std::pin::Pin;
14use std::sync::Arc;
15use std::time::Duration;
16
17use arrow_flight::flight_service_server::FlightServiceServer;
18use datafusion::execution::runtime_env::RuntimeEnv;
19use tonic::transport::Server;
20use tower::ServiceBuilder;
21use tower::layer::layer_fn;
22
23use super::connect_info_layer::ConnectedIncoming;
24use super::flight_sql_service_impl::FlightSqlServiceImpl;
25use super::grpc_health_service::GrpcHealthService;
26use super::log_uri_service::LogUriService;
27
28type ViewFactoryFn = Box<
29    dyn FnOnce(
30            Arc<RuntimeEnv>,
31            Arc<DataLakeConnection>,
32        ) -> Pin<Box<dyn Future<Output = Result<ViewFactory>> + Send>>
33        + Send,
34>;
35
36/// Builder for assembling and running a FlightSQL server.
37///
38/// Encapsulates the full setup sequence: data lake connection, lakehouse migration,
39/// runtime env, view factory, partition provider, session configurator, auth, and
40/// the gRPC tower layer stack.
41///
42/// # Example
43///
44/// ```rust,no_run
45/// use micromegas::servers::flight_sql_server::FlightSqlServer;
46///
47/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
48/// FlightSqlServer::builder()
49///     .with_default_auth()
50///     .build_and_serve()
51///     .await?;
52/// # Ok(())
53/// # }
54/// ```
55pub struct FlightSqlServer;
56
57impl FlightSqlServer {
58    pub fn builder() -> FlightSqlServerBuilder {
59        FlightSqlServerBuilder::default()
60    }
61}
62
63pub struct FlightSqlServerBuilder {
64    view_factory_fn: Option<ViewFactoryFn>,
65    session_configurator: Option<Arc<dyn SessionConfigurator>>,
66    auth_provider: Option<Arc<dyn AuthProvider>>,
67    use_default_auth: bool,
68    max_decoding_message_size: usize,
69    listen_addr: SocketAddr,
70    shutdown_grace: Duration,
71    injected_lakehouse: Option<Arc<LakehouseContext>>,
72    injected_shutdown: Option<Pin<Box<dyn Future<Output = ()> + Send + 'static>>>,
73    health_listen_addr: Option<SocketAddr>,
74}
75
76impl Default for FlightSqlServerBuilder {
77    fn default() -> Self {
78        Self {
79            view_factory_fn: None,
80            session_configurator: None,
81            auth_provider: None,
82            use_default_auth: false,
83            max_decoding_message_size: 100 * 1024 * 1024,
84            listen_addr: "0.0.0.0:50051"
85                .parse()
86                .expect("valid default listen address"),
87            shutdown_grace: Duration::from_secs(25),
88            injected_lakehouse: None,
89            injected_shutdown: None,
90            health_listen_addr: None,
91        }
92    }
93}
94
95impl FlightSqlServerBuilder {
96    /// Override the default view factory with a custom closure.
97    ///
98    /// The closure receives the runtime and data lake created by the builder.
99    pub fn with_view_factory_fn<F, Fut>(mut self, f: F) -> Self
100    where
101        F: FnOnce(Arc<RuntimeEnv>, Arc<DataLakeConnection>) -> Fut + Send + 'static,
102        Fut: Future<Output = Result<ViewFactory>> + Send + 'static,
103    {
104        self.view_factory_fn = Some(Box::new(move |runtime, lake| Box::pin(f(runtime, lake))));
105        self
106    }
107
108    /// Override the default session configurator.
109    ///
110    /// By default the builder loads static tables from `MICROMEGAS_STATIC_TABLES_URL`.
111    /// Use this to replace that behavior entirely.
112    pub fn with_session_configurator(mut self, cfg: Arc<dyn SessionConfigurator>) -> Self {
113        self.session_configurator = Some(cfg);
114        self
115    }
116
117    /// Set an explicit auth provider.
118    pub fn with_auth_provider(mut self, provider: Arc<dyn AuthProvider>) -> Self {
119        self.auth_provider = Some(provider);
120        self.use_default_auth = false;
121        self
122    }
123
124    /// Use the default auth provider from env vars during build.
125    ///
126    /// Errors if no auth providers are configured (fail-fast).
127    pub fn with_default_auth(mut self) -> Self {
128        self.use_default_auth = true;
129        self.auth_provider = None;
130        self
131    }
132
133    /// Set the max decoding message size (default: 100 MB).
134    pub fn with_max_decoding_message_size(mut self, bytes: usize) -> Self {
135        self.max_decoding_message_size = bytes;
136        self
137    }
138
139    /// Set the listen address (default: `0.0.0.0:50051`).
140    pub fn with_listen_addr(mut self, addr: SocketAddr) -> Self {
141        self.listen_addr = addr;
142        self
143    }
144
145    /// Set the grace period for graceful shutdown on SIGTERM (default: 25s).
146    pub fn with_shutdown_grace(mut self, grace: Duration) -> Self {
147        self.shutdown_grace = grace;
148        self
149    }
150
151    /// Inject a pre-built `LakehouseContext` instead of calling `LakehouseContext::from_env`.
152    ///
153    /// Useful for the monolith, which constructs one shared context for all lake-backed roles.
154    pub fn with_lakehouse(mut self, lakehouse: Arc<LakehouseContext>) -> Self {
155        self.injected_lakehouse = Some(lakehouse);
156        self
157    }
158
159    /// Inject a custom shutdown future instead of the default `wait_for_sigterm()`.
160    ///
161    /// The monolith passes `fanout.subscribe()` here so all roles shut down from one signal.
162    pub fn with_shutdown(mut self, shutdown: impl Future<Output = ()> + Send + 'static) -> Self {
163        self.injected_shutdown = Some(Box::pin(shutdown));
164        self
165    }
166
167    /// Spawn a lightweight HTTP sidecar (`/health`, `/ready`) on `addr`.
168    ///
169    /// Enables plain-HTTP ALB health checks without changing the gRPC protocol.
170    /// If not set, no sidecar is started.
171    pub fn with_health_addr(mut self, addr: SocketAddr) -> Self {
172        self.health_listen_addr = Some(addr);
173        self
174    }
175
176    /// Build and run the FlightSQL server.
177    ///
178    /// Runs the full setup sequence and blocks until the server shuts down.
179    pub async fn build_and_serve(self) -> Result<()> {
180        // Use injected lakehouse or build one from environment
181        let lakehouse = if let Some(lh) = self.injected_lakehouse {
182            lh
183        } else {
184            LakehouseContext::from_env().await?
185        };
186        let data_lake = lakehouse.lake().clone();
187        let probe_lake = lakehouse.lake().clone();
188        info!(
189            "created lakehouse context with metadata cache: {:?}",
190            lakehouse.metadata_cache()
191        );
192
193        let view_factory = if let Some(factory_fn) = self.view_factory_fn {
194            Arc::new(factory_fn(lakehouse.runtime().clone(), data_lake).await?)
195        } else {
196            Arc::new(default_view_factory(lakehouse.runtime().clone(), data_lake).await?)
197        };
198
199        let partition_provider =
200            Arc::new(LivePartitionProvider::new(lakehouse.lake().db_pool.clone()));
201
202        let session_configurator: Arc<dyn SessionConfigurator> =
203            if let Some(cfg) = self.session_configurator {
204                cfg
205            } else {
206                StaticTablesConfigurator::from_env(
207                    "MICROMEGAS_STATIC_TABLES_URL",
208                    lakehouse.runtime().clone(),
209                )
210                .await?
211            };
212
213        let svc = FlightServiceServer::new(FlightSqlServiceImpl::new(
214            lakehouse,
215            partition_provider,
216            view_factory,
217            session_configurator,
218        ))
219        .max_decoding_message_size(self.max_decoding_message_size);
220
221        let auth_provider: Option<Arc<dyn AuthProvider>> = if let Some(provider) =
222            self.auth_provider
223        {
224            Some(provider)
225        } else if self.use_default_auth {
226            match micromegas_auth::default_provider::provider().await? {
227                Some(provider) => Some(provider),
228                None => {
229                    anyhow::bail!(
230                        "Authentication required but no auth providers configured. Set MICROMEGAS_API_KEYS or MICROMEGAS_OIDC_CONFIG"
231                    );
232                }
233            }
234        } else {
235            info!("Authentication disabled");
236            None
237        };
238
239        let layer = ServiceBuilder::new()
240            .layer(layer_fn(GrpcHealthService::new))
241            .layer(layer_fn(|service| LogUriService { service }))
242            .layer(layer_fn(move |inner| AuthService {
243                inner,
244                auth_provider: auth_provider.clone(),
245            }))
246            .into_inner();
247
248        use super::shutdown::{ShutdownFanout, wait_for_sigterm};
249
250        info!("Listening on {:?}", self.listen_addr);
251        let listener = std::net::TcpListener::bind(self.listen_addr)?;
252        let incoming = ConnectedIncoming::from_std_listener(listener)?;
253
254        // Use injected shutdown future or default to SIGTERM
255        let shutdown_future: Pin<Box<dyn Future<Output = ()> + Send + 'static>> = self
256            .injected_shutdown
257            .unwrap_or_else(|| Box::pin(wait_for_sigterm()));
258        let fanout = ShutdownFanout::new(shutdown_future);
259        let grace_secs = self.shutdown_grace.as_secs();
260        let grace = self.shutdown_grace;
261
262        if let Some(health_addr) = self.health_listen_addr {
263            use super::readiness::ReadinessProbe;
264            use axum::Extension;
265            use axum::Router;
266            use axum::routing::get;
267            use tokio::net::TcpListener;
268
269            let probe = std::sync::Arc::new(ReadinessProbe::new(probe_lake));
270            let sidecar_listener = TcpListener::bind(health_addr).await?;
271            let shutdown_rx = fanout.subscribe();
272            tokio::spawn(async move {
273                async fn sidecar_ready(
274                    Extension(p): Extension<std::sync::Arc<ReadinessProbe>>,
275                ) -> axum::http::StatusCode {
276                    if p.check_ready().await {
277                        axum::http::StatusCode::OK
278                    } else {
279                        axum::http::StatusCode::SERVICE_UNAVAILABLE
280                    }
281                }
282                let sidecar_app = Router::new()
283                    .route("/health", get(|| async { axum::http::StatusCode::OK }))
284                    .route("/ready", get(sidecar_ready))
285                    .layer(Extension(probe));
286                let _ = axum::serve(sidecar_listener, sidecar_app)
287                    .with_graceful_shutdown(shutdown_rx)
288                    .await;
289                info!("FlightSQL health sidecar stopped");
290            });
291            info!("FlightSQL health sidecar listening on {health_addr}");
292        }
293
294        let serve = Server::builder()
295            .layer(layer)
296            .add_service(svc)
297            .serve_with_incoming_shutdown(incoming, fanout.subscribe());
298
299        let deadline = {
300            let d = fanout.subscribe();
301            async move {
302                d.await;
303                tokio::time::sleep(grace).await;
304            }
305        };
306
307        tokio::select! {
308            res = serve => {
309                info!("drain completed");
310                res?;
311            }
312            _ = deadline => {
313                warn!("grace period of {grace_secs}s elapsed with work still in flight");
314            }
315        }
316
317        info!("bye");
318        Ok(())
319    }
320}
micromegas/servers/flight_sql_server.rs

micromegas/servers/
flight_sql_server.rs
