oxigdal-workflow 0.1.7

//! Prefect integration.

use crate::dag::{ResourceRequirements, RetryPolicy, TaskEdge, TaskNode, WorkflowDag};
use crate::engine::WorkflowDefinition;
use crate::error::{Result, WorkflowError};
use regex::Regex;
use std::collections::HashMap;

/// Header comment prefix identifying Python source generated by
/// [`PrefectIntegration::export_workflow`].
///
/// [`PrefectIntegration::import_workflow`] only supports round-tripping this crate's own
/// generated format; this marker is what it checks for, alongside the `workflow_id` metadata
/// comment and the `from prefect import flow, task` import line.
const GENERATED_MARKER: &str = "# Code generated by oxigdal-workflow PrefectIntegration.";

/// Prefect integration.
pub struct PrefectIntegration;

impl PrefectIntegration {
    /// Export workflow to Prefect flow format (Python).
    pub fn export_workflow(workflow: &WorkflowDefinition) -> Result<String> {
        let mut python_code = String::new();

        // Metadata header. `workflow.name` is already recoverable from the `@flow(name=...)`
        // line below, but the flow/task function names are derived from a sanitized version of
        // `workflow.id` (hyphens and spaces both collapse to `_`), which is not invertible in
        // general. Embedding the raw id lets `import_workflow` recover it losslessly.
        python_code.push_str(GENERATED_MARKER);
        python_code.push_str(" Do not edit.\n");
        python_code.push_str(&format!("# workflow_id: {}\n", workflow.id));

        // Add imports
        python_code.push_str("from prefect import flow, task\n");
        python_code.push_str("from datetime import timedelta\n\n");

        // Define tasks
        for (idx, _task) in workflow.dag.tasks().iter().enumerate() {
            python_code.push_str(&format!("@task(name='task_{}')\n", idx));
            python_code.push_str(&format!("def task_{}():\n", idx));
            python_code.push_str("    print('Task executed')\n");
            python_code.push_str("    return True\n\n");
        }

        // Define flow
        python_code.push_str(&format!("@flow(name='{}')\n", workflow.name));
        python_code.push_str(&format!("def {}():\n", Self::sanitize_id(&workflow.id)));

        for (idx, _task) in workflow.dag.tasks().iter().enumerate() {
            python_code.push_str(&format!("    result_{} = task_{}()\n", idx, idx));
        }

        python_code.push('\n');
        python_code.push_str("if __name__ == '__main__':\n");
        python_code.push_str(&format!("    {}()\n", Self::sanitize_id(&workflow.id)));

        Ok(python_code)
    }

    /// Import workflow from Prefect flow.
    ///
    /// Only Python source previously produced by [`Self::export_workflow`] is supported: this
    /// searches for the `oxigdal-workflow` generated-code marker, the `workflow_id` metadata
    /// comment, the `@flow(name='...')` decorator, and `@task(name='task_{N}')` decorators, then
    /// rebuilds a [`WorkflowDefinition`] whose DAG is a simple sequential chain
    /// (`task-0 -> task-1 -> ...`) over as many tasks as there were `@task` decorators in the
    /// source. Arbitrary hand-written Prefect flows are not a supported input: this function
    /// does not implement a Python parser and returns a [`WorkflowError::Integration`]
    /// describing the supported subset instead of attempting (and likely failing) to interpret
    /// unfamiliar code.
    pub fn import_workflow(flow_code: &str) -> Result<WorkflowDefinition> {
        if !flow_code.contains(GENERATED_MARKER)
            || !flow_code.contains("from prefect import flow, task")
        {
            return Err(WorkflowError::integration(
                "prefect",
                "Import only supports Python source produced by \
                 PrefectIntegration::export_workflow (it must contain the \
                 'oxigdal-workflow PrefectIntegration' generated-code marker and \
                 'from prefect import flow, task'); arbitrary hand-written Prefect flows are \
                 not a supported input format",
            ));
        }

        let id_regex = Regex::new(r"(?m)^# workflow_id:\s*(.+)$").map_err(|e| {
            WorkflowError::integration("prefect", format!("Internal regex error: {}", e))
        })?;
        let name_regex = Regex::new(r"(?m)^@flow\(name='(.*)'\)$").map_err(|e| {
            WorkflowError::integration("prefect", format!("Internal regex error: {}", e))
        })?;
        let task_regex = Regex::new(r"(?m)^@task\(name='task_(\d+)'\)$").map_err(|e| {
            WorkflowError::integration("prefect", format!("Internal regex error: {}", e))
        })?;

        let id = id_regex
            .captures(flow_code)
            .and_then(|c| c.get(1))
            .map(|m| m.as_str().trim().to_string())
            .ok_or_else(|| {
                WorkflowError::integration(
                    "prefect",
                    "Missing '# workflow_id: <id>' metadata comment; this is not a workflow \
                     exported by PrefectIntegration::export_workflow",
                )
            })?;

        let name = name_regex
            .captures(flow_code)
            .and_then(|c| c.get(1))
            .map(|m| m.as_str().trim().to_string())
            .unwrap_or_else(|| id.clone());

        let mut task_indices: Vec<usize> = task_regex
            .captures_iter(flow_code)
            .filter_map(|c| c.get(1).and_then(|m| m.as_str().parse::<usize>().ok()))
            .collect();
        task_indices.sort_unstable();
        task_indices.dedup();

        let mut dag = WorkflowDag::new();
        for (position, _original_idx) in task_indices.iter().enumerate() {
            dag.add_task(Self::sequential_task(position))?;
        }
        for position in 1..task_indices.len() {
            dag.add_dependency(
                &Self::sequential_task_id(position - 1),
                &Self::sequential_task_id(position),
                TaskEdge::default(),
            )?;
        }

        Ok(WorkflowDefinition {
            id,
            name,
            description: None,
            version: "1.0.0".to_string(),
            dag,
        })
    }

