nordisk/knut

crates/knut-bifrost/tests/common/mod.rs

//! Shared by the live Iceberg → Doris tests (`live_iceberg_doris.rs`, and the
//! container harness): Doris SQL over the FE's HTTP API, and the people rows.
//!
//! The FE runs a statement posted to `/api/query/default_cluster/<db>` with
//! HTTP Basic, so DDL and checks need no MySQL client and no shell.

#![allow(dead_code)]

pub mod rig;

use std::sync::Arc;

use arrow_array::{Int64Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema};

/// Doris SQL over the FE's HTTP API.
pub struct DorisSql {
    pub fe: String,
    pub user: String,
    pub password: String,
    client: reqwest::Client,
}

impl DorisSql {
    pub fn new(fe: &str, user: &str, password: &str) -> Self {
        Self {
            fe: fe.trim_end_matches('/').to_string(),
            user: user.into(),
            password: password.into(),
            client: tls_roots(reqwest::Client::builder())
                .timeout(std::time::Duration::from_secs(600))
                .build()
                .expect("http client"),
        }
    }

    /// From `DORIS_USER` / `DORIS_PASSWORD`.
    pub fn from_env(fe: &str) -> Self {
        Self::new(
            fe,
            &std::env::var("DORIS_USER").expect("DORIS_USER"),
            &std::env::var("DORIS_PASSWORD").unwrap_or_default(),
        )
    }

    /// Run one statement; the result rows (empty for DDL). Panics with the
    /// FE's own message on failure — a test wants the reason, not a code.
    pub async fn exec(&self, stmt: &str) -> Vec<Vec<serde_json::Value>> {
        let url = format!("{}/api/query/default_cluster/information_schema", self.fe);
        let resp = self
            .client
            .post(&url)
            .basic_auth(&self.user, Some(&self.password))
            .json(&serde_json::json!({ "stmt": stmt }))
            .send()
            .await
            .unwrap_or_else(|e| panic!("doris sql {stmt}: {e}"));
        let v: serde_json::Value = resp.json().await.expect("doris sql answers JSON");
        if v.get("code").and_then(serde_json::Value::as_i64) != Some(0) {
            panic!("doris sql `{stmt}` failed: {v}");
        }
        v.pointer("/data/data")
            .and_then(serde_json::Value::as_array)
            .map(|rows| rows.iter().filter_map(|r| r.as_array().cloned()).collect())
            .unwrap_or_default()
    }

    /// `select <expr> from <table> [where …]` → one integer.
    pub async fn scalar(&self, stmt: &str) -> i64 {
        let rows = self.exec(stmt).await;
        let cell = rows.first().and_then(|r| r.first()).cloned().unwrap_or_default();
        cell.as_i64()
            .or_else(|| cell.as_str().and_then(|s| s.parse().ok()))
            .unwrap_or_else(|| panic!("`{stmt}` gave no integer: {rows:?}"))
    }

    /// `count(*)` that has to read the rows (a bare count may come from metadata).
    pub async fn count(&self, table: &str) -> i64 {
        self.scalar(&format!("select count(*) from {table} where id >= 0")).await
    }

    /// Wait until the FE answers SQL (a fresh container needs a while).
    pub async fn wait_ready(&self, timeout: std::time::Duration) {
        let t0 = std::time::Instant::now();
        loop {
            let url = format!("{}/api/query/default_cluster/information_schema", self.fe);
            let ok = self
                .client
                .post(&url)
                .basic_auth(&self.user, Some(&self.password))
                .json(&serde_json::json!({ "stmt": "select 1" }))
                .send()
                .await
                .ok()
                .map(|r| r.status().is_success())
                .unwrap_or(false);
            if ok {
                return;
            }
            assert!(t0.elapsed() < timeout, "doris FE {} never answered SQL", self.fe);
            tokio::time::sleep(std::time::Duration::from_secs(2)).await;
        }
    }
}

/// The destination DDL: a unique-key merge-on-write table (so deletes apply)
/// with an inverted index on `name` — the native storage the lake cannot give.
pub fn people_ddl(db: &str, table: &str, props: &str) -> String {
    format!(
        "CREATE TABLE IF NOT EXISTS {db}.{table} (\
           id BIGINT NOT NULL, \
           name VARCHAR(64), \
           pnr BIGINT, \
           INDEX idx_name (name) USING INVERTED\
         ) UNIQUE KEY(id) DISTRIBUTED BY HASH(id) BUCKETS 4 \
         PROPERTIES (\"enable_unique_key_merge_on_write\" = \"true\"{props})"
    )
}

pub fn people_schema() -> Arc<Schema> {
    Arc::new(Schema::new(vec![
        Field::new("id", DataType::Int64, false),
        Field::new("name", DataType::Utf8, true),
        Field::new("pnr", DataType::Int64, true),
    ]))
}

/// Rows `ids`, named `p<id>` (or `<tag><id>`, to tell an update from the original).
pub fn people(ids: std::ops::Range<i64>, tag: &str) -> RecordBatch {
    let n: Vec<i64> = ids.collect();
    let names: Vec<String> = n.iter().map(|i| format!("{tag}{i}")).collect();
    let pnr: Vec<i64> = n.iter().map(|i| 19_000_000_0000 + i).collect();
    RecordBatch::try_new(
        people_schema(),
        vec![
            Arc::new(Int64Array::from(n)),
            Arc::new(StringArray::from(names)),
            Arc::new(Int64Array::from(pnr)),
        ],
    )
    .unwrap()
}

/// A name no two runs share (tables are never dropped — MinIO data is never deleted).
pub fn run_tag() -> String {
    let ns = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_secs();
    format!("r{ns}")
}

/// reqwest 0.13 with reqwest 0.12 `rustls-tls` behaviour (2026-10-02): the ring
/// provider and the bundled Mozilla roots (webpki), not 0.13's aws-lc-rs + OS
/// trust store. Roots added later on the builder are merged in, as in 0.12.
#[allow(dead_code)]
fn tls_roots(b: reqwest::ClientBuilder) -> reqwest::ClientBuilder {
    let _ = rustls::crypto::ring::default_provider().install_default();
    b.tls_certs_only(
        webpki_root_certs::TLS_SERVER_ROOT_CERTS
            .iter()
            .filter_map(|c| reqwest::Certificate::from_der(c).ok()),
    )
}