use pg_query::NodeEnum;
use pg_query::protobuf::{self, Node};
pub const PROJECTION_BACKLOG: &str = "projection_backlog";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BackfillStatement {
pub position: usize,
pub target: String,
pub form: &'static str,
}
#[must_use]
pub fn backfill_statements(sql: &str) -> Vec<BackfillStatement> {
let Ok(parsed) = pg_query::parse(sql) else {
return Vec::new();
};
parsed
.protobuf
.stmts
.iter()
.filter_map(|raw| raw.stmt.as_ref().and_then(|s| s.node.as_ref()))
.enumerate()
.filter_map(|(index, node)| {
classify(node).map(|(target, form)| BackfillStatement {
position: index + 1,
target,
form,
})
})
.collect()
}
fn classify(node: &NodeEnum) -> Option<(String, &'static str)> {
match node {
NodeEnum::UpdateStmt(stmt) => Some((stmt.relation.as_ref()?.relname.clone(), "UPDATE")),
NodeEnum::DeleteStmt(stmt) => Some((stmt.relation.as_ref()?.relname.clone(), "DELETE")),
NodeEnum::InsertStmt(stmt) => {
let select = select_of(stmt.select_stmt.as_deref())?;
let table = stmt.relation.as_ref()?.relname.clone();
if table == PROJECTION_BACKLOG && is_keys_only(select) {
return None;
}
Some((table, "INSERT … SELECT"))
},
NodeEnum::SelectStmt(select) => {
refresh_call(select).map(|name| (name, "SELECT refresh_*()"))
},
_ => None,
}
}
fn select_of(node: Option<&Node>) -> Option<&protobuf::SelectStmt> {
match node?.node.as_ref()? {
NodeEnum::SelectStmt(select) if select.values_lists.is_empty() => Some(select),
_ => None,
}
}
fn is_keys_only(select: &protobuf::SelectStmt) -> bool {
let single_relation = select.from_clause.len() == 1
&& matches!(
select.from_clause[0].node.as_ref(),
Some(NodeEnum::RangeVar(_))
);
single_relation
&& select.target_list.len() == 2
&& select.group_clause.is_empty()
&& select.having_clause.is_none()
&& select.with_clause.is_none()
}
fn refresh_call(select: &protobuf::SelectStmt) -> Option<String> {
select.target_list.iter().find_map(|target| {
let NodeEnum::ResTarget(res) = target.node.as_ref()? else {
return None;
};
let NodeEnum::FuncCall(call) = res.val.as_ref()?.node.as_ref()? else {
return None;
};
let name = call.funcname.last().and_then(|n| match n.node.as_ref()? {
NodeEnum::String(s) => Some(s.sval.clone()),
_ => None,
})?;
name.starts_with("refresh_").then_some(name)
})
}