omni_dev/

snowflake.rs

//! Account-agnostic Snowflake query engine hosted by the daemon.
//!
//! Each `(account, user)` keeps a **bounded pool** of authenticated sessions
//! (see [`session`]). A query checks one out, applies any per-query context with
//! `USE` (skipping `USE`s already in effect), runs concurrently with other
//! checkouts, and returns it. This gives **concurrent queries on a single
//! authentication identity** while still honoring per-query
//! `warehouse`/`role`/`database`/`schema`, with the number of browser auths
//! capped at the pool size and grown lazily.
//!
//! A background keep-alive heartbeat ([`SnowflakeEngine::start_heartbeat`])
//! periodically heartbeats idle sessions so their master tokens stay valid and
//! an idle pool never re-prompts browser SSO.
//!
//! This is the standalone engine, analogous to [`crate::browser`]; the daemon
//! adapter lives in [`crate::daemon::services::snowflake`].

pub mod client;
pub mod session;

use std::sync::{Mutex as StdMutex, MutexGuard};
use std::time::Duration;

use anyhow::{anyhow, bail, Result};
use chrono::TimeDelta;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;

use crate::utils::env::EnvSource;
use crate::utils::settings::Settings;
use client::{Error as ClientError, Row, SnowflakeClient, SnowflakeClientConfig, SnowflakeSession};
use session::{PoolRegistry, QueryContext, SessionInfo, SessionKey};

/// Env var (with `~/.omni-dev/settings.json` fallback) for the default account.
const ENV_ACCOUNT: &str = "SNOWFLAKE_ACCOUNT";
/// Env var for the default user.
const ENV_USER: &str = "SNOWFLAKE_USER";
/// Env var for the default warehouse.
const ENV_WAREHOUSE: &str = "SNOWFLAKE_WAREHOUSE";
/// Env var for the default role.
const ENV_ROLE: &str = "SNOWFLAKE_ROLE";
/// Env var for the default database.
const ENV_DATABASE: &str = "SNOWFLAKE_DATABASE";
/// Env var for the default schema.
const ENV_SCHEMA: &str = "SNOWFLAKE_SCHEMA";
/// Env var for the per-`(account, user)` pool size (max concurrent sessions).
const ENV_POOL_SIZE: &str = "SNOWFLAKE_POOL_SIZE";
/// Env var for the per-request HTTP timeout (seconds).
const ENV_HTTP_TIMEOUT: &str = "SNOWFLAKE_HTTP_TIMEOUT";
/// Env var for the overall sign-in deadline (seconds).
const ENV_AUTH_TIMEOUT: &str = "SNOWFLAKE_AUTH_TIMEOUT";
/// Env var for the overall per-query deadline incl. async polling (seconds).
const ENV_QUERY_TIMEOUT: &str = "SNOWFLAKE_QUERY_TIMEOUT";
/// Env var for the idle-session keep-alive heartbeat interval (seconds; `0`
/// disables the heartbeat).
const ENV_HEARTBEAT_INTERVAL: &str = "SNOWFLAKE_HEARTBEAT_INTERVAL";

/// Default pool size when `SNOWFLAKE_POOL_SIZE` is unset: the max concurrent
/// queries (and max browser auths) per `(account, user)`.
const DEFAULT_POOL_SIZE: usize = 4;
/// Default overall sign-in deadline: comfortably over the SSO callback wait so a
/// genuine sign-in completes, but bounded so a hung auth can't hold the gate.
const DEFAULT_AUTH_TIMEOUT: Duration = Duration::from_secs(150);
/// Max characters of SQL shown in the "running" preview for a busy session.
const SQL_PREVIEW_MAX: usize = 60;
/// Default keep-alive heartbeat interval: a quarter of the default 3600s
/// session-token validity (Snowflake's own drivers clamp their heartbeat
/// frequency to 900–3600s), so an idle session renews comfortably before
/// either token lapses.
const DEFAULT_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(900);
/// Extra margin past the heartbeat interval when deciding to proactively renew
/// an idle session's token before it would lapse mid-cycle.
const HEARTBEAT_RENEW_MARGIN_SECS: i64 = 60;

/// Engine defaults, resolved from environment variables and then
/// `~/.omni-dev/settings.json` (the Atlassian credential-resolution pattern).
///
/// Account/user/context are optional; a request supplies its own and falls back
/// to these. There is **no** hardcoded account list or alias map.
#[derive(Clone, Debug)]
pub struct SnowflakeEngineConfig {
    /// Default account when a request omits `--account`.
    pub default_account: Option<String>,
    /// Default user when a request omits `--user`.
    pub default_user: Option<String>,
    /// Default warehouse applied at session creation.
    pub default_warehouse: Option<String>,
    /// Default role applied at session creation.
    pub default_role: Option<String>,
    /// Default database applied at session creation.
    pub default_database: Option<String>,
    /// Default schema applied at session creation.
    pub default_schema: Option<String>,
    /// Max concurrent sessions (and browser auths) per `(account, user)`.
    pub pool_size: usize,
    /// Per-request HTTP timeout for REST calls.
    pub http_timeout: Duration,
    /// Overall deadline for one sign-in (SSO + login) so a hung auth can't hold
    /// the shared auth gate indefinitely.
    pub auth_timeout: Duration,
    /// Overall deadline for one query (submit + async-result polling).
    pub query_timeout: Duration,
    /// How often the background task heartbeats idle sessions to keep their
    /// master token alive. Zero disables the heartbeat.
    pub heartbeat_interval: Duration,
}

