1use pg_query::protobuf::Node;
10use pg_query::NodeEnum;
11
12use crate::ast::{ColumnType, RangeSubtype};
13use crate::error::{Result, SQLError};
14
15use super::tree::extract_string;
16
17#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct ParsedRegtypeName {
20 pub names: Vec<String>,
21 pub array_dimensions: usize,
22}
23
24#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct ParsedRegprocedureName {
27 pub names: Vec<String>,
28 pub argument_types: Option<Vec<ParsedRegtypeName>>,
29}
30
31const POSTGRES_IDENTIFIER_MAX_BYTES: usize = 63;
32const POSTGRES_FUNCTION_MAX_ARGUMENTS: usize = 100;
33
34fn scanner_isspace(byte: u8) -> bool {
35 matches!(byte, b' ' | b'\t' | b'\n' | b'\r' | 0x0b | 0x0c)
36}
37
38fn truncate_postgres_identifier(mut identifier: String) -> String {
39 if identifier.len() <= POSTGRES_IDENTIFIER_MAX_BYTES {
40 return identifier;
41 }
42 let mut end = POSTGRES_IDENTIFIER_MAX_BYTES;
43 while !identifier.is_char_boundary(end) {
44 end -= 1;
45 }
46 identifier.truncate(end);
47 identifier
48}
49
50#[must_use]
52pub fn parse_regobject_name(input: &str) -> Option<Vec<String>> {
53 let bytes = input.as_bytes();
54 let mut offset = 0usize;
55 while bytes.get(offset).is_some_and(|byte| scanner_isspace(*byte)) {
56 offset += 1;
57 }
58 if offset == bytes.len() {
59 return None;
60 }
61
62 let mut names = Vec::new();
63 loop {
64 let component = if bytes[offset] == b'"' {
65 offset += 1;
66 let mut quoted = String::new();
67 loop {
68 let relative = bytes[offset..].iter().position(|byte| *byte == b'"')?;
69 let quote = offset + relative;
70 quoted.push_str(&input[offset..quote]);
71 offset = quote + 1;
72 if bytes.get(offset) == Some(&b'"') {
73 quoted.push('"');
74 offset += 1;
75 continue;
76 }
77 break;
78 }
79 quoted
80 } else {
81 let start = offset;
82 while bytes
83 .get(offset)
84 .is_some_and(|byte| *byte != b'.' && !scanner_isspace(*byte))
85 {
86 offset += 1;
87 }
88 if offset == start {
89 return None;
90 }
91 input[start..offset].to_ascii_lowercase()
92 };
93 names.push(truncate_postgres_identifier(component));
94
95 while bytes.get(offset).is_some_and(|byte| scanner_isspace(*byte)) {
96 offset += 1;
97 }
98 match bytes.get(offset) {
99 None => return Some(names),
100 Some(b'.') => {
101 offset += 1;
102 while bytes.get(offset).is_some_and(|byte| scanner_isspace(*byte)) {
103 offset += 1;
104 }
105 if offset == bytes.len() {
106 return None;
107 }
108 }
109 Some(_) => return None,
110 }
111 }
112}
113
114pub fn parse_regtype_name(input: &str) -> Result<Option<ParsedRegtypeName>> {
116 if input.bytes().all(scanner_isspace) {
117 return Ok(None);
118 }
119 let parsed = pg_query::parse_with_mode(input, pg_query::ParseMode::TypeName)?;
120 let [raw] = parsed.protobuf.stmts.as_slice() else {
121 return Ok(None);
122 };
123 let Some(NodeEnum::List(names)) = raw.stmt.as_ref().and_then(|node| node.node.as_ref()) else {
124 return Ok(None);
125 };
126 let names = names
127 .items
128 .iter()
129 .map(extract_string)
130 .collect::<Result<Vec<_>>>()?;
131 if names.is_empty() {
132 return Ok(None);
133 }
134 let scanned = pg_query::scan(input)?;
135 let tokens = scanned
136 .tokens
137 .iter()
138 .filter_map(|token| pg_query::protobuf::Token::try_from(token.token).ok())
139 .collect::<Vec<_>>();
140 if tokens.contains(&pg_query::protobuf::Token::Setof) {
141 return Ok(None);
142 }
143 let bracket_dimensions = tokens
144 .iter()
145 .filter(|token| **token == pg_query::protobuf::Token::Ascii91)
146 .count();
147 let array_dimensions = bracket_dimensions.max(usize::from(
148 tokens.contains(&pg_query::protobuf::Token::Array),
149 ));
150 Ok(Some(ParsedRegtypeName {
151 names,
152 array_dimensions,
153 }))
154}
155
156pub fn parse_regprocedure_name(input: &str) -> Result<Option<ParsedRegprocedureName>> {
158 let mut in_quote = false;
159 let left_parenthesis = input.bytes().enumerate().find_map(|(offset, byte)| {
160 if byte == b'"' {
161 in_quote = !in_quote;
162 None
163 } else if byte == b'(' && !in_quote {
164 Some(offset)
165 } else {
166 None
167 }
168 });
169 let Some(left_parenthesis) = left_parenthesis else {
170 return Ok(
171 parse_regobject_name(input).map(|names| ParsedRegprocedureName {
172 names,
173 argument_types: None,
174 }),
175 );
176 };
177 let Some(names) = parse_regobject_name(&input[..left_parenthesis]) else {
178 return Ok(None);
179 };
180
181 let bytes = input.as_bytes();
182 let mut end = bytes.len();
183 while end > left_parenthesis + 1 && scanner_isspace(bytes[end - 1]) {
184 end -= 1;
185 }
186 if end <= left_parenthesis + 1 || bytes[end - 1] != b')' {
187 return Err(SQLError::Parse(format!(
188 "expected a right parenthesis in routine identity \"{input}\""
189 )));
190 }
191 let arguments = &input[left_parenthesis + 1..end - 1];
192 let argument_bytes = arguments.as_bytes();
193 let mut argument_types = Vec::new();
194 let mut offset = 0usize;
195 let mut had_comma = false;
196 loop {
197 while argument_bytes
198 .get(offset)
199 .is_some_and(|byte| scanner_isspace(*byte))
200 {
201 offset += 1;
202 }
203 if offset == argument_bytes.len() {
204 if had_comma {
205 return Err(SQLError::Parse(format!(
206 "expected a type name in routine identity \"{input}\""
207 )));
208 }
209 break;
210 }
211
212 let start = offset;
213 let mut quoted = false;
214 let mut nesting = 0i32;
215 while let Some(byte) = argument_bytes.get(offset).copied() {
216 if byte == b'"' {
217 quoted = !quoted;
218 } else if byte == b',' && !quoted && nesting == 0 {
219 break;
220 } else if !quoted {
221 match byte {
222 b'(' | b'[' => nesting += 1,
223 b')' | b']' => nesting -= 1,
224 _ => {}
225 }
226 }
227 offset += 1;
228 }
229 if quoted || nesting != 0 {
230 return Err(SQLError::Parse(format!(
231 "improper type name in routine identity \"{input}\""
232 )));
233 }
234 let mut type_end = offset;
235 while type_end > start && scanner_isspace(argument_bytes[type_end - 1]) {
236 type_end -= 1;
237 }
238 let Some(type_name) = parse_regtype_name(&arguments[start..type_end])? else {
239 return Ok(None);
240 };
241 if argument_types.len() == POSTGRES_FUNCTION_MAX_ARGUMENTS {
242 return Err(SQLError::Parse(format!(
243 "too many arguments in routine identity \"{input}\""
244 )));
245 }
246 argument_types.push(type_name);
247 had_comma = argument_bytes.get(offset) == Some(&b',');
248 if had_comma {
249 offset += 1;
250 }
251 }
252
253 Ok(Some(ParsedRegprocedureName {
254 names,
255 argument_types: Some(argument_types),
256 }))
257}
258
259pub(super) fn compile_foreign_key_action(raw: &str) -> Result<crate::ast::ForeignKeyAction> {
260 use crate::ast::ForeignKeyAction;
261 match raw.as_bytes().first().copied() {
262 None | Some(0) | Some(b'a') => Ok(ForeignKeyAction::NoAction),
263 Some(b'r') => Ok(ForeignKeyAction::Restrict),
264 Some(b'c') => Ok(ForeignKeyAction::Cascade),
265 Some(b'n') => Ok(ForeignKeyAction::SetNull),
266 Some(b'd') => Ok(ForeignKeyAction::SetDefault),
267 Some(other) => Err(SQLError::Unsupported(format!(
268 "unsupported FOREIGN KEY action byte {other:?}"
269 ))),
270 }
271}
272
273pub(super) fn compile_foreign_key_match(raw: &str) -> Result<crate::ast::ForeignKeyMatch> {
274 use crate::ast::ForeignKeyMatch;
275 match raw.as_bytes().first().copied() {
276 None | Some(0) | Some(b's') => Ok(ForeignKeyMatch::Simple),
277 Some(b'f') => Ok(ForeignKeyMatch::Full),
278 Some(b'p') => Err(SQLError::Unsupported(
279 "FOREIGN KEY MATCH PARTIAL is not implemented by PostgreSQL".into(),
280 )),
281 Some(other) => Err(SQLError::Unsupported(format!(
282 "unsupported FOREIGN KEY match byte {other:?}"
283 ))),
284 }
285}
286
287pub(super) fn validate_foreign_key_set_columns(
288 local_columns: &[String],
289 set_columns: &[String],
290 raw_delete_action: &str,
291) -> Result<()> {
292 if set_columns.is_empty() {
293 return Ok(());
294 }
295 let action = compile_foreign_key_action(raw_delete_action)?;
296 if !matches!(
297 action,
298 crate::ast::ForeignKeyAction::SetNull | crate::ast::ForeignKeyAction::SetDefault
299 ) {
300 return Err(SQLError::Unsupported(
301 "FOREIGN KEY column lists are only valid for ON DELETE SET NULL/DEFAULT".into(),
302 ));
303 }
304 for col in set_columns {
305 if !local_columns.iter().any(|local| local == col) {
306 return Err(SQLError::Unsupported(format!(
307 "FOREIGN KEY SET column `{col}` is not part of the local key"
308 )));
309 }
310 }
311 Ok(())
312}
313
314pub(super) fn raw_type_name(col: &pg_query::protobuf::ColumnDef) -> Result<Option<String>> {
315 let Some(type_name) = col.type_name.as_ref() else {
316 return Ok(None);
317 };
318 let names = type_name
319 .names
320 .iter()
321 .map(extract_string)
322 .collect::<Result<Vec<_>>>()?;
323 Ok(names.last().map(|name| name.to_lowercase()))
324}
325
326pub(super) fn compile_type_name(col: &pg_query::protobuf::ColumnDef) -> Result<ColumnType> {
327 let Some(type_name) = col.type_name.as_ref() else {
328 return Err(SQLError::Internal(format!(
329 "column `{}` has no type",
330 col.colname
331 )));
332 };
333 compile_pg_type_name(type_name, &col.colname)
334}
335
336#[expect(
337 clippy::too_many_lines,
338 reason = "ordered PostgreSQL lowering preserves syntax and error precedence"
339)]
340pub(super) fn compile_pg_type_name(
341 type_name: &pg_query::protobuf::TypeName,
342 column_name: &str,
343) -> Result<ColumnType> {
344 let names = type_name
345 .names
346 .iter()
347 .map(extract_string)
348 .collect::<Result<Vec<_>>>()?;
349 let raw = names
350 .last()
351 .ok_or_else(|| {
352 SQLError::Internal(format!(
353 "type name for `{column_name}` has no name components"
354 ))
355 })?
356 .to_lowercase();
357 let base = match raw.as_str() {
358 "smallint" | "int2" | "smallserial" | "serial2" => Ok(ColumnType::SmallInteger),
359 "int" | "int4" | "integer" | "serial" | "serial4" => Ok(ColumnType::Integer),
360 "bigint" | "int8" | "bigserial" | "serial8" => Ok(ColumnType::BigInteger),
361 "oid" => Ok(ColumnType::Oid),
362 "xid" => Ok(ColumnType::Xid),
363 "void" => Ok(ColumnType::Void),
364 "text" => Ok(ColumnType::Text),
365 "name" => Ok(ColumnType::Name),
366 "uuid" => Ok(ColumnType::Uuid),
367 "varchar" | "character varying" => {
368 if type_name.typmods.len() > 1 {
369 return Err(SQLError::TypeMismatch(format!(
370 "CHARACTER VARYING accepts at most one length modifier, got {}",
371 type_name.typmods.len()
372 )));
373 }
374 let length = type_name
375 .typmods
376 .first()
377 .map(expect_positive_character_length)
378 .transpose()?;
379 Ok(ColumnType::Varchar(length))
380 }
381 "character" | "char" | "bpchar" => {
382 if type_name.typmods.len() > 1 {
383 return Err(SQLError::TypeMismatch(format!(
384 "CHARACTER accepts at most one length modifier, got {}",
385 type_name.typmods.len()
386 )));
387 }
388 let length = type_name
389 .typmods
390 .first()
391 .map(expect_positive_character_length)
392 .transpose()?
393 .unwrap_or(1);
394 Ok(ColumnType::Character(length))
395 }
396 "bool" | "boolean" => Ok(ColumnType::Boolean),
397 "real" | "float4" => Ok(ColumnType::Real),
398 "float8" | "double" | "double precision" => Ok(ColumnType::DoublePrecision),
399 "numeric" | "decimal" => {
400 if type_name.typmods.len() > 2 {
401 return Err(SQLError::TypeMismatch(format!(
402 "NUMERIC accepts at most precision and scale, got {} modifiers",
403 type_name.typmods.len()
404 )));
405 }
406 let mut typmods_iter = type_name.typmods.iter();
407 let precision = typmods_iter
408 .next()
409 .map(|n| {
410 let value = expect_integer_const(n)?;
411 if !(1..=1000).contains(&value) {
412 return Err(SQLError::TypeMismatch(format!(
413 "NUMERIC precision must be between 1 and 1000, got {value}"
414 )));
415 }
416 Ok(value as u32)
417 })
418 .transpose()?;
419 let scale = typmods_iter
420 .next()
421 .map(|n| {
422 let value = expect_integer_const(n)?;
423 if !(-1000..=1000).contains(&value) {
424 return Err(SQLError::TypeMismatch(format!(
425 "NUMERIC scale must be between -1000 and 1000, got {value}"
426 )));
427 }
428 Ok(value as i32)
429 })
430 .transpose()?;
431 let scale = scale.or(precision.map(|_| 0));
434 Ok(ColumnType::Numeric { precision, scale })
435 }
436 "date" => Ok(ColumnType::Date),
437 "time" | "time without time zone" => Ok(ColumnType::Time),
438 "timetz" | "time with time zone" => Ok(ColumnType::TimeTz),
439 "timestamp" | "datetime" | "timestamp without time zone" => Ok(ColumnType::Timestamp),
440 "timestamptz" | "timestamp with time zone" => Ok(ColumnType::TimestampTz),
441 "interval" => Ok(ColumnType::Interval),
442 "int4range" => Ok(ColumnType::Range(RangeSubtype::Integer)),
443 "int8range" => Ok(ColumnType::Range(RangeSubtype::BigInteger)),
444 "numrange" => Ok(ColumnType::Range(RangeSubtype::Numeric)),
445 "daterange" => Ok(ColumnType::Range(RangeSubtype::Date)),
446 "tsrange" => Ok(ColumnType::Range(RangeSubtype::Timestamp)),
447 "tstzrange" => Ok(ColumnType::Range(RangeSubtype::TimestampTz)),
448 "int4multirange" => Ok(ColumnType::Multirange(RangeSubtype::Integer)),
449 "int8multirange" => Ok(ColumnType::Multirange(RangeSubtype::BigInteger)),
450 "nummultirange" => Ok(ColumnType::Multirange(RangeSubtype::Numeric)),
451 "datemultirange" => Ok(ColumnType::Multirange(RangeSubtype::Date)),
452 "tsmultirange" => Ok(ColumnType::Multirange(RangeSubtype::Timestamp)),
453 "tstzmultirange" => Ok(ColumnType::Multirange(RangeSubtype::TimestampTz)),
454 "json" => Ok(ColumnType::Json),
455 "jsonb" => Ok(ColumnType::JsonB),
456 "bytea" => Ok(ColumnType::Bytea),
457 "regproc" => Ok(ColumnType::Regproc),
458 "regprocedure" => Ok(ColumnType::Regprocedure),
459 "regclass" => Ok(ColumnType::Regclass),
460 "regnamespace" => Ok(ColumnType::Regnamespace),
461 "regrole" => Ok(ColumnType::Regrole),
462 "regtype" => Ok(ColumnType::Regtype),
463 "pg_node_tree" => Ok(ColumnType::PgNodeTree),
464 "aclitem" => Ok(ColumnType::AclItem),
465 "int2vector" => Ok(ColumnType::Int2Vector),
466 "oidvector" => Ok(ColumnType::OidVector),
467 "anyarray" => Ok(ColumnType::AnyArray),
468 "record" => Ok(ColumnType::Record),
469 "vector" => {
470 let [arg] = type_name.typmods.as_slice() else {
472 return Err(SQLError::Unsupported(
473 "VECTOR requires exactly one dimension".into(),
474 ));
475 };
476 let raw_dim = expect_integer_const(arg)?;
477 let dim = u32::try_from(raw_dim).map_err(|_| {
478 SQLError::TypeMismatch(format!(
479 "VECTOR dimension must be between 1 and {}, got {raw_dim}",
480 u32::MAX
481 ))
482 })?;
483 if dim == 0 {
484 return Err(SQLError::TypeMismatch(
485 "VECTOR dimension must be greater than zero".into(),
486 ));
487 }
488 Ok(ColumnType::Vector(dim))
489 }
490 "tensor" => {
491 let [arg] = type_name.typmods.as_slice() else {
493 return Err(SQLError::Unsupported(
494 "TENSOR requires exactly one dimension".into(),
495 ));
496 };
497 let raw_dim = expect_integer_const(arg)?;
498 let dim = u32::try_from(raw_dim).map_err(|_| {
499 SQLError::TypeMismatch(format!(
500 "TENSOR dimension must be between 1 and {}, got {raw_dim}",
501 u32::MAX
502 ))
503 })?;
504 if dim == 0 {
505 return Err(SQLError::TypeMismatch(
506 "TENSOR dimension must be greater than zero".into(),
507 ));
508 }
509 Ok(ColumnType::Tensor(dim))
510 }
511 other => Err(SQLError::Unsupported(format!(
512 "column `{column_name}` type `{other}` is not supported"
513 ))),
514 }?;
515 if matches!(base, ColumnType::Void) && !type_name.array_bounds.is_empty() {
516 return Err(SQLError::Routine {
517 sqlstate: "42704".into(),
518 message: "type \"void[]\" does not exist".into(),
519 });
520 }
521 Ok(type_name
522 .array_bounds
523 .iter()
524 .fold(base, |element, _| ColumnType::Array(Box::new(element))))
525}
526
527fn expect_positive_character_length(node: &Node) -> Result<u32> {
528 let length = expect_integer_const(node)?;
529 u32::try_from(length)
530 .ok()
531 .filter(|length| *length > 0)
532 .ok_or_else(|| {
533 SQLError::TypeMismatch(format!(
534 "character length must be greater than zero, got {length}"
535 ))
536 })
537}
538
539fn expect_integer_const(node: &Node) -> Result<i64> {
540 let Some(inner) = node.node.as_ref() else {
541 return Err(SQLError::Internal("missing const node".into()));
542 };
543 match inner {
544 NodeEnum::AConst(c) => match &c.val {
545 Some(pg_query::protobuf::a_const::Val::Ival(i)) => Ok(i64::from(i.ival)),
546 Some(pg_query::protobuf::a_const::Val::Fval(f)) => {
547 f.fval.parse::<i64>().map_err(|_| {
548 SQLError::TypeMismatch(format!(
549 "type modifier must be an integer, got `{}`",
550 f.fval
551 ))
552 })
553 }
554 other => Err(SQLError::Internal(format!(
555 "expected integer constant, got {other:?}"
556 ))),
557 },
558 _ => Err(SQLError::Internal(format!(
559 "expected A_Const, got {inner:?}"
560 ))),
561 }
562}