micromegas/servers/

flight_sql_server.rs

use crate::config::DEFAULT_FLIGHTSQL_LISTEN_ADDR;
use anyhow::{Context, Result};
use micromegas_analytics::lakehouse::lakehouse_context::LakehouseContext;
use micromegas_analytics::lakehouse::partition_cache::LivePartitionProvider;
use micromegas_analytics::lakehouse::read_scope::IsolationConfig;
use micromegas_analytics::lakehouse::session_configurator::SessionConfigurator;
use micromegas_analytics::lakehouse::static_tables_configurator::StaticTablesConfigurator;
use micromegas_analytics::lakehouse::view_definition_store::PgViewDefinitionStore;
use micromegas_analytics::lakehouse::view_factory::{ViewFactory, default_view_factory};
use micromegas_analytics::lakehouse::view_registry::ViewRegistry;
use micromegas_auth::db_api_key::{ApiKeyTable, dedicated_key_store_pool};
use micromegas_auth::db_audience_grants::{DbAudienceGrantsConfig, DbAudienceGrantsSource};
use micromegas_auth::default_provider::ProviderBuilder;
use micromegas_auth::policy::{AudienceGrants, AudienceReadPolicy, ReadPolicy};
use micromegas_auth::tower::AuthService;
use micromegas_auth::types::AuthProvider;
use micromegas_ingestion::data_lake_connection::{DataLakeConnection, WritablePolicy};
use micromegas_ingestion::sql_migration::warn_if_data_lake_schema_stale;
use micromegas_tracing::prelude::*;
use std::future::Future;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;

use arrow_flight::flight_service_server::FlightServiceServer;
use datafusion::execution::runtime_env::RuntimeEnv;
use tonic::transport::Server;
use tower::ServiceBuilder;
use tower::layer::layer_fn;

use super::connect_info_layer::ConnectedIncoming;
use super::flight_sql_service_impl::FlightSqlServiceImpl;
use super::grpc_health_service::GrpcHealthService;
use super::log_uri_service::LogUriService;

type ViewFactoryFn = Box<
    dyn FnOnce(
            Arc<RuntimeEnv>,
            Arc<DataLakeConnection>,
        ) -> Pin<Box<dyn Future<Output = Result<ViewFactory>> + Send>>
        + Send,
>;

/// The auth provider plus the default `ReadPolicy` `build_and_serve` resolves per branch, before
/// `self.read_policy` (if set) overrides it.
type AuthAndDefaults = (Option<Arc<dyn AuthProvider>>, Arc<dyn ReadPolicy>);

/// The env is not read at all when the caller supplied a config: `with_isolation_config` wins
/// regardless, so reading anyway would invent a startup failure on a variable the caller
/// deliberately overrode.
fn resolve_isolation_config(
    explicit: Option<Arc<IsolationConfig>>,
) -> Result<Arc<IsolationConfig>> {
    match explicit {
        Some(config) => Ok(config),
        None => Ok(Arc::new(IsolationConfig::from_env()?)),
    }
}

/// Builder for assembling and running a FlightSQL server.
///
/// Encapsulates the full setup sequence: data lake connection, lakehouse migration,
/// runtime env, view factory, partition provider, session configurator, auth, and
/// the gRPC tower layer stack.
///
/// # Example
///
/// ```rust,no_run
/// use micromegas::servers::flight_sql_server::FlightSqlServer;
///
/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
/// FlightSqlServer::builder()
///     .with_default_auth()
///     .build_and_serve()
///     .await?;
/// # Ok(())
/// # }
/// ```
pub struct FlightSqlServer;

impl FlightSqlServer {
    pub fn builder() -> FlightSqlServerBuilder {
        FlightSqlServerBuilder::default()
    }
}

pub struct FlightSqlServerBuilder {
    view_factory_fn: Option<ViewFactoryFn>,
    session_configurator: Option<Arc<dyn SessionConfigurator>>,
    auth_provider: Option<Arc<dyn AuthProvider>>,
    read_policy: Option<Arc<dyn ReadPolicy>>,
    isolation_config: Option<Arc<IsolationConfig>>,
    use_default_auth: bool,
    max_decoding_message_size: usize,
    listen_addr: SocketAddr,
    shutdown_grace: Duration,
    injected_lakehouse: Option<Arc<LakehouseContext>>,
    injected_shutdown: Option<Pin<Box<dyn Future<Output = ()> + Send + 'static>>>,
    health_listen_addr: Option<SocketAddr>,
}

