1use crate::error::OpError;
2use crate::op::Op;
3use crate::prelude::*;
4use crate::{DryContext, OpMetadata, WetContext};
5use async_trait::async_trait;
6use jsonschema::{Draft, JSONSchema};
7
8pub struct ValidatingWrapper<T> {
9 wrapped_op: Box<dyn Op<T>>,
10 validate_input: bool,
11 validate_output: bool,
12}
13
14impl<T> ValidatingWrapper<T>
15where
16 T: Send + Sync + 'static + serde::Serialize,
17{
18 pub fn new(op: Box<dyn Op<T>>) -> Self {
20 Self {
21 wrapped_op: op,
22 validate_input: true,
23 validate_output: true,
24 }
25 }
26
27 pub fn input_only(op: Box<dyn Op<T>>) -> Self {
29 Self {
30 wrapped_op: op,
31 validate_input: true,
32 validate_output: false,
33 }
34 }
35
36 pub fn output_only(op: Box<dyn Op<T>>) -> Self {
38 Self {
39 wrapped_op: op,
40 validate_input: false,
41 validate_output: true,
42 }
43 }
44
45 fn validate_input_schema(&self, dry: &DryContext, metadata: &OpMetadata) -> OpResult<()> {
47 if !self.validate_input {
48 return Ok(());
49 }
50
51 if let Some(ref schema) = metadata.input_schema {
52 let compiled = JSONSchema::options()
54 .with_draft(Draft::Draft7)
55 .compile(schema)
56 .map_err(|e| {
57 OpError::Context(format!("Invalid input schema for {}: {}", metadata.name, e))
58 })?;
59
60 let context_json = serde_json::json!(dry.values());
62
63 let validation_result = compiled.validate(&context_json);
65 if let Err(errors) = validation_result {
66 let error_messages: Vec<String> = errors
67 .map(|e| format!("{}: {}", e.instance_path, e))
68 .collect();
69
70 return Err(OpError::Context(format!(
71 "Input validation failed for {}: {}",
72 metadata.name,
73 error_messages.join(", ")
74 )));
75 }
76 }
77 Ok(())
78 }
79
80 fn validate_references_schema(&self, wet: &WetContext, metadata: &OpMetadata) -> OpResult<()> {
82 if let Some(ref schema) = metadata.reference_schema {
86 if let Some(required_refs) = schema.get("required") {
89 if let Some(required_array) = required_refs.as_array() {
90 for required_ref in required_array {
91 if let Some(ref_name) = required_ref.as_str() {
92 if !wet.contains(ref_name) {
93 return Err(OpError::Context(format!(
94 "Required reference '{}' not found in WetContext for op '{}'",
95 ref_name, metadata.name
96 )));
97 }
98 }
99 }
100 }
101 }
102
103 if let Some(properties) = schema.get("properties") {
105 if let Some(props_obj) = properties.as_object() {
106 for (ref_name, ref_schema) in props_obj {
107 if wet.contains(ref_name) {
108 if let Some(ref_type) = ref_schema.get("type") {
111 if let Some(type_str) = ref_type.as_str() {
112 tracing::debug!(
114 "Reference '{}' exists but type validation ('{}') is skipped for runtime safety",
115 ref_name, type_str
116 );
117 }
118 }
119 }
120 }
121 }
122 }
123 }
124 Ok(())
125 }
126
127 fn validate_output_schema(&self, output: &T, metadata: &OpMetadata) -> OpResult<()> {
129 if !self.validate_output {
130 return Ok(());
131 }
132
133 if let Some(ref schema) = metadata.output_schema {
134 let compiled = JSONSchema::options()
136 .with_draft(Draft::Draft7)
137 .compile(schema)
138 .map_err(|e| {
139 OpError::Context(format!(
140 "Invalid output schema for {}: {}",
141 metadata.name, e
142 ))
143 })?;
144
145 let output_json = serde_json::to_value(output).map_err(|e| {
147 OpError::Context(format!("Failed to serialize output for validation: {}", e))
148 })?;
149
150 let validation_result = compiled.validate(&output_json);
152 if let Err(errors) = validation_result {
153 let error_messages: Vec<String> = errors
154 .map(|e| format!("{}: {}", e.instance_path, e))
155 .collect();
156
157 return Err(OpError::Context(format!(
158 "Output validation failed for {}: {}",
159 metadata.name,
160 error_messages.join(", ")
161 )));
162 }
163 }
164 Ok(())
165 }
166}
167
168#[async_trait]
169impl<T> Op<T> for ValidatingWrapper<T>
170where
171 T: Send + Sync + 'static + serde::Serialize,
172{
173 async fn perform(&self, dry: &mut DryContext, wet: &mut WetContext) -> OpResult<T> {
174 let metadata = self.wrapped_op.metadata();
175
176 self.validate_input_schema(dry, &metadata)?;
178
179 self.validate_references_schema(wet, &metadata)?;
181
182 let result = self.wrapped_op.perform(dry, wet).await?;
184
185 self.validate_output_schema(&result, &metadata)?;
187
188 Ok(result)
189 }
190
191 fn metadata(&self) -> OpMetadata {
192 self.wrapped_op.metadata()
194 }
195}
196
197#[cfg(test)]
198mod tests {
199 use super::*;
200 use serde_json::json;
201
202 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
203 struct TestOutput {
204 value: i32,
205 }
206
207 struct ValidatedOp;
208
209 #[async_trait]
210 impl Op<TestOutput> for ValidatedOp {
211 async fn perform(
212 &self,
213 dry: &mut DryContext,
214 _wet: &mut WetContext,
215 ) -> OpResult<TestOutput> {
216 let value = dry.get_required::<i32>("value")?;
217 Ok(TestOutput { value })
218 }
219
220 fn metadata(&self) -> OpMetadata {
221 OpMetadata::builder("ValidatedOp")
222 .description("Op with schema validation")
223 .input_schema(json!({
224 "type": "object",
225 "properties": {
226 "value": { "type": "integer", "minimum": 0, "maximum": 100 }
227 },
228 "required": ["value"]
229 }))
230 .output_schema(json!({
231 "type": "object",
232 "properties": {
233 "value": { "type": "integer" }
234 },
235 "required": ["value"]
236 }))
237 .build()
238 }
239 }
240
241 #[tokio::test]
243 async fn test0038_valid_input_output() {
244 let validator = ValidatingWrapper::new(Box::new(ValidatedOp));
245
246 let mut dry = DryContext::new();
247 dry.insert("value", 42);
248 let mut wet = WetContext::new();
249
250 let result = validator.perform(&mut dry, &mut wet).await;
251 assert!(result.is_ok());
252 assert_eq!(result.unwrap().value, 42);
253 }
254
255 #[tokio::test]
257 async fn test0039_invalid_input_missing_required() {
258 let validator = ValidatingWrapper::new(Box::new(ValidatedOp));
259
260 let mut dry = DryContext::new();
261 let mut wet = WetContext::new();
263
264 let result = validator.perform(&mut dry, &mut wet).await;
265 assert!(result.is_err());
266 let err = result.unwrap_err();
267 match err {
268 OpError::Context(msg) => assert!(msg.contains("Input validation failed")),
269 _ => panic!("Expected Context error"),
270 }
271 }
272
273 #[tokio::test]
275 async fn test0040_invalid_input_out_of_range() {
276 let validator = ValidatingWrapper::new(Box::new(ValidatedOp));
277
278 let mut dry = DryContext::new();
279 dry.insert("value", 150); let mut wet = WetContext::new();
281
282 let result = validator.perform(&mut dry, &mut wet).await;
283 assert!(result.is_err());
284 let err = result.unwrap_err();
285 match err {
286 OpError::Context(msg) => assert!(msg.contains("maximum")),
287 _ => panic!("Expected Context error"),
288 }
289 }
290
291 #[tokio::test]
293 async fn test0041_input_only_validation() {
294 struct NoOutputSchemaOp;
295
296 #[async_trait]
297 impl Op<i32> for NoOutputSchemaOp {
298 async fn perform(&self, dry: &mut DryContext, _wet: &mut WetContext) -> OpResult<i32> {
299 dry.get_required::<i32>("value")
300 }
301
302 fn metadata(&self) -> OpMetadata {
303 OpMetadata::builder("NoOutputSchemaOp")
304 .input_schema(json!({
305 "type": "object",
306 "properties": {
307 "value": { "type": "integer" }
308 },
309 "required": ["value"]
310 }))
311 .build()
312 }
313 }
314
315 let validator = ValidatingWrapper::input_only(Box::new(NoOutputSchemaOp));
316
317 let mut dry = DryContext::new();
318 dry.insert("value", 42);
319 let mut wet = WetContext::new();
320
321 let result = validator.perform(&mut dry, &mut wet).await;
322 assert!(result.is_ok());
323 assert_eq!(result.unwrap(), 42);
324 }
325
326 #[tokio::test]
328 async fn test0042_output_only_validation() {
329 struct NoInputSchemaOp;
330
331 #[async_trait]
332 impl Op<TestOutput> for NoInputSchemaOp {
333 async fn perform(
334 &self,
335 _dry: &mut DryContext,
336 _wet: &mut WetContext,
337 ) -> OpResult<TestOutput> {
338 Ok(TestOutput { value: 99 })
339 }
340
341 fn metadata(&self) -> OpMetadata {
342 OpMetadata::builder("NoInputSchemaOp")
343 .output_schema(json!({
344 "type": "object",
345 "properties": {
346 "value": { "type": "integer", "maximum": 100 }
347 },
348 "required": ["value"]
349 }))
350 .build()
351 }
352 }
353
354 let validator = ValidatingWrapper::output_only(Box::new(NoInputSchemaOp));
355
356 let mut dry = DryContext::new();
357 let mut wet = WetContext::new();
358
359 let result = validator.perform(&mut dry, &mut wet).await;
360 assert!(result.is_ok());
361 assert_eq!(result.unwrap().value, 99);
362 }
363
364 #[tokio::test]
366 async fn test0043_no_schema_validation() {
367 struct NoSchemaOp;
368
369 #[async_trait]
370 impl Op<i32> for NoSchemaOp {
371 async fn perform(&self, _dry: &mut DryContext, _wet: &mut WetContext) -> OpResult<i32> {
372 Ok(123)
373 }
374
375 fn metadata(&self) -> OpMetadata {
376 OpMetadata::builder("NoSchemaOp").build()
377 }
378 }
379
380 let validator = ValidatingWrapper::new(Box::new(NoSchemaOp));
381
382 let mut dry = DryContext::new();
383 let mut wet = WetContext::new();
384
385 let result = validator.perform(&mut dry, &mut wet).await;
387 assert!(result.is_ok());
388 assert_eq!(result.unwrap(), 123);
389 }
390
391 #[tokio::test]
393 async fn test0044_metadata_transparency() {
394 let validator = ValidatingWrapper::new(Box::new(ValidatedOp));
395 let metadata = validator.metadata();
396
397 assert_eq!(metadata.name, "ValidatedOp");
398 assert_eq!(
399 metadata.description,
400 Some("Op with schema validation".to_string())
401 );
402 assert!(metadata.input_schema.is_some());
403 assert!(metadata.output_schema.is_some());
404 }
405
406 #[tokio::test]
408 async fn test0045_reference_validation() {
409 struct ServiceRequiringOp;
410
411 #[async_trait]
412 impl Op<String> for ServiceRequiringOp {
413 async fn perform(
414 &self,
415 _dry: &mut DryContext,
416 wet: &mut WetContext,
417 ) -> OpResult<String> {
418 let service = wet.get_required::<String>("database")?;
419 Ok(format!("Used service: {}", service))
420 }
421
422 fn metadata(&self) -> OpMetadata {
423 OpMetadata::builder("ServiceRequiringOp")
424 .reference_schema(json!({
425 "type": "object",
426 "required": ["database", "cache"],
427 "properties": {
428 "database": { "type": "string" },
429 "cache": { "type": "string" }
430 }
431 }))
432 .build()
433 }
434 }
435
436 let validator = ValidatingWrapper::new(Box::new(ServiceRequiringOp));
437
438 let mut dry = DryContext::new();
439 let mut wet = WetContext::new();
440
441 let result = validator.perform(&mut dry, &mut wet).await;
443 assert!(result.is_err());
444 let err = result.unwrap_err();
445 match err {
446 OpError::Context(msg) => {
447 assert!(msg.contains("Required reference 'database' not found"))
448 }
449 _ => panic!("Expected Context error"),
450 }
451
452 wet.insert_ref("database", "postgresql".to_string());
454 let result = validator.perform(&mut dry, &mut wet).await;
455 assert!(result.is_err());
456 let err = result.unwrap_err();
457 match err {
458 OpError::Context(msg) => assert!(msg.contains("Required reference 'cache' not found")),
459 _ => panic!("Expected Context error"),
460 }
461
462 wet.insert_ref("cache", "redis".to_string());
464 let result = validator.perform(&mut dry, &mut wet).await;
465 assert!(result.is_ok());
466 assert_eq!(result.unwrap(), "Used service: postgresql");
467 }
468
469 #[tokio::test]
471 async fn test0046_no_reference_schema() {
472 struct NoRefSchemaOp;
473
474 #[async_trait]
475 impl Op<i32> for NoRefSchemaOp {
476 async fn perform(&self, _dry: &mut DryContext, _wet: &mut WetContext) -> OpResult<i32> {
477 Ok(456)
478 }
479
480 fn metadata(&self) -> OpMetadata {
481 OpMetadata::builder("NoRefSchemaOp").build()
482 }
483 }
484
485 let validator = ValidatingWrapper::new(Box::new(NoRefSchemaOp));
486
487 let mut dry = DryContext::new();
488 let mut wet = WetContext::new();
489
490 let result = validator.perform(&mut dry, &mut wet).await;
492 assert!(result.is_ok());
493 assert_eq!(result.unwrap(), 456);
494 }
495
496 #[tokio::test]
498 async fn test0112_output_only_still_validates_references() {
499 struct RefRequiringOp;
500
501 #[async_trait]
502 impl Op<i32> for RefRequiringOp {
503 async fn perform(&self, _dry: &mut DryContext, _wet: &mut WetContext) -> OpResult<i32> {
504 Ok(42)
505 }
506 fn metadata(&self) -> OpMetadata {
507 OpMetadata::builder("RefRequiringOp")
508 .reference_schema(json!({
509 "type": "object",
510 "required": ["database"],
511 "properties": {
512 "database": { "type": "string" }
513 }
514 }))
515 .build()
516 }
517 }
518
519 let validator = ValidatingWrapper::output_only(Box::new(RefRequiringOp));
520 let mut dry = DryContext::new();
521 let mut wet = WetContext::new();
522
523 let result = validator.perform(&mut dry, &mut wet).await;
525 assert!(
526 result.is_err(),
527 "output_only must still validate references"
528 );
529 match result.unwrap_err() {
530 OpError::Context(msg) => assert!(msg.contains("database")),
531 e => panic!("expected Context error, got {:?}", e),
532 }
533
534 wet.insert_ref("database", "postgres://localhost".to_string());
536 let result = validator.perform(&mut dry, &mut wet).await;
537 assert!(
538 result.is_ok(),
539 "should succeed when required reference is present"
540 );
541 }
542}