arkflow-plugin 0.1.0

High-performance Rust flow processing engine
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

//! Drop output component
//!
//! This component discards all messages without performing any operations.
//! It's useful for testing or when you want to intentionally discard data.

use arkflow_core::output::{register_output_builder, Output, OutputBuilder};
use arkflow_core::{Error, MessageBatch};
use async_trait::async_trait;
use std::sync::Arc;

/// Drop output component that discards all messages
///
/// This component implements the `Output` trait but doesn't perform any actual
/// output operations. All messages sent to this output are simply discarded.
pub struct DropOutput;

#[async_trait]
impl Output for DropOutput {
    async fn connect(&self) -> Result<(), Error> {
        Ok(())
    }

    async fn write(&self, _: &MessageBatch) -> Result<(), Error> {
        Ok(())
    }

    async fn close(&self) -> Result<(), Error> {
        Ok(())
    }
}

pub(crate) struct DropOutputBuilder;
impl OutputBuilder for DropOutputBuilder {
    fn build(&self, _: &Option<serde_json::Value>) -> Result<Arc<dyn Output>, Error> {
        Ok(Arc::new(DropOutput))
    }
}

pub fn init() {
    register_output_builder("drop", Arc::new(DropOutputBuilder));
}

#[cfg(test)]
mod tests {
    use crate::output::drop::DropOutput;

    use arkflow_core::output::Output;
    use arkflow_core::MessageBatch;
    use datafusion::arrow::array::{Int32Array, StringArray};
    use datafusion::arrow::datatypes::{DataType, Field, Schema};
    use datafusion::arrow::record_batch::RecordBatch;
    use std::sync::Arc;

    #[tokio::test]
    async fn test_drop_output_connect() {
        // Create a DropOutput instance
        let drop_output = DropOutput;

        // Test connect method
        let result = drop_output.connect().await;
        assert!(result.is_ok(), "connect() should return Ok(())");
    }

    #[tokio::test]
    async fn test_drop_output_write_binary() {
        // Create a DropOutput instance
        let drop_output = DropOutput;

        // Create a binary message batch
        let binary_data = vec![b"test message".to_vec()];
        let message_batch = MessageBatch::new_binary(binary_data);

        // Test write method with binary data
        let result = drop_output.write(&message_batch).await;
        assert!(
            result.is_ok(),
            "write() should return Ok(()) for binary data"
        );
    }

    #[tokio::test]
    async fn test_drop_output_write_arrow() {
        // Create a DropOutput instance
        let drop_output = DropOutput;

        // Create an Arrow message batch
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int32, false),
            Field::new("name", DataType::Utf8, false),
        ]));

        let id_array = Int32Array::from(vec![1, 2, 3]);
        let name_array = StringArray::from(vec!["Alice", "Bob", "Charlie"]);

        let batch =
            RecordBatch::try_new(schema, vec![Arc::new(id_array), Arc::new(name_array)]).unwrap();

        let message_batch = MessageBatch::new_arrow(batch);

        // Test write method with Arrow data
        let result = drop_output.write(&message_batch).await;
        assert!(
            result.is_ok(),
            "write() should return Ok(()) for Arrow data"
        );
    }

    #[tokio::test]
    async fn test_drop_output_close() {
        // Create a DropOutput instance
        let drop_output = DropOutput;

        // Test close method
        let result = drop_output.close().await;
        assert!(result.is_ok(), "close() should return Ok(())");
    }

    #[tokio::test]
    async fn test_drop_output_full_lifecycle() {
        // Create a DropOutput instance
        let drop_output = DropOutput;

        // Test the full lifecycle: connect -> write -> close
        let connect_result = drop_output.connect().await;
        assert!(connect_result.is_ok(), "connect() should return Ok(())");

        // Create a binary message batch
        let binary_data = vec![b"test message".to_vec()];
        let message_batch = MessageBatch::new_binary(binary_data);

        // Write multiple messages
        for _ in 0..5 {
            let write_result = drop_output.write(&message_batch).await;
            assert!(write_result.is_ok(), "write() should return Ok(())");
        }

        let close_result = drop_output.close().await;
        assert!(close_result.is_ok(), "close() should return Ok(())");
    }
}