impl Default for SnowflakeEngineConfig {
    fn default() -> Self {
        Self {
            default_account: None,
            default_user: None,
            default_warehouse: None,
            default_role: None,
            default_database: None,
            default_schema: None,
            pool_size: DEFAULT_POOL_SIZE,
            http_timeout: client::config::DEFAULT_HTTP_TIMEOUT,
            auth_timeout: DEFAULT_AUTH_TIMEOUT,
            query_timeout: client::config::DEFAULT_QUERY_TIMEOUT,
            heartbeat_interval: DEFAULT_HEARTBEAT_INTERVAL,
        }
    }
}

impl SnowflakeEngineConfig {
    /// Resolves defaults from env vars (then settings.json). Cheap and
    /// side-effect-free; never authenticates.
    ///
    /// # Errors
    ///
    /// Currently infallible, but returns `Result` so the daemon registry wiring
    /// can `?` it and future validation can surface errors.
    pub fn from_env_and_settings() -> Result<Self> {
        let settings = Settings::load().unwrap_or_default();
        let pool_size = settings
            .get_env_var(ENV_POOL_SIZE)
            .and_then(|s| s.trim().parse::<usize>().ok())
            .filter(|&n| n >= 1)
            .unwrap_or(DEFAULT_POOL_SIZE);
        let secs = |key: &str| {
            settings
                .get_env_var(key)
                .and_then(|s| s.trim().parse::<u64>().ok())
                .filter(|&n| n >= 1)
                .map(Duration::from_secs)
        };
        Ok(Self {
            default_account: settings.get_env_var(ENV_ACCOUNT),
            default_user: settings.get_env_var(ENV_USER),
            default_warehouse: settings.get_env_var(ENV_WAREHOUSE),
            default_role: settings.get_env_var(ENV_ROLE),
            default_database: settings.get_env_var(ENV_DATABASE),
            default_schema: settings.get_env_var(ENV_SCHEMA),
            pool_size,
            http_timeout: secs(ENV_HTTP_TIMEOUT).unwrap_or(client::config::DEFAULT_HTTP_TIMEOUT),
            auth_timeout: secs(ENV_AUTH_TIMEOUT).unwrap_or(DEFAULT_AUTH_TIMEOUT),
            query_timeout: secs(ENV_QUERY_TIMEOUT).unwrap_or(client::config::DEFAULT_QUERY_TIMEOUT),
            heartbeat_interval: heartbeat_interval_from(
                settings.get_env_var(ENV_HEARTBEAT_INTERVAL),
            ),
        })
    }
}

/// Parses the heartbeat-interval setting: seconds, with `0` meaning disabled.
/// Unset or unparseable values fall back to the default. (The `secs` helper in
/// [`SnowflakeEngineConfig::from_env_and_settings`] rejects `0`, which here is
/// a meaningful value.)
fn heartbeat_interval_from(raw: Option<String>) -> Duration {
    raw.and_then(|s| s.trim().parse::<u64>().ok())
        .map_or(DEFAULT_HEARTBEAT_INTERVAL, Duration::from_secs)
}

/// A single arbitrary-SQL query request routed to the engine.
///
/// `account`/`user` and the per-query context default to the engine config when
/// omitted. Serialized by the CLI client (after [`fill_defaults_from`]
/// resolution) and deserialized from the daemon `query` op payload.
///
/// [`fill_defaults_from`]: QueryRequest::fill_defaults_from
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
#[serde(default)]
pub struct QueryRequest {
    /// Target account; falls back to `SNOWFLAKE_ACCOUNT`.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub account: Option<String>,
    /// Authenticating user; falls back to `SNOWFLAKE_USER`.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub user: Option<String>,
    /// Per-query `USE WAREHOUSE` override.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub warehouse: Option<String>,
    /// Per-query `USE ROLE` override.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub role: Option<String>,
    /// Per-query `USE DATABASE` override.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub database: Option<String>,
    /// Per-query `USE SCHEMA` override.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub schema: Option<String>,
    /// The SQL to execute.
    pub sql: String,
}

impl QueryRequest {
    /// Fills each unset identity/context field from `env`.
    ///
    /// Called by the CLI client with a profile-aware source
    /// ([`SettingsEnv`](crate::utils::settings::SettingsEnv)) so that
    /// `--profile` / `OMNI_DEV_PROFILE` resolve in the invoking process —
    /// the daemon's startup defaults then only back-fill requests that still
    /// omit a field (e.g. bare socket clients). Explicit values are never
    /// overwritten.
    pub fn fill_defaults_from(&mut self, env: &impl EnvSource) {
        let fill = |slot: &mut Option<String>, key: &str| {
            if slot.is_none() {
                *slot = env.var(key);
            }
        };
        fill(&mut self.account, ENV_ACCOUNT);
        fill(&mut self.user, ENV_USER);
        fill(&mut self.warehouse, ENV_WAREHOUSE);
        fill(&mut self.role, ENV_ROLE);
        fill(&mut self.database, ENV_DATABASE);
        fill(&mut self.schema, ENV_SCHEMA);
    }
    /// The per-query context overrides (the `Some` fields override the session
    /// base context).
    fn overrides(&self) -> QueryContext {
        QueryContext {
            warehouse: self.warehouse.clone(),
            role: self.role.clone(),
            database: self.database.clone(),
            schema: self.schema.clone(),
        }
    }
}

/// The running keep-alive heartbeat loop: cancelled and awaited on shutdown.
struct HeartbeatTask {
    token: CancellationToken,
    handle: JoinHandle<()>,
}

