faucet_source_databricks/

stream.rs

//! Databricks SQL query source over the Statement Execution API.
//!
//! Submits `sql` to `POST /api/2.0/sql/statements`, polls
//! `GET /api/2.0/sql/statements/{id}` until the statement is terminal, then
//! streams the result chunks (INLINE + JSON_ARRAY) as typed JSON rows,
//! following `result.next_chunk_internal_link` across chunks. No SDK — plain
//! `reqwest` over the shared-auth bearer token.

use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Mutex;
use std::time::Duration;

use async_trait::async_trait;
use faucet_common_databricks::{
    ErrorSide, StatementClient, StatementOptions, resolve_authorization, value_to_param_string,
};
use faucet_core::replication::cursor::{CursorBookmark, CursorTracker};
use faucet_core::{FaucetError, SharedAuthProvider, Source, Stream, StreamPage};
use reqwest::Client;
use serde::Deserialize;
use serde_json::{Value, json};

use crate::config::{
    BOOKMARK_TOKEN, DatabricksReplication, DatabricksSourceConfig, ResultDisposition, shared_auth,
};
use crate::convert::{ColumnInfo, row_to_json};

/// Databricks SQL query source.
pub struct DatabricksSource {
    config: DatabricksSourceConfig,
    client: Client,
    /// Base URL override (up to but not including `/api/2.0/...`). `None` uses
    /// `config.workspace_url`. Set by tests to point at a mock server.
    endpoint_base: Option<String>,
    /// Shared auth provider; when set, supplies the bearer token (takes
    /// precedence over inline auth).
    auth_provider: Option<SharedAuthProvider>,
    /// Bookmark applied via [`Source::apply_start_bookmark`].
    start_bookmark: Mutex<Option<CursorBookmark>>,
    /// Backend-call context (#880) every request runs through.
    calls: faucet_core::CallSlot,
}

#[derive(Debug, Deserialize)]
struct ResultChunk {
    #[serde(default)]
    data_array: Option<Vec<Vec<Value>>>,
    #[serde(default)]
    next_chunk_internal_link: Option<String>,
    /// Present under `EXTERNAL_LINKS` disposition: each entry carries a
    /// presigned URL to a chunk (a JSON array of rows, or an Arrow IPC stream
    /// under `arrow_native`) plus its own next-chunk link.
    #[serde(default)]
    external_links: Option<Vec<ExternalLink>>,
}

/// One `result.external_links[]` entry (EXTERNAL_LINKS disposition).
#[derive(Debug, Deserialize)]
struct ExternalLink {
    #[serde(default)]
    external_link: Option<String>,
    #[serde(default)]
    next_chunk_internal_link: Option<String>,
}

/// The internal link to the next result chunk, checked at both the result
/// level and (for EXTERNAL_LINKS) on the last external link of the response.
fn next_chunk_link(chunk: &ResultChunk) -> Option<String> {
    chunk.next_chunk_internal_link.clone().or_else(|| {
        chunk
            .external_links
            .as_ref()
            .and_then(|l| l.last())
            .and_then(|e| e.next_chunk_internal_link.clone())
    })
}

/// A terminal statement: its result columns and first chunk.
struct Executed {
    columns: Vec<ColumnInfo>,
    result: Option<ResultChunk>,
}

/// Incremental-read context (column + starting position).
struct IncrementalCtx {
    column: String,
    start: CursorBookmark,
}

/// A Databricks SQL identifier in backticks (a backtick inside is doubled).
fn quote_ident(name: &str) -> String {
    format!("`{}`", name.replace('`', "``"))
}

impl DatabricksSource {
    /// Create a new source. Validates config; does no I/O.
    pub fn new(config: DatabricksSourceConfig) -> Result<Self, FaucetError> {
        config.validate()?;
        Ok(Self {
            config,
            client: faucet_common_databricks::http_client()?,
            endpoint_base: None,
            auth_provider: None,
            start_bookmark: Mutex::new(None),
            calls: faucet_core::CallSlot::new(faucet_core::CallContext::new().without_timeout()),
        })
    }

