1use rupg_common::{Error, Result, SqlState};
4use rupg_sql::nodes::{ColumnRef, List, Node, RangeVar};
5use rupg_types::oid;
6
7use crate::coerce::{AtOpt, Context};
8use crate::expr::{Expr, ExprKind};
9use crate::select::{Query, Target};
10use crate::{Analyzer, Env, Params, types};
11
12pub(crate) struct Assign {
14 pub(crate) ty: u32,
15 pub(crate) typmod: i32,
16 pub(crate) name: String,
18 pub(crate) indirect: Option<(Expr, List)>,
20}
21
22#[derive(Clone, Copy, Debug, PartialEq, Eq)]
24pub struct Variable {
25 pub param: usize,
27 pub ty: u32,
28 pub typmod: i32,
29}
30
31pub trait Variables {
33 fn column(&self, names: &[&str], at: Option<usize>) -> Result<Option<Variable>>;
39 fn param(&self, number: usize) -> Option<Variable>;
41}
42
43pub fn analyze_with(
49 stmt: &Node,
50 env: &dyn Env,
51 params: &Params,
52 variables: &dyn Variables,
53) -> Result<Query> {
54 crate::analyze_hooked(stmt, env, params, Some(variables))
55}
56
57pub fn analyze_assign(
63 stmt: &Node,
64 env: &dyn Env,
65 variables: &dyn Variables,
66 (ty, typmod): (u32, i32),
67 name: &str,
68) -> Result<Query> {
69 let assign = Assign { ty, typmod, name: name.to_owned(), indirect: None };
70 crate::analyze_assigned(Analyzer::new(env, &Params::default()), stmt, variables, assign)
71}
72
73pub fn analyze_assign_part(stmt: &Node, env: &dyn Env, variables: &dyn Variables) -> Result<Query> {
79 let Node::PLAssignStmt(s) = stmt else {
80 return Err(Error::internal("an assignment is not a PLAssignStmt"));
81 };
82 let Some(val) = &s.val else { return Err(Error::internal("an assignment has no value")) };
83 let name = s.name.as_deref().unwrap_or_default();
84 let split = usize::try_from(s.nnames).unwrap_or(1).saturating_sub(1).min(s.indirection.len());
86 let (names, indirection) = s.indirection.split_at(split);
87 let mut fields = vec![Some(Node::String(name.into()))];
88 fields.extend(names.iter().cloned());
89 let cref = Node::ColumnRef(Box::new(ColumnRef { fields, location: s.location }));
90 let mut analyzer = Analyzer::new(env, &Params::default());
91 analyzer.variables = Some(variables);
92 let target = analyzer.transform(Some(&cref))?;
93 let (ty, typmod) = (target.ty, target.typmod);
94 let indirect = Some((target, indirection.to_vec()));
95 let assign = Assign { ty, typmod, name: name.to_owned(), indirect };
96 crate::analyze_assigned(analyzer, &Node::SelectStmt(val.clone()), variables, assign)
97}
98
99impl Analyzer<'_> {
100 pub(crate) fn assign_target(&mut self, targets: &mut [Target], assign: Assign) -> Result<()> {
102 let Assign { ty, typmod, name, indirect } = assign;
103 let [target] = targets else {
104 let n = targets.len();
105 let word = if n == 1 { "column" } else { "columns" };
106 return Err(Error::new(
107 SqlState::SYNTAX_ERROR,
108 format!("assignment source returned {n} {word}"),
109 ));
110 };
111 if let Some((base, indirection)) = indirect {
112 let rhs = std::mem::replace(&mut target.expr, Expr::new(ExprKind::CaseTest, 0));
113 let at = base.place();
114 let parts = (base, &name[..], (ty, typmod));
115 target.expr =
116 self.assignment_indirection(parts, &indirection, rhs, Context::Plpgsql, at)?;
117 return Ok(());
118 }
119 let source = target.expr.ty;
120 let composite = |t: u32| t == oid::RECORD || types::is_complex(t);
121 if ty != source && composite(ty) && composite(source) {
122 return Ok(());
123 }
124 let expr = std::mem::replace(&mut target.expr, Expr::new(ExprKind::CaseTest, 0));
125 let at = expr.location;
126 let Some(expr) = self.coerce_to_target(expr, ty, typmod, Context::Plpgsql, None)? else {
127 return Err(Error::new(
128 SqlState::DATATYPE_MISMATCH,
129 format!(
130 "variable \"{name}\" is of type {} but expression is of type {}",
131 types::name(ty),
132 types::name(source)
133 ),
134 )
135 .with_hint("You will need to rewrite or cast the expression.")
136 .at_opt(at));
137 };
138 target.expr = expr;
139 Ok(())
140 }
141}
142
143pub fn plpgsql_cast(env: &dyn Env, source: (u32, i32), target: u32, typmod: i32) -> Result<Query> {
149 let mut analyzer = Analyzer::new(env, &Params { types: vec![source.0], variable: false });
150 let param = Expr { kind: ExprKind::Param(1), ty: source.0, typmod: source.1, location: None };
151 let Some(expr) = analyzer.coerce_to_target(param, target, typmod, Context::Plpgsql, None)?
152 else {
153 return Err(Error::new(
154 SqlState::CANNOT_COERCE,
155 format!("cannot cast type {} to {}", types::name(source.0), types::name(target)),
156 ));
157 };
158 let target = Target { name: "?column?".into(), expr, origin: None, junk: false };
159 Ok(Query { targets: vec![target], params: vec![source.0], ..Query::default() })
160}
161
162pub fn row_type_fields(env: &dyn Env, ty: u32) -> Option<Vec<(String, u32, i32)>> {
164 let analyzer = Analyzer::new(env, &Params::default());
165 let columns = analyzer.row_columns(ty)?;
166 Some(columns.into_iter().map(|c| (c.name, c.ty, c.typmod)).collect())
167}
168
169pub fn relation_by_names(env: &dyn Env, names: &[&str]) -> Result<Option<u32>> {
175 let mut rv = RangeVar::default();
176 match names {
177 [rel] => rv.relname = Some((*rel).into()),
178 [schema, rel] => {
179 rv.schemaname = Some((*schema).into());
180 rv.relname = Some((*rel).into());
181 }
182 [database, schema, rel] => {
183 rv.catalogname = Some((*database).into());
184 rv.schemaname = Some((*schema).into());
185 rv.relname = Some((*rel).into());
186 }
187 _ => return Ok(None),
188 }
189 Analyzer::new(env, &Params::default()).lookup_relation(&rv)
190}
191
192pub fn column_type(env: &dyn Env, names: &[&str], column: &str) -> Result<Option<(u32, i32)>> {
198 let Some(relid) = relation_by_names(env, names)? else { return Ok(None) };
199 let analyzer = Analyzer::new(env, &Params::default());
200 let Ok(columns) = analyzer.open_relation(relid, "", None) else { return Ok(None) };
201 Ok(columns.into_iter().find(|c| c.name == column).map(|c| (c.ty, c.typmod)))
202}
203
204pub fn relation_row_type(env: &dyn Env, names: &[&str]) -> Result<Option<u32>> {
210 let Some(relid) = relation_by_names(env, names)? else { return Ok(None) };
211 if let Some(rel) = env.catalog().and_then(|c| c.relation(relid)) {
212 return Ok(Some(rel.row_type).filter(|&ty| ty != 0));
213 }
214 Ok(rupg_pgcatalog::builtin::types().iter().find(|t| t.relid == relid).map(|t| t.oid))
215}
216
217impl Analyzer<'_> {
218 pub(crate) fn column_or_variable(
220 &mut self,
221 c: &ColumnRef,
222 variables: &dyn Variables,
223 ) -> Result<Expr> {
224 let at = crate::typename::place(c.location);
225 let names: Option<Vec<&str>> = c
226 .fields
227 .iter()
228 .map(|f| match f {
229 Some(Node::String(s)) => Some(&**s),
230 _ => None,
231 })
232 .collect();
233 let core = self.column_ref(c);
234 let missing =
235 |e: &Error| matches!(e.state(), SqlState::UNDEFINED_COLUMN | SqlState::UNDEFINED_TABLE);
236 let Some(names) = names else {
237 if let Some((Some(Node::A_Star(_)), prefix)) = c.fields.split_last()
239 && let Err(error) = &core
240 && missing(error)
241 {
242 let prefix: Option<Vec<&str>> = prefix
243 .iter()
244 .map(|f| match f {
245 Some(Node::String(s)) => Some(&**s),
246 _ => None,
247 })
248 .collect();
249 if let Some(prefix) = prefix.filter(|p| !p.is_empty())
250 && let Some(variable) = self.composite_variable(&prefix, at)?
251 {
252 return Ok(self.variable_param(variable, at));
253 }
254 }
255 return core;
256 };
257 match core {
258 Ok(expr) => match variables.column(&names, at)? {
259 Some(_) => Err(Error::new(
260 SqlState::AMBIGUOUS_COLUMN,
261 format!("column reference \"{}\" is ambiguous", names.join(".")),
262 )
263 .with_detail("It could refer to either a PL/pgSQL variable or a table column.")
264 .at_opt(at)),
265 None => Ok(expr),
266 },
267 Err(error) if missing(&error) => match variables.column(&names, at)? {
268 Some(variable) => Ok(self.variable_param(variable, at)),
269 None => Err(error),
270 },
271 Err(error) => Err(error),
272 }
273 }
274
275 pub(crate) fn composite_variable(
277 &self,
278 names: &[&str],
279 at: Option<usize>,
280 ) -> Result<Option<Variable>> {
281 let Some(variables) = self.variables else { return Ok(None) };
282 let Some(variable) = variables.column(names, at)? else { return Ok(None) };
283 let composite =
284 variable.ty == oid::RECORD || self.row_columns(types::base(variable.ty)).is_some();
285 Ok(composite.then_some(variable))
286 }
287
288 pub(crate) fn variable_param(&mut self, variable: Variable, at: Option<usize>) -> Expr {
290 if self.params.len() < variable.param {
291 self.params.resize(variable.param, 0);
292 }
293 self.params[variable.param - 1] = variable.ty;
294 Expr {
295 kind: ExprKind::Param(variable.param),
296 ty: variable.ty,
297 typmod: variable.typmod,
298 location: at,
299 }
300 }
301}