faucet_source_oracle/

stream.rs

//! The Oracle [`Source`] — the only module that talks to the driver. Rows are
//! fetched on a blocking thread and handed to the async stream one page at a
//! time over a bounded channel, so memory stays at O(batch_size).

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

use async_trait::async_trait;
use faucet_common_oracle::oracle::sql_type::ToSql;
use faucet_common_oracle::{
    OraclePool, Side, blocking, call_timeout, checkout, connect_pool, ora_err, quote_ident_oracle,
    row_to_json, shard_bounds_query, shard_wrap, trim_statement,
};
use faucet_core::check::{CheckContext, CheckReport, Probe};
use faucet_core::replication::cursor::CursorBookmark;
use faucet_core::shard::{PkShardBounds, ShardSpec, parse_pk_shard, pk_shards_from_bounds};
use faucet_core::{FaucetError, Source, StreamPage};
use futures::Stream;
use serde_json::Value;

use crate::config::{OracleReplication, OracleSourceConfig};
use crate::query::{
    CatalogRow, IncrementalCtx, OwnedBind, PlannedQuery, cursor_wrap, default_state_key,
    descriptors_from_catalog, is_datetime_conversion_error, owned_binds, plan_query, resolve_binds,
    text_binds,
};

/// Oracle Database query source.
pub struct OracleSource {
    config: OracleSourceConfig,
    pool: OraclePool,
    start_bookmark: Mutex<Option<CursorBookmark>>,
    applied_shard: Mutex<Option<PkShardBounds>>,
    /// Backend-call context (#880): the run's deadline, cancel token and
    /// round-trip recorder.
    calls: faucet_core::CallSlot,
}

static NULL_TEXT: Option<String> = None;

fn as_tosql(b: &OwnedBind) -> &dyn ToSql {
    match b {
        OwnedBind::Null => &NULL_TEXT,
        OwnedBind::Int(i) => i,
        OwnedBind::Float(f) => f,
        OwnedBind::Text(s) => s,
        OwnedBind::Timestamp(t) => t,
        OwnedBind::TimestampTz(t) => t,
    }
}

/// Everything a blocking fetch needs, owned so it can cross into
/// `spawn_blocking`.
struct QueryJob {
    pool: OraclePool,
    sql: String,
    params: Vec<Value>,
    bookmark: Option<Value>,
    array_size: u32,
    chunk: usize,
    json_columns: Vec<String>,
    timeout: Option<Duration>,
}

/// Run a query, handing each full page to `emit`; stops early when `emit`
/// returns `false` (the consumer went away).
///
/// An ISO-shaped `:bookmark` binds as a timestamp first (SQL-41). When the
/// cursor column is text holding ISO strings, Oracle converts the *column* to a
/// timestamp and rejects the `T` before any row is returned; the query then
/// runs once more with the bookmark bound as text (#862 M-067).
fn run_query(job: QueryJob, mut emit: impl FnMut(Vec<Value>) -> bool) -> Result<(), FaucetError> {
    let conn = checkout(&job.pool, Side::Source, job.timeout)?;
    let mut emitted = false;
    match run_query_on(&conn, &job, BookmarkBind::Typed, &mut emit, &mut emitted) {
        Err(Failure {
            message,
            retry_as_text: true,
        }) if !emitted => {
            tracing::debug!(error = %message, "oracle: retrying with the bookmark bound as text");
            run_query_on(&conn, &job, BookmarkBind::Text, &mut emit, &mut emitted)
                .map_err(|f| f.into_error())
        }
        other => other.map_err(|f| f.into_error()),
    }
}

/// How `:bookmark` is bound.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum BookmarkBind {
    Typed,
    Text,
}

/// A failed attempt, and whether it is the text-cursor conversion failure a
/// text-bound bookmark avoids.
struct Failure {
    message: FaucetError,
    retry_as_text: bool,
}

impl Failure {
    fn into_error(self) -> FaucetError {
        self.message
    }
}