/// The account-agnostic Snowflake query engine: lazy multiplexed auth, bounded
/// per-identity session pools, and concurrent arbitrary-SQL execution.
///
/// A background keep-alive heartbeat for idle sessions is started by
/// [`start_heartbeat`](Self::start_heartbeat) and stopped by
/// [`shutdown`](Self::shutdown).
pub struct SnowflakeEngine {
    config: SnowflakeEngineConfig,
    registry: PoolRegistry,
    heartbeat: StdMutex<Option<HeartbeatTask>>,
}

impl SnowflakeEngine {
    /// Builds an engine. Cheap — no eager auth or I/O.
    #[must_use]
    pub fn new(config: SnowflakeEngineConfig) -> Self {
        Self {
            config,
            registry: PoolRegistry::new(),
            heartbeat: StdMutex::new(None),
        }
    }

    /// Starts the background keep-alive heartbeat loop (idempotent).
    ///
    /// Every `heartbeat_interval` the loop heartbeats each pool's idle sessions
    /// — renewing a session token about to lapse — so the server-side
    /// `CLIENT_SESSION_KEEP_ALIVE` actually extends the master token and an
    /// idle pool survives past the token TTL without a new browser SSO (#1107).
    /// Busy sessions are skipped; the query path keeps them alive inline.
    ///
    /// No-op when the interval is zero (disabled) or when called outside a
    /// tokio runtime. Stopped by [`shutdown`](Self::shutdown).
    pub fn start_heartbeat(&self) {
        let interval = self.config.heartbeat_interval;
        if interval.is_zero() {
            return;
        }
        if tokio::runtime::Handle::try_current().is_err() {
            tracing::debug!("no tokio runtime; Snowflake keep-alive heartbeat not started");
            return;
        }
        let mut guard = self.lock_heartbeat();
        if guard.is_some() {
            return;
        }
        let token = CancellationToken::new();
        let loop_token = token.clone();
        let registry = self.registry.clone();
        let handle = tokio::spawn(async move {
            loop {
                tokio::select! {
                    () = loop_token.cancelled() => break,
                    () = tokio::time::sleep(interval) => {
                        heartbeat_all_pools(&registry, interval).await;
                    }
                }
            }
        });
        *guard = Some(HeartbeatTask { token, handle });
    }

    fn lock_heartbeat(&self) -> MutexGuard<'_, Option<HeartbeatTask>> {
        self.heartbeat
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
    }

    /// A snapshot of every active pool.
    #[must_use]
    pub fn sessions(&self) -> Vec<SessionInfo> {
        self.registry.snapshot()
    }

    /// The number of active pools (`(account, user)` identities).
    #[must_use]
    pub fn pool_count(&self) -> usize {
        self.registry.len()
    }

    /// Evicts the pool for `(account, user)`. Returns whether one existed.
    pub fn disconnect(&self, account: &str, user: &str) -> bool {
        let key = SessionKey::new(normalize_account(account), user.trim());
        self.registry.remove(&key).is_some()
    }

    /// Evicts the pool with the given id. Returns whether one existed.
    pub fn disconnect_by_id(&self, id: u64) -> bool {
        self.registry.remove_by_id(id).is_some()
    }

    /// Evicts every pool. Returns how many were removed.
    pub fn disconnect_all(&self) -> usize {
        self.registry.take_all().len()
    }

    /// Stops the keep-alive heartbeat, then drops every pool (and its sessions).
    pub async fn shutdown(&self) {
        let task = self.lock_heartbeat().take();
        if let Some(task) = task {
            task.token.cancel();
            let _ = task.handle.await;
        }
        let pools = self.registry.take_all();
        drop(pools);
    }

    /// Runs arbitrary SQL against the `(account, user)` pool, authenticating a
    /// session on first use, and returns a self-describing `{ columns, rows }`
    /// payload. Concurrent calls run on separate pooled sessions (up to the pool
    /// size).
    ///
    /// # Errors
    ///
    /// Returns an error if no account/user can be resolved, a context flag is not
    /// a valid identifier, authentication fails, or the query fails. On a
    /// session-expiry error that session is discarded and the next query
    /// re-authenticates.
    pub async fn query(&self, req: QueryRequest) -> Result<Value> {
        let account = normalize_account(
            req.account
                .as_deref()
                .or(self.config.default_account.as_deref())
                .ok_or_else(|| {
                    anyhow!("no Snowflake account: pass --account or set SNOWFLAKE_ACCOUNT")
                })?,
        );
        let user = req
            .user
            .as_deref()
            .or(self.config.default_user.as_deref())
            .ok_or_else(|| anyhow!("no Snowflake user: pass --user or set SNOWFLAKE_USER"))?
            .trim()
            .to_string();
        validate_context(&req)?;

        let key = SessionKey::new(account, user);
        let overrides = req.overrides();
        let pool = self.registry.get_or_create(&key, self.config.pool_size);

        // Check out a session. The pool reuses an idle one when available
        // (re-checking after the auth gate so a session freed mid-auth is reused),
        // and only authenticates a new one — serialized to one browser at a time
        // by the pool's shared auth gate — when none is idle and it is under
        // capacity. The permit inside the checkout caps concurrency at pool_size.
        let cfg = self.config.clone();
        let create_key = key.clone();
        let auth_timeout = self.config.auth_timeout;
        let checkout = pool
            .checkout(move || async move {
                // Overall sign-in deadline so a hung auth releases the gate.
                match tokio::time::timeout(
                    auth_timeout,
                    create_session_with_base(&create_key, &cfg),
                )
                .await
                {
                    Ok(result) => result,
                    Err(_) => Err(ClientError::Auth(format!(
                        "Snowflake sign-in timed out after {auth_timeout:?}"
                    ))),
                }
            })
            .await
            .map_err(|e| {
                anyhow::Error::new(e).context(format!(
                    "failed to authenticate Snowflake session for {} / {}",
                    key.account, key.user
                ))
            })?;

        // Proactively renew a session whose token is about to expire, before use.
        if checkout
            .session()
            .session_expiring_within(TimeDelta::seconds(120))
            && checkout.session().renew().await.is_err()
        {
            pool.discard(checkout);
            return Err(anyhow!(
                "Snowflake session expired and was discarded — re-run the query to re-authenticate"
            ));
        }

        // Record what this member is now running, so menus/status show it.
        pool.start_query(checkout.id(), sql_preview(&req.sql, SQL_PREVIEW_MAX));

        // Apply the requested context and run the query, transparently renewing
        // the token and retrying once if it expires mid-flight.
        let target = checkout.base().overlay(&overrides);
        match run_with_renew(checkout.session(), checkout.current(), &target, &req.sql).await {
            Ok(rows) => {
                pool.touch();
                pool.checkin(checkout, target);
                Ok(client::rows_to_payload(&rows))
            }
            Err(e) if e.is_session_expired() => {
                pool.discard(checkout);
                Err(anyhow!(
                    "Snowflake session expired and was discarded — re-run the query to re-authenticate"
                ))
            }
            Err(e) => {
                // Log the underlying cause server-side and surface it to the
                // client (the daemon reply uses the full anyhow chain).
                tracing::warn!("Snowflake query failed: {e}");
                // The session's context is uncertain after a failure; check in
                // with an empty context so the next reuse re-applies every dimension.
                pool.checkin(checkout, QueryContext::default());
                Err(anyhow::Error::new(e).context("Snowflake query failed"))
            }
        }
    }
}

