orion-server 1.0.0

use std::sync::Arc;

use async_trait::async_trait;
use dataflow_rs::engine::error::DataflowError;
use dataflow_rs::engine::functions::AsyncFunctionHandler;
use dataflow_rs::engine::task_context::TaskContext;
use dataflow_rs::engine::task_outcome::TaskOutcome;
use serde_json::Value;
use sqlx::any::{AnyRow, AnyTypeInfoKind};
use sqlx::{Column, Row, ValueRef};

use super::connector_helpers::{
    ConnectorCall, apply_output, bind_json_params, reject_mongo_connector, require_db_connector,
    resolve_bind_params, timed_query, to_connect_error,
};
use super::schema::{FieldKind, FieldSchema};
use crate::connector::ConnectorRegistry;
use crate::connector::pool_cache::SqlPoolCache;

/// This handler's name in metrics, profiles and error messages (F48).
const NAME: &str = "db_read";

/// Executes SQL SELECT queries against external databases configured via connectors.
pub struct DbReadHandler {
    pub pool_cache: Arc<SqlPoolCache>,
    pub registry: Arc<ConnectorRegistry>,
    /// Hard row cap, from `query.max_limit` (F10). Raw SQL can't have a
    /// LIMIT injected reliably, so rows are streamed and counted — one
    /// `SELECT * FROM big_table` must not OOM the process.
    pub max_rows: usize,
}

#[async_trait]
impl AsyncFunctionHandler for DbReadHandler {
    type Input = Value;

    async fn execute(
        &self,
        ctx: &mut TaskContext<'_>,
        input: &Value,
    ) -> dataflow_rs::Result<TaskOutcome> {
        // F48/F58: the literal prologue first — `connector` and `query` are
        // both literal keys, so a task missing either must be told before
        // anything about the message is consulted.
        let call = ConnectorCall::begin(NAME, input, ctx)?;
        let query = call.require_str(input, "query")?;

        // Bind values are resolved against the message context; the SQL text
        // itself is never message-derived, so parameters stay the only
        // request-controlled part of the statement.
        let params = resolve_bind_params(input, call.name, ctx)?;

        call.run(&self.registry, async {
            let connector_config = call.resolve(&self.registry, Some("read")).await?;
            let db_config = require_db_connector(&connector_config, call.connector)?;
            reject_mongo_connector(call.name, call.connector, db_config)?;

            let pool = self
                .pool_cache
                .get_pool(call.connector, db_config)
                .await
                .map_err(to_connect_error)?;

            let sqlx_query = bind_json_params(sqlx::query(query), &params);

            let max_rows = self.max_rows;
            let rows: Vec<AnyRow> = timed_query(db_config.query_timeout_ms, call.name, async {
                use futures::TryStreamExt;
                let mut stream = sqlx_query.fetch(&pool);
                let mut rows: Vec<AnyRow> = Vec::new();
                while let Some(row) = stream.try_next().await.map_err(|e| e.to_string())? {
                    if rows.len() >= max_rows {
                        // F42: marked so `timed_query` reports it as a 400
                        // with the text intact rather than a 500 with the
                        // guidance sanitised away.
                        return Err(format!(
                            "{}{NAME} result exceeds query.max_limit ({max_rows} rows) — \
                             add a LIMIT to the query or raise the cap",
                            crate::engine::functions::connector_helpers::LIMIT_MARKER
                        ));
                    }
                    rows.push(row);
                }
                Ok(rows)
            })
            .await?;

            apply_output(ctx, call.output, Value::Array(rows_to_json(&rows)?));
            Ok(TaskOutcome::Success)
        })
        .await
    }
}

/// Convert AnyRow results to a JSON array of objects.
///
/// Column names are collected once from the first row and reused for all
/// subsequent rows, eliminating O(rows × columns) string allocations.
///
/// Only a genuine SQL `NULL` becomes `Value::Null`. A value the driver cannot
/// represent is an error, never a silent null — the two must stay
/// distinguishable to the workflow reading the result.
pub fn rows_to_json(rows: &[AnyRow]) -> Result<Vec<Value>, DataflowError> {
    if rows.is_empty() {
        return Ok(Vec::new());
    }

    let col_names: Vec<String> = rows[0]
        .columns()
        .iter()
        .map(|col| col.name().to_string())
        .collect();

    let mut result = Vec::with_capacity(rows.len());
    for row in rows {
        let mut obj = serde_json::Map::with_capacity(col_names.len());
        for (i, name) in col_names.iter().enumerate() {
            obj.insert(name.clone(), column_to_json(row, i, name)?);
        }
        result.push(Value::Object(obj));
    }
    Ok(result)
}