fn run_query_on(
    conn: &faucet_common_oracle::oracle::Connection,
    job: &QueryJob,
    mode: BookmarkBind,
    emit: &mut impl FnMut(Vec<Value>) -> bool,
    emitted: &mut bool,
) -> Result<(), Failure> {
    let plain = |message: FaucetError| Failure {
        message,
        retry_as_text: false,
    };
    let mut stmt = conn
        .statement(&job.sql)
        .fetch_array_size(job.array_size)
        .build()
        .map_err(|e| plain(ora_err(Side::Source, "prepare", &e)))?;
    let names: Vec<String> = stmt.bind_names().iter().map(|s| s.to_string()).collect();
    let binds = resolve_binds(&names, &job.params, job.bookmark.as_ref()).map_err(plain)?;
    let owned = match mode {
        BookmarkBind::Typed => owned_binds(binds),
        BookmarkBind::Text => text_binds(binds),
    };
    let typed_bookmark = owned
        .iter()
        .any(|(_, b)| matches!(b, OwnedBind::Timestamp(_) | OwnedBind::TimestampTz(_)));
    let classify = |message: FaucetError, raw: &str| Failure {
        retry_as_text: typed_bookmark && is_datetime_conversion_error(raw),
        message,
    };
    let named: Vec<(&str, &dyn ToSql)> = owned
        .iter()
        .map(|(n, b)| (n.as_str(), as_tosql(b)))
        .collect();
    let rows = stmt.query_named(&named).map_err(|e| {
        let raw = e.to_string();
        classify(query_error(&raw), &raw)
    })?;
    let mut buf = Vec::with_capacity(job.chunk.min(job.array_size as usize));
    for row in rows {
        let row = row.map_err(|e| {
            let raw = e.to_string();
            classify(ora_err(Side::Source, "fetch", &e), &raw)
        })?;
        buf.push(row_to_json(&row, &job.json_columns, Side::Source).map_err(plain)?);
        if buf.len() >= job.chunk {
            *emitted = true;
            if !emit(std::mem::take(&mut buf)) {
                return Ok(());
            }
        }
    }
    if !buf.is_empty() {
        *emitted = true;
        emit(buf);
    }
    Ok(())
}

/// A query failure, with the fix spelled out when the driver refused a native
/// `JSON` column.
fn query_error(message: &str) -> FaucetError {
    if message.contains("unsupported Oracle type JSON") {
        FaucetError::Source(format!(
            "oracle query: {message}; native JSON columns cannot be fetched directly — select \
             JSON_SERIALIZE(<col> RETURNING CLOB) AS <col> and list the column under `json_columns`"
        ))
    } else {
        FaucetError::Source(format!("oracle query: {message}"))
    }
}

impl OracleSource {
    /// Validate the config and connect (fails fast on bad credentials).
    pub async fn new(config: OracleSourceConfig) -> Result<Self, FaucetError> {
        config.validate()?;
        let pool = connect_pool(&config.connection, config.max_connections).await?;
        Ok(Self {
            config,
            pool,
            start_bookmark: Mutex::new(None),
            applied_shard: Mutex::new(None),
            calls: faucet_core::CallSlot::new(faucet_core::CallContext::new().without_timeout()),
        })
    }

    fn current_start(&self) -> Option<CursorBookmark> {
        self.start_bookmark.lock().expect("bookmark mutex").clone()
    }

    fn job(&self, planned: &PlannedQuery) -> QueryJob {
        let sql = trim_statement(&planned.sql).to_string();
        let sql = match &*self.applied_shard.lock().expect("shard mutex") {
            Some(bounds) => shard_wrap(bounds, &sql),
            None => sql,
        };
        let sql = cursor_wrap(&sql, planned.incremental.as_ref(), true);
        let batch = self.config.batch_size;
        QueryJob {
            pool: self.pool.clone(),
            sql,
            params: planned.params.clone(),
            bookmark: planned.bookmark.clone(),
            array_size: if batch == 0 {
                1000
            } else {
                batch.min(u32::MAX as usize) as u32
            },
            chunk: if batch == 0 { usize::MAX } else { batch },
            json_columns: self.config.json_columns.clone(),
            timeout: call_timeout(self.config.statement_timeout_secs),
        }
    }

