1use crate::{
16 BranchSqlSelector, BranchSqlStatement, BranchSqlValue, DeleteStatement, ExprKind,
17 InsertStatement, SelectProjection, SelectStatement, SqlArithmeticOperand, SqlAssignment,
18 SqlAssignmentValue, SqlBound, SqlColumnDefault, SqlExpression, SqlStatement, SqlValue,
19 UpdateStatement,
20};
21use hawdb_core::{HawDBError, Result};
22use std::collections::BTreeSet;
23
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct PreparedPostgresStatement {
26 pub statement: SqlStatement,
27 pub parameters: Vec<PostgresParameterMetadata>,
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub struct PostgresParameterMetadata {
32 pub position: usize,
33}
34
35pub fn prepare_postgres_sql(input: &str) -> Result<PreparedPostgresStatement> {
36 let statement = crate::parse_postgres_sql(input)?;
37 let mut positions = BTreeSet::new();
38 collect_statement_parameters(&statement, &mut positions);
39 if !positions.is_empty()
40 && matches!(
41 &statement,
42 SqlStatement::CreateTable(_)
43 | SqlStatement::CreateIndex(_)
44 | SqlStatement::AlterTableAddColumn(_)
45 )
46 {
47 return Err(HawDBError::Semantic(
48 "PostgreSQL schema statements do not accept parameters".to_string(),
49 ));
50 }
51 let maximum = positions.last().copied().unwrap_or(0);
52 let expected = (1..=maximum).collect::<BTreeSet<_>>();
53 if positions != expected {
54 return Err(HawDBError::Semantic(format!(
55 "PostgreSQL parameter positions must be dense from $1 through ${maximum}; found {positions:?}"
56 )));
57 }
58 Ok(PreparedPostgresStatement {
59 statement,
60 parameters: positions
61 .into_iter()
62 .map(|position| PostgresParameterMetadata { position })
63 .collect(),
64 })
65}
66
67fn collect_statement_parameters(statement: &SqlStatement, positions: &mut BTreeSet<usize>) {
68 match statement {
69 SqlStatement::Explain(explain) => {
70 collect_statement_parameters(&explain.statement, positions)
71 }
72 SqlStatement::Select(select) => collect_select_parameters(select, positions),
73 SqlStatement::Insert(insert) => collect_insert_parameters(insert, positions),
74 SqlStatement::Update(update) => collect_update_parameters(update, positions),
75 SqlStatement::Delete(delete) => collect_delete_parameters(delete, positions),
76 SqlStatement::CreateTable(create) => {
77 for column in &create.columns {
78 if let Some(value) = &column.default {
79 collect_column_default_parameter(value, positions);
80 }
81 }
82 }
83 SqlStatement::AlterTableAddColumn(alter) => {
84 if let Some(value) = &alter.column.default {
85 collect_column_default_parameter(value, positions);
86 }
87 }
88 SqlStatement::CreateIndex(_) => {}
89 SqlStatement::Branch(branch) => collect_branch_parameters(branch, positions),
90 }
91}
92
93fn collect_branch_parameters(branch: &BranchSqlStatement, positions: &mut BTreeSet<usize>) {
94 match branch {
95 BranchSqlStatement::ShowBranches(statement) => {
96 collect_bound_parameter(Some(statement.limit), positions);
97 collect_bound_parameter(statement.offset, positions);
98 }
99 BranchSqlStatement::ShowBranch(statement) => {
100 collect_branch_selector(&statement.selector, positions)
101 }
102 BranchSqlStatement::UseBranch(selector) => collect_branch_selector(selector, positions),
103 BranchSqlStatement::ShowCurrentBranch => {}
104 BranchSqlStatement::CreateBranch(statement) => {
105 if let Some(name) = &statement.name {
106 collect_branch_value(name, positions);
107 }
108 collect_branch_selector(&statement.source, positions);
109 collect_bound_parameter(Some(statement.expected_source_revision), positions);
110 collect_branch_value(&statement.request_key, positions);
111 if let Some(owner) = &statement.owner {
112 collect_branch_value(owner, positions);
113 }
114 }
115 BranchSqlStatement::DropBranch(statement) => {
116 collect_branch_value(&statement.id, positions);
117 collect_bound_parameter(Some(statement.expected_metadata_revision), positions);
118 }
119 }
120}
121
122fn collect_branch_selector(selector: &BranchSqlSelector, positions: &mut BTreeSet<usize>) {
123 let (BranchSqlSelector::Name(value) | BranchSqlSelector::Id(value)) = selector;
124 collect_branch_value(value, positions);
125}
126
127fn collect_branch_value(value: &BranchSqlValue, positions: &mut BTreeSet<usize>) {
128 if let BranchSqlValue::Parameter(position) = value {
129 positions.insert(*position);
130 }
131}
132
133fn collect_column_default_parameter(default: &SqlColumnDefault, positions: &mut BTreeSet<usize>) {
134 if let SqlColumnDefault::Literal(value) = default {
135 collect_value_parameter(value, positions);
136 }
137}
138
139fn collect_select_parameters(select: &SelectStatement, positions: &mut BTreeSet<usize>) {
140 if let Some(search) = &select.vector_index {
141 for value in [&search.index, &search.embedding, &search.top_k] {
142 collect_value_parameter(value, positions);
143 }
144 if let Some(value) = &search.max_work {
145 collect_value_parameter(value, positions);
146 }
147 if let Some(value) = &search.candidates {
148 collect_value_parameter(value, positions);
149 }
150 }
151 for projection in &select.projection {
152 if let SelectProjection::Expression { expression, .. } = projection {
153 collect_expression_parameters(expression, positions);
154 }
155 }
156 for join in &select.joins {
157 collect_expression_parameters(&join.on, positions);
158 }
159 if let Some(predicate) = &select.selection {
160 collect_expression_parameters(predicate, positions);
161 }
162 if let Some(predicate) = &select.having {
163 collect_expression_parameters(predicate, positions);
164 }
165 for order in &select.order_by {
166 collect_expression_parameters(&order.expression, positions);
167 }
168 collect_bound_parameter(select.limit, positions);
169 collect_bound_parameter(select.offset, positions);
170}
171
172fn collect_insert_parameters(insert: &InsertStatement, positions: &mut BTreeSet<usize>) {
173 for row in &insert.rows {
174 for value in row {
175 collect_value_parameter(value, positions);
176 }
177 }
178 if let Some(conflict) = &insert.on_conflict
179 && let crate::SqlConflictAction::DoUpdate(assignments) = &conflict.action
180 {
181 collect_assignment_parameters(assignments, positions);
182 }
183}
184
185fn collect_update_parameters(update: &UpdateStatement, positions: &mut BTreeSet<usize>) {
186 collect_assignment_parameters(&update.assignments, positions);
187 if let Some(predicate) = &update.selection {
188 collect_expression_parameters(predicate, positions);
189 }
190}
191
192fn collect_delete_parameters(delete: &DeleteStatement, positions: &mut BTreeSet<usize>) {
193 if let Some(predicate) = &delete.selection {
194 collect_expression_parameters(predicate, positions);
195 }
196}
197
198fn collect_assignment_parameters(assignments: &[SqlAssignment], positions: &mut BTreeSet<usize>) {
199 for assignment in assignments {
200 match &assignment.value {
201 SqlAssignmentValue::Value(value) => collect_value_parameter(value, positions),
202 SqlAssignmentValue::Column(_) => {}
203 SqlAssignmentValue::Arithmetic { left, right, .. } => {
204 collect_arithmetic_operand_parameter(left, positions);
205 collect_arithmetic_operand_parameter(right, positions);
206 }
207 }
208 }
209}
210
211fn collect_arithmetic_operand_parameter(
212 operand: &SqlArithmeticOperand,
213 positions: &mut BTreeSet<usize>,
214) {
215 if let SqlArithmeticOperand::Value(value) = operand {
216 collect_value_parameter(value, positions);
217 }
218}
219
220fn collect_expression_parameters(expression: &SqlExpression, positions: &mut BTreeSet<usize>) {
221 expression.visit(&mut |expression| {
222 if let ExprKind::Value(value) = &expression.kind {
223 collect_value_parameter(value, positions);
224 }
225 });
226}
227
228fn collect_value_parameter(value: &SqlValue, positions: &mut BTreeSet<usize>) {
229 if let SqlValue::Parameter(position) = value {
230 positions.insert(*position);
231 }
232}
233
234fn collect_bound_parameter(bound: Option<SqlBound>, positions: &mut BTreeSet<usize>) {
235 if let Some(SqlBound::Parameter(position)) = bound {
236 positions.insert(position);
237 }
238}