use super::CsvImportLimits;
use hawdb_core::{HawDBError, Result};
use std::io::Read;
pub const CSV_READ_BUFFER_BYTES: usize = 64 * 1024;
const UTF8_BOM: [u8; 3] = [0xEF, 0xBB, 0xBF];
pub type CsvField = Option<String>;
pub struct CsvRecordReader<R> {
input: R,
buffer: Box<[u8]>,
position: usize,
filled: usize,
eof: bool,
started: bool,
terminator: u8,
limits: CsvImportLimits,
bytes_read: u64,
line: u64,
records: u64,
}
impl<R: Read> CsvRecordReader<R> {
pub fn new(input: R, terminator: u8, limits: CsvImportLimits) -> Result<Self> {
if !terminator.is_ascii() || matches!(terminator, b'"' | b'\r' | b'\n') {
return Err(HawDBError::Semantic(
"FIELDTERMINATOR must be one ASCII character other than a double quote or line break"
.to_string(),
));
}
Ok(Self {
input,
buffer: vec![0; CSV_READ_BUFFER_BYTES].into_boxed_slice(),
position: 0,
filled: 0,
eof: false,
started: false,
terminator,
limits,
bytes_read: 0,
line: 1,
records: 0,
})
}
pub fn records(&self) -> u64 {
self.records
}
pub fn next_record(&mut self) -> Result<Option<Vec<CsvField>>> {
if !self.started {
self.started = true;
self.skip_bom()?;
}
loop {
let Some(first) = self.peek()? else {
return Ok(None);
};
if matches!(first, b'\n' | b'\r') {
self.consume_line_break(first)?;
continue;
}
let start_line = self.line;
let record = self.read_record(start_line)?;
self.records += 1;
return Ok(Some(record));
}
}
fn read_record(&mut self, start_line: u64) -> Result<Vec<CsvField>> {
let mut fields = Vec::new();
let mut record_bytes = 0usize;
loop {
if fields.len() == self.limits.max_fields_per_record {
return Err(self.error(
start_line,
format!(
"has more than max_fields_per_record {} fields",
self.limits.max_fields_per_record
),
));
}
let mut field = Vec::new();
let quoted = self.peek()? == Some(b'"');
let end = if quoted {
self.bump();
self.read_quoted(&mut field, &mut record_bytes, start_line)?
} else {
self.read_unquoted(&mut field, &mut record_bytes, start_line)?
};
let text = String::from_utf8(field).map_err(|_| {
self.error(
start_line,
format!("field {} is not valid UTF-8", fields.len() + 1),
)
})?;
fields.push((quoted || !text.is_empty()).then_some(text));
if end == FieldEnd::Record {
return Ok(fields);
}
}
}
fn read_unquoted(
&mut self,
field: &mut Vec<u8>,
record_bytes: &mut usize,
start_line: u64,
) -> Result<FieldEnd> {
loop {
let Some(byte) = self.peek()? else {
return Ok(FieldEnd::Record);
};
self.bump();
match byte {
byte if byte == self.terminator => return Ok(FieldEnd::Field),
b'\n' | b'\r' => {
self.finish_line_break(byte)?;
return Ok(FieldEnd::Record);
}
byte => self.push(field, byte, record_bytes, start_line)?,
}
}
}
fn read_quoted(
&mut self,
field: &mut Vec<u8>,
record_bytes: &mut usize,
start_line: u64,
) -> Result<FieldEnd> {
loop {
let Some(byte) = self.peek()? else {
return Err(self.error(start_line, "has an unterminated quoted field".to_string()));
};
self.bump();
match byte {
b'"' if self.peek()? == Some(b'"') => {
self.bump();
self.push(field, b'"', record_bytes, start_line)?;
}
b'"' => break,
b'\n' => {
self.line += 1;
self.push(field, byte, record_bytes, start_line)?;
}
b'\r' => {
if self.peek()? != Some(b'\n') {
self.line += 1;
}
self.push(field, byte, record_bytes, start_line)?;
}
byte => self.push(field, byte, record_bytes, start_line)?,
}
}
match self.peek()? {
None => Ok(FieldEnd::Record),
Some(byte) if byte == self.terminator => {
self.bump();
Ok(FieldEnd::Field)
}
Some(byte @ (b'\n' | b'\r')) => {
self.consume_line_break(byte)?;
Ok(FieldEnd::Record)
}
Some(_) => Err(self.error(
start_line,
"has a character after a closing quote; quote the whole field and double embedded quotes"
.to_string(),
)),
}
}
fn push(
&self,
field: &mut Vec<u8>,
byte: u8,
record_bytes: &mut usize,
start_line: u64,
) -> Result<()> {
*record_bytes += 1;
if *record_bytes > self.limits.max_record_bytes {
return Err(self.error(
start_line,
format!("exceeds max_record_bytes {}", self.limits.max_record_bytes),
));
}
field.push(byte);
Ok(())
}
fn consume_line_break(&mut self, byte: u8) -> Result<()> {
self.bump();
self.finish_line_break(byte)
}
fn finish_line_break(&mut self, byte: u8) -> Result<()> {
if byte == b'\r' && self.peek()? == Some(b'\n') {
self.bump();
}
self.line += 1;
Ok(())
}
fn skip_bom(&mut self) -> Result<()> {
while self.filled - self.position < UTF8_BOM.len() && !self.eof {
self.fill()?;
}
if self.buffer[self.position..self.filled].starts_with(&UTF8_BOM) {
self.position += UTF8_BOM.len();
}
Ok(())
}
fn peek(&mut self) -> Result<Option<u8>> {
if self.position == self.filled {
if self.eof {
return Ok(None);
}
self.position = 0;
self.filled = 0;
self.fill()?;
if self.filled == 0 {
return Ok(None);
}
}
Ok(Some(self.buffer[self.position]))
}
fn bump(&mut self) {
self.position += 1;
}
fn fill(&mut self) -> Result<()> {
if self.position > 0 {
self.buffer.copy_within(self.position..self.filled, 0);
self.filled -= self.position;
self.position = 0;
}
let read = loop {
match self.input.read(&mut self.buffer[self.filled..]) {
Ok(read) => break read,
Err(error) if error.kind() == std::io::ErrorKind::Interrupted => continue,
Err(error) => {
return Err(HawDBError::Execution(format!(
"LOAD CSV failed to read the source: {error}"
)))
}
}
};
if read == 0 {
self.eof = true;
return Ok(());
}
self.bytes_read = self.bytes_read.saturating_add(read as u64);
if self.bytes_read > self.limits.max_file_bytes {
return Err(HawDBError::Execution(format!(
"LOAD CSV source exceeds max_file_bytes {}",
self.limits.max_file_bytes
)));
}
self.filled += read;
Ok(())
}
fn error(&self, start_line: u64, detail: String) -> HawDBError {
HawDBError::Execution(format!(
"LOAD CSV record {} starting at line {start_line} {detail}",
self.records + 1
))
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum FieldEnd {
Field,
Record,
}