komp_ac/server/src/steel/server/execution.rs

// src/steel/server/execution.rs
use std::fmt;
use std::collections::HashMap;
use sqlx::{PgPool, Row};

#[derive(Debug)]
pub enum ScriptOperation {
    SetToLocalColumn { source: String },
    SetToExternalColumn { table: String, column: String },
}

#[derive(Debug)]
pub enum ScriptExecutionError {
    ParseError(String),
    MissingSourceColumn(String),
    Mismatch { expected: String, actual: String },
    InvalidReference(String),
    MissingLinkKey(String),
    DatabaseError(String),
    MissingExternalData(String),
}

impl fmt::Display for ScriptExecutionError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::ParseError(msg) => write!(f, "Parse error: {}", msg),
            Self::MissingSourceColumn(col) => write!(f, "Missing source column: {}", col),
            Self::Mismatch { expected, actual } => write!(
                f,
                "Value does not match script expectation. Expected: {}, Actual: {}",
                expected, actual
            ),
            Self::InvalidReference(msg) => write!(f, "Invalid reference: {}", msg),
            Self::MissingLinkKey(key) => write!(f, "Missing link key: {}", key),
            Self::DatabaseError(msg) => write!(f, "Database error: {}", msg),
            Self::MissingExternalData(msg) => write!(f, "External data not found: {}", msg),
        }
    }
}

pub fn parse_script(script: &str, expected_target: &str) -> Result<ScriptOperation, ScriptExecutionError> {
    let script = script.trim();

    if !script.starts_with("(set! ") || !script.ends_with(')') {
        return Err(ScriptExecutionError::ParseError(
            "Script must be in the form (set! target source)".to_string(),
        ));
    }

    let content = &script[6..script.len() - 1].trim();
    let mut parts = content.split_whitespace();

    let target = parts.next().ok_or_else(|| {
        ScriptExecutionError::ParseError("Missing target in set! expression".to_string())
    })?;
    let source = parts.next().ok_or_else(|| {
        ScriptExecutionError::ParseError("Missing source in set! expression".to_string())
    })?;

    if parts.next().is_some() {
        return Err(ScriptExecutionError::ParseError(
            "Too many arguments in set! expression".to_string(),
        ));
    }

    if target != expected_target {
        return Err(ScriptExecutionError::ParseError(format!(
            "Script target '{}' does not match expected '{}'",
            target, expected_target
        )));
    }

    // Check if source is an external reference
    if source.starts_with('@') {
        let (table, column) = source.split_once('.')
            .ok_or_else(|| ScriptExecutionError::InvalidReference(
                format!("Invalid external reference format: {}", source)
            ))?;
        Ok(ScriptOperation::SetToExternalColumn {
            table: table.to_string(),
            column: column.to_string(),
        })
    } else {
        Ok(ScriptOperation::SetToLocalColumn {
            source: source.to_string(),
        })
    }
}

pub async fn resolve_value(
    db_pool: &PgPool,
    profile_id: i64,
    current_table: &str,
    current_data: &HashMap<String, String>,
    source: &str,
) -> Result<String, ScriptExecutionError> {
    if let Some((table, column)) = source.split_once('.') {
        // External table reference
        let external_table = table.strip_prefix('@')
            .ok_or_else(|| ScriptExecutionError::InvalidReference(format!("Invalid external reference: {}", source)))?;

        // Get foreign key relationship info
        let (fk_column, referenced_column) = get_relationship_info(db_pool, current_table, external_table).await?;

        // Get foreign key value from current data
        let fk_value = current_data.get(&fk_column)
            .ok_or_else(|| ScriptExecutionError::MissingLinkKey(fk_column.clone()))?;

        // Build and execute query
        let query = format!(
            "SELECT {} FROM \"{}\" WHERE {} = $1",
            column, external_table, referenced_column
        );

        let result: Option<String> = sqlx::query(&query)
            .bind(fk_value)
            .fetch_optional(db_pool)
            .await
            .map_err(|e| ScriptExecutionError::DatabaseError(e.to_string()))?
            .and_then(|row| row.try_get(0).ok());

        result.ok_or_else(|| ScriptExecutionError::MissingExternalData(
            format!("No data found for {} in {}", column, external_table)
        ))
    } else {
        // Local column reference remains the same
        current_data.get(source)
            .cloned()
            .ok_or_else(|| ScriptExecutionError::MissingSourceColumn(source.into()))
    }
}

async fn get_relationship_info(
    db_pool: &PgPool,
    current_table: &str,
    external_table: &str,
) -> Result<(String, String), ScriptExecutionError> {
    let query = r#"
        SELECT
            kcu.column_name AS fk_column,
            ccu.column_name AS referenced_column
        FROM
            information_schema.table_constraints AS tc
            JOIN information_schema.key_column_usage AS kcu
              ON tc.constraint_name = kcu.constraint_name
              AND tc.table_schema = kcu.table_schema
            JOIN information_schema.constraint_column_usage AS ccu
              ON tc.constraint_name = ccu.constraint_name
              AND tc.table_schema = ccu.table_schema
        WHERE
            tc.constraint_type = 'FOREIGN KEY'
            AND tc.table_name = $1
            AND ccu.table_name = $2
    "#;

    let rows = sqlx::query(query)
        .bind(current_table)
        .bind(external_table)
        .fetch_all(db_pool)
        .await
        .map_err(|e| ScriptExecutionError::DatabaseError(e.to_string()))?;

    match rows.len() {
        0 => Err(ScriptExecutionError::InvalidReference(format!(
            "No foreign key relationship found from {} to {}",
            current_table, external_table
        ))),
        1 => {
            let row = &rows[0];
            Ok((
                row.get::<String, _>("fk_column"),
                row.get::<String, _>("referenced_column"),
            ))
        },
        _ => Err(ScriptExecutionError::InvalidReference(format!(
            "Multiple foreign keys between {} and {} - cannot resolve ambiguity",
            current_table, external_table
        )))
    }
}