diff --git a/datafusion-pg-catalog/src/sql/parser.rs b/datafusion-pg-catalog/src/sql/parser.rs index 56be16a..ac7ce0d 100644 --- a/datafusion-pg-catalog/src/sql/parser.rs +++ b/datafusion-pg-catalog/src/sql/parser.rs @@ -11,17 +11,53 @@ use datafusion::sql::sqlparser::tokenizer::TokenWithSpan; use super::rules::AliasDuplicatedProjectionRewrite; use super::rules::CurrentUserVariableToSessionUserFunctionCall; use super::rules::FixArrayLiteral; -use super::rules::FixCollate; use super::rules::FixVersionColumnName; use super::rules::PrependUnqualifiedPgTableName; -use super::rules::RemoveQualifier; use super::rules::RemoveSubqueryFromProjection; -use super::rules::ResolveUnqualifiedIdentifer; +use super::rules::ResolveUnqualifiedIdentifier; use super::rules::RewriteArrayAnyAllOperation; +use super::rules::RewritePgCatalogOperator; use super::rules::SqlStatementRewriteRule; +use super::rules::StripCallableQualifier; +use super::rules::StripCollate; +/// Last-resort replacements for whole client queries that DataFusion cannot +/// execute and cannot be fixed by a SQL rewrite rule. +/// +/// Each entry is a `(client_sql, replacement_sql)` pair. The parser tokenizes +/// both at construction time (see [`PostgresCompatibilityParser::new`]) and the +/// input token stream is scanned for the client pattern; on a match the matched +/// tokens are replaced wholesale, before any rewrite rule runs. +/// +/// # How matching works +/// +/// * **Token-level, whitespace-insensitive:** only non-whitespace, non-`;` +/// tokens are compared, so formatting differences (indentation, newlines, +/// spacing) do not matter. +/// * **Substring match:** patterns are scanned at *every* token position, so a +/// pattern matches as a contiguous sub-sequence of the input. Most entries +/// are whole statements and therefore only match at a statement boundary; +/// the grafana entry below is intentionally a *partial* replacement (it +/// substitutes just an inner subquery). +/// * **`$N` placeholders are wildcards:** a `Token::Placeholder` in the pattern +/// (e.g. `$1`) matches any single non-`;` token in the input, so the psql +/// `\d` entries match regardless of whether the client binds a parameter or +/// sends a literal oid string. +/// +/// # When to add or remove an entry +/// +/// These entries exist solely because of a specific DataFusion limitation. +/// Each entry's comment names the limitation and a removal condition. Before +/// adding a new entry, check whether a general SQL rewrite rule (or better, an +/// analyzer rule / type planner, as was done for oid resolution) can handle it +/// instead -- those are preferred because they are metadata-aware and +/// order-independent. const BLACKLIST_SQL_MAPPING: &[(&str, &str)] = &[ - // pgcli startup query + // pgcli -- foreign-key column lookup + // + // Blocked by: `unnest()` -- DataFusion cannot unnest the + // result of a subquery (here, `array_agg(...)` over a correlated join). + // Remove when: DF supports unnesting arbitrary subquery results. ( "SELECT s_p.nspname AS parentschema, t_p.relname AS parenttable, @@ -58,7 +94,12 @@ const BLACKLIST_SQL_MAPPING: &[(&str, &str)] = &[ NULL::TEXT AS childcolumn WHERE false"), - // pgcli startup query + // pgcli -- user-defined type lookup + // + // Blocked by: correlated scalar subquery in the WHERE clause -- + // `SELECT c.relkind = 'c' FROM pg_catalog.pg_class c + // WHERE c.oid = t.typrelid` references the outer `t`. + // Remove when: DF plans correlated scalar subqueries in predicates. ( "SELECT n.nspname schema_name, t.typname type_name @@ -83,7 +124,12 @@ const BLACKLIST_SQL_MAPPING: &[(&str, &str)] = &[ "SELECT NULL::TEXT AS schema_name, NULL::TEXT AS type_name WHERE false" ), -// psql \d queries + // psql \d
-- row policies (polrelid bound via the $1 wildcard) + // + // Blocked by: the `array(SELECT ...)` array constructor (here over a + // join on `any(pol.polroles)`). DataFusion lacks the array-from-subquery + // constructor syntax. + // Remove when: DF supports the `array(SELECT ...)` constructor. ( "SELECT pol.polname, pol.polpermissive, CASE WHEN pol.polroles = '{0}' THEN NULL ELSE pg_catalog.array_to_string(array(select rolname from pg_catalog.pg_roles where oid = any (pol.polroles) order by 1),',') END, @@ -107,6 +153,15 @@ const BLACKLIST_SQL_MAPPING: &[(&str, &str)] = &[ WHERE false" ), + // psql \d
-- extended statistics (stxrelid bound via $1) + // + // Blocked by: ` = any()` -- `stxkind` is a + // `char[2]`; DataFusion's `array_has` expects matching element types and + // rejects the char-vs-string comparison, so all three `any(stxkind)` + // predicates fail to plan. + // Remove when: DF coerces element types in `any(array_col)` comparisons, + // or an analyzer rule lowers `lit = ANY(array_col)` to a proper element + // comparison. ( "SELECT oid, stxrelid::pg_catalog.regclass, stxnamespace::pg_catalog.regnamespace::pg_catalog.text AS nsp, stxname, pg_catalog.pg_get_statisticsobjdef_columns(oid) AS columns, @@ -130,6 +185,14 @@ const BLACKLIST_SQL_MAPPING: &[(&str, &str)] = &[ WHERE false" ), + // psql \d
-- publications (prrelid / oid bound via $1) + // + // Blocked by: the three UNION branches each project two unnamed `NULL` + // columns, yielding duplicate/empty projection names that DataFusion + // rejects. (The `generate_series(0, array_upper(...))` over a correlated + // `int2[]` would also fail, but the duplicate-name error fires first.) + // Remove when: DF tolerates duplicate/unnamed NULL projection columns, + // or an aliasing rule can assign synthetic names. ( "SELECT pubname , NULL @@ -164,7 +227,16 @@ const BLACKLIST_SQL_MAPPING: &[(&str, &str)] = &[ WHERE false" ), - // grafana array index magic + // grafana -- search_path membership lookup. NOTE: this is the only entry + // that is a *partial* replacement -- it matches just the inner subquery and + // substitutes the empty string `''`, so the surrounding `IN (...)` simply + // matches nothing and the query returns zero rows. + // + // Blocked by: missing `current_setting()` UDF -- the query resolves the + // session `search_path` at runtime via `current_setting('search_path')`. + // (generate_series over dynamic array bounds now works via the + // CoerceIntArgsToBigInt wrapper; current_setting is the remaining gap.) + // Remove when: a `current_setting()` UDF is registered. (r#"SELECT CASE WHEN trim(s[i]) = '"$user"' THEN user ELSE trim(s[i]) END FROM @@ -179,8 +251,9 @@ const BLACKLIST_SQL_MAPPING: &[(&str, &str)] = &[ /// A parser with Postgres Compatibility for Datafusion /// /// This parser will try its best to rewrite postgres SQL into a form that -/// datafuiosn supports. It also maintains a blacklist that will transform the -/// statement to a similar version if rewrite doesn't worth the effort for now. +/// DataFusion supports. It also maintains a blacklist that will transform the +/// statement to a similar version if rewriting it is not worth the effort for +/// now. #[derive(Debug)] pub struct PostgresCompatibilityParser { blacklist: Vec<(Vec, Vec)>, @@ -221,16 +294,18 @@ impl PostgresCompatibilityParser { Self { blacklist: mapping, rewrite_rules: vec![ - // make sure blacklist based rewriter it on the top to prevent sql - // being rewritten from other rewriters + // The blacklist substitution in `parse()` runs before any of + // these rules, so by the time they see the statement any + // blacklisted fragment has already been replaced. Arc::new(AliasDuplicatedProjectionRewrite), - Arc::new(ResolveUnqualifiedIdentifer), + Arc::new(ResolveUnqualifiedIdentifier), Arc::new(RewriteArrayAnyAllOperation), Arc::new(PrependUnqualifiedPgTableName), - Arc::new(RemoveQualifier), + Arc::new(StripCallableQualifier), Arc::new(FixArrayLiteral), Arc::new(CurrentUserVariableToSessionUserFunctionCall), - Arc::new(FixCollate), + Arc::new(StripCollate), + Arc::new(RewritePgCatalogOperator), Arc::new(RemoveSubqueryFromProjection), Arc::new(FixVersionColumnName), ], @@ -242,23 +317,22 @@ impl PostgresCompatibilityParser { let parser = Parser::new(&PostgreSqlDialect {}); let tokens = parser.try_with_sql(input)?.into_tokens(); - // Get token values (without spans) and filter out only whitespace - // Keep semicolons as they separate statements - // Also rewrite ABORT to ROLLBACK for postgres compatibility - // remove this when https://github.com/apache/datafusion-sqlparser-rs/pull/2332 is ready + // Drop whitespace tokens; keep semicolons as they separate statements. let filtered_tokens: Vec = tokens .iter() .map(|t| t.token.clone()) .filter(|t| !matches!(t, Token::Whitespace(_))) - .map(|t| { - if matches!(&t, Token::Word(w) if w.keyword == Keyword::ABORT) { - Token::make_keyword("ROLLBACK") - } else { - t - } - }) .collect(); + // Rewrite the Postgres `ABORT` statement to `ROLLBACK` (they are + // synonyms). This must happen at the token level because sqlparser does + // not parse `ABORT` as a statement, so it cannot be a post-parse + // [`SqlStatementRewriteRule`]. + // + // TODO: remove when https://github.com/apache/datafusion-sqlparser-rs/pull/2332 + // is released in the version of sqlparser we depend on. + let filtered_tokens = Self::rewrite_abort_to_rollback(filtered_tokens); + // Handle empty input if filtered_tokens.is_empty() { return Ok(Vec::new()); @@ -330,6 +404,21 @@ impl PostgresCompatibilityParser { Ok(result) } + /// Replace every `ABORT` keyword token with `ROLLBACK` (see + /// [`Self::maybe_replace_tokens`] for rationale). + fn rewrite_abort_to_rollback(tokens: Vec) -> Vec { + tokens + .into_iter() + .map(|t| { + if matches!(&t, Token::Word(w) if w.keyword == Keyword::ABORT) { + Token::make_keyword("ROLLBACK") + } else { + t + } + }) + .collect() + } + fn parse_tokens(&self, tokens: Vec) -> Result, ParserError> { let parser = Parser::new(&PostgreSqlDialect {}); // Convert tokens to TokenWithSpan with dummy spans @@ -607,4 +696,26 @@ mod tests { assert_eq!(actual_tokens, expected_token_values); } + + #[test] + fn test_abort_rewritten_to_rollback() { + // ABORT is a Postgres synonym for ROLLBACK; sqlparser does not parse it + // as a statement, so the token-level rewrite_abort_to_rollback turns + // it into ROLLBACK before parsing. + let parser = PostgresCompatibilityParser::new(); + let stmts = parser.parse("ABORT").expect("failed to parse"); + assert_eq!(stmts.len(), 1); + assert_eq!(stmts[0].to_string(), "ROLLBACK"); + + // `ABORT AND CHAIN` -> `ROLLBACK AND CHAIN`. + let stmts = parser.parse("ABORT AND CHAIN").expect("failed to parse"); + assert_eq!(stmts[0].to_string(), "ROLLBACK AND CHAIN"); + + // A column named `abort` (quoted) must NOT be rewritten -- only the + // bare keyword form is a statement. + let stmts = parser + .parse("SELECT \"abort\" FROM t") + .expect("failed to parse"); + assert!(stmts[0].to_string().contains("\"abort\"")); + } } diff --git a/datafusion-pg-catalog/src/sql/rules.rs b/datafusion-pg-catalog/src/sql/rules.rs index b1417f5..50ae11d 100644 --- a/datafusion-pg-catalog/src/sql/rules.rs +++ b/datafusion-pg-catalog/src/sql/rules.rs @@ -150,9 +150,9 @@ impl SqlStatementRewriteRule for AliasDuplicatedProjectionRewrite { /// Postgres allows unqualified identifier in ORDER BY and FILTER but it's not /// accepted by datafusion. #[derive(Debug)] -pub struct ResolveUnqualifiedIdentifer; +pub struct ResolveUnqualifiedIdentifier; -impl ResolveUnqualifiedIdentifer { +impl ResolveUnqualifiedIdentifier { fn rewrite_unqualified_identifiers(query: &mut Box) { if let SetExpr::Select(select) = query.body.as_mut() { // Step 1: Find all table aliases from FROM and JOIN clauses. @@ -288,7 +288,7 @@ impl ResolveUnqualifiedIdentifer { } } -impl SqlStatementRewriteRule for ResolveUnqualifiedIdentifer { +impl SqlStatementRewriteRule for ResolveUnqualifiedIdentifier { fn rewrite(&self, mut statement: Statement) -> Statement { if let Statement::Query(query) = &mut statement { Self::rewrite_unqualified_identifiers(query); @@ -298,14 +298,46 @@ impl SqlStatementRewriteRule for ResolveUnqualifiedIdentifer { } } -/// Rewrite Postgres's ANY operator to array_contains +/// Rewrite the Postgres `ANY`/`ALL` quantified comparison operators over +/// arrays into DataFusion's `array_contains` function. +/// +/// Postgres `x ANY (arr)` / `x ALL (arr)` have no direct DataFusion +/// equivalent, so this rule lowers the two cases that map cleanly onto +/// membership: +/// +/// * `x = ANY(arr)` -> `array_contains(arr, x)` (x equals some element) +/// * `x <> ALL(arr)` -> `NOT array_contains(arr, x)` (x differs from every +/// element, i.e. is not a member) +/// +/// The other two combinations are **not** handled, because `array_contains` +/// only expresses membership, not universal quantification: +/// +/// * `x <> ANY(arr)` (x differs from *some* element) is true for essentially +/// every non-empty array and has no faithful lowering to membership. +/// * `x = ALL(arr)` (x equals *every* element) needs "all elements are x", +/// which membership cannot express without an additional length/count check. +/// +/// These are left untouched; if a client query uses them, DataFusion reports +/// the unsupported operator rather than silently producing wrong results. +/// +/// # String-literal array form +/// +/// Postgres allows an array as a single-quoted string literal, e.g. +/// `x = ANY('{r, l, e}')`. When the right operand is such a literal this rule +/// parses it into an `ARRAY[...]` expression first. The parsing is **naive**: +/// it splits on commas and treats every element as a string. That is correct +/// for the string-array case clients actually send (privilege/role char arrays, +/// `contyp`/`stxkind` style columns), but it does not handle quoted elements +/// containing commas, and it assumes string element types for non-string array +/// literals. A non-string-literal right operand (a column, `ARRAY[...]`, or a +/// function call) is passed through unchanged. #[derive(Debug)] pub struct RewriteArrayAnyAllOperation; struct RewriteArrayAnyAllOperationVisitor; impl RewriteArrayAnyAllOperationVisitor { - fn any_to_array_cofntains(&self, left: &Expr, right: &Expr) -> Expr { + fn any_to_array_contains(&self, left: &Expr, right: &Expr) -> Expr { let array = if let Expr::Value(ValueWithSpan { value: Value::SingleQuotedString(array_literal), .. @@ -316,7 +348,10 @@ impl RewriteArrayAnyAllOperationVisitor { let items = array_literal.trim_matches(|c| c == '{' || c == '}' || c == ' '); let items = items.split(',').map(|s| s.trim()).filter(|s| !s.is_empty()); - // For now, we assume the data type is string + // Elements are treated as strings -- see the rule-level doc + // comment for why this naive split is sufficient for the array + // literals clients actually send (string-typed arrays) but does + // not handle quoted elements containing commas. let elems = items .map(|s| { Expr::Value(Value::SingleQuotedString(s.to_string()).with_empty_span()) @@ -365,10 +400,13 @@ impl VisitorMut for RewriteArrayAnyAllOperationVisitor { .. } => match compare_op { BinaryOperator::Eq => { - *expr = self.any_to_array_cofntains(left.as_ref(), right.as_ref()); + *expr = self.any_to_array_contains(left.as_ref(), right.as_ref()); } BinaryOperator::NotEq => { - // TODO:left not equals to any element in array + // x <> ANY(arr): true if x differs from SOME element. That + // is true for almost every non-empty array and cannot be + // lowered to `array_contains`; leave it for DataFusion to + // report as unsupported. } _ => {} }, @@ -378,12 +416,14 @@ impl VisitorMut for RewriteArrayAnyAllOperationVisitor { right, } => match compare_op { BinaryOperator::Eq => { - // TODO: left equals to every element in array + // x = ALL(arr): true if EVERY element equals x. Membership + // alone cannot express this; leave it for DataFusion to + // report as unsupported. } BinaryOperator::NotEq => { *expr = Expr::UnaryOp { op: UnaryOperator::Not, - expr: Box::new(self.any_to_array_cofntains(left.as_ref(), right.as_ref())), + expr: Box::new(self.any_to_array_contains(left.as_ref(), right.as_ref())), } } _ => {} @@ -529,11 +569,15 @@ impl SqlStatementRewriteRule for FixArrayLiteral { } } -/// Remove qualifier from unsupported items +/// Strip the schema qualifier from callable names. /// -/// This rewriter removes the `pg_catalog.` qualifier from: -/// 1. function names, for example `pg_catalog.array_to_string` -/// 2. table function names +/// DataFusion resolves functions (including table functions) by bare name in +/// a flat namespace -- it has no notion of a schema-qualified function path. +/// Postgres clients routinely write the qualifier explicitly, most often +/// `pg_catalog.` (e.g. `pg_catalog.array_to_string(...)`, +/// `pg_catalog.pg_get_keywords()`), so this rule drops the leading +/// identifier(s) from any multi-part function or table-function name and +/// keeps only the last segment. /// /// Type-cast qualifiers such as `pg_catalog.text` / `pg_catalog.int2[]` used /// to be handled here too, but are now resolved by the [`PgOidTypePlanner`] @@ -542,22 +586,23 @@ impl SqlStatementRewriteRule for FixArrayLiteral { /// /// [`PgOidTypePlanner`]: crate::pg_catalog::oid_type_planner::PgOidTypePlanner #[derive(Debug)] -pub struct RemoveQualifier; +pub struct StripCallableQualifier; -struct RemoveQualifierVisitor; +struct StripCallableQualifierVisitor; -impl VisitorMut for RemoveQualifierVisitor { +impl VisitorMut for StripCallableQualifierVisitor { type Break = (); fn pre_visit_table_factor( &mut self, table_factor: &mut TableFactor, ) -> ControlFlow { - // remove table function qualifier + // Strip qualifier from table function names (any schema, not just + // pg_catalog): DF resolves table functions by bare name. if let TableFactor::Table { name, args, .. } = table_factor && args.is_some() { - // multiple idents in name, which means it's a qualified table name + // multiple idents in name means it's a qualified table name if name.0.len() > 1 && let Some(last_ident) = name.0.pop() { @@ -568,7 +613,8 @@ impl VisitorMut for RemoveQualifierVisitor { } fn pre_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow { - // remove qualifier from pg_catalog.function + // Strip qualifier from function names (any schema, not just + // pg_catalog): DF resolves functions by bare name. if let Expr::Function(function) = expr { let name = &mut function.name; if name.0.len() > 1 @@ -581,9 +627,9 @@ impl VisitorMut for RemoveQualifierVisitor { } } -impl SqlStatementRewriteRule for RemoveQualifier { +impl SqlStatementRewriteRule for StripCallableQualifier { fn rewrite(&self, mut s: Statement) -> Statement { - let mut visitor = RemoveQualifierVisitor; + let mut visitor = StripCallableQualifierVisitor; let _ = s.visit(&mut visitor); s @@ -719,38 +765,89 @@ impl SqlStatementRewriteRule for CurrentUserVariableToSessionUserFunctionCall { } } -/// Fix collate and regex calls +/// Strip Postgres `COLLATE` clauses. +/// +/// Postgres allows trailing `COLLATE .` on expressions +/// (e.g. `c.relname COLLATE pg_catalog.default`). DataFusion has no notion of +/// collation and rejects this syntax at parse time, so strip the clause and +/// keep the inner expression. The collation is irrelevant to DataFusion's +/// byte-wise string comparison anyway. #[derive(Debug)] -pub struct FixCollate; +pub struct StripCollate; -struct FixCollateVisitor; +struct StripCollateVisitor; -impl VisitorMut for FixCollateVisitor { +impl VisitorMut for StripCollateVisitor { type Break = (); fn pre_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow { - match expr { - Expr::Collate { expr: inner, .. } => { - *expr = inner.as_ref().clone(); - } - Expr::BinaryOp { op, .. } => { - if let BinaryOperator::PGCustomBinaryOperator(ops) = op - && *ops == ["pg_catalog", "~"] - { - *op = BinaryOperator::PGRegexMatch; - } - } - _ => {} + if let Expr::Collate { expr: inner, .. } = expr { + *expr = inner.as_ref().clone(); } - ControlFlow::Continue(()) } } -impl SqlStatementRewriteRule for FixCollate { +impl SqlStatementRewriteRule for StripCollate { fn rewrite(&self, mut s: Statement) -> Statement { - let mut visitor = FixCollateVisitor; + let mut visitor = StripCollateVisitor; + let _ = s.visit(&mut visitor); + s + } +} +/// Rewrite Postgres `OPERATOR(schema.name)` call syntax to the bare operator. +/// +/// Postgres lets clients force an operator's schema with +/// `lhs OPERATOR(pg_catalog.~) rhs`. sqlparser represents this as a +/// [`BinaryOperator::PGCustomBinaryOperator`]; DataFusion only understands the +/// built-in operator variants, so map every `pg_catalog.` we see +/// back to its builtin form: +/// +/// | pg_catalog name | builtin operator | +/// | --------------- | --------------------------- | +/// | `~` | `PGRegexMatch` (`~`) | +/// | `~*` | `PGRegexIMatch` (`~*`) | +/// | `!~` | `PGRegexNotMatch` (`!~`) | +/// | `!~*` | `PGRegexNotIMatch`(`!~*`) | +/// +/// Only the four regex operators are mapped, because those are the ones +/// clients emit via `OPERATOR(...)` in practice (psql's `\d` queries use +/// `OPERATOR(pg_catalog.~)`). Any other operator name is left untouched for +/// DataFusion to report. +#[derive(Debug)] +pub struct RewritePgCatalogOperator; + +struct RewritePgCatalogOperatorVisitor; + +impl VisitorMut for RewritePgCatalogOperatorVisitor { + type Break = (); + + fn pre_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow { + if let Expr::BinaryOp { op, .. } = expr + && let BinaryOperator::PGCustomBinaryOperator(ops) = op + && ops.first().is_some_and(|s| s == "pg_catalog") + && let Some(name) = ops.last().map(String::as_str) + { + // Map every `OPERATOR(pg_catalog.)` to its builtin form. + // Match on the operator name (last segment); the schema is checked + // above. Only the four regex ops have a builtin equivalent + // DataFusion understands, so any other name is left untouched. + *op = match name { + "~" => BinaryOperator::PGRegexMatch, + "~*" => BinaryOperator::PGRegexIMatch, + "!~" => BinaryOperator::PGRegexNotMatch, + "!~*" => BinaryOperator::PGRegexNotIMatch, + _ => return ControlFlow::Continue(()), + }; + } + ControlFlow::Continue(()) + } +} + +impl SqlStatementRewriteRule for RewritePgCatalogOperator { + fn rewrite(&self, mut s: Statement) -> Statement { + let mut visitor = RewritePgCatalogOperatorVisitor; let _ = s.visit(&mut visitor); s } @@ -994,7 +1091,7 @@ mod tests { #[test] fn test_qualifier_prepend() { let rules: Vec> = - vec![Arc::new(ResolveUnqualifiedIdentifer)]; + vec![Arc::new(ResolveUnqualifiedIdentifier)]; assert_rewrite!( &rules, @@ -1117,8 +1214,8 @@ mod tests { } #[test] - fn test_remove_qualifier_from_table_function() { - let rules: Vec> = vec![Arc::new(RemoveQualifier)]; + fn test_strip_callable_qualifier_from_table_function() { + let rules: Vec> = vec![Arc::new(StripCallableQualifier)]; assert_rewrite!( &rules, @@ -1162,13 +1259,56 @@ mod tests { } #[test] - fn test_collate_fix() { - let rules: Vec> = vec![Arc::new(FixCollate)]; + fn test_strip_collate() { + let rules: Vec> = vec![Arc::new(StripCollate)]; + + assert_rewrite!( + &rules, + "SELECT c.oid, c.relname FROM pg_catalog.pg_class c WHERE c.relname COLLATE pg_catalog.default = 'foo'", + "SELECT c.oid, c.relname FROM pg_catalog.pg_class c WHERE c.relname = 'foo'" + ); + + // Combined with OPERATOR(...) syntax: only COLLATE is stripped here; + // the operator is left for RewritePgCatalogOperator (next test). + assert_rewrite!( + &rules, + "SELECT c.oid, c.relname FROM pg_catalog.pg_class c WHERE c.relname OPERATOR(pg_catalog.~) '^(tablename)$' COLLATE pg_catalog.default ORDER BY 2, 3;", + "SELECT c.oid, c.relname FROM pg_catalog.pg_class c WHERE c.relname OPERATOR(pg_catalog.~) '^(tablename)$' ORDER BY 2, 3" + ); + } + + #[test] + fn test_rewrite_pg_catalog_operator() { + let rules: Vec> = vec![Arc::new(RewritePgCatalogOperator)]; + + // All four pg_catalog regex operators map to their builtin forms. + assert_rewrite!( + &rules, + "SELECT 'a' OPERATOR(pg_catalog.~) 'b'", + "SELECT 'a' ~ 'b'" + ); + assert_rewrite!( + &rules, + "SELECT 'a' OPERATOR(pg_catalog.~*) 'b'", + "SELECT 'a' ~* 'b'" + ); + assert_rewrite!( + &rules, + "SELECT 'a' OPERATOR(pg_catalog.!~) 'b'", + "SELECT 'a' !~ 'b'" + ); + assert_rewrite!( + &rules, + "SELECT 'a' OPERATOR(pg_catalog.!~*) 'b'", + "SELECT 'a' !~* 'b'" + ); + // The original psql \d shape: OPERATOR + COLLATE together. The operator + // arm fires; COLLATE is left for StripCollate. assert_rewrite!( &rules, - "SELECT c.oid, c.relname FROM pg_catalog.pg_class c WHERE c.relname OPERATOR(pg_catalog.~) '^(tablename)$' COLLATE pg_catalog.default AND pg_catalog.pg_table_is_visible(c.oid) ORDER BY 2, 3;", - "SELECT c.oid, c.relname FROM pg_catalog.pg_class c WHERE c.relname ~ '^(tablename)$' AND pg_catalog.pg_table_is_visible(c.oid) ORDER BY 2, 3" + "SELECT c.oid, c.relname FROM pg_catalog.pg_class c WHERE c.relname OPERATOR(pg_catalog.~) '^(tablename)$' COLLATE pg_catalog.default ORDER BY 2, 3;", + "SELECT c.oid, c.relname FROM pg_catalog.pg_class c WHERE c.relname ~ '^(tablename)$' COLLATE pg_catalog.default ORDER BY 2, 3" ); } diff --git a/datafusion-postgres/tests/array_bounds.rs b/datafusion-postgres/tests/array_bounds.rs index 73155d6..cfd6131 100644 --- a/datafusion-postgres/tests/array_bounds.rs +++ b/datafusion-postgres/tests/array_bounds.rs @@ -12,7 +12,7 @@ use datafusion_postgres::testing::*; const QUERIES: &[&str] = &[ // Bare UDFs over an array literal. "SELECT array_upper(ARRAY[1, 2, 3], 1), array_lower(ARRAY[1, 2, 3], 1)", - // pg_catalog-qualified form (qualifier stripped by RemoveQualifier rule). + // pg_catalog-qualified form (qualifier stripped by StripCallableQualifier rule). "SELECT pg_catalog.array_upper(ARRAY['a', 'b'], 1)", "SELECT pg_catalog.array_lower(ARRAY['a', 'b'], 1)", // Typical usage pattern: generate_series over array bounds, mirroring the