1#![allow(missing_docs)]
17
18use crate::error::{IoError, Result};
19use serde_json::Value;
20use std::fs::File;
21use std::io::{BufRead, BufReader};
22use std::path::{Path, PathBuf};
23use std::sync::{Arc, Mutex};
24
25pub trait DataSource: Send {
31 fn next_batch(&mut self, batch_size: usize) -> Result<Option<Vec<Value>>>;
33
34 fn reset(&mut self) -> Result<()> {
36 Err(IoError::Other(
37 "reset not supported by this source".to_string(),
38 ))
39 }
40
41 fn name(&self) -> &str {
43 "unknown"
44 }
45
46 fn estimated_len(&self) -> Option<usize> {
48 None
49 }
50}
51
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub enum FileSourceFormat {
59 Csv,
61 JsonLines,
63 Json,
65 Auto,
67}
68
69pub struct FileSource {
71 path: PathBuf,
72 format: FileSourceFormat,
73 reader: Option<FileSourceReader>,
74 exhausted: bool,
75 csv_headers: Option<Vec<String>>,
77 json_records: Option<std::collections::VecDeque<Value>>,
79}
80
81enum FileSourceReader {
82 Lines(BufReader<File>),
83}
84
85impl FileSource {
86 pub fn new(path: impl AsRef<Path>) -> Self {
88 Self::with_format(path, FileSourceFormat::Auto)
89 }
90
91 pub fn with_format(path: impl AsRef<Path>, format: FileSourceFormat) -> Self {
93 Self {
94 path: path.as_ref().to_path_buf(),
95 format,
96 reader: None,
97 exhausted: false,
98 csv_headers: None,
99 json_records: None,
100 }
101 }
102
103 fn resolve_format(&self) -> FileSourceFormat {
104 if self.format != FileSourceFormat::Auto {
105 return self.format;
106 }
107 match self
108 .path
109 .extension()
110 .and_then(|e| e.to_str())
111 .map(|e| e.to_ascii_lowercase())
112 .as_deref()
113 {
114 Some("csv") => FileSourceFormat::Csv,
115 Some("jsonl") | Some("ndjson") => FileSourceFormat::JsonLines,
116 Some("json") => FileSourceFormat::Json,
117 _ => FileSourceFormat::JsonLines,
118 }
119 }
120
121 fn open(&mut self) -> Result<()> {
122 let fmt = self.resolve_format();
123 match fmt {
124 FileSourceFormat::Json => {
125 let file = File::open(&self.path).map_err(IoError::Io)?;
127 let value: Value = serde_json::from_reader(BufReader::new(file))
128 .map_err(|e| IoError::ParseError(e.to_string()))?;
129 let records: Vec<Value> = match value {
130 Value::Array(arr) => arr,
131 other => vec![other],
132 };
133 self.json_records = Some(records.into());
134 }
135 _ => {
136 let file = File::open(&self.path).map_err(IoError::Io)?;
138 let reader = BufReader::new(file);
139 self.reader = Some(FileSourceReader::Lines(reader));
140 }
141 }
142 Ok(())
143 }
144
145 fn next_batch_csv(&mut self, batch_size: usize) -> Result<Option<Vec<Value>>> {
146 let reader = match &mut self.reader {
147 Some(FileSourceReader::Lines(r)) => r,
148 None => return Ok(None),
149 };
150
151 if self.csv_headers.is_none() {
153 let mut line = String::new();
154 if reader.read_line(&mut line).map_err(IoError::Io)? == 0 {
155 return Ok(None);
156 }
157 let headers: Vec<String> = line
158 .trim_end_matches(['\n', '\r'])
159 .split(',')
160 .map(|h| h.trim().trim_matches('"').to_string())
161 .collect();
162 self.csv_headers = Some(headers);
163 }
164
165 let headers = self.csv_headers.clone().unwrap_or_default();
166 let mut records = Vec::with_capacity(batch_size);
167 let mut line = String::new();
168
169 while records.len() < batch_size {
170 line.clear();
171 let n = reader.read_line(&mut line).map_err(IoError::Io)?;
172 if n == 0 {
173 break;
174 }
175 let row = line.trim_end_matches(['\n', '\r']);
176 if row.is_empty() {
177 continue;
178 }
179 let fields: Vec<&str> = row.split(',').collect();
180 let obj: serde_json::Map<String, Value> = headers
181 .iter()
182 .enumerate()
183 .map(|(i, h)| {
184 let raw = fields.get(i).copied().unwrap_or("").trim();
185 let val = raw.trim_matches('"');
186 let json_val: Value = if let Ok(n) = val.parse::<i64>() {
188 Value::Number(n.into())
189 } else if let Ok(f) = val.parse::<f64>() {
190 Value::Number(
191 serde_json::Number::from_f64(f).unwrap_or(serde_json::Number::from(0)),
192 )
193 } else if val == "true" {
194 Value::Bool(true)
195 } else if val == "false" {
196 Value::Bool(false)
197 } else {
198 Value::String(val.to_string())
199 };
200 (h.clone(), json_val)
201 })
202 .collect();
203 records.push(Value::Object(obj));
204 }
205
206 if records.is_empty() {
207 Ok(None)
208 } else {
209 Ok(Some(records))
210 }
211 }
212
213 fn next_batch_jsonl(&mut self, batch_size: usize) -> Result<Option<Vec<Value>>> {
214 let reader = match &mut self.reader {
215 Some(FileSourceReader::Lines(r)) => r,
216 None => return Ok(None),
217 };
218
219 let mut records = Vec::with_capacity(batch_size);
220 let mut line = String::new();
221
222 while records.len() < batch_size {
223 line.clear();
224 let n = reader.read_line(&mut line).map_err(IoError::Io)?;
225 if n == 0 {
226 break;
227 }
228 let trimmed = line.trim();
229 if trimmed.is_empty() {
230 continue;
231 }
232 let val: Value = serde_json::from_str(trimmed)
233 .map_err(|e| IoError::ParseError(format!("JSONL parse: {e}")))?;
234 records.push(val);
235 }
236
237 if records.is_empty() {
238 Ok(None)
239 } else {
240 Ok(Some(records))
241 }
242 }
243}
244
245impl DataSource for FileSource {
246 fn next_batch(&mut self, batch_size: usize) -> Result<Option<Vec<Value>>> {
247 if self.exhausted {
248 return Ok(None);
249 }
250
251 if self.reader.is_none() && self.json_records.is_none() {
252 self.open()?;
253 }
254
255 let result = if let Some(queue) = &mut self.json_records {
256 if queue.is_empty() {
258 Ok(None)
259 } else {
260 let batch: Vec<Value> = queue.drain(..batch_size.min(queue.len())).collect();
261 Ok(Some(batch))
262 }
263 } else {
264 match self.resolve_format() {
265 FileSourceFormat::Csv => self.next_batch_csv(batch_size),
266 _ => self.next_batch_jsonl(batch_size),
267 }
268 };
269
270 if matches!(result, Ok(None)) {
271 self.exhausted = true;
272 }
273 result
274 }
275
276 fn reset(&mut self) -> Result<()> {
277 self.reader = None;
278 self.json_records = None;
279 self.exhausted = false;
280 self.csv_headers = None;
281 Ok(())
282 }
283
284 fn name(&self) -> &str {
285 "file_source"
286 }
287}
288
289pub struct DatabaseSource {
298 connection_string: String,
300 query: String,
302 rows: Option<std::collections::VecDeque<Value>>,
304 exhausted: bool,
305}
306
307impl DatabaseSource {
308 pub fn new(connection_string: impl Into<String>, query: impl Into<String>) -> Self {
313 Self {
314 connection_string: connection_string.into(),
315 query: query.into(),
316 rows: None,
317 exhausted: false,
318 }
319 }
320
321 #[cfg(feature = "sqlite")]
322 fn load_rows(&mut self) -> Result<()> {
323 use oxisql_core::{Connection as OxiConnection, Value as OxiValue};
324 use oxisql_sqlite_compat::SqliteConnection;
325 use std::future::Future;
326
327 fn run_sync_local<F, T, E>(fut: F) -> std::result::Result<T, E>
329 where
330 F: Future<Output = std::result::Result<T, E>>,
331 {
332 match tokio::runtime::Handle::try_current() {
333 Ok(handle) => tokio::task::block_in_place(|| handle.block_on(fut)),
334 Err(_) => tokio::runtime::Builder::new_current_thread()
335 .enable_all()
336 .build()
337 .expect("tokio runtime for sqlite source")
338 .block_on(fut),
339 }
340 }
341
342 let conn = run_sync_local(SqliteConnection::open(&self.connection_string))
343 .map_err(|e| IoError::DatabaseError(e.to_string()))?;
344
345 let oxi_rows = run_sync_local(conn.query(&self.query, &[]))
346 .map_err(|e| IoError::DatabaseError(e.to_string()))?;
347
348 let mut rows_vec = Vec::with_capacity(oxi_rows.len());
349
350 for row in &oxi_rows {
351 let mut obj = serde_json::Map::new();
352 let col_count = row.column_count();
353 let col_names = row.columns().to_vec();
354
355 for i in 0..col_count {
356 let col_name = col_names
357 .get(i)
358 .cloned()
359 .unwrap_or_else(|| format!("column_{i}"));
360 let json_val = match row.get_by_index(i) {
361 Some(OxiValue::Null) | None => Value::Null,
362 Some(OxiValue::Bool(b)) => Value::Bool(*b),
363 Some(OxiValue::I64(n)) => Value::Number((*n).into()),
364 Some(OxiValue::F64(f)) => Value::Number(
365 serde_json::Number::from_f64(*f).unwrap_or(serde_json::Number::from(0)),
366 ),
367 Some(OxiValue::Text(t)) => Value::String(t.clone()),
368 Some(OxiValue::Blob(b)) => Value::String(format!("<blob {} bytes>", b.len())),
369 Some(OxiValue::Json(j)) => {
370 serde_json::from_str(j).unwrap_or_else(|_| Value::String(j.clone()))
371 }
372 Some(OxiValue::Decimal(d)) => Value::String(d.clone()),
373 Some(OxiValue::Timestamp(ts)) => Value::Number((*ts).into()),
374 Some(OxiValue::Date(d)) => Value::Number((*d).into()),
375 Some(OxiValue::Time(t)) => Value::Number((*t).into()),
376 Some(OxiValue::Uuid(_)) => {
377 Value::String(format!(
379 "{}",
380 row.get_by_index(i)
381 .expect("index checked above in outer match")
382 ))
383 }
384 Some(OxiValue::Array(_)) => Value::String(format!("{:?}", row.get_by_index(i))),
385 Some(OxiValue::TypedArray { .. }) => {
388 Value::String(format!("{:?}", row.get_by_index(i)))
389 }
390 };
391 obj.insert(col_name, json_val);
392 }
393 rows_vec.push(Value::Object(obj));
394 }
395
396 self.rows = Some(rows_vec.into());
397 Ok(())
398 }
399
400 #[cfg(not(feature = "sqlite"))]
401 fn load_rows(&mut self) -> Result<()> {
402 Err(IoError::Other(
403 "DatabaseSource requires the `sqlite` feature".to_string(),
404 ))
405 }
406}
407
408impl DataSource for DatabaseSource {
409 fn next_batch(&mut self, batch_size: usize) -> Result<Option<Vec<Value>>> {
410 if self.exhausted {
411 return Ok(None);
412 }
413 if self.rows.is_none() {
414 self.load_rows()?;
415 }
416 let queue = self
417 .rows
418 .as_mut()
419 .expect("rows must be populated after load_rows");
420 if queue.is_empty() {
421 self.exhausted = true;
422 return Ok(None);
423 }
424 let batch: Vec<Value> = queue.drain(..batch_size.min(queue.len())).collect();
425 Ok(Some(batch))
426 }
427
428 fn reset(&mut self) -> Result<()> {
429 self.rows = None;
430 self.exhausted = false;
431 Ok(())
432 }
433
434 fn name(&self) -> &str {
435 "database_source"
436 }
437}
438
439pub struct StreamSource {
448 receiver: crossbeam_channel::Receiver<Value>,
449 timeout: std::time::Duration,
450 label: String,
451}
452
453pub type StreamSender = crossbeam_channel::Sender<Value>;
455
456impl StreamSource {
457 pub fn new_unbounded() -> (Self, StreamSender) {
462 let (tx, rx) = crossbeam_channel::unbounded();
463 let source = Self {
464 receiver: rx,
465 timeout: std::time::Duration::from_millis(50),
466 label: "stream_source".to_string(),
467 };
468 (source, tx)
469 }
470
471 pub fn new_bounded(capacity: usize) -> (Self, StreamSender) {
475 let (tx, rx) = crossbeam_channel::bounded(capacity);
476 let source = Self {
477 receiver: rx,
478 timeout: std::time::Duration::from_millis(50),
479 label: "stream_source_bounded".to_string(),
480 };
481 (source, tx)
482 }
483
484 pub fn with_timeout(mut self, timeout: std::time::Duration) -> Self {
486 self.timeout = timeout;
487 self
488 }
489
490 pub fn with_name(mut self, name: impl Into<String>) -> Self {
492 self.label = name.into();
493 self
494 }
495}
496
497impl DataSource for StreamSource {
498 fn next_batch(&mut self, batch_size: usize) -> Result<Option<Vec<Value>>> {
501 use crossbeam_channel::RecvTimeoutError;
502
503 let mut batch = Vec::with_capacity(batch_size);
504
505 match self.receiver.recv_timeout(self.timeout) {
507 Ok(msg) => batch.push(msg),
508 Err(RecvTimeoutError::Disconnected) => return Ok(None),
509 Err(RecvTimeoutError::Timeout) => return Ok(Some(batch)), }
511
512 while batch.len() < batch_size {
514 match self.receiver.try_recv() {
515 Ok(msg) => batch.push(msg),
516 Err(_) => break,
517 }
518 }
519
520 Ok(Some(batch))
521 }
522
523 fn name(&self) -> &str {
524 &self.label
525 }
526}
527
528pub struct GeneratorSource<F>
537where
538 F: FnMut() -> Option<Value> + Send,
539{
540 generator: F,
541 label: String,
542 exhausted: bool,
543}
544
545impl<F> GeneratorSource<F>
546where
547 F: FnMut() -> Option<Value> + Send,
548{
549 pub fn new(generator: F) -> Self {
551 Self {
552 generator,
553 label: "generator_source".to_string(),
554 exhausted: false,
555 }
556 }
557
558 pub fn with_name(mut self, name: impl Into<String>) -> Self {
560 self.label = name.into();
561 self
562 }
563}
564
565impl<F> DataSource for GeneratorSource<F>
566where
567 F: FnMut() -> Option<Value> + Send,
568{
569 fn next_batch(&mut self, batch_size: usize) -> Result<Option<Vec<Value>>> {
570 if self.exhausted {
571 return Ok(None);
572 }
573 let mut batch = Vec::with_capacity(batch_size);
574 while batch.len() < batch_size {
575 match (self.generator)() {
576 Some(v) => batch.push(v),
577 None => {
578 self.exhausted = true;
579 break;
580 }
581 }
582 }
583 if batch.is_empty() {
584 Ok(None)
585 } else {
586 Ok(Some(batch))
587 }
588 }
589
590 fn reset(&mut self) -> Result<()> {
591 Err(IoError::Other(
593 "GeneratorSource does not support reset".to_string(),
594 ))
595 }
596
597 fn name(&self) -> &str {
598 &self.label
599 }
600}
601
602pub fn drain_source(source: &mut dyn DataSource, batch_size: usize) -> Result<Vec<Value>> {
608 let mut all = Vec::new();
609 while let Some(batch) = source.next_batch(batch_size)? {
610 all.extend(batch);
611 }
612 Ok(all)
613}
614
615#[cfg(test)]
620mod tests {
621 use super::*;
622 use serde_json::json;
623 use std::io::Write;
624
625 fn temp_path(suffix: &str) -> PathBuf {
626 let mut p = std::env::temp_dir();
627 p.push(format!("scirs2_io_test_{suffix}_{}", std::process::id()));
628 p
629 }
630
631 #[test]
632 fn test_generator_source() {
633 let mut counter = 0u64;
634 let mut src = GeneratorSource::new(move || {
635 if counter < 5 {
636 let v = json!(counter);
637 counter += 1;
638 Some(v)
639 } else {
640 None
641 }
642 });
643
644 let batch = src.next_batch(3).unwrap().unwrap();
645 assert_eq!(batch.len(), 3);
646 let rest = drain_source(&mut src, 10).unwrap();
647 assert_eq!(rest.len(), 2);
648 assert!(src.next_batch(10).unwrap().is_none());
649 }
650
651 #[test]
652 fn test_stream_source_basic() {
653 let (mut src, tx) = StreamSource::new_unbounded();
654 tx.send(json!({"id": 1})).unwrap();
655 tx.send(json!({"id": 2})).unwrap();
656 drop(tx); let all = drain_source(&mut src, 10).unwrap();
659 let total: usize = all.len() + drain_source(&mut src, 10).unwrap_or_default().len();
661 assert!(total >= 2 || all.len() >= 2);
662 }
663
664 #[test]
665 fn test_file_source_jsonl() {
666 let path = temp_path("jsonl");
667 {
668 let mut f = std::fs::File::create(&path).unwrap();
669 writeln!(f, r#"{{"a":1,"b":"x"}}"#).unwrap();
670 writeln!(f, r#"{{"a":2,"b":"y"}}"#).unwrap();
671 writeln!(f, r#"{{"a":3,"b":"z"}}"#).unwrap();
672 }
673 let mut src = FileSource::with_format(&path, FileSourceFormat::JsonLines);
674 let all = drain_source(&mut src, 10).unwrap();
675 assert_eq!(all.len(), 3);
676 assert_eq!(all[0]["a"], json!(1));
677 std::fs::remove_file(&path).ok();
678 }
679
680 #[test]
681 fn test_file_source_csv() {
682 let path = temp_path("csv");
683 {
684 let mut f = std::fs::File::create(&path).unwrap();
685 writeln!(f, "name,score").unwrap();
686 writeln!(f, "alice,90").unwrap();
687 writeln!(f, "bob,85").unwrap();
688 }
689 let mut src = FileSource::with_format(&path, FileSourceFormat::Csv);
690 let all = drain_source(&mut src, 10).unwrap();
691 assert_eq!(all.len(), 2);
692 assert_eq!(all[0]["name"], json!("alice"));
693 assert_eq!(all[0]["score"], json!(90));
694 std::fs::remove_file(&path).ok();
695 }
696
697 #[test]
698 fn test_file_source_json_array() {
699 let path = temp_path("json");
700 {
701 let mut f = std::fs::File::create(&path).unwrap();
702 write!(f, r#"[{{"x":1}},{{"x":2}},{{"x":3}}]"#).unwrap();
703 }
704 let mut src = FileSource::with_format(&path, FileSourceFormat::Json);
705 let b1 = src.next_batch(2).unwrap().unwrap();
707 assert_eq!(b1.len(), 2);
708 let b2 = src.next_batch(2).unwrap().unwrap();
709 assert_eq!(b2.len(), 1);
710 assert!(src.next_batch(2).unwrap().is_none());
711 std::fs::remove_file(&path).ok();
712 }
713
714 #[test]
715 fn test_file_source_batching() {
716 let path = temp_path("jsonl2");
717 {
718 let mut f = std::fs::File::create(&path).unwrap();
719 for i in 0..10 {
720 writeln!(f, r#"{{"i":{i}}}"#).unwrap();
721 }
722 }
723 let mut src = FileSource::with_format(&path, FileSourceFormat::JsonLines);
724 let b1 = src.next_batch(4).unwrap().unwrap();
725 assert_eq!(b1.len(), 4);
726 let b2 = src.next_batch(4).unwrap().unwrap();
727 assert_eq!(b2.len(), 4);
728 let b3 = src.next_batch(4).unwrap().unwrap();
729 assert_eq!(b3.len(), 2);
730 assert!(src.next_batch(4).unwrap().is_none());
731 std::fs::remove_file(&path).ok();
732 }
733}