uqa-sql 0.5.4

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Result rows returned by `Engine::sql`.

use std::collections::BTreeMap;

use uqa_core::Value;

use crate::ast::ColumnType;
use crate::RowSchema;

mod labels;
mod text;
pub use labels::render_result_enum_labels;
pub use text::{format_postgres_text, format_postgres_text_with_catalog};

pub type ResultRow = BTreeMap<String, Value>;

/// Whether execution produced a row descriptor, independently of its column or row count.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum SQLResultKind {
    #[default]
    Command,
    Rows,
    /// An external result source did not provide descriptor information.
    Unknown,
}

#[derive(Debug, Clone, Default)]
pub struct SQLResult {
    /// Descriptor presence, including zero-column and empty row results.
    pub kind: SQLResultKind,
    /// `PostgreSQL` command completion, including its command-specific row count. Execution sets this from the command that actually ran; row constructors leave it absent because rows alone do not identify a SQL command.
    pub command_tag: Option<String>,
    /// Column order as the SELECT clause specified.
    pub columns: Vec<String>,
    /// Statically bound SQL type for each output position. A missing entry
    /// represents a type that has not yet been resolved, never a type inferred
    /// from the first runtime value.
    pub column_types: Vec<Option<ColumnType>>,
    /// Complete, densely ordered native output metadata, including selected collations and nested record descriptors. External result sources may omit it.
    pub row_schema: Option<RowSchema>,
    /// One row per result document, with the named columns in
    /// `columns`. Extra columns from `_score` etc. are included here
    /// too.
    pub rows: Vec<ResultRow>,
    /// Positional values for result sets whose output contains repeated column
    /// labels. `rows` remains available for named lookup, while this carrier
    /// preserves values that cannot be represented by a string-keyed map.
    #[doc(hidden)]
    pub positional_rows: Option<Vec<Vec<Value>>>,
    /// Number of rows touched by an INSERT / UPDATE / DELETE.
    pub affected_rows: u64,
}

impl SQLResult {
    /// Attach the completion chosen by the executing SQL command.
    pub fn with_command_tag(mut self, tag: impl Into<String>) -> Self {
        self.command_tag = Some(tag.into());
        self
    }

    pub fn empty() -> Self {
        Self::default()
    }

    pub fn from_rows(columns: Vec<String>, rows: Vec<ResultRow>) -> Self {
        let column_types = vec![None; columns.len()];
        Self {
            kind: SQLResultKind::Rows,
            command_tag: None,
            columns,
            column_types,
            row_schema: None,
            rows,
            positional_rows: None,
            affected_rows: 0,
        }
    }

    pub fn from_rows_with_positions(
        columns: Vec<String>,
        rows: Vec<ResultRow>,
        positional_rows: Option<Vec<Vec<Value>>>,
    ) -> Self {
        let column_types = vec![None; columns.len()];
        Self::from_typed_rows_with_positions(columns, column_types, rows, positional_rows)
    }

    pub fn from_typed_rows_with_positions(
        columns: Vec<String>,
        column_types: Vec<Option<ColumnType>>,
        mut rows: Vec<ResultRow>,
        positional_rows: Option<Vec<Vec<Value>>>,
    ) -> Self {
        debug_assert_eq!(columns.len(), column_types.len());
        debug_assert!(positional_rows.as_ref().is_none_or(|values| {
            values.len() == rows.len() && values.iter().all(|row| row.len() == columns.len())
        }));
        if let Some(positional_rows) = positional_rows.as_ref() {
            let compatibility_labels = unique_compatibility_labels(&columns);
            for (row, positional) in rows.iter_mut().zip(positional_rows) {
                for (label, value) in compatibility_labels.iter().zip(positional) {
                    row.insert(label.clone(), value.clone());
                }
            }
        }
        Self {
            kind: SQLResultKind::Rows,
            command_tag: None,
            columns,
            column_types,
            row_schema: None,
            rows,
            positional_rows,
            affected_rows: 0,
        }
    }

    /// Retain a native output descriptor independently of whether execution returned any rows.
    pub fn from_schema_rows_with_positions(
        schema: RowSchema,
        rows: Vec<ResultRow>,
        positional_rows: Option<Vec<Vec<Value>>>,
    ) -> Self {
        let mut result = Self::from_typed_rows_with_positions(
            schema.columns().to_vec(),
            schema.column_types().to_vec(),
            rows,
            positional_rows,
        );
        result.row_schema = Some(schema);
        result
    }