    /// Every page of an incremental read, and its last bookmark.
    async fn collect_incremental(
        &self,
        context: &HashMap<String, Value>,
    ) -> Result<(Vec<Value>, Option<Value>), FaucetError> {
        use futures::TryStreamExt as _;
        let mut pages = self.stream_pages(context, 0);
        let mut records = Vec::new();
        let mut bookmark = None;
        while let Some(page) = pages.try_next().await? {
            records.extend(page.records);
            bookmark = page.bookmark.or(bookmark);
        }
        Ok((records, bookmark))
    }

    /// Run a full-replication query in one call.
    async fn collect_all(
        &self,
        context: &HashMap<String, Value>,
    ) -> Result<Vec<Value>, FaucetError> {
        let planned = plan_query(&self.config, context, None);
        let mut job = self.job(&planned);
        job.chunk = usize::MAX;
        let records = self
            .calls
            .call_once(
                "query",
                blocking(move || {
                    let mut all = Vec::new();
                    run_query(job, |page| {
                        all.extend(page);
                        true
                    })?;
                    Ok(all)
                }),
            )
            .await?;
        Ok(records)
    }
}

#[async_trait]
impl Source for OracleSource {
    fn set_call_context(&self, ctx: faucet_core::CallContext) {
        self.calls.install(ctx);
    }

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

    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> {
        if matches!(self.config.replication, OracleReplication::Full) {
            return Ok((self.collect_all(context).await?, None));
        }
        self.collect_incremental(context).await
    }

