diff --git a/crates/pgls_pretty_print/src/emitter.rs b/crates/pgls_pretty_print/src/emitter.rs index 018ecf090..9d1c5f758 100644 --- a/crates/pgls_pretty_print/src/emitter.rs +++ b/crates/pgls_pretty_print/src/emitter.rs @@ -61,6 +61,8 @@ pub struct EventEmitter { comments_after_separator: Vec, /// `config.layout` as configured, while `with_layout` overrides it for part of a statement. configured_layout: crate::Layout, + /// Identifiers the parser truncated, by their truncated name, as written in the source. + truncated_identifiers: HashMap, } impl EventEmitter { @@ -71,6 +73,7 @@ impl EventEmitter { config, leading_comments: HashMap::new(), trailing_comments: HashMap::new(), + truncated_identifiers: HashMap::new(), leading_line_comments_require_break: false, comments_after_separator: Vec::new(), } @@ -105,11 +108,25 @@ impl EventEmitter { config, leading_comments, trailing_comments, + truncated_identifiers: HashMap::new(), leading_line_comments_require_break: false, comments_after_separator: Vec::new(), } } + /// Identifiers of the source that the parser truncated, by their truncated name. The emitter + /// writes them as written rather than as the AST has them. + pub fn set_truncated_identifiers(&mut self, identifiers: HashMap) { + self.truncated_identifiers = identifiers; + } + + /// `name` as written in the source, when the parser truncated it. + pub fn source_identifier<'a>(&'a self, name: &'a str) -> &'a str { + self.truncated_identifiers + .get(name) + .map_or(name, String::as_str) + } + pub fn token(&mut self, token: TokenKind) { // A comment held for a comma follows it. Other tokens closing the item before the comma, // a parenthesis or an alias for instance, are written first. A statement never ends on a diff --git a/crates/pgls_pretty_print/src/identifiers.rs b/crates/pgls_pretty_print/src/identifiers.rs new file mode 100644 index 000000000..6974f517a --- /dev/null +++ b/crates/pgls_pretty_print/src/identifiers.rs @@ -0,0 +1,128 @@ +use std::collections::HashMap; + +use pgls_query::protobuf::Token; + +/// The longest identifier Postgres keeps, `NAMEDATALEN - 1` bytes. The parser cuts longer ones +/// down, so the AST only has their first 63 bytes. +const MAX_IDENTIFIER_LENGTH: usize = 63; + +/// The identifiers of `sql` that are longer than Postgres keeps, by the name the parser cut them +/// down to. +/// +/// The formatter writes names from the AST. Without this, it would rewrite a long identifier to its +/// truncated form, which means the same to Postgres but is not what was written. When two +/// different spellings are cut to the same name, neither is restored. +pub(crate) fn truncated_identifiers(sql: &str) -> HashMap { + let Ok(scan) = pgls_query::scan(sql) else { + return HashMap::new(); + }; + + let mut by_truncated: HashMap> = HashMap::new(); + for token in &scan.tokens { + if Token::try_from(token.token) != Ok(Token::Ident) { + continue; + } + let Some(text) = sql.get(token.start as usize..token.end as usize) else { + continue; + }; + let name = identifier_name(text); + // A name of exactly the maximum length takes part too: written next to a longer one with + // the same prefix, it is not clear which spelling to restore. + if name.len() < MAX_IDENTIFIER_LENGTH { + continue; + } + + let truncated = truncate(&name).to_string(); + by_truncated + .entry(truncated) + .and_modify(|existing| { + if existing.as_deref() != Some(name.as_str()) { + *existing = None; + } + }) + .or_insert(Some(name)); + } + + by_truncated + .into_iter() + .filter_map(|(truncated, name)| Some((truncated, name?))) + .filter(|(truncated, name)| truncated != name) + .collect() +} + +/// The name an identifier token stands for: a quoted identifier without its quotes, an unquoted +/// one in lower case, as `downcase_identifier` does for ASCII letters. +fn identifier_name(text: &str) -> String { + match text + .strip_prefix('"') + .and_then(|rest| rest.strip_suffix('"')) + { + Some(quoted) => quoted.replace("\"\"", "\""), + None => text.to_ascii_lowercase(), + } +} + +/// The longest prefix of whole characters that fits in `MAX_IDENTIFIER_LENGTH` bytes, as +/// `truncate_identifier` keeps it. +fn truncate(name: &str) -> &str { + let mut end = MAX_IDENTIFIER_LENGTH; + while !name.is_char_boundary(end) { + end -= 1; + } + &name[..end] +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_long_identifier_is_found_by_its_truncated_name() { + let long = "a_role_name_that_is_definitely_way_too_long_for_anything_postgres_keeps"; + let identifiers = truncated_identifiers(&format!("grant r to {long}")); + + assert_eq!(identifiers.get(&long[..63]).map(String::as_str), Some(long)); + assert_eq!(identifiers.len(), 1); + } + + #[test] + fn a_quoted_identifier_keeps_its_case_and_spaces() { + let long = "A Quoted Name That Is Definitely Way Too Long For Anything Postgres Keeps"; + let identifiers = truncated_identifiers(&format!("select 1 as \"{long}\"")); + + assert_eq!(identifiers.get(&long[..63]).map(String::as_str), Some(long)); + } + + #[test] + fn a_multibyte_character_is_not_split() { + // 62 ASCII bytes followed by a two-byte character: Postgres keeps the 62 bytes. + let long = format!("{}é_and_more", "a".repeat(62)); + let identifiers = truncated_identifiers(&format!("select 1 as {long}")); + + assert_eq!( + identifiers.get(&"a".repeat(62)).map(String::as_str), + Some(long.as_str()) + ); + } + + #[test] + fn two_identifiers_cut_to_the_same_name_are_not_restored() { + let prefix = "x".repeat(63); + let identifiers = truncated_identifiers(&format!("select {prefix}_one, {prefix}_two")); + + assert!(identifiers.is_empty()); + } + + #[test] + fn a_name_written_both_truncated_and_in_full_is_not_restored() { + let prefix = "x".repeat(63); + let identifiers = truncated_identifiers(&format!("select {prefix}, {prefix}_long")); + + assert!(identifiers.is_empty()); + } + + #[test] + fn short_identifiers_are_ignored() { + assert!(truncated_identifiers("select a, \"B\" from t").is_empty()); + } +} diff --git a/crates/pgls_pretty_print/src/lib.rs b/crates/pgls_pretty_print/src/lib.rs index 469715287..29aa10da2 100644 --- a/crates/pgls_pretty_print/src/lib.rs +++ b/crates/pgls_pretty_print/src/lib.rs @@ -1,6 +1,7 @@ mod codegen; pub mod comments; pub mod emitter; +mod identifiers; pub mod nodes; pub mod normalize; pub mod renderer; @@ -222,7 +223,7 @@ pub fn format_statement( || !attached.unattached.is_empty(); if !has_comments { - return format_statement_once(ast, config, attached); + return format_statement_once(ast, sql, config, attached); } let mut current_sql = sql.to_string(); @@ -231,7 +232,7 @@ pub fn format_statement( let mut seen_layouts = HashSet::from([current_sql.clone()]); for pass in 1..=MAX_COMMENT_FORMAT_PASSES { - let result = format_statement_once(¤t_ast, config, attached)?; + let result = format_statement_once(¤t_ast, ¤t_sql, config, attached)?; if result.formatted == current_sql { return Ok(result); } @@ -260,6 +261,7 @@ pub fn format_statement( /// Formats a statement exactly once, including semantic verification. fn format_statement_once( ast: &NodeEnum, + sql: &str, config: &FormatConfig, attached: AttachedComments, ) -> Result { @@ -276,6 +278,7 @@ fn format_statement_once( attached.leading_by_location, attached.trailing_by_location, ); + emitter.set_truncated_identifiers(identifiers::truncated_identifiers(sql)); nodes::emit_node_enum(ast, &mut emitter); let pending = emitter.pending_comments(); if pending > 0 { diff --git a/crates/pgls_pretty_print/src/nodes/string.rs b/crates/pgls_pretty_print/src/nodes/string.rs index d4425b2e9..732f617d8 100644 --- a/crates/pgls_pretty_print/src/nodes/string.rs +++ b/crates/pgls_pretty_print/src/nodes/string.rs @@ -25,6 +25,7 @@ pub(super) fn emit_string_identifier(e: &mut EventEmitter, n: &PgString) { } pub(super) fn emit_identifier(e: &mut EventEmitter, value: &str) { + let value = e.source_identifier(value); let escaped = value.replace('"', "\"\""); e.token(TokenKind::IDENT(format!("\"{escaped}\""))); } @@ -42,10 +43,11 @@ pub(super) fn emit_identifier_maybe_quoted(e: &mut EventEmitter, value: &str) { return; } - if needs_quoting_as_column_name(value) { - emit_identifier(e, value); + let value = e.source_identifier(value).to_string(); + if needs_quoting_as_column_name(&value) { + emit_identifier(e, &value); } else { - e.token(TokenKind::IDENT(value.to_string())); + e.token(TokenKind::IDENT(value)); } } @@ -58,10 +60,11 @@ pub(super) fn emit_type_identifier_maybe_quoted(e: &mut EventEmitter, value: &st return; } - if needs_quoting(value) { - emit_identifier(e, value); + let value = e.source_identifier(value).to_string(); + if needs_quoting(&value) { + emit_identifier(e, &value); } else { - e.token(TokenKind::TYPE_IDENT(value.to_string())); + e.token(TokenKind::TYPE_IDENT(value)); } } diff --git a/crates/pgls_pretty_print/tests/source_fidelity.rs b/crates/pgls_pretty_print/tests/source_fidelity.rs index d34a58d13..443d022ed 100644 --- a/crates/pgls_pretty_print/tests/source_fidelity.rs +++ b/crates/pgls_pretty_print/tests/source_fidelity.rs @@ -325,3 +325,26 @@ fn grant_role_with_admin_option_keeps_the_portable_spelling() { "grant r1 to u1 with inherit false, admin false, set true;" ); } + +/// Postgres keeps the first 63 bytes of an identifier, and so does the AST. The formatter writes +/// the name as it was written, not as the parser cut it down. +#[test] +fn identifiers_longer_than_postgres_keeps_are_written_in_full() { + let role = "a_role_name_that_is_definitely_way_too_long_for_anything_postgres_keeps"; + assert_eq!( + format(&format!("GRANT SELECT ON t TO {role}")), + format!("grant select on table t to {role};") + ); + + let column = "A Quoted Column Name That Is Definitely Way Too Long For Anything Postgres Keeps"; + assert_eq!( + format(&format!("SELECT \"{column}\" FROM t")), + format!("select \"{column}\" from t;") + ); + + let table = "an_unquoted_table_name_that_is_definitely_way_too_long_for_anything"; + assert_eq!( + format(&format!("CREATE TABLE {table} (id int)")), + format!("create table {table} (id int);") + ); +}