trino_rust_client/spooling/
segment.rs

1use serde::Deserialize;
2use serde::Serialize;
3use std::collections::HashMap;
4
5// Data attributes for a segment
6#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
7#[serde(rename_all = "camelCase")]
8pub struct DataAttributes {
9    #[serde(flatten)]
10    attributes: HashMap<String, serde_json::Value>,
11}
12
13// Data attributes for a segment
14impl DataAttributes {
15    // Get the row offset for a segment
16    pub fn row_offset(&self) -> Option<u64> {
17        self.attributes.get("rowOffset")?.as_u64()
18    }
19
20    // Get the number of rows for a segment
21    pub fn rows_count(&self) -> Option<u64> {
22        self.attributes.get("rowsCount")?.as_u64()
23    }
24
25    // Get the size of a segment
26    pub fn segment_size(&self) -> Option<u64> {
27        self.attributes.get("segmentSize")?.as_u64()
28    }
29}
30
31// Segment is a part of a query result when using the spooling protocol
32#[derive(Debug, PartialEq, Deserialize, Serialize)]
33#[serde(rename_all = "camelCase", untagged)]
34pub enum Segment {
35    // Inlined segment
36    Inlined {
37        #[serde(rename = "type")]
38        segment_type: String,
39        data: String,
40        metadata: DataAttributes,
41    },
42    // Spooled segment
43    Spooled {
44        #[serde(rename = "type")]
45        segment_type: String,
46        uri: String,
47        #[serde(rename = "ackUri")]
48        #[serde(skip_serializing_if = "Option::is_none")]
49        ack_uri: Option<String>,
50        #[serde(skip_serializing_if = "Option::is_none")]
51        headers: Option<HashMap<String, Vec<String>>>,
52        metadata: DataAttributes,
53    },
54}
55
56#[cfg(test)]
57mod tests {
58    use super::*;
59
60    #[test]
61    fn test_deserialize_inlined_segment() {
62        let json = r#"{
63            "type": "inline",
64            "data": "SGVsbG8gV29ybGQ=",
65            "metadata": {
66                "rowOffset": 0,
67                "rowsCount": 1,
68                "segmentSize": 1024
69            }
70        }"#;
71        let segment: Segment = serde_json::from_str(json).unwrap();
72        match segment {
73            Segment::Inlined {
74                segment_type,
75                data,
76                metadata,
77            } => {
78                assert_eq!(segment_type, "inline");
79                assert_eq!(data, "SGVsbG8gV29ybGQ=");
80                assert_eq!(metadata.row_offset(), Some(0));
81                assert_eq!(metadata.rows_count(), Some(1));
82                assert_eq!(metadata.segment_size(), Some(1024));
83            }
84            _ => panic!("Expected Inlined segment"),
85        }
86    }
87
88    #[test]
89    fn test_deserialize_spooled_segment_minimal() {
90        let json = r#"{
91            "type": "spooled",
92            "uri": "http://minio:9000/bucket/segment.json?signature=abc123",
93            "metadata": {
94                "rowOffset": 0,
95                "rowsCount": 1000,
96                "segmentSize": 1048576
97            }
98        }"#;
99        let segment: Segment = serde_json::from_str(json).unwrap();
100        match segment {
101            Segment::Spooled {
102                segment_type,
103                uri,
104                ack_uri,
105                headers,
106                metadata,
107            } => {
108                assert_eq!(segment_type, "spooled");
109                assert_eq!(
110                    uri,
111                    "http://minio:9000/bucket/segment.json?signature=abc123"
112                );
113                assert_eq!(ack_uri, None);
114                assert_eq!(headers, None);
115                assert_eq!(metadata.row_offset(), Some(0));
116                assert_eq!(metadata.rows_count(), Some(1000));
117                assert_eq!(metadata.segment_size(), Some(1048576));
118            }
119            _ => panic!("Expected Spooled segment"),
120        }
121    }
122
123    #[test]
124    fn test_deserialize_spooled_segment_with_ack() {
125        let json = r#"{
126            "type": "spooled",
127            "uri": "http://minio:9000/bucket/segment.json",
128            "ackUri": "http://minio:9000/bucket/segment.ack",
129            "headers": {
130                "X-Custom": ["value1"]
131            },
132            "metadata": {
133                "rowOffset": 0,
134                "rowsCount": 100
135            }
136        }"#;
137        let segment: Segment = serde_json::from_str(json).unwrap();
138        match segment {
139            Segment::Spooled {
140                segment_type,
141                uri,
142                ack_uri,
143                headers,
144                metadata,
145            } => {
146                assert_eq!(segment_type, "spooled");
147                assert_eq!(uri, "http://minio:9000/bucket/segment.json");
148                assert_eq!(
149                    ack_uri,
150                    Some("http://minio:9000/bucket/segment.ack".to_string())
151                );
152                assert!(headers.is_some());
153                assert_eq!(metadata.row_offset(), Some(0));
154                assert_eq!(metadata.rows_count(), Some(100));
155            }
156            _ => panic!("Expected Spooled segment"),
157        }
158    }
159
160    #[test]
161    fn test_deserialize_inlined_segment_minimal() {
162        let json = r#"{
163            "type": "inline",
164            "data": "YWJjZGVmZw==",
165            "metadata": {}
166        }"#;
167        let segment: Segment = serde_json::from_str(json).unwrap();
168        match segment {
169            Segment::Inlined {
170                segment_type,
171                data,
172                metadata,
173            } => {
174                assert_eq!(segment_type, "inline");
175                assert_eq!(data, "YWJjZGVmZw==");
176                assert_eq!(metadata.row_offset(), None);
177                assert_eq!(metadata.rows_count(), None);
178                assert_eq!(metadata.segment_size(), None);
179            }
180            _ => panic!("Expected Inlined segment"),
181        }
182    }
183
184    #[test]
185    fn test_data_attributes_row_offset() {
186        let json = r#"{
187            "rowOffset": 42
188        }"#;
189        let data_attributes: DataAttributes = serde_json::from_str(json).unwrap();
190        assert_eq!(data_attributes.row_offset(), Some(42));
191        assert_eq!(data_attributes.rows_count(), None);
192        assert_eq!(data_attributes.segment_size(), None);
193    }
194
195    #[test]
196    fn test_data_attributes_rows_count() {
197        let json = r#"{
198            "rowsCount": 100
199        }"#;
200        let data_attributes: DataAttributes = serde_json::from_str(json).unwrap();
201        assert_eq!(data_attributes.row_offset(), None);
202        assert_eq!(data_attributes.rows_count(), Some(100));
203        assert_eq!(data_attributes.segment_size(), None);
204    }
205
206    #[test]
207    fn test_data_attributes_segment_size() {
208        let json = r#"{
209            "segmentSize": 4096
210        }"#;
211        let data_attributes: DataAttributes = serde_json::from_str(json).unwrap();
212        assert_eq!(data_attributes.row_offset(), None);
213        assert_eq!(data_attributes.rows_count(), None);
214        assert_eq!(data_attributes.segment_size(), Some(4096));
215    }
216
217    #[test]
218    fn test_data_attributes_all_fields() {
219        let json = r#"{
220            "rowOffset": 0,
221            "rowsCount": 10,
222            "segmentSize": 512
223        }"#;
224        let data_attributes: DataAttributes = serde_json::from_str(json).unwrap();
225        assert_eq!(data_attributes.row_offset(), Some(0));
226        assert_eq!(data_attributes.rows_count(), Some(10));
227        assert_eq!(data_attributes.segment_size(), Some(512));
228    }
229
230    #[test]
231    fn test_spooled_segment_with_multiple_header_values() {
232        let json = r#"{
233            "type": "spooled",
234            "uri": "http://storage/segment.json",
235            "ackUri": "http://storage/segment.ack",
236            "headers": {
237                "Authorization": ["Bearer token123"],
238                "X-Custom": ["value1", "value2"]
239            },
240            "metadata": {
241                "rowOffset": 100,
242                "rowsCount": 50
243            }
244        }"#;
245        let segment: Segment = serde_json::from_str(json).unwrap();
246
247        match segment {
248            Segment::Spooled {
249                uri,
250                ack_uri,
251                headers,
252                metadata,
253                ..
254            } => {
255                assert_eq!(uri, "http://storage/segment.json");
256                assert_eq!(ack_uri, Some("http://storage/segment.ack".to_string()));
257                assert!(headers.is_some());
258                let headers_map = headers.unwrap();
259                assert_eq!(
260                    headers_map.get("Authorization"),
261                    Some(&vec!["Bearer token123".to_string()])
262                );
263                assert_eq!(
264                    headers_map.get("X-Custom"),
265                    Some(&vec!["value1".to_string(), "value2".to_string()])
266                );
267                assert_eq!(metadata.row_offset(), Some(100));
268                assert_eq!(metadata.rows_count(), Some(50));
269            }
270            _ => panic!("Expected Spooled segment"),
271        }
272    }
273}
trino_rust_client/spooling/segment.rs

trino_rust_client/spooling/
segment.rs