impl Default for FlightSqlServerBuilder {
    fn default() -> Self {
        Self {
            view_factory_fn: None,
            session_configurator: None,
            auth_provider: None,
            read_policy: None,
            isolation_config: None,
            use_default_auth: false,
            max_decoding_message_size: 100 * 1024 * 1024,
            listen_addr: DEFAULT_FLIGHTSQL_LISTEN_ADDR,
            shutdown_grace: Duration::from_secs(25),
            injected_lakehouse: None,
            injected_shutdown: None,
            health_listen_addr: None,
        }
    }
}

impl FlightSqlServerBuilder {
    /// Override the default view factory with a custom closure.
    ///
    /// The closure receives the runtime and data lake created by the builder.
    pub fn with_view_factory_fn<F, Fut>(mut self, f: F) -> Self
    where
        F: FnOnce(Arc<RuntimeEnv>, Arc<DataLakeConnection>) -> Fut + Send + 'static,
        Fut: Future<Output = Result<ViewFactory>> + Send + 'static,
    {
        self.view_factory_fn = Some(Box::new(move |runtime, lake| Box::pin(f(runtime, lake))));
        self
    }

    /// Override the default session configurator.
    ///
    /// By default the builder loads static tables from `MICROMEGAS_STATIC_TABLES_URL`.
    /// Use this to replace that behavior entirely.
    pub fn with_session_configurator(mut self, cfg: Arc<dyn SessionConfigurator>) -> Self {
        self.session_configurator = Some(cfg);
        self
    }

    /// Set an explicit auth provider.
    pub fn with_auth_provider(mut self, provider: Arc<dyn AuthProvider>) -> Self {
        self.auth_provider = Some(provider);
        self.use_default_auth = false;
        self
    }

    /// Set an explicit `ReadPolicy`, resolved once per request against the caller's
    /// `AuthContext`. This policy wins on every `build_and_serve` branch,
    /// overriding that branch's own default. When never called, the default depends on how auth
    /// is configured: with `with_default_auth()` or `with_auth_provider(..)`,
    /// `AudienceReadPolicy::default()` backed by the DB grant store (so the seeded `public`
    /// read row still resolves); with auth left disabled, `AudienceReadPolicy` with an empty
    /// grant map, though that branch never actually resolves it. **Not** `ReadScope::All`: the
    /// absent-`AuthContext`-extension convention already supplies `All` when no provider is
    /// configured, so this default must not duplicate
    /// that decision -- it only ever resolves a scope when an `AuthContext` is present to resolve
    /// one from.
    pub fn with_read_policy(mut self, policy: Arc<dyn ReadPolicy>) -> Self {
        self.read_policy = Some(policy);
        self
    }

    /// Set an explicit `IsolationConfig` -- the data-isolation deployment knob
    /// (`MICROMEGAS_PUBLIC_VIEW_SETS`) consumed by the row-level filter
    /// (`OwnershipRewrite`). Mirrors `with_read_policy`: wins regardless of how auth is
    /// configured, overriding the single default `build_and_serve` would otherwise resolve via
    /// `IsolationConfig::from_env()`.
    pub fn with_isolation_config(mut self, config: Arc<IsolationConfig>) -> Self {
        self.isolation_config = Some(config);
        self
    }

    /// Use the default auth provider from env vars during build.
    ///
    /// Errors if no auth providers are configured (fail-fast).
    pub fn with_default_auth(mut self) -> Self {
        self.use_default_auth = true;
        self.auth_provider = None;
        self
    }

    /// Set the max decoding message size (default: 100 MB).
    pub fn with_max_decoding_message_size(mut self, bytes: usize) -> Self {
        self.max_decoding_message_size = bytes;
        self
    }

    /// Set the listen address (default: [`DEFAULT_FLIGHTSQL_LISTEN_ADDR`], `0.0.0.0:50051`).
    pub fn with_listen_addr(mut self, addr: SocketAddr) -> Self {
        self.listen_addr = addr;
        self
    }