    /// Attach a shared auth provider (yields the bearer token). Takes
    /// precedence over inline auth — lets several connectors share one token.
    pub fn with_auth_provider(mut self, provider: SharedAuthProvider) -> Self {
        self.auth_provider = Some(provider);
        self
    }

    /// Override the base URL (e.g. a wiremock server). Requests go to
    /// `{base}/api/2.0/sql/statements`.
    pub fn with_endpoint_base(mut self, base: impl Into<String>) -> Self {
        self.endpoint_base = Some(base.into());
        self
    }

    fn base_url(&self) -> String {
        match &self.endpoint_base {
            Some(b) => b.trim_end_matches('/').to_owned(),
            None => self.config.workspace_url.trim_end_matches('/').to_owned(),
        }
    }

    fn statements_url(&self) -> String {
        format!("{}/api/2.0/sql/statements", self.base_url())
    }

    /// Resolve the `Authorization` header value: shared provider first, else
    /// inline auth.
    async fn auth_header(&self) -> Result<String, FaucetError> {
        resolve_authorization(&shared_auth(&self.config.auth), self.auth_provider.as_ref()).await
    }

    /// A Statement Execution API client for this source's warehouse.
    fn statement_client(&self) -> StatementClient {
        StatementClient::new(
            self.client.clone(),
            self.base_url(),
            self.config.warehouse_id.clone(),
            shared_auth(&self.config.auth),
            self.auth_provider.clone(),
            StatementOptions {
                wait_timeout_secs: self.config.wait_timeout_secs,
                poll_interval: Duration::from_secs(self.config.poll_interval_secs.max(1)),
                statement_timeout: (self.config.statement_timeout_secs > 0)
                    .then(|| Duration::from_secs(self.config.statement_timeout_secs)),
                ..StatementOptions::default()
            },
            ErrorSide::Source,
        )
        .with_call_context(self.calls.get())
    }

    /// The effective incremental start bookmark (persisted bookmark, else the
    /// configured `initial_value`).
    fn incremental_ctx(&self) -> Option<IncrementalCtx> {
        match &self.config.replication {
            DatabricksReplication::Full => None,
            DatabricksReplication::Incremental {
                column,
                initial_value,
            } => {
                let start = self
                    .start_bookmark
                    .lock()
                    .expect("start_bookmark mutex poisoned")
                    .clone()
                    .unwrap_or_else(|| CursorBookmark::initial(initial_value.clone()));
                Some(IncrementalCtx {
                    column: column.clone(),
                    start,
                })
            }
        }
    }

