forge_runtime/function/

executor.rs

use std::sync::Arc;
use std::time::Duration;

use forge_core::{AuthContext, ForgeError, JobDispatch, RequestMetadata, Result, WorkflowDispatch};
use serde_json::Value;
use tokio::time::timeout;
use tracing::{debug, error, info, trace, warn};

use super::registry::FunctionRegistry;
use super::router::{FunctionRouter, RouteResult};

/// Executes functions with timeout and error handling.
pub struct FunctionExecutor {
    router: FunctionRouter,
    registry: Arc<FunctionRegistry>,
    default_timeout: Duration,
}

impl FunctionExecutor {
    /// Create a new function executor.
    pub fn new(registry: Arc<FunctionRegistry>, db_pool: sqlx::PgPool) -> Self {
        Self {
            router: FunctionRouter::new(Arc::clone(&registry), db_pool),
            registry,
            default_timeout: Duration::from_secs(30),
        }
    }

    /// Create a new function executor with custom timeout.
    pub fn with_timeout(
        registry: Arc<FunctionRegistry>,
        db_pool: sqlx::PgPool,
        default_timeout: Duration,
    ) -> Self {
        Self {
            router: FunctionRouter::new(Arc::clone(&registry), db_pool),
            registry,
            default_timeout,
        }
    }

    /// Create a new function executor with dispatch capabilities.
    pub fn with_dispatch(
        registry: Arc<FunctionRegistry>,
        db_pool: sqlx::PgPool,
        job_dispatcher: Option<Arc<dyn JobDispatch>>,
        workflow_dispatcher: Option<Arc<dyn WorkflowDispatch>>,
    ) -> Self {
        let mut router = FunctionRouter::new(Arc::clone(&registry), db_pool);
        if let Some(jd) = job_dispatcher {
            router = router.with_job_dispatcher(jd);
        }
        if let Some(wd) = workflow_dispatcher {
            router = router.with_workflow_dispatcher(wd);
        }
        Self {
            router,
            registry,
            default_timeout: Duration::from_secs(30),
        }
    }

    /// Execute a function call.
    pub async fn execute(
        &self,
        function_name: &str,
        args: Value,
        auth: AuthContext,
        request: RequestMetadata,
    ) -> Result<ExecutionResult> {
        let start = std::time::Instant::now();

        // Get function-specific timeout or use default
        let fn_timeout = self.get_function_timeout(function_name);

        // Get log level for this function (default to trace)
        let log_level = self.get_function_log_level(function_name);

        // Execute with timeout
        let result = match timeout(
            fn_timeout,
            self.router
                .route(function_name, args.clone(), auth, request),
        )
        .await
        {
            Ok(result) => result,
            Err(_) => {
                let duration = start.elapsed();
                self.log_execution(
                    log_level,
                    function_name,
                    "unknown",
                    &args,
                    duration,
                    false,
                    Some(&format!("Timeout after {:?}", fn_timeout)),
                );
                return Err(ForgeError::Timeout(format!(
                    "Function '{}' timed out after {:?}",
                    function_name, fn_timeout
                )));
            }
        };

        let duration = start.elapsed();

        match result {
            Ok(route_result) => {
                let (kind, value) = match route_result {
                    RouteResult::Query(v) => ("query", v),
                    RouteResult::Mutation(v) => ("mutation", v),
                    RouteResult::Job(v) => ("job", v),
                    RouteResult::Workflow(v) => ("workflow", v),
                };

                self.log_execution(log_level, function_name, kind, &args, duration, true, None);

                Ok(ExecutionResult {
                    function_name: function_name.to_string(),
                    function_kind: kind.to_string(),
                    result: value,
                    duration,
                    success: true,
                    error: None,
                })
            }
            Err(e) => {
                let kind = self
                    .router
                    .get_function_kind(function_name)
                    .map(|k| k.to_string())
                    .unwrap_or_else(|| "unknown".to_string());

                self.log_execution(
                    log_level,
                    function_name,
                    &kind,
                    &args,
                    duration,
                    false,
                    Some(&e.to_string()),
                );

                Err(e)
            }
        }
    }

    /// Log function execution at the configured level.
    #[allow(clippy::too_many_arguments)]
    fn log_execution(
        &self,
        log_level: &str,
        function_name: &str,
        kind: &str,
        input: &Value,
        duration: Duration,
        success: bool,
        error: Option<&str>,
    ) {
        macro_rules! log_fn {
            ($level:ident) => {
                if success {
                    $level!(
                        function = function_name,
                        kind = kind,
                        input = %input,
                        duration_ms = duration.as_millis() as u64,
                        success = success,
                        "Function executed"
                    );
                } else {
                    $level!(
                        function = function_name,
                        kind = kind,
                        input = %input,
                        duration_ms = duration.as_millis() as u64,
                        success = success,
                        error = error,
                        "Function failed"
                    );
                }
            };
        }

        match log_level {
            "off" => {}
            "error" => log_fn!(error),
            "warn" => log_fn!(warn),
            "info" => log_fn!(info),
            "debug" => log_fn!(debug),
            _ => log_fn!(trace),
        }
    }

    /// Get the log level for a specific function.
    fn get_function_log_level(&self, function_name: &str) -> &'static str {
        self.registry
            .get(function_name)
            .and_then(|entry| entry.info().log_level)
            .unwrap_or("trace")
    }

    /// Get the timeout for a specific function.
    fn get_function_timeout(&self, function_name: &str) -> Duration {
        self.registry
            .get(function_name)
            .and_then(|entry| entry.info().timeout)
            .map(Duration::from_secs)
            .unwrap_or(self.default_timeout)
    }

    /// Check if a function exists.
    pub fn has_function(&self, function_name: &str) -> bool {
        self.router.has_function(function_name)
    }
}

/// Result of executing a function.
#[derive(Debug, Clone, serde::Serialize)]
pub struct ExecutionResult {
    /// Function name that was executed.
    pub function_name: String,
    /// Kind of function (query, mutation).
    pub function_kind: String,
    /// The result value (or null on error).
    pub result: Value,
    /// Execution duration.
    #[serde(with = "duration_millis")]
    pub duration: Duration,
    /// Whether execution succeeded.
    pub success: bool,
    /// Error message if failed.
    pub error: Option<String>,
}

mod duration_millis {
    use serde::{Deserialize, Deserializer, Serializer};
    use std::time::Duration;

    pub fn serialize<S>(duration: &Duration, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.serialize_u64(duration.as_millis() as u64)
    }

    #[allow(dead_code)]
    pub fn deserialize<'de, D>(deserializer: D) -> Result<Duration, D::Error>
    where
        D: Deserializer<'de>,
    {
        let millis = u64::deserialize(deserializer)?;
        Ok(Duration::from_millis(millis))
    }
}

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

    #[test]
    fn test_execution_result_serialization() {
        let result = ExecutionResult {
            function_name: "get_user".to_string(),
            function_kind: "query".to_string(),
            result: serde_json::json!({"id": "123"}),
            duration: Duration::from_millis(42),
            success: true,
            error: None,
        };

        let json = serde_json::to_string(&result).unwrap();
        assert!(json.contains("\"duration\":42"));
        assert!(json.contains("\"success\":true"));
    }
}