impl Drop for SnowflakeEngine {
    fn drop(&mut self) {
        // Best-effort: an engine dropped without `shutdown()` must not leave the
        // heartbeat loop running forever. Cancellation is sync; the task itself
        // is detached and exits on its next select.
        let task = self.lock_heartbeat().take();
        if let Some(task) = task {
            task.token.cancel();
        }
    }
}

/// Sends one keep-alive round to every pool's currently-idle sessions,
/// discarding any session that is dead beyond renewal.
async fn heartbeat_all_pools(registry: &PoolRegistry, interval: Duration) {
    for pool in registry.pools() {
        let checkouts = pool.checkout_all_idle();
        if checkouts.is_empty() {
            continue;
        }
        let total = checkouts.len();
        let mut kept = 0usize;
        for checkout in checkouts {
            if keep_session_alive(checkout.session(), interval).await {
                pool.restore(checkout);
                kept += 1;
            } else {
                pool.discard(checkout);
            }
        }
        tracing::debug!(
            pool = pool.id(),
            kept,
            total,
            "Snowflake keep-alive heartbeat round"
        );
    }
}

/// Keeps one idle session alive: proactively renews a session token that would
/// lapse before the next tick (the heartbeat itself is authorized by the
/// session token), then heartbeats so the server extends the master token.
///
/// Returns whether the session is still usable: `false` only when the master
/// token has expired (a full re-auth is unavoidable), so the caller discards
/// it. Transient errors keep the session for the next tick.
async fn keep_session_alive(session: &SnowflakeSession, interval: Duration) -> bool {
    let margin_secs = i64::try_from(interval.as_secs())
        .unwrap_or(i64::MAX)
        .saturating_add(HEARTBEAT_RENEW_MARGIN_SECS);
    let margin = TimeDelta::try_seconds(margin_secs).unwrap_or(TimeDelta::MAX);
    if session.session_expiring_within(margin) {
        match session.renew().await {
            Ok(()) => {}
            Err(e) if e.is_session_expired() => {
                tracing::warn!("Snowflake keep-alive: master token expired; discarding session");
                return false;
            }
            Err(e) => {
                // Transient: keep the session and try again next tick.
                tracing::warn!("Snowflake keep-alive renew failed: {e}");
                return true;
            }
        }
    }
    match session.heartbeat().await {
        Ok(()) => true,
        Err(e) if e.is_session_expired() => match session.renew().await {
            Ok(()) => true,
            Err(renew_err) if renew_err.is_session_expired() => {
                tracing::warn!("Snowflake keep-alive: master token expired; discarding session");
                false
            }
            Err(renew_err) => {
                tracing::warn!("Snowflake keep-alive renew failed: {renew_err}");
                true
            }
        },
        Err(e) => {
            tracing::warn!("Snowflake keep-alive heartbeat failed: {e}");
            true
        }
    }
}

/// Authenticates a session (external-browser SSO), enables keep-alive, and
/// captures its base (account/user default) context.
async fn create_session_with_base(
    key: &SessionKey,
    config: &SnowflakeEngineConfig,
) -> client::Result<(SnowflakeSession, QueryContext)> {
    let mut cfg = SnowflakeClientConfig::external_browser(&key.account, &key.user);
    cfg.warehouse = config.default_warehouse.clone();
    cfg.role = config.default_role.clone();
    cfg.database = config.default_database.clone();
    cfg.schema = config.default_schema.clone();
    cfg.http_timeout = config.http_timeout;
    cfg.query_timeout = config.query_timeout;

    let client = SnowflakeClient::new(cfg)?;
    let session = client.create_session().await?;
    session
        .query("ALTER SESSION SET CLIENT_SESSION_KEEP_ALIVE = true")
        .await?;
    let base = capture_base_context(&session).await?;
    Ok((session, base))
}