    /// Build the request body: the SQL (with `${bookmark}` and `{ctx}` tokens
    /// rewritten to named params, and for an incremental read wrapped as
    /// `SELECT * FROM (<sql>) AS q WHERE q.<column> >= :_faucet_bookmark
    /// ORDER BY q.<column>`) plus the `parameters` array.
    fn build_body(&self, context: &HashMap<String, Value>, incr: Option<&IncrementalCtx>) -> Value {
        let mut sql = self.config.sql.clone();
        // Named parameters: start with the user's static ones.
        let mut params: Vec<Value> = self
            .config
            .parameters
            .iter()
            .map(|p| {
                json!({
                    "name": p.name,
                    "value": value_to_param_string(&p.value),
                    "type": p.param_type.clone().unwrap_or_else(|| "STRING".into()),
                })
            })
            .collect();

        // Parent-context `{key}` tokens → `:_faucet_ctN` named params.
        if !context.is_empty() {
            let (rewritten, ctx_values) =
                faucet_core::util::substitute_context_bind_params(&sql, context, 0, |i| {
                    format!(":_faucet_ct{i}")
                });
            sql = rewritten;
            for (i, v) in ctx_values.into_iter().enumerate() {
                params.push(json!({
                    "name": format!("_faucet_ct{i}"),
                    "value": value_to_param_string(&v),
                }));
            }
        }

        // Incremental: `${bookmark}` → the `:_faucet_bookmark` named param,
        // and the outer `>=` cursor filter + order on the same param.
        if let Some(ctx) = incr {
            sql = sql.replace(BOOKMARK_TOKEN, ":_faucet_bookmark");
            let c = quote_ident(&ctx.column);
            let inner = sql.trim_end_matches(|ch: char| ch.is_whitespace() || ch == ';');
            sql = format!(
                "SELECT * FROM ({inner}\n) AS q WHERE q.{c} >= :_faucet_bookmark ORDER BY q.{c}"
            );
            params.push(json!({
                "name": "_faucet_bookmark",
                "value": value_to_param_string(&ctx.start.value),
            }));
        }

        // ARROW_STREAM is only valid with EXTERNAL_LINKS disposition; JSON_ARRAY
        // stays INLINE. `arrow_native` is validated to require the `arrow`
        // feature at config load, so requesting Arrow here is always decodable.
        // INLINE results are capped at 25 MiB server-side, so the row path
        // reads presigned JSON chunks by default (SQL-45).
        let (disposition, format) = if self.config.arrow_native {
            ("EXTERNAL_LINKS", "ARROW_STREAM")
        } else {
            match self.config.result_disposition {
                ResultDisposition::ExternalLinks => ("EXTERNAL_LINKS", "JSON_ARRAY"),
                ResultDisposition::Inline => ("INLINE", "JSON_ARRAY"),
            }
        };
        let mut body = json!({
            "statement": sql,
            "warehouse_id": self.config.warehouse_id,
            "wait_timeout": format!("{}s", self.config.wait_timeout_secs),
            "on_wait_timeout": "CONTINUE",
            "disposition": disposition,
            "format": format,
        });
        if let Some(c) = &self.config.catalog {
            body["catalog"] = json!(c);
        }
        if let Some(s) = &self.config.schema {
            body["schema"] = json!(s);
        }
        if !params.is_empty() {
            body["parameters"] = Value::Array(params);
        }
        body
    }

    /// Submit the statement and poll until it reaches a terminal state.
    async fn run_statement(
        &self,
        context: &HashMap<String, Value>,
        incr: Option<&IncrementalCtx>,
    ) -> Result<Executed, FaucetError> {
        let body = self.build_body(context, incr);
        let resp = self.statement_client().execute_body(&body).await?;
        let columns = resp.columns().to_vec();
        crate::convert::check_unique_columns(&columns)?;
        let result = resp.result.map(decode_chunk).transpose()?;
        Ok(Executed { columns, result })
    }

    /// Fetch a follow-up chunk by its `next_chunk_internal_link` (a full API path).
    async fn fetch_chunk(&self, link: &str) -> Result<ResultChunk, FaucetError> {
        decode_chunk(self.statement_client().fetch_chunk(link).await?)
    }

    /// Fetch a presigned external link and decode its body as an Arrow IPC
    /// stream. The link is a pre-signed cloud-storage URL, so it is fetched
    /// **without** an `Authorization` header (adding one breaks the signature).
    #[cfg(feature = "arrow")]
    async fn fetch_arrow_link(
        &self,
        url: &str,
    ) -> Result<Vec<arrow::array::RecordBatch>, FaucetError> {
        let data = bytes::Bytes::from(self.get_external(url).await?);
        tokio::task::spawn_blocking(move || decode_arrow_ipc(data))
            .await
            .map_err(|e| {
                FaucetError::Source(format!("databricks: arrow decode task panicked: {e}"))
            })?
    }
}

/// A presigned link must be an `http(s)` URL whose host is not a link-local
/// or cloud-metadata address.
fn check_external_link(url: &str) -> Result<(), FaucetError> {
    let refused = |why: &str| {
        FaucetError::Source(format!(
            "databricks: external link `{}` refused: {why}",
            faucet_core::util::redact_uri_credentials(url)
        ))
    };
    let parsed =
        reqwest::Url::parse(url).map_err(|e| refused(&format!("it does not parse: {e}")))?;
    if !matches!(parsed.scheme(), "http" | "https") {
        return Err(refused("it is not an http(s) URL"));
    }
    match parsed.host_str() {
        None => Err(refused("it names no host")),
        Some(h) if faucet_core::resolve::is_metadata_host(h) => {
            Err(refused("the host is a link-local / cloud-metadata address"))
        }
        Some(_) => Ok(()),
    }
}

