uqa_client/
sql_stream.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use crate::{HttpEngineError, SQLStreamFrame};
8
9const MAX_STREAM_FRAME_BYTES: usize = 64 * 1024 * 1024;
10
11/// Incremental reader for one authenticated UQA NDJSON SQL response.
12pub struct SQLStream {
13    response: reqwest::Response,
14    request_id: String,
15    buffer: Vec<u8>,
16    scan_start: usize,
17    phase: StreamPhase,
18    body_finished: bool,
19}
20
21#[derive(Clone, Copy, Eq, PartialEq)]
22enum StreamPhase {
23    AwaitingMetadata,
24    Rows,
25    Terminal,
26    Finished,
27}
28
29impl SQLStream {
30    pub(crate) fn new(response: reqwest::Response, request_id: String) -> Self {
31        Self {
32            response,
33            request_id,
34            buffer: Vec::new(),
35            scan_start: 0,
36            phase: StreamPhase::AwaitingMetadata,
37            body_finished: false,
38        }
39    }
40
41    pub fn request_id(&self) -> &str {
42        &self.request_id
43    }
44
45    pub async fn next_frame(&mut self) -> Result<Option<SQLStreamFrame>, HttpEngineError> {
46        loop {
47            if self.phase == StreamPhase::Finished {
48                return Ok(None);
49            }
50            if self.phase == StreamPhase::Terminal {
51                return self.finish().await;
52            }
53            if let Some(newline) = self.next_newline() {
54                if newline > MAX_STREAM_FRAME_BYTES {
55                    return Err(HttpEngineError::StreamFrameTooLarge);
56                }
57                let line = take_line(&mut self.buffer, newline);
58                self.scan_start = 0;
59                if line_is_empty(&line) {
60                    continue;
61                }
62                return self.decode_frame(&line).map(Some);
63            }
64            if self.buffer.len() > MAX_STREAM_FRAME_BYTES {
65                return Err(HttpEngineError::StreamFrameTooLarge);
66            }
67            if self.body_finished {
68                if self.buffer.is_empty() {
69                    return Err(HttpEngineError::TruncatedStream);
70                }
71                if self.buffer.len() > MAX_STREAM_FRAME_BYTES {
72                    return Err(HttpEngineError::StreamFrameTooLarge);
73                }
74                let line = std::mem::take(&mut self.buffer);
75                return self.decode_frame(&line).map(Some);
76            }
77            match self
78                .response
79                .chunk()
80                .await
81                .map_err(HttpEngineError::transport)?
82            {
83                Some(chunk) => {
84                    self.buffer.extend_from_slice(&chunk);
85                }
86                None => self.body_finished = true,
87            }
88        }
89    }
90
91    fn decode_frame(&mut self, line: &[u8]) -> Result<SQLStreamFrame, HttpEngineError> {
92        let line = line.strip_suffix(b"\r").unwrap_or(line);
93        let mut frame = serde_json::from_slice::<SQLStreamFrame>(line)
94            .map_err(HttpEngineError::InvalidResponse)?;
95        if let SQLStreamFrame::Error {
96            code, diagnostic, ..
97        } = &mut frame
98        {
99            if diagnostic
100                .as_ref()
101                .is_some_and(|value| !value.matches_code(code))
102            {
103                *diagnostic = None;
104            }
105        }
106        if frame
107            .request_id()
108            .is_some_and(|request_id| request_id != self.request_id)
109        {
110            return Err(HttpEngineError::StreamRequestIdMismatch);
111        }
112        self.phase = match (&self.phase, &frame) {
113            (StreamPhase::AwaitingMetadata, SQLStreamFrame::Metadata { .. })
114            | (StreamPhase::Rows, SQLStreamFrame::Row { .. }) => StreamPhase::Rows,
115            (StreamPhase::AwaitingMetadata | StreamPhase::Rows, SQLStreamFrame::Error { .. })
116            | (StreamPhase::Rows, SQLStreamFrame::Complete { .. }) => StreamPhase::Terminal,
117            _ => return Err(HttpEngineError::InvalidStreamSequence),
118        };
119        Ok(frame)
120    }
121
122    async fn finish(&mut self) -> Result<Option<SQLStreamFrame>, HttpEngineError> {
123        loop {
124            if let Some(newline) = self.next_newline() {
125                if newline > MAX_STREAM_FRAME_BYTES {
126                    return Err(HttpEngineError::StreamFrameTooLarge);
127                }
128                let line = take_line(&mut self.buffer, newline);
129                self.scan_start = 0;
130                if !line_is_empty(&line) {
131                    return Err(HttpEngineError::InvalidStreamSequence);
132                }
133                continue;
134            }
135            if self.buffer.len() > MAX_STREAM_FRAME_BYTES {
136                return Err(HttpEngineError::StreamFrameTooLarge);
137            }
138            if self.body_finished {
139                if !line_is_empty(&self.buffer) {
140                    return Err(HttpEngineError::InvalidStreamSequence);
141                }
142                self.buffer.clear();
143                self.phase = StreamPhase::Finished;
144                return Ok(None);
145            }
146            match self
147                .response
148                .chunk()
149                .await
150                .map_err(HttpEngineError::transport)?
151            {
152                Some(chunk) => {
153                    self.buffer.extend_from_slice(&chunk);
154                }
155                None => self.body_finished = true,
156            }
157        }
158    }
159
160    fn next_newline(&mut self) -> Option<usize> {
161        let start = self.scan_start.min(self.buffer.len());
162        let newline = self.buffer[start..]
163            .iter()
164            .position(|byte| *byte == b'\n')
165            .map(|offset| start + offset);
166        if newline.is_none() {
167            self.scan_start = self.buffer.len();
168        }
169        newline
170    }
171}
172
173fn line_is_empty(line: &[u8]) -> bool {
174    line.is_empty() || line == b"\r"
175}
176
177fn take_line(buffer: &mut Vec<u8>, newline: usize) -> Vec<u8> {
178    let remainder = buffer.split_off(newline + 1);
179    let mut line = std::mem::replace(buffer, remainder);
180    line.truncate(newline);
181    line
182}
uqa_client/sql_stream.rs

uqa_client/
sql_stream.rs