    /// Set the grace period for graceful shutdown on SIGTERM (default: 25s).
    pub fn with_shutdown_grace(mut self, grace: Duration) -> Self {
        self.shutdown_grace = grace;
        self
    }

    /// Inject a pre-built `LakehouseContext` instead of calling `LakehouseContext::from_env`.
    ///
    /// Useful for the monolith, which constructs one shared context for all lake-backed roles.
    pub fn with_lakehouse(mut self, lakehouse: Arc<LakehouseContext>) -> Self {
        self.injected_lakehouse = Some(lakehouse);
        self
    }

    /// Inject a custom shutdown future instead of the default `wait_for_sigterm()`.
    ///
    /// The monolith passes `fanout.subscribe()` here so all roles shut down from one signal.
    pub fn with_shutdown(mut self, shutdown: impl Future<Output = ()> + Send + 'static) -> Self {
        self.injected_shutdown = Some(Box::pin(shutdown));
        self
    }

    /// Spawn a lightweight HTTP sidecar (`/health`, `/ready`) on `addr`.
    ///
    /// Enables plain-HTTP ALB health checks without changing the gRPC protocol.
    /// If not set, no sidecar is started.
    pub fn with_health_addr(mut self, addr: SocketAddr) -> Self {
        self.health_listen_addr = Some(addr);
        self
    }