impl DatabricksSource {
    /// GET a presigned external link. It is signed for the storage service, so
    /// it is fetched **without** an `Authorization` header (adding one breaks
    /// the signature).
    async fn get_external(&self, url: &str) -> Result<Vec<u8>, FaucetError> {
        check_external_link(url)?;
        self.calls
            .call("chunk", || async {
                let resp = self.client.get(url).send().await.map_err(|e| {
                    FaucetError::Source(format!("databricks: external-link request failed: {e}"))
                })?;
                let status = resp.status();
                if !status.is_success() {
                    let body = faucet_core::http::read_body_limited(
                        resp,
                        Some(faucet_core::util::DEFAULT_ERROR_BODY_MAX_LEN as u64),
                    )
                    .await
                    .map(|b| String::from_utf8_lossy(&b).into_owned())
                    .unwrap_or_default();
                    return Err(FaucetError::Source(format!(
                        "databricks: external link HTTP {status}: {body}"
                    )));
                }
                faucet_core::http::read_body_limited(
                    resp,
                    Some(faucet_core::http::DEFAULT_MAX_RESPONSE_BYTES),
                )
                .await
                .map_err(|e| {
                    FaucetError::Source(format!(
                        "databricks: reading external-link body failed: {e}"
                    ))
                })
            })
            .await
    }

    /// Fetch a presigned `JSON_ARRAY` chunk: a JSON array of row arrays.
    async fn fetch_json_link(&self, url: &str) -> Result<Vec<Vec<Value>>, FaucetError> {
        let body = self.get_external(url).await?;
        serde_json::from_slice(&body).map_err(|e| {
            FaucetError::Source(format!(
                "databricks: could not parse an external-link chunk: {e}"
            ))
        })
    }

    /// Every row of one result chunk: its inline `data_array` plus the rows of
    /// each presigned external link.
    async fn chunk_rows(&self, chunk: &mut ResultChunk) -> Result<Vec<Vec<Value>>, FaucetError> {
        let mut rows = chunk.data_array.take().unwrap_or_default();
        for link in chunk.external_links.iter().flatten() {
            if let Some(url) = link.external_link.as_deref() {
                rows.extend(self.fetch_json_link(url).await?);
            }
        }
        Ok(rows)
    }
}

/// Decode an Arrow IPC **stream** (the ARROW_STREAM chunk body) into its
/// `RecordBatch`es. Synchronous (runs inside `spawn_blocking`).
#[cfg(feature = "arrow")]
fn decode_arrow_ipc(data: bytes::Bytes) -> Result<Vec<arrow::array::RecordBatch>, FaucetError> {
    use arrow::ipc::reader::StreamReader;

    let reader = StreamReader::try_new(std::io::Cursor::new(data), None).map_err(|e| {
        FaucetError::Source(format!("databricks: arrow IPC reader init failed: {e}"))
    })?;
    let mut batches = Vec::new();
    for batch in reader {
        batches.push(batch.map_err(|e| {
            FaucetError::Source(format!("databricks: arrow IPC decode failed: {e}"))
        })?);
    }
    Ok(batches)
}

/// Derive a stable state key from the workspace, warehouse, and query.
fn default_state_key(config: &DatabricksSourceConfig) -> String {
    // Fixed FNV-1a: the key is durable and `DefaultHasher` is not stable across Rust releases.
    let material = format!(
        "{}\u{1f}{}\u{1f}{}",
        config.workspace_url, config.warehouse_id, config.sql
    );
    format!(
        "databricks:{:016x}",
        faucet_core::shard::shard_hash(&material)
    )
}