    fn stream_pages<'a>(
        &'a self,
        context: &'a HashMap<String, Value>,
        _batch_size: usize,
    ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
        let start = self.current_start();
        let planned = plan_query(&self.config, context, start.as_ref());
        let job = self.job(&planned);
        Box::pin(async_stream::try_stream! {
            let (tx, mut rx) = tokio::sync::mpsc::channel::<Vec<Value>>(2);
            let handle = tokio::task::spawn_blocking(move || {
                run_query(job, |page| tx.blocking_send(page).is_ok())
            });
            let calls = self.calls.get();
            calls.record("query");
            // A page of the running query is not a round trip of its own; the
            // run's deadline and cancel token end the wait for it (#880).
            let steps = calls.with_roundtrip_metering(false);
            let mut tracker = planned.incremental.as_ref().map(IncrementalCtx::tracker);
            let mut total = 0usize;
            while let Some(page) = steps.call_once("query", async { Ok(rx.recv().await) }).await? {
                let kept: Vec<Value> = match tracker.as_mut() {
                    None => page,
                    Some(t) => page.into_iter().filter(|r| t.admit(r)).collect(),
                };
                total += kept.len();
                let bookmark = tracker.as_ref().map(|t| t.bookmark().to_value());
                if !kept.is_empty() || bookmark.is_some() {
                    yield StreamPage { records: kept, bookmark };
                }
            }
            handle
                .await
                .map_err(|e| FaucetError::Source(format!("oracle fetch task failed: {e}")))??;
            if let Some(t) = tracker.as_ref() {
                // The final page carries the position even when it is empty.
                yield StreamPage { records: Vec::new(), bookmark: Some(t.bookmark().to_value()) };
                if t.missing() > 0 {
                    tracing::warn!(
                        missing = t.missing(),
                        "oracle incremental read: rows without a cursor value were emitted \
                         without moving the bookmark"
                    );
                }
            }
            tracing::info!(rows = total, query = %self.config.query, "oracle source stream complete");
        })
    }

    fn config_schema(&self) -> Value {
        serde_json::to_value(faucet_core::schema_for!(OracleSourceConfig))
            .expect("schema serialization")
    }

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

    fn dataset_uri(&self) -> String {
        format!(
            "{}?query={}",
            self.config.connection.display_uri(),
            self.config.query
        )
    }

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

    fn legacy_state_keys(&self) -> Vec<String> {
        match (&self.config.replication, &self.config.state_key) {
            (OracleReplication::Incremental { .. }, None) => {
                crate::query::legacy_state_keys(&self.config)
            }
            _ => Vec::new(),
        }
    }

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

    async fn check(&self, ctx: &CheckContext) -> Result<CheckReport, FaucetError> {
        let started = std::time::Instant::now();
        let pool = self.pool.clone();
        let probe = blocking(move || {
            let conn = checkout(&pool, Side::Source, None)?;
            conn.query_row("SELECT 1 FROM DUAL", &[])
                .map_err(|e| ora_err(Side::Source, "probe", &e))?;
            Ok(())
        });
        let hint = "check the connection settings / credentials and that the listener is reachable";
        let probe = match tokio::time::timeout(ctx.timeout, probe).await {
            Ok(Ok(())) => Probe::pass("connect", started.elapsed()),
            Ok(Err(e)) => Probe::fail_hint("connect", started.elapsed(), e.to_string(), hint),
            Err(_) => Probe::fail_hint("connect", started.elapsed(), "timed out", hint),
        };
        Ok(CheckReport::single(probe))
    }

    fn supports_discover(&self) -> bool {
        true
    }

    /// Enumerate tables in non-Oracle-maintained schemas from `ALL_TAB_COLUMNS`,
    /// with row estimates from optimizer statistics (`ALL_TABLES.NUM_ROWS`).
    async fn discover(&self) -> Result<Vec<faucet_core::DatasetDescriptor>, FaucetError> {
        const SQL: &str = "SELECT c.OWNER, c.TABLE_NAME, c.COLUMN_NAME, c.DATA_TYPE, \
                c.DATA_SCALE, c.NULLABLE, t.NUM_ROWS \
           FROM ALL_TAB_COLUMNS c \
           JOIN ALL_TABLES t ON t.OWNER = c.OWNER AND t.TABLE_NAME = c.TABLE_NAME \
          WHERE c.OWNER IN (SELECT USERNAME FROM ALL_USERS WHERE ORACLE_MAINTAINED = 'N') \
            AND t.TEMPORARY = 'N' AND t.NESTED = 'NO' AND t.SECONDARY = 'N' \
          ORDER BY c.OWNER, c.TABLE_NAME, c.COLUMN_ID";
        const KEY_SQL: &str = "SELECT c.OWNER || '.' || c.TABLE_NAME, cc.COLUMN_NAME \
           FROM ALL_CONSTRAINTS c \
           JOIN ALL_CONS_COLUMNS cc \
             ON cc.OWNER = c.OWNER AND cc.CONSTRAINT_NAME = c.CONSTRAINT_NAME \
          WHERE c.CONSTRAINT_TYPE = 'P' \
            AND c.OWNER IN (SELECT USERNAME FROM ALL_USERS WHERE ORACLE_MAINTAINED = 'N') \
          ORDER BY c.OWNER, c.TABLE_NAME, cc.POSITION";
        let pool = self.pool.clone();
        let timeout = call_timeout(self.config.statement_timeout_secs);
        let (rows, keys) = self
            .calls
            .call_once(
                "discover",
                blocking(move || {
                    let conn = checkout(&pool, Side::Source, timeout)?;
                    let mut keys: Vec<(String, String)> = Vec::new();
                    for r in conn
                        .query_as::<(String, String)>(KEY_SQL, &[])
                        .map_err(|e| ora_err(Side::Source, "primary-key discovery", &e))?
                    {
                        keys.push(
                            r.map_err(|e| ora_err(Side::Source, "primary-key discovery", &e))?,
                        );
                    }
                    let rs = conn
                        .query_as::<(
                            String,
                            String,
                            String,
                            String,
                            Option<i64>,
                            String,
                            Option<i64>,
                        )>(SQL, &[])
                        .map_err(|e| ora_err(Side::Source, "catalog discovery", &e))?;
                    let mut out = Vec::new();
                    for r in rs {
                        let (owner, table, column, data_type, scale, nullable, num_rows) =
                            r.map_err(|e| ora_err(Side::Source, "catalog discovery", &e))?;
                        out.push(CatalogRow {
                            owner,
                            table,
                            column,
                            data_type,
                            scale,
                            nullable: nullable != "N",
                            num_rows,
                        });
                    }
                    Ok((out, keys))
                }),
            )
            .await?;
        Ok(faucet_core::attach_primary_keys(
            descriptors_from_catalog(rows)?,
            keys,
        ))
    }

    fn is_shardable(&self) -> bool {
        self.config.shard.is_some()
    }

    async fn enumerate_shards(&self, target: usize) -> Result<Vec<ShardSpec>, FaucetError> {
        let Some(shard_cfg) = self.config.shard.clone() else {
            return Ok(vec![ShardSpec::whole()]);
        };
        let start = self.current_start();
        let planned = plan_query(&self.config, &HashMap::new(), start.as_ref());
        let key = quote_ident_oracle(&shard_cfg.key)?;
        let base = cursor_wrap(
            trim_statement(&planned.sql),
            planned.incremental.as_ref(),
            false,
        );
        let sql = shard_bounds_query(&base, &key);
        let pool = self.pool.clone();
        let timeout = call_timeout(self.config.statement_timeout_secs);
        let (lo, hi) = self.calls.call_once("query", blocking(move || {
            let conn = checkout(&pool, Side::Source, timeout)?;
            let stmt_err = |e: &faucet_common_oracle::oracle::Error| {
                FaucetError::Source(format!(
                    "oracle: shard bounds for key {:?} failed (it must be an integer-valued \
                     output column): {e}",
                    shard_cfg.key
                ))
            };
            let mut stmt = conn.statement(&sql).build().map_err(|e| stmt_err(&e))?;
            let names: Vec<String> = stmt.bind_names().iter().map(|s| s.to_string()).collect();
            let binds = resolve_binds(&names, &planned.params, planned.bookmark.as_ref())?;
            let owned = owned_binds(binds);
            let named: Vec<(&str, &dyn ToSql)> = owned
                .iter()
                .map(|(n, b)| (n.as_str(), as_tosql(b)))
                .collect();
            stmt.query_row_as_named::<(Option<i64>, Option<i64>)>(&named)
                .map_err(|e| stmt_err(&e))
        }))
        .await?;
        Ok(pk_shards_from_bounds(
            &self.config.shard.as_ref().expect("checked").key,
            lo,
            hi,
            target,
        ))
    }

    async fn apply_shard(&self, shard: &ShardSpec) -> Result<(), FaucetError> {
        let bounds = parse_pk_shard(shard, "oracle")?;
        if let Some(b) = &bounds {
            quote_ident_oracle(&b.key)?;
        }
        *self.applied_shard.lock().expect("shard mutex") = bounds;
        Ok(())
    }
}

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

    #[test]
    fn query_error_hints_native_json() {
        let e = query_error("DPI-xxxx: unsupported Oracle type JSON");
        assert!(e.to_string().contains("JSON_SERIALIZE"), "{e}");
        let e = query_error("ORA-00942: table or view does not exist");
        assert!(!e.to_string().contains("JSON_SERIALIZE"), "{e}");
    }

    #[test]
    fn owned_binds_convert_to_driver_values() {
        for b in [
            OwnedBind::Null,
            OwnedBind::Int(1),
            OwnedBind::Float(1.5),
            OwnedBind::Text("x".into()),
            OwnedBind::Timestamp(chrono::DateTime::UNIX_EPOCH.naive_utc()),
            OwnedBind::TimestampTz(chrono::DateTime::UNIX_EPOCH.fixed_offset()),
        ] {
            let _ = as_tosql(&b);
        }
    }
}