    /// Build and run the FlightSQL server.
    ///
    /// Runs the full setup sequence and blocks until the server shuts down.
    pub async fn build_and_serve(self) -> Result<()> {
        // Resolved first, ahead of the Postgres/object-store setup below, so a malformed
        // `MICROMEGAS_PUBLIC_VIEW_SETS` fails fast at no cost.
        let isolation_config = resolve_isolation_config(self.isolation_config)?;

        // Use injected lakehouse or build one from environment. The non-injected path is
        // standalone FlightSQL, the one role that prefers a writable primary but falls back to
        // a read-only replica rather than refuse to start (see `high-availability.md`); an
        // injected lakehouse (the monolith) shares its lake pool with the ingestion role, which
        // stays strict, so it is never `Prefer` here.
        let lakehouse = if let Some(lh) = self.injected_lakehouse {
            lh
        } else {
            LakehouseContext::from_env(WritablePolicy::Prefer).await?
        };
        let data_lake = lakehouse.lake().clone();
        let probe_lake = lakehouse.lake().clone();
        // Cloned here, before `lakehouse` is moved into `FlightSqlServiceImpl::new` below, so the
        // refresh task (spawned further down, once the shutdown fanout exists) has its own
        // handle. Only the FlightSQL server builder spawns this -- other `LakehouseContext`
        // holders (maintenance daemon, tests) keep an empty snapshot and never deny anything.
        let query_denials = lakehouse.query_denials().clone();
        info!(
            "created lakehouse context with metadata cache: {:?}",
            lakehouse.metadata_cache()
        );

        let base_view_factory: Arc<ViewFactory> = if let Some(factory_fn) = self.view_factory_fn {
            Arc::new(factory_fn(lakehouse.runtime().clone(), data_lake.clone()).await?)
        } else {
            Arc::new(
                default_view_factory(
                    lakehouse.runtime().clone(),
                    data_lake.clone(),
                    lakehouse.default_audience(),
                )
                .await?,
            )
        };

        let partition_provider =
            Arc::new(LivePartitionProvider::new(lakehouse.lake().db_pool.clone()));
        // Cloned here, before `lakehouse` is moved into `FlightSqlServiceImpl::new`
        // below, so `dedicated_key_store_pool` has a pool to read pool options
        // from.
        let lake_pool_for_keys = lakehouse.lake().db_pool.clone();

        let session_configurator: Arc<dyn SessionConfigurator> =
            if let Some(cfg) = self.session_configurator {
                cfg
            } else {
                StaticTablesConfigurator::from_env(
                    "MICROMEGAS_STATIC_TABLES_URL",
                    lakehouse.runtime().clone(),
                )
                .await?
            };

        // The shared registry: `base_view_factory` (code-driven views) plus every DDL-defined
        // global view (`lakehouse_view_set_definitions`), rebuilt on
        // `MICROMEGAS_VIEW_DEFINITION_REFRESH_SECONDS`. Loaded once, synchronously, before
        // `serve()` -- failing startup on error, the way `migrate_lakehouse` already does --
        // so there is never a window after startup where a DDL-defined view (or the now-seeded
        // `log_stats`) is absent from `current()`.
        let view_registry = Arc::new(ViewRegistry::new(
            base_view_factory,
            Arc::new(PgViewDefinitionStore::new(lakehouse.lake().db_pool.clone())),
            lakehouse.runtime().clone(),
            data_lake,
            session_configurator.clone(),
        ));
        view_registry
            .reload()
            .await
            .with_context(|| "initial view_registry reload")?;

        // `FlightSqlServiceImpl::new` takes the resolved `ReadPolicy` as a constructor argument,
        // so auth/policy resolution happens here, before construction.
        //
        // `self.read_policy` (set via `with_read_policy`) must win on every branch below --
        // each branch only computes its own *default*, used solely when the caller never set an
        // explicit policy. This keeps `with_read_policy` order-independent with respect to
        // `with_auth_provider` / `with_default_auth`.
        let (auth_provider, default_policy): AuthAndDefaults =
            if let Some(provider) = self.auth_provider {
                // Injected-provider path (the monolith's `with_auth_provider` call): the caller is
                // expected to pair `with_auth_provider` with its own `with_read_policy` call, but
                // when it didn't, this default must still resolve the seeded `public` read grant
                // rather than the empty set -- an empty-grant-map default was only ever safe
                // while `resolve` had a built-in `PUBLIC_AUDIENCE` arm; with that arm removed it
                // would silently deny every query for an embedder that forgot `with_read_policy`.
                // So this builds the same store-backed policy the `use_default_auth` branch
                // below does.
                //
                // Built only when `self.read_policy` is unset -- `with_read_policy` always wins
                // below, and the only in-repo `with_auth_provider` caller (the monolith) always
                // pairs the two, so skipping this avoids a DB round trip whenever it would be
                // discarded anyway.
                let policy: Arc<dyn ReadPolicy> = if self.read_policy.is_none() {
                    let key_store_pool = dedicated_key_store_pool(&lake_pool_for_keys);
                    warn_if_data_lake_schema_stale(&lake_pool_for_keys).await;
                    let audience_grants_config = DbAudienceGrantsConfig::from_env_with_prefix("");
                    let audience_grants_store = Arc::new(DbAudienceGrantsSource::new(
                        key_store_pool,
                        Duration::from_secs(audience_grants_config.cache_ttl_secs),
                    ));
                    Arc::new(AudienceReadPolicy::default().with_store(audience_grants_store))
                } else {
                    Arc::new(AudienceReadPolicy::new(AudienceGrants::empty()))
                };
                (Some(provider), policy)
            } else if self.use_default_auth {
                let key_store_pool = dedicated_key_store_pool(&lake_pool_for_keys);
                let group_store_pool = dedicated_key_store_pool(&lake_pool_for_keys);
                warn_if_data_lake_schema_stale(&lake_pool_for_keys).await;
                let provider = match ProviderBuilder::new("")
                    .with_db_key_store(key_store_pool, ApiKeyTable::Analytics)
                    .with_group_store(group_store_pool)
                    .build()
                    .await?
                {
                    Some(provider) => provider,
                    None => {
                        anyhow::bail!(
                            "Authentication required but no auth providers configured. Set \
                         MICROMEGAS_OIDC_CONFIG, or populate the analytics_api_keys DB table"
                        );
                    }
                };
                // One shared snapshot cache for this process, built from
                // its own dedicated pool -- not `key_store_pool` above, which
                // `DbApiKeyAuthProvider` already owns -- via the same `dedicated_key_store_pool`
                // convention. The cache TTL is the flat, unprefixed
                // `MICROMEGAS_AUTH_CACHE_TTL_SECONDS`; the `prefix` argument below is passed only
                // for call-site symmetry -- `from_env_with_prefix` never actually consults it.
                let audience_grants_pool = dedicated_key_store_pool(&lake_pool_for_keys);
                let audience_grants_config = DbAudienceGrantsConfig::from_env_with_prefix("");
                let audience_grants_store = Arc::new(DbAudienceGrantsSource::new(
                    audience_grants_pool,
                    Duration::from_secs(audience_grants_config.cache_ttl_secs),
                ));
                let policy: Arc<dyn ReadPolicy> =
                    Arc::new(AudienceReadPolicy::default().with_store(audience_grants_store));
                (Some(provider), policy)
            } else {
                info!("Authentication disabled");
                // No `AuthContext` extension is ever inserted on this path (no `AuthService`
                // provider configured), so the absent-extension convention already supplies
                // `ReadScope::All` -- this default policy is never actually resolved against a
                // real caller, but must still exist since `FlightSqlServiceImpl::new` requires
                // one. This never-resolved property is now load-bearing: with the built-in
                // `PUBLIC_AUDIENCE` read arm removed, an empty grant map resolves to the empty
                // set, not `{public}` -- fine here only because `resolve` is never actually
                // called on this branch.
                (
                    None,
                    Arc::new(AudienceReadPolicy::new(AudienceGrants::empty())),
                )
            };
        let read_policy = self.read_policy.unwrap_or(default_policy);

        let svc = FlightServiceServer::new(FlightSqlServiceImpl::new(
            lakehouse,
            partition_provider,
            view_registry.clone(),
            session_configurator,
            read_policy,
            isolation_config,
        ))
        .max_decoding_message_size(self.max_decoding_message_size);

        let layer = ServiceBuilder::new()
            .layer(layer_fn(GrpcHealthService::new))
            .layer(layer_fn(|service| LogUriService { service }))
            .layer(layer_fn(move |inner| AuthService {
                inner,
                auth_provider: auth_provider.clone(),
            }))
            .into_inner();

        use super::shutdown::{ShutdownFanout, wait_for_sigterm};

        info!("Listening on {:?}", self.listen_addr);
        let listener = std::net::TcpListener::bind(self.listen_addr)?;
        let incoming = ConnectedIncoming::from_std_listener(listener)?;

        // Use injected shutdown future or default to SIGTERM
        let shutdown_future: Pin<Box<dyn Future<Output = ()> + Send + 'static>> = self
            .injected_shutdown
            .unwrap_or_else(|| Box::pin(wait_for_sigterm()));
        let fanout = ShutdownFanout::new(shutdown_future);
        let grace_secs = self.shutdown_grace.as_secs();
        let grace = self.shutdown_grace;

        // Admin-managed query deny list: refreshes the shared snapshot every
        // `MICROMEGAS_QUERY_DENY_REFRESH_SECONDS`, once immediately and then on a tick, until
        // this replica shuts down.
        query_denials.spawn_refresh_task(fanout.subscribe());

        // DDL-defined view sets: refreshes `view_registry.current()` every
        // `MICROMEGAS_VIEW_DEFINITION_REFRESH_SECONDS`, until this replica shuts down. A
        // same-node `CREATE`/`DROP MATERIALIZED VIEW` reloads inline instead (see
        // `flight_sql_service_impl.rs::execute_view_ddl`); this bounds how long a *different*
        // node lags a definition change.
        view_registry.spawn_refresh_task(fanout.subscribe());

        if let Some(health_addr) = self.health_listen_addr {
            use super::readiness::ReadinessProbe;
            use axum::Extension;
            use axum::Router;
            use axum::routing::get;
            use tokio::net::TcpListener;

            let probe = std::sync::Arc::new(ReadinessProbe::new(probe_lake));
            let sidecar_listener = TcpListener::bind(health_addr).await?;
            let shutdown_rx = fanout.subscribe();
            tokio::spawn(async move {
                async fn sidecar_ready(
                    Extension(p): Extension<std::sync::Arc<ReadinessProbe>>,
                ) -> axum::http::StatusCode {
                    if p.check_ready().await {
                        axum::http::StatusCode::OK
                    } else {
                        axum::http::StatusCode::SERVICE_UNAVAILABLE
                    }
                }
                let sidecar_app = Router::new()
                    .route("/health", get(|| async { axum::http::StatusCode::OK }))
                    .route("/ready", get(sidecar_ready))
                    .layer(Extension(probe));
                let _ = axum::serve(sidecar_listener, sidecar_app)
                    .with_graceful_shutdown(shutdown_rx)
                    .await;
                info!("FlightSQL health sidecar stopped");
            });
            info!("FlightSQL health sidecar listening on {health_addr}");
        }

        let serve = Server::builder()
            .layer(layer)
            .add_service(svc)
            .serve_with_incoming_shutdown(incoming, fanout.subscribe());

        let deadline = {
            let d = fanout.subscribe();
            async move {
                d.await;
                tokio::time::sleep(grace).await;
            }
        };

        tokio::select! {
            res = serve => {
                info!("drain completed");
                res?;
            }
            _ = deadline => {
                warn!("grace period of {grace_secs}s elapsed with work still in flight");
            }
        }

        info!("bye");
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serial_test::serial;

    const PUBLIC_VIEW_SETS_VAR: &str = "MICROMEGAS_PUBLIC_VIEW_SETS";

    /// Clears the var on drop so a failing assertion in one test can't leak state into the next.
    struct EnvGuard;

    impl Drop for EnvGuard {
        fn drop(&mut self) {
            // SAFETY: tests are serialized with `#[serial]`.
            unsafe {
                std::env::remove_var(PUBLIC_VIEW_SETS_VAR);
            }
        }
    }

    #[test]
    #[serial]
    fn env_set_is_parsed() {
        let _guard = EnvGuard;
        // SAFETY: serialized via `#[serial]`.
        unsafe {
            std::env::set_var(PUBLIC_VIEW_SETS_VAR, "log_stats,images");
        }
        let config = resolve_isolation_config(None).expect("resolve_isolation_config");
        assert_eq!(
            config.public_view_sets,
            vec!["log_stats".to_string(), "images".to_string()]
        );
    }

    #[test]
    #[serial]
    fn env_unset_is_empty() {
        let _guard = EnvGuard;
        // SAFETY: serialized via `#[serial]`.
        unsafe {
            std::env::remove_var(PUBLIC_VIEW_SETS_VAR);
        }
        let config = resolve_isolation_config(None).expect("resolve_isolation_config");
        assert!(
            config.public_view_sets.is_empty(),
            "an unset var must resolve to no public view sets"
        );
    }

    #[test]
    #[serial]
    fn env_malformed_is_err() {
        let _guard = EnvGuard;
        // SAFETY: serialized via `#[serial]`.
        unsafe {
            std::env::set_var(PUBLIC_VIEW_SETS_VAR, r#"["log_stats"]"#);
        }
        let err =
            resolve_isolation_config(None).expect_err("a JSON-array-shaped value must be rejected");
        assert!(err.to_string().contains(PUBLIC_VIEW_SETS_VAR));
    }

    /// An explicit config wins even when the env is set to a different value -- pinning that
    /// `with_isolation_config` skips the env read entirely on this path, not merely that it
    /// takes precedence over a matching value.
    #[test]
    #[serial]
    fn explicit_config_wins_over_a_different_env_value() {
        let _guard = EnvGuard;
        // SAFETY: serialized via `#[serial]`.
        unsafe {
            std::env::set_var(PUBLIC_VIEW_SETS_VAR, "log_stats");
        }
        let explicit = Arc::new(IsolationConfig {
            public_view_sets: vec!["images".to_string()],
        });
        let config =
            resolve_isolation_config(Some(explicit.clone())).expect("resolve_isolation_config");
        assert_eq!(config.public_view_sets, explicit.public_view_sets);
    }

    /// Same as above, but with the env set to a value that would fail to parse -- proving the
    /// explicit config really does skip the env read, rather than merely winning after a
    /// successful-but-discarded parse.
    #[test]
    #[serial]
    fn explicit_config_wins_over_a_malformed_env_value() {
        let _guard = EnvGuard;
        // SAFETY: serialized via `#[serial]`.
        unsafe {
            std::env::set_var(PUBLIC_VIEW_SETS_VAR, r#"["log_stats"]"#);
        }
        let explicit = Arc::new(IsolationConfig {
            public_view_sets: vec!["images".to_string()],
        });
        let config =
            resolve_isolation_config(Some(explicit.clone())).expect("resolve_isolation_config");
        assert_eq!(config.public_view_sets, explicit.public_view_sets);
    }
}