1use rupg_common::{Error, Result, SqlState};
4use rupg_pgcatalog::builtin;
5use rupg_sql::nodes::{
6 Alias, ColumnRef, CopyStmt, InsertStmt, Node, OverridingKind, ParamRef, RangeSubselect,
7 RangeVar, ResTarget, SelectStmt,
8};
9
10use super::lock::{lookup, relkind};
11use super::relkind_detail;
12use crate::walk::{self, Walker};
13use crate::{Analyzer, Env, Params};
14
15#[derive(Debug)]
17pub struct CopyRelation {
18 pub name: String,
20 pub columns: Vec<String>,
22 pub selected: Vec<usize>,
24 pub query: Node,
26 pub partitions: Option<Vec<u32>>,
28}
29
30#[derive(Debug, Clone)]
32pub struct CopyTarget {
33 pub relid: u32,
35 pub name: String,
37 pub kind: u8,
39 pub columns: Vec<String>,
41 pub selected: Vec<usize>,
43}
44
45pub fn copy_target(stmt: &CopyStmt, rv: &RangeVar, env: &dyn Env) -> Result<CopyTarget> {
51 let params = Params::default();
52 let an = Analyzer::new(env, ¶ms);
53 let relid = lookup(&an, rv)?;
54 let name = rv.relname.as_deref().unwrap_or_default().to_string();
55 let kind = relkind(env, relid);
56 if matches!(kind, b'i' | b'I' | b'c') {
57 return Err(Error::new(
58 SqlState::WRONG_OBJECT_TYPE,
59 format!("cannot open relation \"{name}\""),
60 )
61 .with_detail(relkind_detail(kind)));
62 }
63 let columns = relation_columns(env, relid);
64 let names: Vec<String> =
65 stmt.attlist.iter().flatten().filter_map(string).map(str::to_string).collect();
66 let mut selected = attnums(&columns, &names, Some(&name))?;
67 let generated: Vec<bool> =
69 env.catalog().and_then(|c| c.relation(relid)).map_or_else(Vec::new, |rel| {
70 rel.columns.iter().filter(|c| !c.dropped).map(|c| c.generated != 0).collect()
71 });
72 let is_generated = |i: usize| generated.get(i).copied().unwrap_or(false);
73 if let Some(condition) = &stmt.whereClause {
74 let mut refs = WhereColumns { relname: &name, columns: &columns, used: Vec::new() };
75 walk::one(&mut refs, Some(condition));
76 let mut used = refs.used;
77 used.sort_unstable();
78 if let Some(&i) = used.iter().find(|&&i| is_generated(i)) {
79 return Err(Error::new(
80 SqlState::INVALID_COLUMN_REFERENCE,
81 "generated columns are not supported in COPY FROM WHERE conditions",
82 )
83 .with_detail(format!("Column \"{}\" is a generated column.", columns[i])));
84 }
85 }
86 if names.is_empty() {
87 selected.retain(|&i| !is_generated(i));
88 } else if let Some(&i) = selected.iter().find(|&&i| is_generated(i)) {
89 return Err(Error::new(
90 SqlState::INVALID_COLUMN_REFERENCE,
91 format!("column \"{}\" is a generated column", columns[i]),
92 )
93 .with_detail("Generated columns cannot be used in COPY."));
94 }
95 Ok(CopyTarget { relid, name, kind, columns, selected })
96}
97
98struct WhereColumns<'a> {
100 relname: &'a str,
101 columns: &'a [String],
102 used: Vec<usize>,
103}
104
105impl Walker for WhereColumns<'_> {
106 fn node(&mut self, node: &Node) -> bool {
107 let Node::ColumnRef(cref) = node else { return walk::walk(self, node) };
108 let fields: Vec<Option<&str>> = cref.fields.iter().flatten().map(string).collect();
109 let position = |name: &str| self.columns.iter().position(|c| c == name);
110 match fields.as_slice() {
111 [Some(name)] if position(name).is_some() => self.used.extend(position(name)),
112 [Some(rel)] | [Some(rel), None] if *rel == self.relname => {
113 self.used.extend(0..self.columns.len());
114 }
115 [Some(rel), Some(name)] if *rel == self.relname => self.used.extend(position(name)),
116 _ => {}
117 }
118 false
119 }
120
121 fn select(&mut self, stmt: &SelectStmt) -> bool {
122 walk::select(self, stmt)
123 }
124}
125
126pub fn copy_relation(stmt: &CopyStmt, rv: &RangeVar, env: &dyn Env) -> Result<CopyRelation> {
132 let CopyTarget { relid, name, kind, columns, selected } = copy_target(stmt, rv, env)?;
133 let mut partitions = None;
134 match kind {
135 b'r' => {}
136 b'p' => partitions = Some(copy_partitions(env, relid, &name)?),
137 b'm' => {
138 let rel = env.catalog().and_then(|c| c.relation(relid));
139 if rel.and_then(|r| r.view.as_ref()).is_some_and(|v| !v.populated) {
140 return Err(Error::new(
141 SqlState::FEATURE_NOT_SUPPORTED,
142 format!("cannot copy from unpopulated materialized view \"{name}\""),
143 )
144 .with_hint("Use the REFRESH MATERIALIZED VIEW command."));
145 }
146 }
147 b'f' => {
148 return Err(Error::new(
149 SqlState::WRONG_OBJECT_TYPE,
150 format!("cannot copy from foreign table \"{name}\""),
151 )
152 .with_hint("Try the COPY (SELECT ...) TO variant."));
153 }
154 b'v' => {
155 return Err(Error::new(
156 SqlState::WRONG_OBJECT_TYPE,
157 format!("cannot copy from view \"{name}\""),
158 )
159 .with_hint("Try the COPY (SELECT ...) TO variant."));
160 }
161 b'S' => {
162 return Err(Error::new(
163 SqlState::WRONG_OBJECT_TYPE,
164 format!("cannot copy from sequence \"{name}\""),
165 ));
166 }
167 _ => {
168 return Err(Error::new(
169 SqlState::WRONG_OBJECT_TYPE,
170 format!("cannot copy from non-table relation \"{name}\""),
171 ));
172 }
173 }
174 let tableoid = partitions.as_ref().map(|_| "tableoid");
175 let targets = selected
176 .iter()
177 .map(|&i| columns[i].as_str())
178 .chain(tableoid)
179 .map(|column| {
180 let field = Some(Node::String(column.into()));
181 let column = ColumnRef { fields: vec![field], location: -1 };
182 let target = ResTarget { val: Some(column.into()), location: -1, ..Default::default() };
183 Some(target.into())
184 })
185 .collect();
186 let from = RangeVar { inh: partitions.is_some(), alias: None, location: -1, ..rv.clone() };
187 let select = SelectStmt {
188 targetList: targets,
189 fromClause: vec![Some(from.into())],
190 ..Default::default()
191 };
192 Ok(CopyRelation { name, columns, selected, query: select.into(), partitions })
193}
194
195fn copy_partitions(env: &dyn Env, root: u32, name: &str) -> Result<Vec<u32>> {
201 let Some(catalog) = env.catalog() else { return Ok(Vec::new()) };
202 let mut all = vec![root];
203 let mut at = 0;
204 while let Some(&oid) = all.get(at) {
205 at += 1;
206 all.extend(catalog.partitions(oid).iter().map(|rel| rel.oid));
207 }
208 let mut out = Vec::new();
209 for oid in all {
210 match relkind(env, oid) {
211 b'p' => {}
212 b'f' => {
213 let child = catalog.relation(oid).map_or("", |rel| rel.name.as_str());
214 return Err(Error::new(
215 SqlState::WRONG_OBJECT_TYPE,
216 format!("cannot copy from foreign table \"{child}\""),
217 )
218 .with_detail(format!(
219 "Partition \"{child}\" is a foreign table in partitioned table \"{name}\"."
220 ))
221 .with_hint("Try the COPY (SELECT ...) TO variant."));
222 }
223 _ => out.push(oid),
224 }
225 }
226 Ok(out)
227}
228
229#[must_use]
231pub fn copy_insert(target: &CopyTarget, rv: &RangeVar, defaults: &[bool]) -> Node {
232 let cols: Vec<_> = target
233 .selected
234 .iter()
235 .enumerate()
236 .filter(|&(k, _)| !defaults.get(k).copied().unwrap_or(false))
237 .map(|(_, &i)| {
238 let name = Some(target.columns[i].as_str().into());
239 Some(ResTarget { name, location: -1, ..Default::default() }.into())
240 })
241 .collect();
242 let select = (!cols.is_empty()).then(|| values(cols.len()).into());
244 let insert = InsertStmt {
245 relation: Some(Box::new(RangeVar { alias: None, location: -1, ..rv.clone() })),
246 cols,
247 selectStmt: select,
248 r#override: OverridingKind::OVERRIDING_SYSTEM_VALUE,
249 ..Default::default()
250 };
251 insert.into()
252}
253
254#[must_use]
256pub fn copy_filter(target: &CopyTarget, rv: &RangeVar, condition: &Node) -> Node {
257 let colnames =
258 target.columns.iter().map(|name| Some(Node::String(name.as_str().into()))).collect();
259 let alias = Alias { aliasname: rv.relname.clone(), colnames };
260 let from = RangeSubselect {
261 lateral: false,
262 subquery: Some(values(target.columns.len()).into()),
263 alias: Some(Box::new(alias)),
264 };
265 let target = ResTarget { val: Some(condition.clone()), location: -1, ..Default::default() };
266 let select = SelectStmt {
267 targetList: vec![Some(target.into())],
268 fromClause: vec![Some(from.into())],
269 ..Default::default()
270 };
271 select.into()
272}
273
274#[must_use]
276pub fn rewrite_insert(target: &CopyTarget, rv: &RangeVar, transforms: &[Option<&Node>]) -> Node {
277 let names: Vec<&str> = target.selected.iter().map(|&i| target.columns[i].as_str()).collect();
278 let cols = names
279 .iter()
280 .map(|&name| {
281 Some(ResTarget { name: Some(name.into()), location: -1, ..Default::default() }.into())
282 })
283 .collect();
284 let targets = names
285 .iter()
286 .enumerate()
287 .map(|(k, &name)| {
288 let value = match transforms.get(k).copied().flatten() {
289 Some(expr) => expr.clone(),
290 None => {
291 let field = Some(Node::String(name.into()));
292 ColumnRef { fields: vec![field], location: -1 }.into()
293 }
294 };
295 Some(ResTarget { val: Some(value), location: -1, ..Default::default() }.into())
296 })
297 .collect();
298 let colnames = names.iter().map(|&name| Some(Node::String(name.into()))).collect();
299 let alias = Alias { aliasname: rv.relname.clone(), colnames };
300 let from = RangeSubselect {
301 lateral: false,
302 subquery: Some(values(names.len()).into()),
303 alias: Some(Box::new(alias)),
304 };
305 let select = SelectStmt {
306 targetList: targets,
307 fromClause: vec![Some(from.into())],
308 ..Default::default()
309 };
310 let insert = InsertStmt {
311 relation: Some(Box::new(RangeVar { alias: None, location: -1, ..rv.clone() })),
312 cols,
313 selectStmt: Some(select.into()),
314 r#override: OverridingKind::OVERRIDING_SYSTEM_VALUE,
315 ..Default::default()
316 };
317 insert.into()
318}
319
320fn values(count: usize) -> SelectStmt {
322 let params = (1..=count)
323 .map(|n| {
324 let number = i32::try_from(n).unwrap_or(i32::MAX);
325 Some(ParamRef { number, location: -1 }.into())
326 })
327 .collect();
328 SelectStmt { valuesLists: vec![Some(Node::List(params))], ..Default::default() }
329}
330
331pub fn attnums(columns: &[String], names: &[String], rel: Option<&str>) -> Result<Vec<usize>> {
337 if names.is_empty() {
338 return Ok((0..columns.len()).collect());
339 }
340 let mut out = Vec::with_capacity(names.len());
341 for name in names {
342 let Some(at) = columns.iter().position(|c| c == name) else {
343 let message = match rel {
344 Some(rel) => format!("column \"{name}\" of relation \"{rel}\" does not exist"),
345 None => format!("column \"{name}\" does not exist"),
346 };
347 return Err(Error::new(SqlState::UNDEFINED_COLUMN, message));
348 };
349 if out.contains(&at) {
350 return Err(Error::new(
351 SqlState::DUPLICATE_COLUMN,
352 format!("column \"{name}\" specified more than once"),
353 ));
354 }
355 out.push(at);
356 }
357 Ok(out)
358}
359
360pub(super) fn relation_columns(env: &dyn Env, relid: u32) -> Vec<String> {
362 if let Some(rel) = env.catalog().and_then(|c| c.relation(relid)) {
363 return rel.columns.iter().filter(|c| !c.dropped).map(|c| c.name.clone()).collect();
364 }
365 builtin::attributes(relid)
366 .iter()
367 .filter(|a| a.num > 0 && !a.dropped)
368 .map(|a| a.name.to_string())
369 .collect()
370}
371
372fn string(node: &Node) -> Option<&str> {
374 match node {
375 Node::String(s) => Some(s),
376 _ => None,
377 }
378}