    /// Borrow the native descriptor, or construct the metadata supplied by a legacy or external result source. Runtime values never determine declared types or collations.
    pub fn row_schema(&self) -> std::borrow::Cow<'_, RowSchema> {
        self.row_schema.as_ref().map_or_else(
            || {
                std::borrow::Cow::Owned(RowSchema::with_types(
                    self.columns.clone(),
                    self.column_types.clone(),
                ))
            },
            std::borrow::Cow::Borrowed,
        )
    }

    /// Return a result value by row and output-column position.
    ///
    /// Positional access is the canonical way to distinguish repeated output
    /// labels. Named rows remain available for compatibility with existing
    /// callers.
    pub fn value_at(&self, row: usize, column: usize) -> Option<&Value> {
        self.positional_rows
            .as_ref()
            .and_then(|rows| rows.get(row))
            .and_then(|row| row.get(column))
            .or_else(|| {
                self.columns
                    .get(column)
                    .and_then(|name| self.rows.get(row)?.get(name))
            })
    }

    pub fn from_affected(affected: u64) -> Self {
        Self {
            affected_rows: affected,
            ..Self::default()
        }
    }
}

fn unique_compatibility_labels(columns: &[String]) -> Vec<String> {
    let mut labels = Vec::with_capacity(columns.len());
    for base in columns {
        let mut label = base.clone();
        let mut suffix = 1usize;
        while labels.contains(&label) {
            label = format!("{base}_{suffix}");
            suffix += 1;
        }
        labels.push(label);
    }
    labels
}

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

    #[test]
    fn result_metadata_survives_empty_rows_and_shared_result_clones() {
        let schema = RowSchema::with_column_metadata(
            None,
            [
                ("value".into(), Some(ColumnType::Text), Some(950)),
                ("value".into(), Some(ColumnType::Text), Some(951)),
            ],
        );
        for rows in [Vec::new(), vec![vec![Value::Null, Value::Null]]] {
            let result = SQLResult::from_schema_rows_with_positions(
                schema.clone(),
                vec![ResultRow::new(); rows.len()],
                Some(rows),
            );
            let retained = result.row_schema();
            assert!(std::ptr::eq(retained.columns(), schema.columns()));
            assert!(matches!(retained, std::borrow::Cow::Borrowed(_)));
            assert!(std::ptr::eq(
                retained.as_ref(),
                result.row_schema.as_ref().unwrap()
            ));
            assert_eq!(result.kind, SQLResultKind::Rows);
            assert_eq!(retained.column_collation(0), Some(950));
            assert_eq!(retained.column_collation(1), Some(951));
            assert_eq!(*result.clone().row_schema(), schema);
        }
    }

    #[test]
    fn repeated_postgresql_labels_keep_unique_named_compatibility_keys() {
        let result = SQLResult::from_rows_with_positions(
            vec!["value".into(), "value".into()],
            vec![ResultRow::from([("value".into(), Value::Int(6))])],
            Some(vec![vec![Value::Int(5), Value::Int(6)]]),
        );

        assert_eq!(result.columns, ["value", "value"]);
        assert_eq!(result.rows[0].get("value"), Some(&Value::Int(5)));
        assert_eq!(result.rows[0].get("value_1"), Some(&Value::Int(6)));
        assert_eq!(result.value_at(0, 0), Some(&Value::Int(5)));
        assert_eq!(result.value_at(0, 1), Some(&Value::Int(6)));
    }
}

pub mod completion;
pub mod vector;

/// Execution measurements supplied to EXPLAIN rendering.
pub struct ExplainAnalysis {
    pub vector_searches: vector::ExplainVectorSearches,
    pub elapsed: std::time::Duration,
    pub rows: u64,
    pub affected_rows: u64,
}

/// Structured physical-plan output captured before execution. Method-specific properties are rendered by the owning planner; execution only carries this diagnostic result to the renderer.
#[derive(Debug, Clone, Default)]
pub struct ExplainPhysicalPlan {
    pub nodes: Vec<serde_json::Value>,
}