alopex_cli/streaming/
writer.rs

1//! StreamingWriter - Streaming output controller
2//!
3//! Manages streaming output with buffer limits for non-streaming formats.
4
5use std::io::Write;
6
7use serde::Serialize;
8
9use crate::cli::RoutingReportFormat;
10use crate::error::{CliError, Result};
11use crate::models::{Column, Row};
12use crate::output::formatter::Formatter;
13/// Default buffer limit for non-streaming formats (table).
14pub const DEFAULT_BUFFER_LIMIT: usize = 10 * 1024 * 1024;
15
16/// A routing/status report kept separate from row output. Fields unavailable
17/// before a server coordinator responds remain null rather than being filled
18/// with local guesses.
19#[derive(Debug, Clone, Serialize)]
20pub struct DistributedReadRoutingReport {
21    pub schema_version: u32,
22    pub requested_mode: String,
23    pub effective_mode: Option<String>,
24    pub decision: String,
25    pub ranges: Vec<String>,
26    pub freshness: Option<String>,
27    pub retry_count: u32,
28    pub failover_count: u32,
29    pub outcome: String,
30    #[serde(skip_serializing_if = "Option::is_none")]
31    pub reason: Option<String>,
32}
33
34impl DistributedReadRoutingReport {
35    pub fn new(
36        requested_mode: impl Into<String>,
37        effective_mode: Option<String>,
38        decision: impl Into<String>,
39        outcome: impl Into<String>,
40        reason: Option<String>,
41    ) -> Self {
42        Self {
43            schema_version: 1,
44            requested_mode: requested_mode.into(),
45            effective_mode,
46            decision: decision.into(),
47            ranges: Vec::new(),
48            freshness: None,
49            retry_count: 0,
50            failover_count: 0,
51            outcome: outcome.into(),
52            reason,
53        }
54    }
55}
56
57/// Write the report to a caller-provided stderr writer. No `StreamingWriter`
58/// or SQL row formatter is involved, so table/JSON/CSV/TSV/JSONL stdout stays
59/// byte-for-byte independent from routing diagnostics.
60pub fn write_distributed_read_routing_report<W: Write>(
61    writer: &mut W,
62    format: RoutingReportFormat,
63    report: &DistributedReadRoutingReport,
64) -> Result<()> {
65    match format {
66        RoutingReportFormat::Json => {
67            serde_json::to_writer(&mut *writer, report)?;
68            writeln!(writer)?;
69        }
70        RoutingReportFormat::Human => {
71            writeln!(writer, "distributed read routing report")?;
72            writeln!(writer, "  requested_mode: {}", report.requested_mode)?;
73            if let Some(mode) = &report.effective_mode {
74                writeln!(writer, "  effective_mode: {mode}")?;
75            }
76            writeln!(writer, "  decision: {}", report.decision)?;
77            writeln!(writer, "  ranges: {}", report.ranges.join(","))?;
78            if let Some(freshness) = &report.freshness {
79                writeln!(writer, "  freshness: {freshness}")?;
80            }
81            writeln!(writer, "  retry_count: {}", report.retry_count)?;
82            writeln!(writer, "  failover_count: {}", report.failover_count)?;
83            writeln!(writer, "  outcome: {}", report.outcome)?;
84            if let Some(reason) = &report.reason {
85                writeln!(writer, "  reason: {reason}")?;
86            }
87        }
88    }
89    writer.flush()?;
90    Ok(())
91}
92
93/// Status returned by write operations.
94#[derive(Debug, Clone, Copy, PartialEq, Eq)]
95pub enum WriteStatus {
96    /// Row was written successfully, continue writing.
97    Continue,
98    /// Limit reached, no more rows will be written.
99    LimitReached,
100}
101
102/// Streaming writer for output.
103///
104/// Controls streaming output with the following behaviors:
105/// - **Streaming formats** (json, jsonl, csv, tsv): Output rows immediately.
106/// - **Non-streaming formats** (table): Buffer rows up to `buffer_limit`.
107///
108/// # Output Boundary
109///
110/// - **Output started**: When `output_started == true` (header has been written).
111/// - **Buffer overflow**: Returns an error prompting `--output json|csv|tsv` or `--limit`.
112pub struct StreamingWriter<W> {
113    /// The underlying writer.
114    writer: W,
115    /// The formatter to use for output.
116    formatter: Box<dyn Formatter>,
117    /// Column definitions (schema).
118    columns: Vec<Column>,
119    /// Optional row limit.
120    limit: Option<usize>,
121    /// Buffer limit for non-streaming formats.
122    buffer_limit: usize,
123    /// Buffer for non-streaming formats (table).
124    buffer: Vec<Row>,
125    /// Approximate buffered size in bytes.
126    buffer_bytes: usize,
127    /// Number of rows written (for limit checking).
128    written_count: usize,
129    /// Whether header has been output (output started).
130    output_started: bool,
131    /// Whether quiet mode is enabled (suppress warnings).
132    quiet: bool,
133}
134
135impl<W: Write> StreamingWriter<W> {
136    /// Create a new StreamingWriter.
137    ///
138    /// # Arguments
139    ///
140    /// * `writer` - The output writer (e.g., stdout).
141    /// * `formatter` - The formatter to use for output.
142    /// * `columns` - Column definitions for the output.
143    /// * `limit` - Optional row limit.
144    pub fn new(
145        writer: W,
146        formatter: Box<dyn Formatter>,
147        columns: Vec<Column>,
148        limit: Option<usize>,
149    ) -> Self {
150        Self {
151            writer,
152            formatter,
153            columns,
154            limit,
155            buffer_limit: DEFAULT_BUFFER_LIMIT,
156            buffer: Vec::new(),
157            buffer_bytes: 0,
158            written_count: 0,
159            output_started: false,
160            quiet: false,
161        }
162    }
163
164    /// Create a new StreamingWriter with a custom buffer limit.
165    #[allow(dead_code)]
166    pub fn with_buffer_limit(mut self, buffer_limit: usize) -> Self {
167        self.buffer_limit = buffer_limit;
168        self
169    }
170
171    /// Enable quiet mode (suppress warnings).
172    pub fn with_quiet(mut self, quiet: bool) -> Self {
173        self.quiet = quiet;
174        self
175    }
176
177    /// Check if quiet mode is enabled.
178    ///
179    /// When quiet mode is enabled, status-only output (OK messages) should be suppressed.
180    pub fn is_quiet(&self) -> bool {
181        self.quiet
182    }
183
184    /// Prepare the writer for output.
185    ///
186    /// For streaming formats (json, jsonl, csv, tsv), immediately outputs the header.
187    /// For non-streaming formats (table), defers header output.
188    ///
189    /// # Arguments
190    ///
191    /// * `row_count_hint` - Optional estimated row count (unused for buffer sizing).
192    pub fn prepare(&mut self, row_count_hint: Option<usize>) -> Result<()> {
193        let _ = row_count_hint;
194
195        // For streaming formats, output header immediately
196        if self.formatter.supports_streaming() {
197            self.formatter
198                .write_header(&mut self.writer, &self.columns)?;
199            self.output_started = true;
200        }
201
202        Ok(())
203    }
204
205    /// Write a row to the output.
206    ///
207    /// For streaming formats, the row is output immediately.
208    /// For non-streaming formats, the row is buffered.
209    ///
210    /// # Returns
211    ///
212    /// * `WriteStatus::Continue` - Row was written, continue writing.
213    /// * `WriteStatus::LimitReached` - Row limit reached, stop writing.
214    ///
215    /// # Note
216    ///
217    /// The row is taken by ownership to avoid unnecessary cloning when buffering.
218    pub fn write_row(&mut self, row: Row) -> Result<WriteStatus> {
219        // Check limit
220        if let Some(limit) = self.limit {
221            if self.written_count >= limit {
222                return Ok(WriteStatus::LimitReached);
223            }
224        }
225
226        if self.formatter.supports_streaming() {
227            // Streaming format: output immediately
228            self.formatter.write_row(&mut self.writer, &row)?;
229            self.written_count += 1;
230        } else {
231            // Non-streaming format: buffer the row
232            let row_bytes = estimate_row_bytes(&row);
233            self.buffer_bytes = self.buffer_bytes.saturating_add(row_bytes);
234            if self.buffer_bytes > self.buffer_limit {
235                return Err(CliError::InvalidArgument(
236                    "Buffer limit exceeded (~10MB). \
237                     Use --output json|csv|tsv or --limit to reduce results."
238                        .into(),
239                ));
240            }
241            self.buffer.push(row);
242            self.written_count += 1;
243        }
244
245        Ok(WriteStatus::Continue)
246    }
247
248    /// Finish output, flushing any buffered rows and writing the footer.
249    pub fn finish(&mut self) -> Result<()> {
250        // For non-streaming formats, output header if not yet done
251        if !self.output_started {
252            self.formatter
253                .write_header(&mut self.writer, &self.columns)?;
254            self.output_started = true;
255
256            // Flush any buffered rows
257            for row in self.buffer.drain(..) {
258                self.formatter.write_row(&mut self.writer, &row)?;
259            }
260        }
261
262        // Write footer
263        self.formatter.write_footer(&mut self.writer)?;
264
265        Ok(())
266    }
267
268    /// Returns the number of rows written.
269    #[allow(dead_code)]
270    pub fn written_count(&self) -> usize {
271        self.written_count
272    }
273
274    /// Returns the remaining configured row limit, if any.
275    pub fn remaining_limit(&self) -> Option<usize> {
276        self.limit
277            .map(|limit| limit.saturating_sub(self.written_count))
278    }
279
280    /// Returns whether output has started.
281    #[allow(dead_code)]
282    pub fn output_started(&self) -> bool {
283        self.output_started
284    }
285
286    /// Returns approximate buffered bytes.
287    #[allow(dead_code)]
288    pub fn buffered_bytes(&self) -> usize {
289        self.buffer_bytes
290    }
291}
292
293fn estimate_row_bytes(row: &Row) -> usize {
294    row.columns
295        .iter()
296        .map(estimate_value_bytes)
297        .sum::<usize>()
298        .saturating_add(row.columns.len() * 8)
299}
300
301fn estimate_value_bytes(value: &crate::models::Value) -> usize {
302    match value {
303        crate::models::Value::Null => 4,
304        crate::models::Value::Bool(_) => 1,
305        crate::models::Value::Int(_) => 8,
306        crate::models::Value::Float(_) => 8,
307        crate::models::Value::Text(text) => text.len(),
308        crate::models::Value::Bytes(bytes) => bytes.len(),
309        crate::models::Value::Vector(values) => values.len() * 4,
310    }
311}
312
313#[cfg(test)]
314mod tests {
315    use super::*;
316    use crate::error::CliError;
317    use crate::models::{DataType, Value};
318    use crate::output::csv::CsvFormatter;
319    use crate::output::json::JsonFormatter;
320    use crate::output::jsonl::JsonlFormatter;
321    use crate::output::table::TableFormatter;
322
323    fn test_columns() -> Vec<Column> {
324        vec![
325            Column::new("id", DataType::Int),
326            Column::new("name", DataType::Text),
327        ]
328    }
329
330    fn test_row(id: i64, name: &str) -> Row {
331        Row::new(vec![Value::Int(id), Value::Text(name.to_string())])
332    }
333
334    #[test]
335    fn test_streaming_format_immediate_output() {
336        let mut output = Vec::new();
337        let formatter = Box::new(JsonlFormatter::new());
338        let columns = test_columns();
339
340        let mut writer = StreamingWriter::new(&mut output, formatter, columns, None);
341
342        writer.prepare(None).unwrap();
343        assert!(writer.output_started());
344
345        let status = writer.write_row(test_row(1, "Alice")).unwrap();
346        assert_eq!(status, WriteStatus::Continue);
347        assert_eq!(writer.written_count(), 1);
348
349        writer.finish().unwrap();
350
351        let result = String::from_utf8(output).unwrap();
352        assert!(result.contains("\"id\":1"));
353        assert!(result.contains("\"name\":\"Alice\""));
354    }
355
356    #[test]
357    fn test_non_streaming_format_buffered_output() {
358        let mut output = Vec::new();
359        let formatter = Box::new(TableFormatter::new());
360        let columns = test_columns();
361
362        let mut writer = StreamingWriter::new(&mut output, formatter, columns, None);
363
364        writer.prepare(None).unwrap();
365        assert!(!writer.output_started()); // Header not output yet
366
367        let status = writer.write_row(test_row(1, "Alice")).unwrap();
368        assert_eq!(status, WriteStatus::Continue);
369        assert!(!writer.output_started()); // Still buffering
370
371        writer.finish().unwrap();
372        assert!(writer.output_started()); // Now output started
373
374        let result = String::from_utf8(output).unwrap();
375        assert!(result.contains("id"));
376        assert!(result.contains("Alice"));
377    }
378
379    #[test]
380    fn test_limit_enforcement() {
381        let mut output = Vec::new();
382        let formatter = Box::new(CsvFormatter::new());
383        let columns = test_columns();
384
385        let mut writer = StreamingWriter::new(&mut output, formatter, columns, Some(2));
386
387        writer.prepare(None).unwrap();
388
389        assert_eq!(
390            writer.write_row(test_row(1, "Alice")).unwrap(),
391            WriteStatus::Continue
392        );
393        assert_eq!(
394            writer.write_row(test_row(2, "Bob")).unwrap(),
395            WriteStatus::Continue
396        );
397        assert_eq!(
398            writer.write_row(test_row(3, "Charlie")).unwrap(),
399            WriteStatus::LimitReached
400        );
401
402        assert_eq!(writer.written_count(), 2);
403
404        writer.finish().unwrap();
405    }
406
407    #[test]
408    fn test_buffer_overflow_errors() {
409        let mut output = Vec::new();
410        let formatter = Box::new(TableFormatter::new());
411        let columns = test_columns();
412
413        let mut writer =
414            StreamingWriter::new(&mut output, formatter, columns, None).with_buffer_limit(40); // Very small buffer
415
416        writer.prepare(None).unwrap();
417        assert!(!writer.output_started());
418
419        // Add rows to buffer
420        assert_eq!(
421            writer.write_row(test_row(1, "Alice")).unwrap(),
422            WriteStatus::Continue
423        );
424        let err = writer.write_row(test_row(2, "Bob")).unwrap_err();
425        assert!(matches!(err, CliError::InvalidArgument(_)));
426    }
427
428    #[test]
429    fn test_empty_output() {
430        let mut output = Vec::new();
431        let formatter = Box::new(JsonFormatter::new());
432        let columns = test_columns();
433
434        let mut writer = StreamingWriter::new(&mut output, formatter, columns, None);
435
436        writer.prepare(None).unwrap();
437        writer.finish().unwrap();
438
439        let result = String::from_utf8(output).unwrap();
440        // JSON array format: should be valid empty array
441        assert!(result.contains('['));
442        assert!(result.contains(']'));
443    }
444
445    #[test]
446    fn test_csv_streaming() {
447        let mut output = Vec::new();
448        let formatter = Box::new(CsvFormatter::new());
449        let columns = test_columns();
450
451        let mut writer = StreamingWriter::new(&mut output, formatter, columns, None);
452
453        writer.prepare(None).unwrap();
454        assert!(writer.output_started()); // CSV is streaming
455
456        writer.write_row(test_row(1, "Alice")).unwrap();
457        writer.write_row(test_row(2, "Bob")).unwrap();
458        writer.finish().unwrap();
459
460        let result = String::from_utf8(output).unwrap();
461        assert_eq!(result, "id,name\n1,Alice\n2,Bob\n");
462    }
463
464    #[test]
465    fn test_written_count() {
466        let mut output = Vec::new();
467        let formatter = Box::new(CsvFormatter::new());
468        let columns = test_columns();
469
470        let mut writer = StreamingWriter::new(&mut output, formatter, columns, None);
471
472        writer.prepare(None).unwrap();
473
474        assert_eq!(writer.written_count(), 0);
475        writer.write_row(test_row(1, "Alice")).unwrap();
476        assert_eq!(writer.written_count(), 1);
477        writer.write_row(test_row(2, "Bob")).unwrap();
478        assert_eq!(writer.written_count(), 2);
479
480        writer.finish().unwrap();
481    }
482
483    #[test]
484    fn test_table_with_small_data() {
485        let mut output = Vec::new();
486        let formatter = Box::new(TableFormatter::new());
487        let columns = test_columns();
488
489        let mut writer = StreamingWriter::new(&mut output, formatter, columns, None);
490
491        writer.prepare(None).unwrap();
492        assert!(!writer.output_started()); // Table is non-streaming
493
494        writer.write_row(test_row(1, "Alice")).unwrap();
495        writer.write_row(test_row(2, "Bob")).unwrap();
496        writer.finish().unwrap();
497
498        let result = String::from_utf8(output).unwrap();
499        // Table output should contain the data
500        assert!(result.contains("id"));
501        assert!(result.contains("name"));
502        assert!(result.contains("Alice"));
503        assert!(result.contains("Bob"));
504    }
505
506    #[test]
507    fn test_streaming_large_row_count_does_not_buffer() {
508        let formatter = Box::new(JsonlFormatter::new());
509        let columns = test_columns();
510
511        let mut writer = StreamingWriter::new(std::io::sink(), formatter, columns, None);
512
513        writer.prepare(None).unwrap();
514        for i in 0..12_000 {
515            writer.write_row(test_row(i, "row")).unwrap();
516        }
517
518        assert_eq!(writer.written_count(), 12_000);
519        assert!(writer.buffer.is_empty());
520
521        writer.finish().unwrap();
522    }
523
524    #[test]
525    fn routing_report_is_serialized_to_its_own_writer_without_row_output() {
526        let report = DistributedReadRoutingReport::new(
527            "strong",
528            Some("strong".into()),
529            "cluster",
530            "retryable_failure",
531            Some("read point expired".into()),
532        );
533        let mut stderr = Vec::new();
534        write_distributed_read_routing_report(&mut stderr, RoutingReportFormat::Json, &report)
535            .unwrap();
536        let value: serde_json::Value = serde_json::from_slice(&stderr).unwrap();
537        assert_eq!(value["requested_mode"], "strong");
538        assert_eq!(value["outcome"], "retryable_failure");
539        assert_eq!(value["reason"], "read point expired");
540    }
541}
alopex_cli/streaming/writer.rs

alopex_cli/streaming/
writer.rs
