feat: cluster variants, ADBC health/reconcile, and distributed runnin…

crates/queryflux-core/src/config_json.rs

@@ -146,6 +146,9 @@ pub fn cluster_config_from_persisted_json(
         tls: None,
         auth,
         query_auth,
+        variants: Vec::new(),
+        health_check_query: json_str(config, "healthCheckQuery"),
+        reconcile_query: json_str(config, "reconcileQuery"),
     }
 }
 

crates/queryflux-e2e-tests/src/harness.rs

@@ -282,6 +282,8 @@ impl TestHarness {
             cluster_manager,
             adapters,
             health_check_targets: vec![],
+            custom_health_queries: HashMap::new(),
+            custom_reconcile_queries: HashMap::new(),
             cluster_configs: HashMap::new(),
             group_members,
             group_order,
@@ -507,6 +509,8 @@ impl WireTestHarness {
             cluster_manager,
             adapters,
             health_check_targets: vec![],
+            custom_health_queries: HashMap::new(),
+            custom_reconcile_queries: HashMap::new(),
             cluster_configs: HashMap::new(),
             group_members: HashMap::from([(GROUP_DUCKDB.to_string(), vec![cluster.0.clone()])]),
             group_order: vec![GROUP_DUCKDB.to_string()],
@@ -640,6 +644,8 @@ impl WireTestHarness {
             cluster_manager,
             adapters,
             health_check_targets: vec![],
+            custom_health_queries: HashMap::new(),
+            custom_reconcile_queries: HashMap::new(),
             cluster_configs: HashMap::new(),
             group_members: HashMap::from([(GROUP_STARROCKS.to_string(), vec![cluster.0.clone()])]),
             group_order: vec![GROUP_STARROCKS.to_string()],

crates/queryflux-engine-adapters/src/adbc/bigquery.rs

@@ -0,0 +1,125 @@
+use async_trait::async_trait;
+use queryflux_core::query::ClusterName;
+
+use super::introspection::AdbcIntrospection;
+use super::sql_helpers;
+use super::AdbcPool;
+
+/// BigQuery introspection via `INFORMATION_SCHEMA.JOBS_BY_PROJECT` (metadata only, no slot cost).
+pub struct BigQueryIntrospection {
+    pool: AdbcPool,
+    project_id: String,
+    region: Option<String>,
+}
+
+pub fn try_from_adbc_config(
+    cluster_name: &ClusterName,
+    uri: &str,
+    db_kwargs: &[(String, String)],
+    pool: AdbcPool,
+) -> Option<BigQueryIntrospection> {
+    let project_id = sql_helpers::db_kwarg(db_kwargs, "project_id")
+        .or_else(|| sql_helpers::db_kwarg(db_kwargs, "project"))
+        .or_else(|| uri.split('/').next().map(str::to_string))
+        .filter(|p| !p.is_empty())?;
+    let region = sql_helpers::db_kwarg(db_kwargs, "location")
+        .or_else(|| sql_helpers::db_kwarg(db_kwargs, "region"));
+
+    tracing::info!(
+        cluster = %cluster_name.0,
+        project_id = %project_id,
+        region = ?region,
+        "BigQuery introspection initialized (INFORMATION_SCHEMA.JOBS_BY_PROJECT)"
+    );
+
+    Some(BigQueryIntrospection {
+        pool,
+        project_id,
+        region,
+    })
+}
+
+impl BigQueryIntrospection {
+    pub(crate) fn build_reconcile_sql(project_id: &str, region: Option<&str>) -> String {
+        if let Some(region) = region {
+            format!(
+                "SELECT COUNT(*) FROM `{region}`.INFORMATION_SCHEMA.JOBS_BY_PROJECT \
+                 WHERE state = 'RUNNING' AND project_id = '{}'",
+                sql_helpers::escape_sql_literal(project_id)
+            )
+        } else {
+            format!(
+                "SELECT COUNT(*) FROM `{project}`.INFORMATION_SCHEMA.JOBS_BY_PROJECT \
+                 WHERE state = 'RUNNING'",
+                project = project_id
+            )
+        }
+    }
+
+    fn reconcile_sql(&self) -> String {
+        Self::build_reconcile_sql(&self.project_id, self.region.as_deref())
+    }
+}
+
+#[async_trait]
+impl AdbcIntrospection for BigQueryIntrospection {
+    async fn health_check(&self) -> bool {
+        // Metadata queries are cheap; no warehouse to wake.
+        true
+    }
+
+    async fn fetch_running_query_count(&self) -> Option<u64> {
+        let pool = self.pool.clone();
+        let sql = self.reconcile_sql();
+        tokio::task::spawn_blocking(move || {
+            let batches = sql_helpers::query_batches(&pool, &sql)?;
+            sql_helpers::first_cell_u64(&batches)
+        })
+        .await
+        .ok()
+        .flatten()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn reconcile_sql_without_region_uses_project_dataset() {
+        let sql = BigQueryIntrospection::build_reconcile_sql("my-proj", None);
+        assert!(sql.contains("`my-proj`.INFORMATION_SCHEMA.JOBS_BY_PROJECT"));
+        assert!(sql.contains("state = 'RUNNING'"));
+        assert!(!sql.contains("project_id ="));
+    }
+
+    #[test]
+    fn reconcile_sql_with_region_filters_project() {
+        let sql = BigQueryIntrospection::build_reconcile_sql("my-proj", Some("region-us"));
+        assert!(sql.contains("`region-us`.INFORMATION_SCHEMA.JOBS_BY_PROJECT"));
+        assert!(sql.contains("project_id = 'my-proj'"));
+    }
+
+    #[test]
+    fn reconcile_sql_escapes_project_id_quotes() {
+        let sql = BigQueryIntrospection::build_reconcile_sql("proj'x", Some("region-us"));
+        assert!(sql.contains("project_id = 'proj''x'"));
+    }
+
+    #[test]
+    fn project_id_from_db_kwargs_or_uri() {
+        let kwargs = vec![("project_id".into(), "from-kw".into())];
+        assert_eq!(
+            sql_helpers::db_kwarg(&kwargs, "project_id").as_deref(),
+            Some("from-kw")
+        );
+        assert_eq!(
+            "my-proj/dataset"
+                .split('/')
+                .next()
+                .map(str::to_string)
+                .as_deref(),
+            Some("my-proj")
+        );
+    }
+}

crates/queryflux-engine-adapters/src/adbc/databricks.rs

@@ -0,0 +1,205 @@
+use async_trait::async_trait;
+use queryflux_core::query::ClusterName;
+
+use super::introspection::AdbcIntrospection;
+
+/// REST-based introspection for Databricks SQL Warehouses.
+/// Avoids waking the warehouse via SQL `SELECT 1` or system tables.
+pub struct DatabricksIntrospection {
+    http: reqwest::Client,
+    workspace_url: String,
+    warehouse_id: String,
+    auth_token: String,
+}
+
+pub fn try_from_adbc_config(
+    cluster_name: &ClusterName,
+    uri: &str,
+    db_kwargs: &[(String, String)],
+) -> Option<DatabricksIntrospection> {
+    let workspace_url = uri.trim_end_matches('/').to_string();
+    if workspace_url.is_empty() {
+        tracing::warn!(cluster = %cluster_name.0, "Databricks REST: no workspace URL in ADBC URI");
+        return None;
+    }
+
+    let find = |key: &str| {
+        db_kwargs
+            .iter()
+            .find(|(k, _)| k == key)
+            .map(|(_, v)| v.clone())
+    };
+
+    let auth_token = find("token")
+        .or_else(|| find("access_token"))
+        .or_else(|| find("databricks_token"))
+        .unwrap_or_default();
+    if auth_token.is_empty() {
+        tracing::warn!(cluster = %cluster_name.0, "Databricks REST: no auth token found in dbKwargs");
+        return None;
+    }
+
+    let http_path = find("http_path").unwrap_or_default();
+    let warehouse_id = http_path.rsplit('/').next().unwrap_or("").to_string();
+    if warehouse_id.is_empty() {
+        tracing::warn!(cluster = %cluster_name.0, "Databricks REST: cannot parse warehouse ID from http_path");
+        return None;
+    }
+
+    tracing::info!(
+        cluster = %cluster_name.0,
+        workspace = %workspace_url,
+        warehouse_id = %warehouse_id,
+        "Databricks REST client initialized for health check and reconciliation"
+    );
+
+    Some(DatabricksIntrospection {
+        http: reqwest::Client::new(),
+        workspace_url,
+        warehouse_id,
+        auth_token,
+    })
+}
+
+#[async_trait]
+impl AdbcIntrospection for DatabricksIntrospection {
+    /// Health check via `GET /api/2.0/sql/warehouses/{warehouse_id}`.
+    /// Healthy = state is RUNNING, STARTING, or STOPPED (warehouse exists and is accessible).
+    async fn health_check(&self) -> bool {
+        let url = format!(
+            "{}/api/2.0/sql/warehouses/{}",
+            self.workspace_url, self.warehouse_id
+        );
+        match self
+            .http
+            .get(&url)
+            .bearer_auth(&self.auth_token)
+            .send()
+            .await
+        {
+            Ok(resp) if resp.status().is_success() => {
+                if let Ok(body) = resp.json::<serde_json::Value>().await {
+                    let state = body.get("state").and_then(|v| v.as_str()).unwrap_or("");
+                    matches!(state, "RUNNING" | "STARTING" | "STOPPED")
+                } else {
+                    false
+                }
+            }
+            Ok(resp) => {
+                tracing::warn!(
+                    warehouse_id = %self.warehouse_id,
+                    status = %resp.status(),
+                    "Databricks REST health check failed"
+                );
+                false
+            }
+            Err(e) => {
+                tracing::warn!(
+                    warehouse_id = %self.warehouse_id,
+                    error = %e,
+                    "Databricks REST health check request error"
+                );
+                false
+            }
+        }
+    }
+
+    /// Reconciliation via `GET /api/2.0/sql/history/queries`.
+    /// Returns the count of RUNNING queries for this warehouse.
+    async fn fetch_running_query_count(&self) -> Option<u64> {
+        let url = format!("{}/api/2.0/sql/history/queries", self.workspace_url);
+        let filter = serde_json::json!({
+            "filter_by": {
+                "statuses": ["RUNNING"],
+                "warehouse_ids": [&self.warehouse_id]
+            }
+        });
+        match self
+            .http
+            .get(&url)
+            .bearer_auth(&self.auth_token)
+            .json(&filter)
+            .send()
+            .await
+        {
+            Ok(resp) if resp.status().is_success() => {
+                if let Ok(body) = resp.json::<serde_json::Value>().await {
+                    body.get("res")
+                        .and_then(|v| v.as_array())
+                        .map(|arr| arr.len() as u64)
+                } else {
+                    None
+                }
+            }
+            Ok(resp) => {
+                tracing::warn!(
+                    warehouse_id = %self.warehouse_id,
+                    status = %resp.status(),
+                    "Databricks REST reconcile query failed"
+                );
+                None
+            }
+            Err(e) => {
+                tracing::warn!(
+                    warehouse_id = %self.warehouse_id,
+                    error = %e,
+                    "Databricks REST reconcile request error"
+                );
+                None
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::adbc::sql_helpers;
+    use queryflux_core::query::ClusterName;
+
+    #[test]
+    fn extracts_warehouse_id_from_http_path() {
+        let kwargs = vec![(
+            "http_path".into(),
+            "/sql/1.0/warehouses/abc123def456".into(),
+        )];
+        let path = sql_helpers::db_kwarg(&kwargs, "http_path").unwrap();
+        let id = path.rsplit('/').next().unwrap();
+        assert_eq!(id, "abc123def456");
+    }
+
+    #[test]
+    fn try_from_requires_token_and_http_path() {
+        let name = ClusterName("wh".into());
+        assert!(try_from_adbc_config(&name, "https://dbc.example.com", &[],).is_none());
+        assert!(try_from_adbc_config(
+            &name,
+            "https://dbc.example.com",
+            &[("token".into(), "tok".into())],
+        )
+        .is_none());
+        assert!(try_from_adbc_config(
+            &name,
+            "",
+            &[
+                ("token".into(), "tok".into(),),
+                ("http_path".into(), "/sql/1.0/warehouses/id1".into(),)
+            ],
+        )
+        .is_none());
+    }
+
+    #[test]
+    fn try_from_accepts_token_and_http_path() {
+        let name = ClusterName("wh".into());
+        let intro = try_from_adbc_config(
+            &name,
+            "https://dbc.example.com/",
+            &[
+                ("token".into(), "tok".into()),
+                ("http_path".into(), "/sql/1.0/warehouses/wh-id".into()),
+            ],
+        );
+        assert!(intro.is_some());
+    }
+}

crates/queryflux-engine-adapters/src/adbc/introspection.rs

@@ -0,0 +1,9 @@
+use async_trait::async_trait;
+
+/// Non-pool introspection for ADBC drivers that must not use default SQL health
+/// checks or reconcile queries (Databricks REST, auto-suspending SaaS warehouses).
+#[async_trait]
+pub trait AdbcIntrospection: Send + Sync {
+    async fn health_check(&self) -> bool;
+    async fn fetch_running_query_count(&self) -> Option<u64>;
+}