    /// Build the `position`-th task of a reconstructed sequential DAG.
    fn sequential_task(position: usize) -> TaskNode {
        TaskNode {
            id: Self::sequential_task_id(position),
            name: format!("Task {}", position),
            description: None,
            config: serde_json::json!({}),
            retry: RetryPolicy::default(),
            timeout_secs: None,
            resources: ResourceRequirements::default(),
            metadata: HashMap::new(),
        }
    }

    /// Deterministic task id used when reconstructing a sequential DAG on import.
    fn sequential_task_id(position: usize) -> String {
        format!("task-{}", position)
    }

    /// Sanitize ID for Prefect compatibility.
    fn sanitize_id(id: &str) -> String {
        id.replace(['-', ' '], "_")
    }

    /// Trigger a Prefect flow via API.
    #[cfg(feature = "integrations")]
    pub async fn trigger_flow(
        base_url: &str,
        flow_id: &str,
        api_key: Option<&str>,
    ) -> Result<String> {
        use reqwest::Client;

        let url = format!("{}/api/flows/{}/runs", base_url, flow_id);
        let client = Client::new();

        let mut request = client
            .post(&url)
            .header("Content-Type", "application/json")
            .json(&serde_json::json!({
                "parameters": {}
            }));

        if let Some(key) = api_key {
            request = request.bearer_auth(key);
        }

        let response = request
            .send()
            .await
            .map_err(|e| WorkflowError::integration("prefect", format!("Request failed: {}", e)))?;

        let status = response.status();

        let body = response.text().await.map_err(|e| {
            WorkflowError::integration("prefect", format!("Failed to read response: {}", e))
        })?;

        if !status.is_success() {
            return Err(WorkflowError::integration(
                "prefect",
                format!("HTTP {}: {}", status, body),
            ));
        }

        Ok(body)
    }
}

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

    fn simple_task(id: &str) -> TaskNode {
        TaskNode {
            id: id.to_string(),
            name: id.to_string(),
            description: None,
            config: serde_json::json!({}),
            retry: RetryPolicy::default(),
            timeout_secs: Some(60),
            resources: ResourceRequirements::default(),
            metadata: HashMap::new(),
        }
    }

    #[test]
    fn test_export_to_prefect() {
        let workflow = WorkflowDefinition {
            id: "test-workflow".to_string(),
            name: "Test Workflow".to_string(),
            description: None,
            version: "1.0.0".to_string(),
            dag: WorkflowDag::new(),
        };

        let result = PrefectIntegration::export_workflow(&workflow);
        assert!(result.is_ok());

        let python_code = result.expect("Failed to export");
        assert!(python_code.contains("from prefect import flow, task"));
        assert!(python_code.contains("@flow"));
    }

    #[test]
    fn test_sanitize_id() {
        assert_eq!(
            PrefectIntegration::sanitize_id("test-workflow-id"),
            "test_workflow_id"
        );
    }

    #[test]
    fn test_import_rejects_foreign_python_code() {
        let foreign = "def main():\n    pass\n";
        let result = PrefectIntegration::import_workflow(foreign);
        assert!(result.is_err());
        let message = result.expect_err("expected error").to_string();
        assert!(
            message.contains("not a supported input format")
                || message.contains("PrefectIntegration::export_workflow")
        );
    }

    #[test]
    fn test_export_import_round_trip() {
        let mut dag = WorkflowDag::new();
        dag.add_task(simple_task("download")).expect("add_task");
        dag.add_task(simple_task("process")).expect("add_task");
        dag.add_task(simple_task("upload")).expect("add_task");

        let original = WorkflowDefinition {
            id: "sentinel-2-pipeline".to_string(),
            name: "Sentinel-2 Pipeline".to_string(),
            description: Some("ignored on round trip".to_string()),
            version: "2.3.1".to_string(),
            dag,
        };

        let exported =
            PrefectIntegration::export_workflow(&original).expect("export_workflow failed");
        let imported =
            PrefectIntegration::import_workflow(&exported).expect("import_workflow failed");

        assert_eq!(imported.id, original.id);
        assert_eq!(imported.dag.task_count(), original.dag.task_count());
    }

    #[test]
    fn test_export_import_round_trip_empty_dag() {
        let original = WorkflowDefinition {
            id: "empty-workflow".to_string(),
            name: "Empty Workflow".to_string(),
            description: None,
            version: "1.0.0".to_string(),
            dag: WorkflowDag::new(),
        };

        let exported =
            PrefectIntegration::export_workflow(&original).expect("export_workflow failed");
        let imported =
            PrefectIntegration::import_workflow(&exported).expect("import_workflow failed");

        assert_eq!(imported.id, original.id);
        assert_eq!(imported.dag.task_count(), original.dag.task_count());
    }
}