/// Decode a raw result chunk.
fn decode_chunk(v: Value) -> Result<ResultChunk, FaucetError> {
    serde_json::from_value(v)
        .map_err(|e| FaucetError::Source(format!("databricks: could not parse response: {e}")))
}

#[async_trait]
impl Source for DatabricksSource {
    fn config_schema(&self) -> Value {
        serde_json::to_value(faucet_core::schema_for!(DatabricksSourceConfig))
            .expect("schema serialization")
    }

    fn connector_name(&self) -> &'static str {
        "databricks"
    }

    fn set_call_context(&self, ctx: faucet_core::io::CallContext) {
        self.calls.install(ctx);
    }

    fn set_roundtrip_recorder(
        &self,
        recorder: std::sync::Arc<faucet_core::observability::RoundtripRecorder>,
    ) {
        self.calls.install_recorder(recorder);
    }

    fn dataset_uri(&self) -> String {
        format!(
            "databricks://{}/warehouses/{}",
            self.config
                .workspace_url
                .trim_start_matches("https://")
                .trim_end_matches('/'),
            self.config.warehouse_id
        )
    }

    fn state_key(&self) -> Option<String> {
        match &self.config.replication {
            DatabricksReplication::Full => None,
            DatabricksReplication::Incremental { .. } => Some(
                self.config
                    .state_key
                    .clone()
                    .unwrap_or_else(|| default_state_key(&self.config)),
            ),
        }
    }

    async fn apply_start_bookmark(&self, bookmark: Value) -> Result<(), FaucetError> {
        let start = CursorBookmark::from_state(bookmark)?;
        *self
            .start_bookmark
            .lock()
            .expect("start_bookmark mutex poisoned") = Some(start);
        Ok(())
    }

    /// The Databricks source advertises the columnar fast path only when
    /// [`arrow_native`](DatabricksSourceConfig::arrow_native) is set — i.e. the
    /// statement is fetched as `ARROW_STREAM` (RFC 0002 / #375).
    #[cfg(feature = "arrow")]
    fn supports_columnar(&self) -> bool {
        self.config.arrow_native
    }

    /// Stream the statement's `ARROW_STREAM` result chunks as Arrow
    /// `RecordBatch`es — one [`ColumnarPage`](faucet_core::columnar::ColumnarPage)
    /// per batch — so a `databricks → parquet`/`delta`/`sql` chain never
    /// materializes `serde_json::Value`. `arrow_native` is Full-replication
    /// only, so every page carries `bookmark: None`. Empty batches are skipped.
    #[cfg(feature = "arrow")]
    fn stream_batches<'a>(
        &'a self,
        context: &'a HashMap<String, Value>,
        _batch_size: usize,
    ) -> Pin<
        Box<
            dyn Stream<Item = Result<faucet_core::columnar::ColumnarPage, FaucetError>> + Send + 'a,
        >,
    > {
        Box::pin(async_stream::try_stream! {
            if !self.config.arrow_native {
                Err(FaucetError::Source(
                    "databricks: stream_batches requires `arrow_native: true`".into(),
                ))?;
            }
            let resp = self.run_statement(context, None).await?;
            let mut chunk = resp.result;
            let mut total_records = 0usize;
            let mut total_pages = 0usize;
            while let Some(c) = chunk {
                if let Some(links) = c.external_links.as_ref() {
                    for link in links {
                        if let Some(url) = link.external_link.as_deref() {
                            let batches = self.fetch_arrow_link(url).await?;
                            for batch in batches {
                                if batch.num_rows() == 0 {
                                    continue;
                                }
                                total_records += batch.num_rows();
                                total_pages += 1;
                                yield faucet_core::columnar::ColumnarPage { batch, bookmark: None };
                            }
                        }
                    }
                }
                chunk = match next_chunk_link(&c) {
                    Some(link) => Some(self.fetch_chunk(&link).await?),
                    None => None,
                };
            }
            tracing::info!(
                pages = total_pages,
                total_records,
                "databricks columnar stream complete",
            );
        })
    }

    fn stream_pages<'a>(
        &'a self,
        context: &'a HashMap<String, Value>,
        _batch_size: usize,
    ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
        Box::pin(async_stream::try_stream! {
            // Arrow-native row path: fetch ARROW_STREAM chunks and decode each
            // RecordBatch to JSON rows (for a non-columnar sink). `arrow_native`
            // is Full-replication only (enforced by config validation), so no
            // incremental filter runs here.
            #[cfg(feature = "arrow")]
            if self.config.arrow_native {
                let resp = self.run_statement(context, None).await?;
                let cap = if self.config.batch_size == 0 { 1024 } else { self.config.batch_size };
                let mut buffer: Vec<Value> = Vec::with_capacity(cap);
                let mut chunk = resp.result;
                while let Some(c) = chunk {
                    if let Some(links) = c.external_links.as_ref() {
                        for link in links {
                            if let Some(url) = link.external_link.as_deref() {
                                let batches = self.fetch_arrow_link(url).await?;
                                for batch in &batches {
                                    for row in faucet_core::columnar::record_batch_to_values(batch)? {
                                        buffer.push(row);
                                        if self.config.batch_size != 0
                                            && buffer.len() >= self.config.batch_size
                                        {
                                            let page = std::mem::replace(
                                                &mut buffer,
                                                Vec::with_capacity(cap),
                                            );
                                            yield StreamPage { records: page, bookmark: None };
                                        }
                                    }
                                }
                            }
                        }
                    }
                    chunk = match next_chunk_link(&c) {
                        Some(link) => Some(self.fetch_chunk(&link).await?),
                        None => None,
                    };
                }
                if !buffer.is_empty() {
                    yield StreamPage { records: buffer, bookmark: None };
                }
                return;
            }

            let incr = self.incremental_ctx();
            let resp = self.run_statement(context, incr.as_ref()).await?;

            let columns: Vec<ColumnInfo> = resp.columns;

            let batch = self.config.batch_size;
            let cap = if batch == 0 { 1024 } else { batch };
            let mut buffer: Vec<Value> = Vec::with_capacity(cap);
            let mut tracker = incr
                .as_ref()
                .map(|ic| CursorTracker::new(ic.column.clone(), ic.start.clone()));

            // Walk chunks: the initial `result`, then follow next_chunk_internal_link.
            let mut chunk = resp.result;
            while let Some(mut c) = chunk {
                let data = self.chunk_rows(&mut c).await?;
                {
                    for row in &data {
                        let obj = row_to_json(row, &columns);
                        if tracker.as_mut().is_none_or(|t| t.admit(&obj)) {
                            buffer.push(obj);
                        }
                        if batch != 0 && buffer.len() >= batch {
                            let page = std::mem::replace(&mut buffer, Vec::with_capacity(cap));
                            let bookmark = tracker.as_ref().map(|t| t.bookmark().to_value());
                            yield StreamPage { records: page, bookmark };
                        }
                    }
                }
                chunk = match next_chunk_link(&c) {
                    Some(link) => Some(self.fetch_chunk(&link).await?),
                    None => None,
                };
            }

            // In incremental mode every page carries the position after its
            // last row (the final page carries one even when empty).
            if !buffer.is_empty() || tracker.is_some() {
                let bookmark = tracker.as_ref().map(|t| t.bookmark().to_value());
                yield StreamPage { records: buffer, bookmark };
            }
            if let Some(missing) = tracker.as_ref().map(CursorTracker::missing).filter(|n| *n > 0) {
                tracing::warn!(
                    missing,
                    "databricks incremental read: rows without a cursor value were emitted \
                     without moving the bookmark"
                );
            }
        })
    }

    async fn fetch_with_context(
        &self,
        context: &HashMap<String, Value>,
    ) -> Result<Vec<Value>, FaucetError> {
        Ok(self.fetch_with_context_incremental(context).await?.0)
    }

    async fn fetch_with_context_incremental(
        &self,
        context: &HashMap<String, Value>,
    ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
        use futures::StreamExt;
        let mut records = Vec::new();
        let mut bookmark = None;
        let mut s = self.stream_pages(context, self.config.batch_size);
        while let Some(page) = s.next().await {
            let page = page?;
            records.extend(page.records);
            bookmark = page.bookmark.or(bookmark);
        }
        Ok((records, bookmark))
    }

    async fn check(
        &self,
        ctx: &faucet_core::check::CheckContext,
    ) -> Result<faucet_core::check::CheckReport, FaucetError> {
        use faucet_core::check::{CheckReport, Probe};
        let started = std::time::Instant::now();
        // Non-scanning probe: run `SELECT 1` on the warehouse.
        let auth = match self.auth_header().await {
            Ok(a) => a,
            Err(e) => {
                return Ok(CheckReport::single(Probe::fail(
                    "auth",
                    started.elapsed(),
                    e.to_string(),
                )));
            }
        };
        let body = json!({
            "statement": "SELECT 1",
            "warehouse_id": self.config.warehouse_id,
            "wait_timeout": "50s",
            "disposition": "INLINE",
            "format": "JSON_ARRAY",
        });
        let calls = self.calls.get().with_timeout(ctx.timeout);
        let fut = calls.call("submit", || async {
            self.client
                .post(self.statements_url())
                .header("Authorization", &auth)
                .header("Content-Type", "application/json")
                .json(&body)
                .send()
                .await
                .map_err(FaucetError::from)
        });
        let probe = match tokio::time::timeout(ctx.timeout, fut).await {
            Ok(Ok(r)) if r.status().is_success() => Probe::pass("warehouse", started.elapsed()),
            Ok(Ok(r)) => Probe::fail_hint(
                "warehouse",
                started.elapsed(),
                format!("databricks probe returned HTTP {}", r.status()),
                "Verify workspace_url, warehouse_id, and token permissions (CAN USE).",
            ),
            Ok(Err(e)) => Probe::fail_hint(
                "warehouse",
                started.elapsed(),
                format!("databricks probe request failed: {e}"),
                "Verify workspace_url and network reachability.",
            ),
            Err(_) => Probe::fail_hint(
                "warehouse",
                started.elapsed(),
                format!("databricks probe timed out after {:?}", ctx.timeout),
                "Check warehouse availability and network reachability.",
            ),
        };
        Ok(CheckReport::single(probe))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::{DatabricksAuth, DatabricksParam};
    use faucet_core::AuthSpec;

    fn cfg() -> DatabricksSourceConfig {
        DatabricksSourceConfig {
            workspace_url: "https://x.cloud.databricks.com".into(),
            warehouse_id: "wh1".into(),
            sql: "SELECT * FROM t WHERE ts >= ${bookmark}".into(),
            auth: AuthSpec::Inline(DatabricksAuth::Pat {
                token: "tok".into(),
            }),
            catalog: Some("main".into()),
            schema: Some("s".into()),
            parameters: vec![DatabricksParam {
                name: "min".into(),
                value: json!(10),
                param_type: Some("INT".into()),
            }],
            wait_timeout_secs: 50,
            poll_interval_secs: 1,
            statement_timeout_secs: 3600,
            batch_size: 1000,
            arrow_native: false,
            result_disposition: ResultDisposition::default(),
            replication: DatabricksReplication::Incremental {
                column: "ts".into(),
                initial_value: json!("2026-01-01"),
            },
            state_key: None,
        }
    }

    fn source(c: DatabricksSourceConfig) -> DatabricksSource {
        DatabricksSource::new(c).unwrap()
    }

    #[test]
    fn body_has_required_fields_and_params() {
        let s = source(cfg());
        let incr = s.incremental_ctx();
        let body = s.build_body(&HashMap::new(), incr.as_ref());
        assert_eq!(body["warehouse_id"], json!("wh1"));
        assert_eq!(body["catalog"], json!("main"));
        // SQL-45: the row path no longer requests INLINE (capped at 25 MiB).
        assert_eq!(body["disposition"], json!("EXTERNAL_LINKS"));
        assert_eq!(body["format"], json!("JSON_ARRAY"));
        assert_eq!(body["wait_timeout"], json!("50s"));
        let mut inline = cfg();
        inline.result_disposition = ResultDisposition::Inline;
        let inline_body = source(inline).build_body(&HashMap::new(), incr.as_ref());
        assert_eq!(inline_body["disposition"], json!("INLINE"));
        assert_eq!(inline_body["format"], json!("JSON_ARRAY"));
        // ${bookmark} rewritten to the named param marker.
        assert!(
            body["statement"]
                .as_str()
                .unwrap()
                .contains(":_faucet_bookmark")
        );
        assert!(!body["statement"].as_str().unwrap().contains("${bookmark}"));
        assert_eq!(
            body["statement"],
            json!(
                "SELECT * FROM (SELECT * FROM t WHERE ts >= :_faucet_bookmark\n) AS q \
                 WHERE q.`ts` >= :_faucet_bookmark ORDER BY q.`ts`"
            )
        );
        assert_eq!(quote_ident("a`b"), "`a``b`");
        let params = body["parameters"].as_array().unwrap();
        // static `min` (typed) + the bookmark param.
        assert!(
            params
                .iter()
                .any(|p| p["name"] == json!("min") && p["type"] == json!("INT"))
        );
        let bm = params
            .iter()
            .find(|p| p["name"] == json!("_faucet_bookmark"))
            .unwrap();
        assert_eq!(bm["value"], json!("2026-01-01"));
    }

    #[test]
    fn full_mode_has_no_state_key_or_bookmark_param() {
        let mut c = cfg();
        c.replication = DatabricksReplication::Full;
        c.sql = "SELECT 1".into();
        let s = source(c);
        assert!(s.state_key().is_none());
        let body = s.build_body(&HashMap::new(), None);
        assert!(
            body.get("parameters").is_none()
                || body["parameters"]
                    .as_array()
                    .unwrap()
                    .iter()
                    .all(|p| p["name"] != json!("_faucet_bookmark"))
        );
    }

    #[test]
    fn incremental_state_key_derived_and_stable() {
        let s = source(cfg());
        let k1 = s.state_key().unwrap();
        let k2 = source(cfg()).state_key().unwrap();
        assert_eq!(k1, k2);
        assert!(k1.starts_with("databricks:"));
        let c = cfg();
        let material = format!("{}\u{1f}{}\u{1f}{}", c.workspace_url, c.warehouse_id, c.sql);
        assert_eq!(
            k1,
            format!(
                "databricks:{:016x}",
                faucet_core::shard::shard_hash(&material)
            )
        );
    }

    #[tokio::test]
    async fn explicit_state_key_wins() {
        let mut c = cfg();
        c.state_key = Some("my-key".into());
        let s = source(c);
        assert_eq!(s.state_key().as_deref(), Some("my-key"));
        // apply_start_bookmark overrides initial_value in the incr ctx.
        s.apply_start_bookmark(json!("2026-06-01")).await.unwrap();
        assert_eq!(
            s.incremental_ctx().unwrap().start.value,
            json!("2026-06-01")
        );
        assert!(
            s.apply_start_bookmark(json!({"value": null}))
                .await
                .is_err()
        );
    }

    #[test]
    fn value_to_param_string_stringifies() {
        assert_eq!(value_to_param_string(&json!(5)), json!("5"));
        assert_eq!(value_to_param_string(&json!(true)), json!("true"));
        assert_eq!(value_to_param_string(&json!("x")), json!("x"));
        assert_eq!(value_to_param_string(&Value::Null), Value::Null);
    }

    #[test]
    fn decode_chunk_rejects_non_objects() {
        assert!(decode_chunk(json!({"data_array": [["1"]]})).is_ok());
        let err = decode_chunk(json!(5)).unwrap_err().to_string();
        assert!(err.contains("could not parse"), "{err}");
    }

    #[test]
    fn dataset_uri_redacts_scheme() {
        let s = source(cfg());
        assert_eq!(
            s.dataset_uri(),
            "databricks://x.cloud.databricks.com/warehouses/wh1"
        );
    }
}