/// Reads the session's effective default context so per-query overrides can
/// later be reset back to it.
async fn capture_base_context(session: &SnowflakeSession) -> client::Result<QueryContext> {
    let rows = session
        .query("SELECT CURRENT_WAREHOUSE(), CURRENT_ROLE(), CURRENT_DATABASE(), CURRENT_SCHEMA()")
        .await?;
    let Some(row) = rows.first() else {
        return Ok(QueryContext::default());
    };
    Ok(QueryContext {
        warehouse: row.raw_at(0).map(str::to_string),
        role: row.raw_at(1).map(str::to_string),
        database: row.raw_at(2).map(str::to_string),
        schema: row.raw_at(3).map(str::to_string),
    })
}

/// Applies the context and runs the SQL, transparently renewing the session
/// token (via the master token) and retrying once if it expired mid-flight.
async fn run_with_renew(
    session: &SnowflakeSession,
    current: &QueryContext,
    target: &QueryContext,
    sql: &str,
) -> client::Result<Vec<Row>> {
    match apply_and_query(session, current, target, sql).await {
        Err(e) if e.is_session_expired() => {
            session.renew().await?;
            // Re-apply the full context on the renewed session, then retry.
            apply_and_query(session, &QueryContext::default(), target, sql).await
        }
        other => other,
    }
}

/// Issues any needed `USE` statements, then runs the SQL.
async fn apply_and_query(
    session: &SnowflakeSession,
    current: &QueryContext,
    target: &QueryContext,
    sql: &str,
) -> client::Result<Vec<Row>> {
    apply_context(session, current, target).await?;
    session.query(sql).await
}

/// Issues `USE` for each context dimension whose target differs from the
/// session's current value. Target names are either validated user overrides or
/// Snowflake-reported base names.
async fn apply_context(
    session: &SnowflakeSession,
    current: &QueryContext,
    target: &QueryContext,
) -> client::Result<()> {
    for (keyword, cur, tgt) in [
        (
            "WAREHOUSE",
            current.warehouse.as_deref(),
            target.warehouse.as_deref(),
        ),
        ("ROLE", current.role.as_deref(), target.role.as_deref()),
        (
            "DATABASE",
            current.database.as_deref(),
            target.database.as_deref(),
        ),
        (
            "SCHEMA",
            current.schema.as_deref(),
            target.schema.as_deref(),
        ),
    ] {
        if let Some(name) = tgt {
            if cur != Some(name) {
                session
                    .query(format!("USE {keyword} {name}").as_str())
                    .await?;
            }
        }
    }
    Ok(())
}

/// Normalizes an account identifier for keying (Snowflake is case-insensitive).
fn normalize_account(account: &str) -> String {
    account.trim().to_ascii_uppercase()
}

/// A single-line, length-bounded preview of SQL for the "running" display
/// (collapses whitespace/newlines; appends `…` when truncated).
fn sql_preview(sql: &str, max: usize) -> String {
    let collapsed = sql.split_whitespace().collect::<Vec<_>>().join(" ");
    if collapsed.chars().count() > max {
        let head: String = collapsed.chars().take(max).collect();
        format!("{}…", head.trim_end())
    } else {
        collapsed
    }
}

/// Validates every present context flag as a safe Snowflake identifier before it
/// is interpolated into a `USE …` statement.
fn validate_context(req: &QueryRequest) -> Result<()> {
    for (name, value) in [
        ("warehouse", req.warehouse.as_deref()),
        ("role", req.role.as_deref()),
        ("database", req.database.as_deref()),
        ("schema", req.schema.as_deref()),
    ] {
        if let Some(value) = value {
            validate_identifier(name, value)?;
        }
    }
    Ok(())
}

