flowcode-core 0.4.1-alpha

//! flowcode-core/src/duckdb_session.rs
//! Manages DuckDB connections and SQL query execution for FlowCode.

use std::error::Error;
use std::fmt;

use crate::ast::ArgValue;
use crate::types::{TypedValue, ValueKind};

// Use DuckDB's re-exported Arrow types to ensure version compatibility
use duckdb::Connection;
use duckdb::types::{Value, ValueRef};

/// Custom error type for DuckDB operations
#[derive(Debug)]
pub enum DuckDBError {
    ConnectionError(String),
    QueryError(String),
    TypeConversionError(String),
}

impl fmt::Display for DuckDBError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DuckDBError::ConnectionError(e) => write!(f, "Connection error: {}", e),
            DuckDBError::QueryError(e) => write!(f, "Query error: {}", e),
            DuckDBError::TypeConversionError(e) => write!(f, "Type conversion error: {}", e),
        }
    }
}

impl Error for DuckDBError {}

/// A session for interacting with a DuckDB database.
#[cfg(feature = "duckdb")]
pub struct DuckSession {
    conn: Connection,
}

#[cfg(feature = "duckdb")]
impl DuckSession {
    /// Creates a new DuckDB session with an in-memory database.
    pub fn new() -> Result<Self, DuckDBError> {
        let conn = Connection::open_in_memory()
            .map_err(|e| DuckDBError::ConnectionError(e.to_string()))?;
        Ok(DuckSession { conn })
    }

    /// Executes an SQL query and returns the result as a TypedValue.
    pub fn execute_sql(&mut self, query: &str) -> Result<TypedValue, DuckDBError> {
        self.execute_with_params(query, &[])
    }
    
    /// Executes a query with parameters and returns the result as a TypedValue.
    pub fn execute_with_params(
        &mut self,
        query: &str,
        params: &[ArgValue],
    ) -> Result<TypedValue, DuckDBError> {
        let mut stmt = self.conn.prepare(query)
            .map_err(|e| DuckDBError::QueryError(e.to_string()))?;

        // Convert parameters to DuckDB-compatible values
        let duck_params: Vec<Value> = params
            .iter()
            .map(convert_arg_to_duckdb_value)
            .collect::<Result<Vec<_>, _>>()?;

        // Create parameter references
        let params: Vec<&dyn duckdb::ToSql> = duck_params
            .iter()
            .map(|p| p as &dyn duckdb::ToSql)
            .collect();

        // Execute query and get results
        let column_count = stmt.column_count();
        let mut rows = stmt.query(&params[..])
            .map_err(|e| DuckDBError::QueryError(e.to_string()))?;

        // Process results
        let mut table_data = Vec::new();
        while let Ok(Some(row)) = rows.next() {
            let mut row_data = Vec::new();
            for i in 0..column_count {
                let val = row.get_ref(i)
                    .map_err(|e| DuckDBError::QueryError(e.to_string()))?;
                let typed_val = self.value_to_typed_value(&val)?;
                row_data.push(typed_val);
            }
            table_data.push(row_data);
        }

        Ok(TypedValue::new(
            ValueKind::Table,
            ArgValue::Table(table_data)
        ))
    }
    
    /// Registers a CSV file as a table in DuckDB for querying.
    pub fn register_csv(&mut self, table_name: &str, file_path: &str) -> Result<(), DuckDBError> {
        let query = format!(
            "CREATE TABLE \"{}\" AS SELECT * FROM read_csv_auto('{}')",
            table_name, file_path
        );
        
        self.conn.execute(&query, [])
            .map_err(|e| DuckDBError::QueryError(e.to_string()))?;
        Ok(())
    }
    
    /// Executes a query that returns a single value.
    pub fn query_scalar(&mut self, query: &str) -> Result<TypedValue, DuckDBError> {
        let mut stmt = self.conn.prepare(query)
            .map_err(|e| DuckDBError::QueryError(e.to_string()))?;
        
        // Execute the query and get the first row
        let mut rows = stmt.query([])
            .map_err(|e| DuckDBError::QueryError(e.to_string()))?;
        
        if let Ok(Some(row)) = rows.next() {
            // Get the first column value
            self.value_to_typed_value(&row.get_ref(0)
                .map_err(|e| DuckDBError::QueryError(e.to_string()))?)
        } else {
            // No rows returned
            Ok(TypedValue::new(ValueKind::Null, ArgValue::Null))
        }
    }
    
    /// Converts a DuckDB value to a TypedValue.
    fn value_to_typed_value(&self, value: &ValueRef) -> Result<TypedValue, DuckDBError> {
        match value {
            ValueRef::Null => Ok(TypedValue::new(ValueKind::Null, ArgValue::Null)),
            ValueRef::Boolean(b) => Ok(TypedValue::new(ValueKind::Bool, ArgValue::Bool(*b))),
            ValueRef::TinyInt(i) => Ok(TypedValue::new(ValueKind::Number, ArgValue::Number(*i as f64))),
            ValueRef::SmallInt(i) => Ok(TypedValue::new(ValueKind::Number, ArgValue::Number(*i as f64))),
            ValueRef::Int(i) => Ok(TypedValue::new(ValueKind::Number, ArgValue::Number(*i as f64))),
            ValueRef::BigInt(i) => Ok(TypedValue::new(ValueKind::Number, ArgValue::Number(*i as f64))),
            ValueRef::HugeInt(i) => Ok(TypedValue::new(ValueKind::Number, ArgValue::Number(*i as f64))),
            ValueRef::Float(f) => Ok(TypedValue::new(ValueKind::Number, ArgValue::Number(*f as f64))),
            ValueRef::Double(d) => Ok(TypedValue::new(ValueKind::Number, ArgValue::Number(*d))),
            ValueRef::Text(s) => {
                let s = String::from_utf8_lossy(s).into_owned();
                Ok(TypedValue::new(ValueKind::String, ArgValue::String(s)))
            },
            _ => Err(DuckDBError::TypeConversionError(
                format!("Unsupported DuckDB value type: {:?}", value)
            )),
        }
    }

    /// Executes a query on a TypedValue Table and returns the result as a new Table.
    pub fn query_table(
        &mut self,
        _table: &TypedValue,
        query: &str,
    ) -> Result<TypedValue, DuckDBError> {
        // For now, just execute the query directly
        // In a real implementation, we would register the table first
        self.execute_sql(query)
    }
    
    /// Convenience wrapper: LOAD CSV via DuckDB's read_csv_auto into a table.
    pub fn load_csv_auto(&mut self, table_name: &str, file_path: &str) -> Result<(), DuckDBError> {
        let query = format!(
            "CREATE TABLE \"{}\" AS SELECT * FROM read_csv_auto('{}')",
            table_name, file_path
        );
        
        self.conn.execute(&query, [])
            .map_err(|e| DuckDBError::QueryError(e.to_string()))?;
        Ok(())
    }
}

/// Helper function to convert ArgValue to DuckDB Value
fn convert_arg_to_duckdb_value(arg: &ArgValue) -> Result<Value, DuckDBError> {
    match arg {
        ArgValue::Number(n) => Ok(Value::Double(*n)),
        ArgValue::String(s) => Ok(Value::Text(s.clone())),
        ArgValue::Bool(b) => Ok(Value::Boolean(*b)),
        ArgValue::Null => Ok(Value::Null),
        _ => Err(DuckDBError::TypeConversionError(
            "Unsupported parameter type".to_string()
        )),
    }
}