use std::str::FromStr;
use datafusion::sql::{
parser::{DFParser, Statement},
sqlparser::{keywords::Keyword, tokenizer::Token},
};
use crate::datasources::ExonFileType;
use super::exon_copy_statement::ExonCopyToStatement;
pub(crate) struct ExonParser<'a> {
df_parser: DFParser<'a>,
}
#[derive(Debug)]
pub(crate) enum ExonStatement {
DFStatement(Box<Statement>),
ExonCopyTo(ExonCopyToStatement),
}
impl ExonParser<'_> {
pub fn new(sql: &str) -> crate::Result<Self> {
let df_parser = DFParser::new(sql)?;
Ok(Self { df_parser })
}
fn is_copy(&self) -> bool {
matches!(
self.df_parser.parser.peek_token().token,
Token::Word(w) if w.keyword == Keyword::COPY
)
}
pub fn parse_statement(&mut self) -> crate::Result<ExonStatement> {
if self.is_copy() {
self.df_parser.parser.next_token(); let df_statement = self.df_parser.parse_copy()?;
if let Statement::CopyTo(s) = &df_statement {
match &s.stored_as {
Some(v) if ExonFileType::from_str(v).is_ok() => Ok(ExonStatement::ExonCopyTo(
ExonCopyToStatement::from(s.clone()),
)),
_ => Ok(ExonStatement::DFStatement(Box::from(df_statement))),
}
} else {
Ok(ExonStatement::DFStatement(Box::from(df_statement)))
}
} else {
let df_statement = self.df_parser.parse_statement()?;
Ok(ExonStatement::DFStatement(Box::from(df_statement)))
}
}
}