/// Rejects context values that are not bare Snowflake identifiers (letters,
/// digits, `_`, `$`, `.`), so a `--warehouse` flag cannot smuggle extra SQL into
/// the `USE …` statement.
fn validate_identifier(field: &str, value: &str) -> Result<()> {
    if value.is_empty() {
        bail!("--{field} must not be empty");
    }
    if !value
        .chars()
        .all(|c| c.is_ascii_alphanumeric() || matches!(c, '_' | '$' | '.'))
    {
        bail!(
            "--{field} '{value}' is not a valid Snowflake identifier \
             (allowed: letters, digits, '_', '$', '.')"
        );
    }
    Ok(())
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;
    use crate::test_support::env::MapEnv;

    #[test]
    fn default_config_has_a_nonzero_pool_size() {
        assert!(SnowflakeEngineConfig::default().pool_size >= 1);
    }

    #[test]
    fn heartbeat_interval_from_parses_seconds_zero_and_garbage() {
        assert_eq!(heartbeat_interval_from(None), DEFAULT_HEARTBEAT_INTERVAL);
        assert_eq!(
            heartbeat_interval_from(Some("300".to_string())),
            Duration::from_secs(300)
        );
        // `0` is meaningful: it disables the heartbeat.
        assert_eq!(
            heartbeat_interval_from(Some(" 0 ".to_string())),
            Duration::ZERO
        );
        assert_eq!(
            heartbeat_interval_from(Some("garbage".to_string())),
            DEFAULT_HEARTBEAT_INTERVAL
        );
    }

    #[tokio::test]
    async fn start_heartbeat_is_a_noop_when_disabled() {
        let engine = SnowflakeEngine::new(SnowflakeEngineConfig {
            heartbeat_interval: Duration::ZERO,
            ..SnowflakeEngineConfig::default()
        });
        engine.start_heartbeat();
        assert!(engine.lock_heartbeat().is_none());
    }

    #[test]
    fn start_heartbeat_is_a_noop_outside_a_runtime() {
        let engine = SnowflakeEngine::new(SnowflakeEngineConfig::default());
        engine.start_heartbeat();
        assert!(engine.lock_heartbeat().is_none());
    }

    #[tokio::test]
    async fn start_heartbeat_is_idempotent_and_shutdown_stops_it() {
        let engine = SnowflakeEngine::new(SnowflakeEngineConfig::default());
        engine.start_heartbeat();
        // Cancelling the running task's token lets a replacement be detected: a
        // second start must keep this task, not spawn (and orphan) a fresh one.
        engine.lock_heartbeat().as_ref().unwrap().token.cancel();
        engine.start_heartbeat();
        assert!(
            engine
                .lock_heartbeat()
                .as_ref()
                .unwrap()
                .token
                .is_cancelled(),
            "second start must not replace the running task"
        );
        engine.shutdown().await;
        assert!(engine.lock_heartbeat().is_none());
    }

    #[test]
    fn normalize_account_uppercases_and_trims() {
        assert_eq!(normalize_account("  my-org.acct  "), "MY-ORG.ACCT");
    }

    #[test]
    fn validate_identifier_accepts_bare_identifiers() {
        for value in ["WH", "my_wh", "DB.SCHEMA", "wh$1"] {
            assert!(validate_identifier("warehouse", value).is_ok(), "{value}");
        }
    }

    #[test]
    fn validate_identifier_rejects_injection_and_empty() {
        for value in ["", "wh; DROP TABLE t", "wh OR 1=1", "wh'", "a b"] {
            assert!(validate_identifier("warehouse", value).is_err(), "{value}");
        }
    }

    #[test]
    fn fill_defaults_from_fills_unset_fields() {
        let env = MapEnv::new()
            .with(ENV_ACCOUNT, "ACCT")
            .with(ENV_USER, "me")
            .with(ENV_WAREHOUSE, "WH")
            .with(ENV_ROLE, "R")
            .with(ENV_DATABASE, "DB")
            .with(ENV_SCHEMA, "S");
        let mut req = QueryRequest {
            sql: "SELECT 1".to_string(),
            ..QueryRequest::default()
        };
        req.fill_defaults_from(&env);
        assert_eq!(req.account.as_deref(), Some("ACCT"));
        assert_eq!(req.user.as_deref(), Some("me"));
        assert_eq!(req.warehouse.as_deref(), Some("WH"));
        assert_eq!(req.role.as_deref(), Some("R"));
        assert_eq!(req.database.as_deref(), Some("DB"));
        assert_eq!(req.schema.as_deref(), Some("S"));
    }

    #[test]
    fn fill_defaults_from_keeps_explicit_values() {
        let env = MapEnv::new()
            .with(ENV_ACCOUNT, "ENV_ACCT")
            .with(ENV_USER, "env_user");
        let mut req = QueryRequest {
            account: Some("FLAG_ACCT".to_string()),
            sql: "SELECT 1".to_string(),
            ..QueryRequest::default()
        };
        req.fill_defaults_from(&env);
        assert_eq!(req.account.as_deref(), Some("FLAG_ACCT"));
        assert_eq!(req.user.as_deref(), Some("env_user"));
    }

    #[test]
    fn fill_defaults_from_leaves_unresolved_fields_none() {
        let mut req = QueryRequest {
            sql: "SELECT 1".to_string(),
            ..QueryRequest::default()
        };
        req.fill_defaults_from(&MapEnv::new());
        assert!(req.account.is_none());
        assert!(req.user.is_none());
        assert!(req.warehouse.is_none());
    }

    #[test]
    fn query_request_serializes_without_none_fields() {
        let req = QueryRequest {
            account: Some("ACCT".to_string()),
            sql: "SELECT 1".to_string(),
            ..QueryRequest::default()
        };
        let value = serde_json::to_value(&req).unwrap();
        assert_eq!(
            value,
            serde_json::json!({ "account": "ACCT", "sql": "SELECT 1" })
        );
    }

    #[test]
    fn validate_context_checks_each_present_flag() {
        let mut req = QueryRequest {
            sql: "SELECT 1".to_string(),
            ..QueryRequest::default()
        };
        assert!(validate_context(&req).is_ok());
        req.role = Some("good_role".to_string());
        assert!(validate_context(&req).is_ok());
        req.database = Some("bad; drop".to_string());
        assert!(validate_context(&req).is_err());
    }

    #[tokio::test]
    async fn query_without_account_errors_without_network() {
        let engine = SnowflakeEngine::new(SnowflakeEngineConfig::default());
        let err = engine
            .query(QueryRequest {
                sql: "SELECT 1".to_string(),
                ..QueryRequest::default()
            })
            .await
            .unwrap_err();
        assert!(err.to_string().contains("account"));
    }

    #[tokio::test]
    async fn query_without_user_errors_without_network() {
        let engine = SnowflakeEngine::new(SnowflakeEngineConfig {
            default_account: Some("ACCT".to_string()),
            ..SnowflakeEngineConfig::default()
        });
        let err = engine
            .query(QueryRequest {
                sql: "SELECT 1".to_string(),
                ..QueryRequest::default()
            })
            .await
            .unwrap_err();
        assert!(err.to_string().contains("user"));
    }

    #[test]
    fn disconnect_and_sessions_on_empty_engine() {
        let engine = SnowflakeEngine::new(SnowflakeEngineConfig::default());
        assert_eq!(engine.pool_count(), 0);
        assert!(engine.sessions().is_empty());
        assert!(!engine.disconnect("ACCT", "user"));
        assert!(!engine.disconnect_by_id(1));
        assert_eq!(engine.disconnect_all(), 0);
    }

    #[test]
    fn sql_preview_collapses_whitespace_and_truncates() {
        // Whitespace/newlines collapse to single spaces; short SQL is unchanged.
        assert_eq!(sql_preview("SELECT   1\n  FROM t", 60), "SELECT 1 FROM t");
        // Over-length SQL is truncated with an ellipsis.
        let long = format!("SELECT {}", "a".repeat(100));
        let preview = sql_preview(&long, 20);
        assert!(preview.ends_with('…'));
        assert!(preview.chars().count() <= 21, "{preview}");
    }

    #[test]
    fn overrides_extracts_only_the_set_dimensions() {
        let req = QueryRequest {
            warehouse: Some("WH".to_string()),
            schema: Some("S".to_string()),
            sql: "SELECT 1".to_string(),
            ..QueryRequest::default()
        };
        let overrides = req.overrides();
        assert_eq!(overrides.warehouse.as_deref(), Some("WH"));
        assert_eq!(overrides.schema.as_deref(), Some("S"));
        assert!(overrides.role.is_none());
        assert!(overrides.database.is_none());
    }

    mod orchestration {
        use super::*;
        use serde_json::json;
        use wiremock::matchers::{method, path};
        use wiremock::{Mock, MockServer, ResponseTemplate};

        /// Mounts a `query-request` handler that returns `data` for every POST.
        async fn mount_query(server: &MockServer, data: serde_json::Value) {
            Mock::given(method("POST"))
                .and(path("/queries/v1/query-request"))
                .respond_with(
                    ResponseTemplate::new(200)
                        .set_body_json(json!({ "success": true, "data": data })),
                )
                .mount(server)
                .await;
        }

        #[tokio::test]
        async fn capture_base_context_reads_current_context() {
            let server = MockServer::start().await;
            mount_query(
                &server,
                json!({
                    "rowtype": [
                        { "name": "CURRENT_WAREHOUSE()", "type": "text" },
                        { "name": "CURRENT_ROLE()", "type": "text" },
                        { "name": "CURRENT_DATABASE()", "type": "text" },
                        { "name": "CURRENT_SCHEMA()", "type": "text" },
                    ],
                    "rowset": [["WH", "R", "DB", "S"]],
                }),
            )
            .await;
            let session = client::test_session(&server.uri(), Duration::from_secs(5));
            let base = capture_base_context(&session).await.unwrap();
            assert_eq!(base.warehouse.as_deref(), Some("WH"));
            assert_eq!(base.role.as_deref(), Some("R"));
            assert_eq!(base.database.as_deref(), Some("DB"));
            assert_eq!(base.schema.as_deref(), Some("S"));
        }

        #[tokio::test]
        async fn capture_base_context_defaults_when_no_rows() {
            let server = MockServer::start().await;
            mount_query(
                &server,
                json!({ "rowtype": [{ "name": "X", "type": "text" }], "rowset": [] }),
            )
            .await;
            let session = client::test_session(&server.uri(), Duration::from_secs(5));
            assert_eq!(
                capture_base_context(&session).await.unwrap(),
                QueryContext::default()
            );
        }

        #[tokio::test]
        async fn apply_context_issues_use_only_for_differing_dimensions() {
            let server = MockServer::start().await;
            mount_query(&server, json!({ "rowtype": [], "rowset": [] })).await;
            let session = client::test_session(&server.uri(), Duration::from_secs(5));

            let current = QueryContext {
                warehouse: Some("WH".to_string()),
                ..QueryContext::default()
            };
            let target = QueryContext {
                warehouse: Some("WH".to_string()), // same → no USE
                role: Some("R2".to_string()),      // differs → one USE
                ..QueryContext::default()
            };
            apply_context(&session, &current, &target).await.unwrap();

            let reqs = server.received_requests().await.unwrap();
            assert_eq!(reqs.len(), 1, "only the differing dimension issues a USE");
            assert!(String::from_utf8_lossy(&reqs[0].body).contains("USE ROLE R2"));
        }

        #[tokio::test]
        async fn run_with_renew_runs_without_renew_when_not_expired() {
            let server = MockServer::start().await;
            mount_query(
                &server,
                json!({ "rowtype": [{ "name": "N", "type": "text" }], "rowset": [["x"]] }),
            )
            .await;
            let session = client::test_session(&server.uri(), Duration::from_secs(5));
            let ctx = QueryContext::default();
            let rows = run_with_renew(&session, &ctx, &ctx, "SELECT 1")
                .await
                .unwrap();
            assert_eq!(rows.len(), 1);
        }

        #[tokio::test]
        async fn run_with_renew_renews_and_retries_once_on_expiry() {
            let server = MockServer::start().await;
            // First query attempt: session expired (then this rule is exhausted).
            Mock::given(method("POST"))
                .and(path("/queries/v1/query-request"))
                .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                    "success": false, "code": "390112", "message": "expired", "data": {}
                })))
                .up_to_n_times(1)
                .with_priority(1)
                .mount(&server)
                .await;
            // Renew succeeds.
            Mock::given(method("POST"))
                .and(path("/session/token-request"))
                .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                    "success": true,
                    "data": { "sessionToken": "fresh", "validityInSecondsST": 3600 }
                })))
                .mount(&server)
                .await;
            // The retried query succeeds.
            Mock::given(method("POST"))
                .and(path("/queries/v1/query-request"))
                .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                    "success": true,
                    "data": { "rowtype": [{ "name": "N", "type": "text" }], "rowset": [["x"]] }
                })))
                .with_priority(2)
                .mount(&server)
                .await;

            let session = client::test_session(&server.uri(), Duration::from_secs(5));
            let ctx = QueryContext::default();
            let rows = run_with_renew(&session, &ctx, &ctx, "SELECT 1")
                .await
                .unwrap();
            assert_eq!(rows.len(), 1, "renewed and retried transparently");
        }
    }

    mod keep_alive {
        use super::*;
        use serde_json::json;
        use wiremock::matchers::{method, path};
        use wiremock::{Mock, MockServer, ResponseTemplate};

        /// A short interval so `session_expiring_within(interval + margin)` is
        /// false for the test session's fresh 3600s token.
        const INTERVAL: Duration = Duration::from_secs(60);

        /// Mounts a `session/heartbeat` handler answering with `body`.
        async fn mount_heartbeat(server: &MockServer, body: serde_json::Value) {
            Mock::given(method("POST"))
                .and(path("/session/heartbeat"))
                .respond_with(ResponseTemplate::new(200).set_body_json(body))
                .mount(server)
                .await;
        }

        /// Mounts a `session/token-request` (renew) handler answering with `body`.
        async fn mount_renew(server: &MockServer, body: serde_json::Value) {
            Mock::given(method("POST"))
                .and(path("/session/token-request"))
                .respond_with(ResponseTemplate::new(200).set_body_json(body))
                .mount(server)
                .await;
        }

        fn ok_body() -> serde_json::Value {
            json!({ "success": true, "data": {} })
        }

        fn renew_ok_body() -> serde_json::Value {
            json!({
                "success": true,
                "data": { "sessionToken": "fresh", "validityInSecondsST": 3600 }
            })
        }

        fn expired_body() -> serde_json::Value {
            json!({ "success": false, "code": "390112", "message": "expired", "data": {} })
        }

        #[tokio::test]
        async fn keep_session_alive_heartbeats_a_healthy_session() {
            let server = MockServer::start().await;
            mount_heartbeat(&server, ok_body()).await;
            let session = client::test_session(&server.uri(), Duration::from_secs(5));
            assert!(keep_session_alive(&session, INTERVAL).await);
            let reqs = server.received_requests().await.unwrap();
            assert_eq!(reqs.len(), 1, "one heartbeat, no renew");
        }

        #[tokio::test]
        async fn keep_session_alive_renews_when_the_heartbeat_reports_expiry() {
            let server = MockServer::start().await;
            mount_heartbeat(&server, expired_body()).await;
            mount_renew(&server, renew_ok_body()).await;
            let session = client::test_session(&server.uri(), Duration::from_secs(5));
            assert!(keep_session_alive(&session, INTERVAL).await);
            let reqs = server.received_requests().await.unwrap();
            assert!(
                reqs.iter()
                    .any(|r| r.url.path() == "/session/token-request"),
                "renewed after the expired heartbeat"
            );
        }

        #[tokio::test]
        async fn keep_session_alive_discards_when_the_master_token_is_dead() {
            let server = MockServer::start().await;
            mount_heartbeat(&server, expired_body()).await;
            mount_renew(&server, expired_body()).await;
            let session = client::test_session(&server.uri(), Duration::from_secs(5));
            assert!(!keep_session_alive(&session, INTERVAL).await);
        }

        #[tokio::test]
        async fn keep_session_alive_keeps_the_session_on_transient_errors() {
            let server = MockServer::start().await;
            mount_heartbeat(
                &server,
                json!({ "success": false, "code": "390001", "message": "hiccup", "data": {} }),
            )
            .await;
            let session = client::test_session(&server.uri(), Duration::from_secs(5));
            assert!(keep_session_alive(&session, INTERVAL).await);
        }

        #[tokio::test]
        async fn keep_session_alive_proactively_renews_a_token_expiring_before_the_next_tick() {
            let server = MockServer::start().await;
            mount_heartbeat(&server, ok_body()).await;
            mount_renew(&server, renew_ok_body()).await;
            let session = client::test_session(&server.uri(), Duration::from_secs(5));
            // interval + margin exceeds the fresh 3600s validity → renew first.
            assert!(keep_session_alive(&session, Duration::from_secs(7200)).await);
            let reqs = server.received_requests().await.unwrap();
            assert_eq!(
                reqs[0].url.path(),
                "/session/token-request",
                "renew ran first"
            );
            assert_eq!(reqs[1].url.path(), "/session/heartbeat");
        }

        #[tokio::test]
        async fn engine_heartbeat_loop_beats_idle_sessions_and_stops_on_shutdown() {
            let server = MockServer::start().await;
            mount_heartbeat(&server, ok_body()).await;

            let engine = SnowflakeEngine::new(SnowflakeEngineConfig {
                heartbeat_interval: Duration::from_millis(50),
                ..SnowflakeEngineConfig::default()
            });
            // Park one idle session in a pool, bypassing the (live-only) SSO.
            let pool = engine
                .registry
                .get_or_create(&SessionKey::new("ACCT", "user"), 2);
            let uri = server.uri();
            let checkout = pool
                .checkout(|| async {
                    Ok::<_, std::convert::Infallible>((
                        client::test_session(&uri, Duration::from_secs(5)),
                        QueryContext::default(),
                    ))
                })
                .await
                .unwrap();
            pool.checkin(checkout, QueryContext::default());

            engine.start_heartbeat();
            let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
            while server.received_requests().await.unwrap().is_empty() {
                assert!(
                    tokio::time::Instant::now() < deadline,
                    "no heartbeat within the deadline"
                );
                tokio::time::sleep(Duration::from_millis(10)).await;
            }
            // The borrowed session was restored, not consumed.
            assert_eq!(pool.live(), 1);

            engine.shutdown().await;
            let after = server.received_requests().await.unwrap().len();
            tokio::time::sleep(Duration::from_millis(150)).await;
            assert_eq!(
                server.received_requests().await.unwrap().len(),
                after,
                "no heartbeats after shutdown"
            );
            assert_eq!(engine.pool_count(), 0, "pools drained on shutdown");
        }
    }
}