/// Decode one column of one row into JSON, dispatching on the value's own type
/// rather than probing candidate Rust types in turn.
///
/// The probe-cascade this replaced fell through to `Value::Null` for anything
/// it did not recognise, so `REAL` and `BLOB` columns — and any future
/// [`AnyTypeInfoKind`] — read back as null even though the query succeeded.
/// Matching the kind exhaustively means a new kind is a compile error instead.
fn column_to_json(row: &AnyRow, index: usize, name: &str) -> Result<Value, DataflowError> {
    let raw = row.try_get_raw(index).map_err(|e| {
        DataflowError::function_execution(
            format!("{NAME}: column '{name}' is unreadable: {e}"),
            None,
        )
    })?;
    if raw.is_null() {
        return Ok(Value::Null);
    }
    let kind = raw.type_info().kind();

    let decode_err = |e: sqlx::Error| {
        DataflowError::function_execution(
            format!("{NAME}: column '{name}' ({kind:?}) failed to decode: {e}"),
            None,
        )
    };

    let value = match kind {
        // Already handled above; a value whose own type info is NULL carries
        // nothing to decode.
        AnyTypeInfoKind::Null => Value::Null,
        AnyTypeInfoKind::Bool => Value::Bool(row.try_get::<bool, _>(index).map_err(decode_err)?),
        AnyTypeInfoKind::SmallInt | AnyTypeInfoKind::Integer | AnyTypeInfoKind::BigInt => {
            Value::Number(row.try_get::<i64, _>(index).map_err(decode_err)?.into())
        }
        AnyTypeInfoKind::Real => float_to_json(
            f64::from(row.try_get::<f32, _>(index).map_err(decode_err)?),
            name,
        )?,
        AnyTypeInfoKind::Double => {
            float_to_json(row.try_get::<f64, _>(index).map_err(decode_err)?, name)?
        }
        AnyTypeInfoKind::Text => {
            Value::String(row.try_get::<String, _>(index).map_err(decode_err)?)
        }
        AnyTypeInfoKind::Blob => {
            blob_to_json(row.try_get::<Vec<u8>, _>(index).map_err(decode_err)?)
        }
    };
    Ok(value)
}

/// JSON has no NaN or infinity. Rather than emit null — indistinguishable from
/// a SQL NULL — say so.
fn float_to_json(v: f64, name: &str) -> Result<Value, DataflowError> {
    serde_json::Number::from_f64(v)
        .map(Value::Number)
        .ok_or_else(|| {
            DataflowError::function_execution(
                format!("{NAME}: column '{name}' holds {v}, which JSON cannot represent"),
                None,
            )
        })
}

/// Binary columns become a string: the UTF-8 text when the bytes are valid
/// UTF-8 (MySQL reports `TEXT`/`JSON` columns as `BLOB`, so this is the common
/// case), otherwise lowercase hex.
fn blob_to_json(bytes: Vec<u8>) -> Value {
    match String::from_utf8(bytes) {
        Ok(s) => Value::String(s),
        Err(e) => Value::String(hex::encode(e.into_bytes())),
    }
}

// -- Input schema (F53) --
//
// The table describing this handler's `function.input` lives next to the
// handler it describes. It used to sit in `schema.rs` with the other nine,
// which is how every schema/handler divergence in the 1.0 audit happened:
// a field was added, renamed or made conditional here and the table saying
// so was in a different file.

pub(super) const DB_READ_FIELDS: &[FieldSchema] = &[
    FieldSchema {
        name: "connector",
        description: "Name of the SQL connector to query.",
        kind: FieldKind::String,
        required: true,
        resolvable: false,
        alias: None,
    },
    FieldSchema {
        name: "query",
        description: "SQL query. Use $1, $2, ... placeholders bound from `params`.",
        kind: FieldKind::String,
        required: true,
        resolvable: false,
        alias: None,
    },
    FieldSchema {
        name: "params",
        description: "Array of values to bind to query placeholders, in order. Accepts {\"var\": \"path\"} to read the value from the message.",
        kind: FieldKind::Array,
        required: false,
        resolvable: true,
        alias: None,
    },
    FieldSchema {
        name: "output",
        description: "Dotted path in the message where rows are written. Defaults to \"data\".",
        kind: FieldKind::String,
        required: false,
        resolvable: false,
        alias: None,
    },
];