Skip to main content

stormchaser_engine/handler/integrations/
jinja.rs

1use crate::handler::{fetch_outputs, fetch_run_context, fetch_step_instance};
2use anyhow::Result;
3use chrono::Utc;
4use serde_json::Value;
5use sqlx::PgPool;
6use stormchaser_model::dsl::JinjaRenderSpec;
7use stormchaser_model::events::{
8    EventSource, EventType, SchemaVersion, StepCompletedEvent, StepEventType,
9};
10use stormchaser_model::nats::NatsSubject;
11use stormchaser_model::RunId;
12use stormchaser_model::StepInstanceId;
13use tracing::info;
14
15/// Handle jinja render.
16pub async fn handle_jinja_render(
17    run_id: RunId,
18    step_id: StepInstanceId,
19    spec: Value,
20    pool: PgPool,
21    nats_client: async_nats::Client,
22) -> Result<()> {
23    let spec: JinjaRenderSpec = serde_json::from_value(spec)?;
24
25    info!("Rendering Jinja template for run {}", run_id);
26
27    // 1. Mark as Running
28    let instance = fetch_step_instance(step_id, &pool).await?;
29    let machine =
30        crate::step_machine::StepMachine::<crate::step_machine::state::Pending>::from_instance(
31            instance,
32        );
33    let _ = machine
34        .start("jinja".to_string(), &mut *pool.acquire().await?)
35        .await?;
36
37    // 2. Prepare Context for Template Rendering
38    let run_context = fetch_run_context(run_id, &pool).await?;
39    let outputs = fetch_outputs(run_id, &pool).await?;
40
41    let template_ctx = prepare_template_context(&spec, run_context, outputs, run_id);
42
43    // 3. Render Template
44    let rendered = render_template(&spec.template, &template_ctx)?;
45
46    // 4. Save Output and Complete Step
47    save_output_and_complete(run_id, step_id, &spec, rendered, pool, nats_client).await
48}
49
50use stormchaser_model::dsl;
51
52fn prepare_template_context(
53    spec: &dsl::JinjaRenderSpec,
54    run_context: crate::handler::RunContext,
55    outputs: Value,
56    run_id: RunId,
57) -> Value {
58    let mut template_ctx = serde_json::json!({
59        "inputs": run_context.inputs,
60        "steps": outputs,
61        "run": {
62            "id": run_id.to_string(),
63        }
64    });
65
66    // Merge custom context if provided
67    if let Some(custom_ctx) = &spec.context {
68        if let Some(obj) = template_ctx.as_object_mut() {
69            if let Some(custom_obj) = custom_ctx.as_object() {
70                for (k, v) in custom_obj {
71                    obj.insert(k.clone(), v.clone());
72                }
73            }
74        }
75    }
76    template_ctx
77}
78
79fn render_template(template: &str, context: &Value) -> Result<String> {
80    use minijinja::Environment;
81    let env = Environment::new();
82    env.render_str(template, context)
83        .map_err(|e| anyhow::anyhow!("Failed to render Jinja template: {:?}", e))
84}
85
86async fn save_output_and_complete(
87    run_id: RunId,
88    step_id: StepInstanceId,
89    spec: &dsl::JinjaRenderSpec,
90    rendered: String,
91    pool: PgPool,
92    nats_client: async_nats::Client,
93) -> Result<()> {
94    let mut tx = pool.begin().await?;
95    let instance = fetch_step_instance(step_id, &mut *tx).await?;
96    let machine =
97        crate::step_machine::StepMachine::<crate::step_machine::state::Running>::from_instance(
98            instance,
99        );
100    let _ = machine.succeed(&mut *tx).await?;
101
102    let output_key = spec
103        .output_key
104        .clone()
105        .unwrap_or_else(|| "result".to_string());
106    crate::db::upsert_step_output(
107        &mut *tx,
108        step_id,
109        &output_key,
110        &Value::String(rendered.clone()),
111    )
112    .await?;
113
114    tx.commit().await?;
115
116    let mut outputs_map = std::collections::HashMap::new();
117    outputs_map.insert(output_key.clone(), serde_json::json!(rendered));
118
119    let event = StepCompletedEvent {
120        run_id,
121        step_id,
122        event_type: EventType::Step(StepEventType::Completed),
123        outputs: Some(outputs_map),
124        exit_code: Some(0),
125        runner_id: None,
126        storage_hashes: None,
127        artifacts: None,
128        test_reports: None,
129        timestamp: Utc::now(),
130    };
131    let js = async_nats::jetstream::new(nats_client);
132    stormchaser_model::nats::publish_cloudevent(
133        &js,
134        NatsSubject::StepCompleted,
135        EventType::Step(StepEventType::Completed),
136        EventSource::System,
137        serde_json::to_value(event).unwrap(),
138        Some(SchemaVersion::new("1.0".to_string())),
139        None,
140    )
141    .await?;
142
143    Ok(())
144}
145
146#[cfg(test)]
147mod tests {
148    use super::*;
149    use serde_json::json;
150    use stormchaser_model::workflow::RunContext;
151
152    #[test]
153    fn test_prepare_template_context() {
154        let run_id = RunId::new_v4();
155        let spec = JinjaRenderSpec {
156            template: "Hello".to_string(),
157            context: Some(json!({"extra": "value"})),
158            output_key: None,
159        };
160        let run_context = RunContext {
161            run_id,
162            dsl_version: "1.0".to_string(),
163            workflow_definition: json!({}),
164            source_code: "".to_string(),
165            inputs: json!({"name": "User"}),
166            secrets: json!({}),
167            sensitive_values: vec![],
168        };
169        let outputs = json!({"prev_step": {"outputs": {"res": "done"}}});
170
171        let ctx = prepare_template_context(&spec, run_context, outputs, run_id);
172
173        assert_eq!(ctx["inputs"]["name"], "User");
174        assert_eq!(ctx["steps"]["prev_step"]["outputs"]["res"], "done");
175        assert_eq!(ctx["run"]["id"], run_id.to_string());
176        assert_eq!(ctx["extra"], "value");
177    }
178
179    #[test]
180    fn test_render_template_happy() {
181        let ctx = json!({"name": "World"});
182        let rendered = render_template("Hello {{ name }}!", &ctx).unwrap();
183        assert_eq!(rendered, "Hello World!");
184    }
185
186    #[test]
187    fn test_render_template_error() {
188        let ctx = json!({});
189        let result = render_template("Hello {{ name", &ctx);
190        assert!(result.is_err());
191    }
192}