From f11d368d325a779082a1805c035e6896c124a076 Mon Sep 17 00:00:00 2001 From: mysticmind Date: Fri, 9 Oct 2026 19:27:26 +0530 Subject: [PATCH 1/2] Add statement classification, resolved references, Format and error rendering; bump to 2.2.0 - Parser.Classify and StatementInfo: kind, read-only, and facts such as a data-modifying CTE, SELECT INTO, a locking clause, EXPLAIN ANALYZE, COPY PROGRAM and the wrapped statement - GetReferences: tables with role and CTE detection, functions, and columns matched to their table where the query alone settles it - Parser.Format: pretty print a query or script, keeping comments, one statement per paragraph - Error.Format and Error.GetLineAndColumn: psql-style error rendering --- README.md | 97 ++++++++ src/PackageSmokeTest/Program.cs | 7 + src/PgSqlParser.Tests/AnalysisTests.cs | 316 +++++++++++++++++++++++++ src/PgSqlParser/Parser.cs | 173 ++++++++++++++ src/PgSqlParser/PgSqlParser.csproj | 2 +- src/PgSqlParser/QueryReferences.cs | 289 ++++++++++++++++++++++ src/PgSqlParser/StatementInfo.cs | 216 +++++++++++++++++ 7 files changed, 1099 insertions(+), 1 deletion(-) create mode 100644 src/PgSqlParser.Tests/AnalysisTests.cs create mode 100644 src/PgSqlParser/QueryReferences.cs create mode 100644 src/PgSqlParser/StatementInfo.cs diff --git a/README.md b/README.md index 3ef354c..71ac31f 100644 --- a/README.md +++ b/README.md @@ -287,6 +287,81 @@ if (result.Error is null) // { "version": 180006, "tokens": [ { "end": 6, "token": "SELECT", "keywordKind": "RESERVED_KEYWORD" }, { "start": 7, "end": 8, "token": "ICONST" } ] } ``` +### Classify + +Get the facts about each statement: its kind, whether it is read-only, and anything that makes a harmless-looking statement write. This is the basis for read/write routing and for checking SQL before running it. + +```csharp +using PgSqlParser; + +foreach (var info in Parser.Classify("SELECT 1; WITH gone AS (DELETE FROM t RETURNING *) SELECT * FROM gone; EXPLAIN ANALYZE DELETE FROM t").GetValueOrThrow()) +{ + Console.WriteLine($"{info.Kind}: read-only {info.IsReadOnly}"); +} +// Output: +// Select: read-only True +// Select: read-only False (HasDataModifyingCte is true) +// Explain: read-only False (ExecutesInner is true, and Inner is a DELETE) +``` + +`StatementInfo` also has `IsUtility`, `HasSelectInto`, `HasLockingClause`, `RunsProgram` and `Inner` (the statement wrapped by EXPLAIN, PREPARE, DECLARE CURSOR, CREATE TABLE AS or COPY). `IsReadOnly` errs on the side of false: CALL, DO and EXECUTE are not read-only because what they run is not visible. A function called from a SELECT can still write, which no parser can see. These are facts, not a security boundary; what to allow is up to you. + +### References + +Find the tables, functions and columns a query refers to, as parse tree nodes with their role: + +```csharp +using PgSqlParser; + +var tree = Parser.Parse( + "WITH recent AS (SELECT * FROM orders) " + + "UPDATE customers c SET total = r.total FROM recent r WHERE r.customer_id = c.id").GetValueOrThrow(); +var references = tree.GetReferences(); + +foreach (var table in references.Tables) +{ + Console.WriteLine($"{table.Name}: {table.Role}{(table.IsCte ? " (a CTE, not a table)" : "")}"); +} +// Output: +// customers: Write +// recent: Read (a CTE, not a table) +// orders: Read + +foreach (var column in references.Columns) +{ + Console.WriteLine($"{column} -> {column.Table?.Name ?? "unknown"}"); +} +// Output: +// r.total -> unknown (recent is a CTE) +// r.customer_id -> unknown +// c.id -> customers +// * -> orders +``` + +A column is matched to its table when the query alone settles it: its qualifier names a table or alias in scope, or it has no qualifier and there is only one table in scope. Otherwise `Table` is null, because it would take the database catalog to tell. `Summary` returns similar information as plain names and is faster; use `GetReferences` when you need the nodes, aliases or column matching. + +### Format + +Format a query across indented lines, keeping its comments. Each statement of a script is formatted on its own. + +```csharp +using PgSqlParser; + +var formatted = Parser.Format("select a, b from t -- pick\n where x = 1; delete from t where a = 1").GetValueOrThrow(); +Console.WriteLine(formatted); +// Output: +// SELECT a, b +// FROM t +// WHERE +// -- pick +// x = 1; +// +// DELETE FROM t +// WHERE a = 1; +``` + +`FormatOptions` sets the indent size, the line length, comma placement, a trailing newline, and whether comments are kept. + ### ParameterRefs Find the parameter references (`$1`, `$2`, ...) in a query, with where each one is in the query string. References inside string literals and comments are not parameters and are left out. @@ -566,6 +641,19 @@ if (result.Error is not null) // Output: syntax error at or near "SELEC" ``` +`error.Format(query)` renders the error the way psql does, with the line it points at, and `error.GetLineAndColumn(query)` gives the position: + +```csharp +var query = "SELECT a,\n b FROM WHERE x"; +var error = Parser.Parse(query).Error!; + +Console.WriteLine(error.Format(query)); +// Output: +// ERROR: syntax error at or near "WHERE" +// LINE 2: b FROM WHERE x +// ^ +``` + The `Error` also carries `CursorPos`, the position of the error in the query (see [Offsets and non-ASCII text](#offsets-and-non-ascii-text)), and the PostgreSQL source location that raised it in `FuncName`, `FileName` and `LineNo`. Two kinds of invalid input are caught before the query reaches libpg_query: @@ -577,6 +665,15 @@ Deeply nested queries, such as a long chain of operators without parentheses (`a Protobuf's own recursive operations on a parse tree, such as `ToString()`, `Clone()` and `Equals()`, run on your thread. On a very deeply nested tree they can still overflow a small stack. +## What's new in 2.2 + +All additions; nothing from earlier versions changes. + +- **Classify statements**: `Parser.Classify` gives each statement's kind, whether it is read-only, and facts such as a data-modifying CTE, `SELECT INTO`, a locking clause or `EXPLAIN ANALYZE`. +- **Resolved references**: `GetReferences` returns the tables, functions and columns of a query as nodes, with table roles (read, write, DDL), CTE detection and columns matched to their table. +- **Format in one call**: `Parser.Format` pretty prints a query or a script and keeps its comments. +- **Readable errors**: `Error.Format` and `Error.GetLineAndColumn`. + ## What's new in 2.1 All additions; nothing from 2.0 changes. diff --git a/src/PackageSmokeTest/Program.cs b/src/PackageSmokeTest/Program.cs index 14af966..021b448 100644 --- a/src/PackageSmokeTest/Program.cs +++ b/src/PackageSmokeTest/Program.cs @@ -65,6 +65,13 @@ void Check(string name, bool ok) Check("ParseExpression", Parser.ParseExpression("(a > 0)").Value?.Deparse().Value == "a > 0"); Check("ParseTypeName", Parser.ParseTypeName("character varying(20)").Value?.Deparse().Value == "varchar(20)"); Check("PgIdentifier", PgIdentifier.Quote("public", "order") == "public.\"order\""); +var classified = Parser.Classify("WITH d AS (DELETE FROM t RETURNING 1) SELECT * FROM d; EXPLAIN SELECT 1").Value!; +Check("Classify", classified is [{ Kind: StatementKind.Select, IsReadOnly: false }, { Kind: StatementKind.Explain, IsReadOnly: true }]); +var refs = Parser.Parse("UPDATE a SET x = b.y FROM b WHERE b.id = a.id").Value!.GetReferences(); +Check("GetReferences", refs.Tables.Select(t => t.Role).SequenceEqual([TableRole.Write, TableRole.Read]) + && refs.Columns.All(c => c.Table is not null)); +Check("Format", Parser.Format("select a from t -- note\n where x = 1").Value == "SELECT a\nFROM t\nWHERE\n -- note\n x = 1"); +Check("Error.Format", Parser.Parse("SELECT FROM WHERE").Error!.Format("SELECT FROM WHERE").EndsWith(" ^")); Check("ParameterRefs", Parser.ParameterRefs("SELECT $1, '$2', $3").Value?.Select(p => p.Number).SequenceEqual([1, 3]) == true); Check("ParsePlpgsqlFunctions", Parser.ParsePlpgsqlFunctions( "CREATE FUNCTION f() RETURNS int AS $$ BEGIN RETURN 1; END; $$ LANGUAGE plpgsql").Value?[0].Queries().SequenceEqual(["1"]) == true); diff --git a/src/PgSqlParser.Tests/AnalysisTests.cs b/src/PgSqlParser.Tests/AnalysisTests.cs new file mode 100644 index 0000000..5bf30d3 --- /dev/null +++ b/src/PgSqlParser.Tests/AnalysisTests.cs @@ -0,0 +1,316 @@ +using Shouldly; +using Xunit; + +namespace PgSqlParser.Tests; + +public class AnalysisTests +{ + private static StatementInfo Classify(string query) => Parser.Classify(query).GetValueOrThrow().ShouldHaveSingleItem(); + + [Theory] + [InlineData("SELECT 1", StatementKind.Select, true)] + [InlineData("VALUES (1)", StatementKind.Select, true)] + [InlineData("TABLE t", StatementKind.Select, true)] + [InlineData("INSERT INTO t SELECT 1", StatementKind.Insert, false)] + [InlineData("UPDATE t SET a = 1", StatementKind.Update, false)] + [InlineData("DELETE FROM t", StatementKind.Delete, false)] + [InlineData("MERGE INTO t USING s ON true WHEN MATCHED THEN DELETE", StatementKind.Merge, false)] + [InlineData("TRUNCATE t", StatementKind.Truncate, false)] + [InlineData("COPY t TO STDOUT", StatementKind.Copy, true)] + [InlineData("COPY t FROM STDIN", StatementKind.Copy, false)] + [InlineData("CREATE TABLE t (a int)", StatementKind.Ddl, false)] + [InlineData("CREATE INDEX i ON t (a)", StatementKind.Ddl, false)] + [InlineData("ALTER TABLE t ADD COLUMN b int", StatementKind.Ddl, false)] + [InlineData("DROP TABLE t", StatementKind.Ddl, false)] + [InlineData("GRANT SELECT ON t TO x", StatementKind.Ddl, false)] + [InlineData("COMMENT ON TABLE t IS 'x'", StatementKind.Ddl, false)] + [InlineData("CREATE TABLE t AS SELECT 1", StatementKind.Ddl, false)] + [InlineData("BEGIN", StatementKind.Transaction, true)] + [InlineData("COMMIT", StatementKind.Transaction, true)] + [InlineData("SET search_path = x", StatementKind.Session, true)] + [InlineData("SHOW all", StatementKind.Session, true)] + [InlineData("DECLARE c CURSOR FOR SELECT 1", StatementKind.Cursor, true)] + [InlineData("DECLARE c CURSOR FOR SELECT 1 FROM t FOR UPDATE", StatementKind.Cursor, false)] + [InlineData("FETCH 1 FROM c", StatementKind.Cursor, true)] + [InlineData("PREPARE p AS SELECT 1", StatementKind.Prepared, true)] + [InlineData("PREPARE p AS DELETE FROM t", StatementKind.Prepared, false)] + [InlineData("EXECUTE p", StatementKind.Prepared, false)] + [InlineData("DEALLOCATE p", StatementKind.Prepared, true)] + [InlineData("CALL p()", StatementKind.Call, false)] + [InlineData("DO $$ BEGIN END $$", StatementKind.Call, false)] + [InlineData("VACUUM t", StatementKind.Maintenance, false)] + [InlineData("ANALYZE t", StatementKind.Maintenance, false)] + [InlineData("REFRESH MATERIALIZED VIEW m", StatementKind.Maintenance, false)] + [InlineData("LOCK TABLE t", StatementKind.Lock, false)] + public void KindAndReadOnly(string query, StatementKind kind, bool readOnly) + { + var info = Classify(query); + + info.Kind.ShouldBe(kind); + info.IsReadOnly.ShouldBe(readOnly); + // IsUtility agrees with libpg_query's own answer. + info.IsUtility.ShouldBe(Parser.IsUtilityStmt(query).GetValueOrThrow()[0]); + } + + [Fact] + public void WritesHiddenInsideASelect() + { + var into = Classify("SELECT * INTO backup FROM t"); + into.HasSelectInto.ShouldBeTrue(); + into.IsReadOnly.ShouldBeFalse(); + + var cte = Classify("WITH gone AS (DELETE FROM t RETURNING *) SELECT * FROM gone"); + cte.Kind.ShouldBe(StatementKind.Select); + cte.HasDataModifyingCte.ShouldBeTrue(); + cte.IsReadOnly.ShouldBeFalse(); + + var nested = Classify("SELECT * FROM (WITH x AS (INSERT INTO t VALUES (1) RETURNING a) SELECT a FROM x) s"); + nested.HasDataModifyingCte.ShouldBeTrue(); + + var locking = Classify("SELECT * FROM t WHERE id IN (SELECT id FROM u FOR SHARE)"); + locking.HasLockingClause.ShouldBeTrue(); + locking.IsReadOnly.ShouldBeFalse(); + + Classify("WITH x AS (SELECT 1) SELECT * FROM x").IsReadOnly.ShouldBeTrue(); + } + + [Theory] + [InlineData("EXPLAIN DELETE FROM t", false, true)] + [InlineData("EXPLAIN ANALYZE DELETE FROM t", true, false)] + [InlineData("EXPLAIN (ANALYZE) DELETE FROM t", true, false)] + [InlineData("EXPLAIN (ANALYZE true, BUFFERS) DELETE FROM t", true, false)] + [InlineData("EXPLAIN (ANALYZE false) DELETE FROM t", false, true)] + [InlineData("EXPLAIN (ANALYZE off) DELETE FROM t", false, true)] + [InlineData("EXPLAIN ANALYZE SELECT 1", true, true)] + public void ExplainRunsItsStatementOnlyWithAnalyze(string query, bool executes, bool readOnly) + { + var info = Classify(query); + + info.Kind.ShouldBe(StatementKind.Explain); + info.ExecutesInner.ShouldBe(executes); + info.IsReadOnly.ShouldBe(readOnly); + info.Inner.ShouldNotBeNull(); + } + + [Fact] + public void InnerStatementsAndPrograms() + { + Classify("EXPLAIN SELECT 1").Inner.ShouldNotBeNull().Kind.ShouldBe(StatementKind.Select); + Classify("CREATE TABLE t AS SELECT 1").Inner.ShouldNotBeNull().Node.ShouldBeOfType(); + Classify("COPY (SELECT 1) TO STDOUT").IsReadOnly.ShouldBeTrue(); + + var program = Classify("COPY (SELECT 1) TO PROGRAM 'cat'"); + program.RunsProgram.ShouldBeTrue(); + program.IsReadOnly.ShouldBeFalse(); + + // A write inside the wrapped statement belongs to Inner, not to the wrapper. + var explain = Classify("EXPLAIN WITH d AS (DELETE FROM t RETURNING 1) SELECT * FROM d"); + explain.HasDataModifyingCte.ShouldBeFalse(); + explain.Inner.ShouldNotBeNull().HasDataModifyingCte.ShouldBeTrue(); + } + + [Fact] + public async Task ClassifyScriptsAndErrors() + { + var infos = Parser.Classify("SELECT 1; DROP TABLE t; BEGIN").GetValueOrThrow(); + infos.Select(info => info.Kind).ShouldBe([StatementKind.Select, StatementKind.Ddl, StatementKind.Transaction]); + + var tree = Parser.Parse("SELECT 1; DELETE FROM t").GetValueOrThrow(); + tree.Classify().Select(info => info.IsReadOnly).ShouldBe([true, false]); + tree.Stmts[1].Classify().Kind.ShouldBe(StatementKind.Delete); + + Parser.Classify("SELECT FROM WHERE").IsSuccess.ShouldBeFalse(); + Parser.Classify("").GetValueOrThrow().ShouldBeEmpty(); + (await Parser.ClassifyAsync("SELECT 1")).GetValueOrThrow().Single().Kind.ShouldBe(StatementKind.Select); + } + + [Fact] + public void TableReferencesWithRolesAliasesAndCtes() + { + const string query = + "WITH recent AS (SELECT * FROM orders WHERE placed > now()) " + + "SELECT c.name, lower(c.email) FROM public.customers c JOIN recent r ON r.customer_id = c.id " + + "WHERE c.id IN (SELECT customer_id FROM vip)"; + + var references = Parser.Parse(query).GetValueOrThrow().GetReferences(); + + references.Tables.Select(table => (table.ToString(), table.Alias, table.Role, table.IsCte)).ShouldBe( + [ + ("public.customers", "c", TableRole.Read, false), + ("recent", "r", TableRole.Read, true), + ("vip", null, TableRole.Read, false), + ("orders", null, TableRole.Read, false) + ]); + references.Tables[0].Schema.ShouldBe("public"); + references.Tables[0].ReferenceName.ShouldBe("c"); + references.Tables[2].ReferenceName.ShouldBe("vip"); + references.Functions.Select(function => function.Name).ShouldBe(["lower", "now"]); + } + + [Fact] + public void TableRolesAcrossStatementKinds() + { + const string script = + "INSERT INTO audit (who) SELECT name FROM users; " + + "UPDATE accounts a SET balance = 0 FROM transfers x WHERE x.account_id = a.id; " + + "DELETE FROM sessions s USING users u WHERE u.id = s.user_id; " + + "TRUNCATE logs; COPY imports FROM STDIN; COPY exports TO STDOUT; " + + "CREATE INDEX i ON items (sku); CREATE VIEW v AS SELECT id FROM items; " + + "SELECT * INTO backup FROM items; ALTER TABLE items ADD COLUMN x int; LOCK TABLE items"; + + var tables = Parser.Parse(script).GetValueOrThrow().GetReferences().Tables; + + tables.Select(table => (table.Name, table.Role, table.StatementIndex)).ShouldBe( + [ + ("audit", TableRole.Write, 0), ("users", TableRole.Read, 0), + ("accounts", TableRole.Write, 1), ("transfers", TableRole.Read, 1), + ("sessions", TableRole.Write, 2), ("users", TableRole.Read, 2), + ("logs", TableRole.Write, 3), + ("imports", TableRole.Write, 4), + ("exports", TableRole.Read, 5), + ("items", TableRole.Ddl, 6), + ("v", TableRole.Ddl, 7), ("items", TableRole.Read, 7), + ("backup", TableRole.Ddl, 8), ("items", TableRole.Read, 8), + ("items", TableRole.Ddl, 9), + ("items", TableRole.Ddl, 10) + ]); + } + + [Fact] + public void ColumnsAreMatchedToTablesWhereTheQuerySettlesIt() + { + const string query = + "SELECT c.name, r.total, status, s.x FROM public.customers c " + + "JOIN recent r ON r.customer_id = c.id, (SELECT x FROM other) s " + + "WHERE c.id IN (SELECT customer_id FROM vip WHERE c.region = region)"; + + var columns = Parser.Parse(query).GetValueOrThrow().GetReferences().Columns; + + columns.Select(column => (column.ToString(), column.Table?.ToString())).ShouldBe( + [ + ("c.name", "public.customers"), + ("r.total", "recent"), + // No qualifier and several tables in scope: it takes the catalog to tell. + ("status", null), + // s is a subquery, not a table. + ("s.x", null), + ("r.customer_id", "recent"), + ("c.id", "public.customers"), + ("x", "other"), + ("c.id", "public.customers"), + // Inside the subquery: one table in scope, and c reaches the outer query. + ("customer_id", "vip"), + ("c.region", "public.customers"), + ("region", "vip") + ]); + columns[0].Qualifier.ShouldBe("c"); + columns[0].Name.ShouldBe("name"); + } + + [Fact] + public void ColumnsOfDmlStarsAndCtes() + { + var update = Parser.Parse("UPDATE accounts a SET balance = a.balance - x.amount FROM transfers x WHERE x.id = a.id RETURNING id") + .GetValueOrThrow().GetReferences(); + update.Columns.Select(column => (column.ToString(), column.Table?.Name)).ShouldBe( + [("a.balance", "accounts"), ("x.amount", "transfers"), ("x.id", "transfers"), ("a.id", "accounts"), ("id", null)]); + + var star = Parser.Parse("SELECT *, t.* FROM t").GetValueOrThrow().GetReferences(); + star.Columns.Select(column => (column.Qualifier, column.Name, column.Table?.Name)).ShouldBe([(null, "*", "t"), ("t", "*", "t")]); + + // A column of a CTE is not attributed to a table. + var cte = Parser.Parse("WITH r AS (SELECT id FROM base) SELECT r.id, id FROM r").GetValueOrThrow().GetReferences(); + cte.Columns.Select(column => (column.ToString(), column.Table?.Name)).ShouldBe([("r.id", null), ("id", null), ("id", "base")]); + } + + [Fact] + public void ReferencesFromANodeAndWithNonAsciiNames() + { + var tree = Parser.Parse("SELECT d.\"é\" FROM \"données\" d WHERE d.\"😀\" = 1").GetValueOrThrow(); + + var references = tree.Stmts[0].Stmt.SelectStmt.GetReferences(); + + references.Tables.Single().Name.ShouldBe("données"); + references.Tables.Single().StatementIndex.ShouldBeNull(); + references.Columns.Select(column => column.Name).ShouldBe(["é", "😀"]); + references.Columns.ShouldAllBe(column => column.Table != null && column.Table.Name == "données"); + } + + [Fact] + public void ReferencesFromASingleStatementNode() + { + var tree = Parser.Parse("WITH r AS (SELECT 1) INSERT INTO target SELECT x FROM r; CREATE INDEX i ON items (sku)").GetValueOrThrow(); + + // Starting from the statement node itself, with nothing above it. + var insert = tree.Stmts[0].Stmt.InsertStmt.GetReferences(); + insert.Tables.Select(table => (table.Name, table.Role, table.IsCte)).ShouldBe( + [("target", TableRole.Write, false), ("r", TableRole.Read, true)]); + + tree.Stmts[1].Stmt.IndexStmt.GetReferences().Tables.Single().Role.ShouldBe(TableRole.Ddl); + // And from the Node wrapper or the RawStmt. + tree.Stmts[1].Stmt.GetReferences().Tables.Single().Role.ShouldBe(TableRole.Ddl); + tree.Stmts[1].GetReferences().Tables.Single().Role.ShouldBe(TableRole.Ddl); + } + + [Fact] + public void FormatKeepsCommentsAndSeparatesStatements() + { + const string query = "select a, b from t -- pick\n where x = 1; -- second\ndelete from t where a = 1; -- done"; + + Parser.Format(query).GetValueOrThrow().ShouldBe( + "SELECT a, b\nFROM t\nWHERE\n -- pick\n x = 1;\n\n-- second\nDELETE FROM t\nWHERE a = 1;\n\n-- done"); + Parser.Format(query, new FormatOptions { KeepComments = false }).GetValueOrThrow().ShouldBe( + "SELECT a, b\nFROM t\nWHERE x = 1;\n\nDELETE FROM t\nWHERE a = 1;"); + } + + [Fact] + public async Task FormatOptionsAndEdgeCases() + { + Parser.Format("select 1").GetValueOrThrow().ShouldBe("SELECT 1"); + Parser.Format("select 1;").GetValueOrThrow().ShouldBe("SELECT 1"); + Parser.Format("select 1", new FormatOptions { TrailingNewline = true }).GetValueOrThrow().ShouldBe("SELECT 1\n"); + Parser.Format("select a from t where x = 1 and y = 2", new FormatOptions { IndentSize = 2 }).GetValueOrThrow() + .ShouldBe("SELECT a\nFROM t\nWHERE\n x = 1\n AND y = 2"); + Parser.Format("-- only a comment").GetValueOrThrow().ShouldBe("-- only a comment"); + Parser.Format("").GetValueOrThrow().ShouldBe(""); + Parser.Format("select '😀' as \"é\" from \"données\"").GetValueOrThrow().ShouldBe("SELECT '😀' AS \"é\"\nFROM \"données\""); + Parser.Format("select from where").IsSuccess.ShouldBeFalse(); + (await Parser.FormatAsync("select 1")).GetValueOrThrow().ShouldBe("SELECT 1"); + + // Formatting what was formatted changes nothing. + var once = Parser.Format("select a,b from t where x in (select 1) order by a").GetValueOrThrow(); + Parser.Format(once).GetValueOrThrow().ShouldBe(once); + } + + [Fact] + public void ErrorLineColumnAndRendering() + { + const string query = "SELECT a,\n\t'😀' AS e,\n\tb FROM WHERE x"; + var error = Parser.Parse(query).Error.ShouldNotBeNull(); + + error.GetLineAndColumn(query).ShouldBe((3, 9)); + error.Format(query).ShouldBe( + "ERROR: syntax error at or near \"WHERE\"\n" + + "LINE 3: b FROM WHERE x\n" + + " ^"); + } + + [Fact] + public void ErrorRenderingEdgeCases() + { + const string single = "SELECT FROM WHERE"; + var error = Parser.Parse(single).Error.ShouldNotBeNull(); + error.GetLineAndColumn(single).ShouldBe((1, 13)); + error.Format(single).ShouldBe("ERROR: syntax error at or near \"WHERE\"\nLINE 1: SELECT FROM WHERE\n ^"); + + // At the end of the input, and with Windows line endings. + const string atEnd = "SELECT 1\r\nFROM"; + var endError = Parser.Parse(atEnd).Error.ShouldNotBeNull(); + endError.GetLineAndColumn(atEnd).ShouldBe((2, 5)); + endError.Format(atEnd).ShouldBe("ERROR: syntax error at end of input\nLINE 2: FROM\n ^"); + + var noPosition = new Error("something failed", null, null, 0, 0, null); + noPosition.GetLineAndColumn(single).ShouldBeNull(); + noPosition.Format(single).ShouldBe("ERROR: something failed"); + } +} diff --git a/src/PgSqlParser/Parser.cs b/src/PgSqlParser/Parser.cs index c0a28a2..aa5ff92 100644 --- a/src/PgSqlParser/Parser.cs +++ b/src/PgSqlParser/Parser.cs @@ -41,6 +41,50 @@ public int GetCursorCharOffset(string query) return charOffset; } + + /// + /// Returns the 1-based line and column of in , or + /// null if the error has no position in it. The column counts UTF-16 chars. + /// + public (int Line, int Column)? GetLineAndColumn(string query) + { + var offset = GetCursorCharOffset(query); + if (offset < 0) + return null; + + var line = 1; + var lineStart = 0; + for (var i = 0; i < offset; i++) + { + if (query[i] == '\n') + { + line++; + lineStart = i + 1; + } + } + + return (line, offset - lineStart + 1); + } + + /// + /// Formats the error the way psql does: the message, then the line of it + /// points at with a caret under the position. + /// + public string Format(string query) + { + var header = $"ERROR: {Message}"; + if (GetLineAndColumn(query) is not var (line, column)) + return header; + + var offset = GetCursorCharOffset(query); + var lineStart = offset - (column - 1); + var lineEnd = query.IndexOf('\n', lineStart); + // Tabs become spaces so the caret lines up whatever the tab width. + var text = query[lineStart..(lineEnd < 0 ? query.Length : lineEnd)].TrimEnd('\r').Replace('\t', ' '); + var prefix = $"LINE {line}: "; + + return $"{header}\n{prefix}{text}\n{new string(' ', prefix.Length + column - 1)}^"; + } } /// @@ -218,6 +262,27 @@ public class DeparseOptions public bool CommasStartOfLine { get; init; } } +/// +/// Options for . +/// +public class FormatOptions +{ + /// Keep the comments of the query. On by default. + public bool KeepComments { get; init; } = true; + + /// Indentation size in spaces. + public int IndentSize { get; init; } = 4; + + /// Restricts the line length of certain lists of items. + public int MaxLineLength { get; init; } = 80; + + /// Add a trailing newline at the end of the output. + public bool TrailingNewline { get; init; } + + /// Place separating commas at the start of the line. + public bool CommasStartOfLine { get; init; } +} + /// /// Flags that control how fingerprints are calculated. /// @@ -903,6 +968,114 @@ public static Result Deparse(ParseResult parseResult, DeparseOptions? op } } + /// + /// Format a query: parse it and print it again across indented lines, keeping its comments. + /// This is , and + /// with pretty printing in one call. + /// + /// + /// + /// + public static Result Format(string query, FormatOptions? options = null) + { + options ??= new FormatOptions(); + + // Each statement is formatted on its own and they are joined one per paragraph. The scanner + // split keeps the comments in front of a statement together with it. + if (!SplitWithScanner(query).TryGetValue(out var split, out var error)) + return Result.Failure(error); + + var formatted = new List(); + var end = 0; + foreach (var statement in split.Statements) + { + if (!FormatStatement(statement.Text.Trim(), options).TryGetValue(out var text, out error)) + return Result.Failure(error); + + if (text.Length > 0) + formatted.Add(text.TrimEnd('\n')); + + end = statement.Location + statement.Length; + } + + var result = string.Join(";\n\n", formatted); + if (formatted.Count > 1) + result += ";"; + + // What follows the last statement can only be comments. + var tail = query[end..].Trim().TrimStart(';').Trim(); + if (options.KeepComments && tail.Length > 0) + result = result.Length == 0 ? tail : $"{result}\n\n{tail}"; + + return Result.Success(options.TrailingNewline && result.Length > 0 ? result + "\n" : result); + } + + private static Result FormatStatement(string statement, FormatOptions options) + { + if (!Parse(statement).TryGetValue(out var tree, out var error)) + return Result.Failure(error); + + // Only comments, or nothing at all. + if (tree.Stmts.Count == 0) + return Result.Success(options.KeepComments ? statement : string.Empty); + + IReadOnlyList comments = []; + if (options.KeepComments) + { + if (!DeparseComments(statement).TryGetValue(out var found, out error)) + return Result.Failure(error); + + comments = found; + } + + return Deparse(tree, new DeparseOptions + { + Comments = comments, + PrettyPrint = true, + IndentSize = options.IndentSize, + MaxLineLength = options.MaxLineLength, + CommasStartOfLine = options.CommasStartOfLine + }); + } + + /// + /// Async format a query + /// + /// + /// + /// + /// + public static Task> FormatAsync(string query, FormatOptions? options = null, + CancellationToken cancellationToken = default) + { + return RunAsync(() => Format(query, options), cancellationToken); + } + + /// + /// Classify each statement of a query: its kind, whether it is read-only, and facts such as a + /// data-modifying CTE or a locking clause. See . + /// + /// + /// + public static Result> Classify(string query) + { + return Parse(query).TryGetValue(out var tree, out var error) + ? Result>.Success(tree.Classify()) + : Result>.Failure(error); + } + + /// + /// Async classify each statement of a query + /// + /// + /// + /// + public static Task>> ClassifyAsync(string query, + CancellationToken cancellationToken = default) + { + return RunAsync(() => Classify(query), cancellationToken); + } + /// /// Deparse a single node of a parse tree back into SQL, for example only a WHERE clause, one /// expression or one table reference. Supported are statements, expressions, items of a FROM diff --git a/src/PgSqlParser/PgSqlParser.csproj b/src/PgSqlParser/PgSqlParser.csproj index f285d3a..232d774 100644 --- a/src/PgSqlParser/PgSqlParser.csproj +++ b/src/PgSqlParser/PgSqlParser.csproj @@ -2,7 +2,7 @@ .NET wrapper on libpg_query to access PostgreSQL parser outside of the server environment - 2.1.0 + 2.2.0 Babu Annamalai net8.0;net9.0;net10.0 enable diff --git a/src/PgSqlParser/QueryReferences.cs b/src/PgSqlParser/QueryReferences.cs new file mode 100644 index 0000000..0ccf70b --- /dev/null +++ b/src/PgSqlParser/QueryReferences.cs @@ -0,0 +1,289 @@ +using Google.Protobuf; + +namespace PgSqlParser; + +/// +/// How a statement uses a table. +/// +public enum TableRole +{ + /// The table is read: in a FROM clause, a join, a subquery and so on. + Read, + + /// The table is the target of INSERT, UPDATE, DELETE, MERGE, TRUNCATE or COPY ... FROM. + Write, + + /// The table is defined, altered, dropped or otherwise managed by a utility statement. + Ddl +} + +/// +/// A table named in a query. +/// +/// The node that names it. +/// How the statement uses it. +/// +/// True if the name refers to a common table expression of the statement and not to a table in the database. +/// +/// The top-level statement it is in, when known. +public sealed record TableReference(RangeVar Node, TableRole Role, bool IsCte, int? StatementIndex) +{ + public string? Schema => Node.Schemaname.Length > 0 ? Node.Schemaname : null; + public string Name => Node.Relname; + public string? Alias => Node.Alias is { Aliasname.Length: > 0 } alias ? alias.Aliasname : null; + + /// The name a column reference uses for this table: its alias, or its name without one. + public string ReferenceName => Alias ?? Name; + + public override string ToString() => Schema is null ? Name : $"{Schema}.{Name}"; +} + +/// +/// A function called in a query. +/// +public sealed record FunctionReference(FuncCall Node, string? Schema, string Name, int? StatementIndex) +{ + public override string ToString() => Schema is null ? Name : $"{Schema}.{Name}"; +} + +/// +/// A column named in a query. +/// +/// The node that names it. +/// What is written in front of the column, such as a table alias, or null. +/// The column name, or *. +/// +/// The table the column belongs to, when the query alone settles it: the qualifier matches a table or +/// alias in scope, or there is no qualifier and only one table in scope. Null when it would take the +/// database catalog to tell, or when the column comes from a subquery or a CTE. +/// +/// The top-level statement it is in, when known. +public sealed record ColumnReference(ColumnRef Node, string? Qualifier, string Name, TableReference? Table, int? StatementIndex) +{ + public override string ToString() => Qualifier is null ? Name : $"{Qualifier}.{Name}"; +} + +/// +/// The tables, functions and columns a query refers to, in the order they appear in its parse tree. +/// +public sealed class QueryReferences +{ + internal QueryReferences(IReadOnlyList tables, IReadOnlyList functions, + IReadOnlyList columns) + { + Tables = tables; + Functions = functions; + Columns = columns; + } + + public IReadOnlyList Tables { get; } + public IReadOnlyList Functions { get; } + public IReadOnlyList Columns { get; } +} + +/// +/// Finds what a query refers to by reading its parse tree. +/// +public static class QueryReferenceExtensions +{ + /// + /// Returns the tables, functions and columns referred to below . Tables carry + /// their role and whether they are really a CTE; columns are matched to their table where the query + /// alone settles it. + /// + public static QueryReferences GetReferences(this IMessage root) + { + ArgumentNullException.ThrowIfNull(root); + + // The root is not among a visit's ancestors, but it is the outermost thing enclosing every node. + var top = root is Node wrapper ? wrapper.Unwrap() ?? root : root; + var tables = new List(); + var functions = new List(); + var columnVisits = new List(); + var tableByNode = new Dictionary(ReferenceEqualityComparer.Instance); + + foreach (var visit in root.Walk()) + { + switch (visit.Node) + { + case RangeVar table: + var reference = new TableReference(table, RoleOf(visit, top), IsCte(table, visit, top), visit.StatementIndex); + tables.Add(reference); + tableByNode[table] = reference; + break; + case FuncCall call when call.Funcname.Count > 0: + var names = call.Funcname.Select(part => part.String?.Sval ?? string.Empty).ToList(); + functions.Add(new FunctionReference(call, names.Count > 1 ? names[^2] : null, names[^1], visit.StatementIndex)); + break; + case ColumnRef: + columnVisits.Add(visit); + break; + } + } + + // Columns are resolved last, since a column can come before the table it belongs to. + var columns = columnVisits.Select(visit => ResolveColumn(visit, top, tableByNode)).ToList(); + return new QueryReferences(tables, functions, columns); + } + + private static IEnumerable Enclosing(NodeVisit visit, IMessage top) + { + foreach (var ancestor in visit.Ancestors) + yield return ancestor; + + yield return top; + } + + private static TableRole RoleOf(NodeVisit visit, IMessage top) + { + switch (visit.Parent ?? top) + { + case InsertStmt or UpdateStmt or DeleteStmt or MergeStmt when visit.FieldName == nameof(InsertStmt.Relation): + case TruncateStmt: + case CopyStmt { IsFrom: true }: + return TableRole.Write; + // SELECT ... INTO and CREATE TABLE AS create the table they name. + case IntoClause: + return TableRole.Ddl; + } + + // Anything inside a query, or elsewhere in a data-changing statement, is read. + foreach (var ancestor in Enclosing(visit, top)) + { + if (ancestor is SelectStmt or InsertStmt or UpdateStmt or DeleteStmt or MergeStmt or CopyStmt) + return TableRole.Read; + } + + var statement = visit.Statement ?? top as RawStmt ?? new RawStmt { Stmt = top is ParseResult ? null : top.AsNode() }; + return statement.Classify().Kind is StatementKind.Ddl or StatementKind.Maintenance or StatementKind.Lock + ? TableRole.Ddl + : TableRole.Read; + } + + // An unqualified name that matches a CTE of an enclosing statement refers to that CTE. + private static bool IsCte(RangeVar table, NodeVisit visit, IMessage top) + { + if (table.Schemaname.Length > 0) + return false; + + foreach (var ancestor in Enclosing(visit, top)) + { + var with = ancestor switch + { + SelectStmt select => select.WithClause, + InsertStmt insert => insert.WithClause, + UpdateStmt update => update.WithClause, + DeleteStmt delete => delete.WithClause, + MergeStmt merge => merge.WithClause, + _ => null + }; + + if (with is not null && with.Ctes.Any(cte => cte.CommonTableExpr?.Ctename == table.Relname)) + return true; + } + + return false; + } + + private static ColumnReference ResolveColumn(NodeVisit visit, IMessage top, Dictionary tableByNode) + { + var column = (ColumnRef)visit.Node; + var parts = column.Fields.Select(field => field.String?.Sval ?? (field.AStar is not null ? "*" : string.Empty)).ToList(); + var name = parts.Count > 0 ? parts[^1] : string.Empty; + var qualifier = parts.Count > 1 ? parts[^2] : null; + + TableReference? table = null; + foreach (var ancestor in Enclosing(visit, top)) + { + var scope = ScopeOf(ancestor); + if (scope is null) + continue; + + if (qualifier is null) + { + // Without a qualifier only a single table in the nearest scope is certain. + if (scope is [{ Table: { } only }]) + table = tableByNode.GetValueOrDefault(only); + break; + } + + var match = scope.FirstOrDefault(item => item.Name == qualifier); + if (match.Name is not null) + { + table = match.Table is null ? null : tableByNode.GetValueOrDefault(match.Table); + break; + } + } + + // A name that is really a CTE is not a table to attribute columns to. + return new ColumnReference(column, qualifier, name, table is { IsCte: true } ? null : table, visit.StatementIndex); + } + + // The names a statement brings into scope for its columns. Table is null for a subquery, function + // or other source that is not a plain table. + private static List<(string Name, RangeVar? Table)>? ScopeOf(IMessage statement) + { + var scope = new List<(string Name, RangeVar? Table)>(); + switch (statement) + { + case SelectStmt select: + foreach (var item in select.FromClause) + AddFromItem(item, scope); + break; + case UpdateStmt update: + AddTable(update.Relation, scope); + foreach (var item in update.FromClause) + AddFromItem(item, scope); + break; + case DeleteStmt delete: + AddTable(delete.Relation, scope); + foreach (var item in delete.UsingClause) + AddFromItem(item, scope); + break; + case InsertStmt insert: + AddTable(insert.Relation, scope); + break; + case MergeStmt merge: + AddTable(merge.Relation, scope); + AddFromItem(merge.SourceRelation, scope); + break; + default: + return null; + } + + return scope; + } + + private static void AddFromItem(Node? item, List<(string Name, RangeVar? Table)> scope) + { + switch (item?.Unwrap()) + { + case RangeVar table: + AddTable(table, scope); + break; + case JoinExpr join: + AddFromItem(join.Larg, scope); + AddFromItem(join.Rarg, scope); + if (join.Alias is { Aliasname.Length: > 0 } joinAlias) + scope.Add((joinAlias.Aliasname, null)); + break; + case RangeSubselect { Alias.Aliasname: { Length: > 0 } alias }: + scope.Add((alias, null)); + break; + case RangeFunction { Alias.Aliasname: { Length: > 0 } alias }: + scope.Add((alias, null)); + break; + case RangeTableFunc { Alias.Aliasname: { Length: > 0 } alias }: + scope.Add((alias, null)); + break; + } + } + + private static void AddTable(RangeVar? table, List<(string Name, RangeVar? Table)> scope) + { + if (table is null) + return; + + scope.Add((table.Alias is { Aliasname.Length: > 0 } alias ? alias.Aliasname : table.Relname, table)); + } +} diff --git a/src/PgSqlParser/StatementInfo.cs b/src/PgSqlParser/StatementInfo.cs new file mode 100644 index 0000000..6789eff --- /dev/null +++ b/src/PgSqlParser/StatementInfo.cs @@ -0,0 +1,216 @@ +using Google.Protobuf; + +namespace PgSqlParser; + +/// +/// The broad kind of a statement. +/// +public enum StatementKind +{ + /// Anything not covered by the other kinds. + Other, + + /// SELECT, VALUES and TABLE. + Select, + Insert, + Update, + Delete, + Merge, + Truncate, + Copy, + + /// + /// Statements that define or change database objects, privileges or roles: CREATE, ALTER, DROP, + /// COMMENT, GRANT, REVOKE and so on. + /// + Ddl, + + /// BEGIN, COMMIT, ROLLBACK, SAVEPOINT, PREPARE TRANSACTION, SET CONSTRAINTS. + Transaction, + + /// SET, RESET, SHOW, DISCARD, LISTEN, UNLISTEN, NOTIFY, LOAD. + Session, + + /// DECLARE, FETCH, MOVE, CLOSE. + Cursor, + + /// PREPARE, EXECUTE, DEALLOCATE. + Prepared, + Explain, + + /// CALL and DO. + Call, + + /// VACUUM, ANALYZE, CLUSTER, REINDEX, CHECKPOINT, REFRESH MATERIALIZED VIEW. + Maintenance, + Lock +} + +/// +/// Facts about one statement, read from its parse tree. These are facts, not a verdict: what to +/// allow is for the caller to decide. +/// +public sealed record StatementInfo +{ + /// The statement node, for example a . + public required IMessage Node { get; init; } + + public required StatementKind Kind { get; init; } + + /// + /// True for everything except SELECT, INSERT, UPDATE, DELETE and MERGE, which are the statements + /// PostgreSQL plans. Matches . + /// + public bool IsUtility => Kind is not (StatementKind.Select or StatementKind.Insert or StatementKind.Update + or StatementKind.Delete or StatementKind.Merge); + + /// True if a WITH clause anywhere in the statement contains an INSERT, UPDATE, DELETE or MERGE. + public bool HasDataModifyingCte { get; init; } + + /// True for SELECT ... INTO, which creates a table. + public bool HasSelectInto { get; init; } + + /// True if a SELECT anywhere in the statement has FOR UPDATE, FOR SHARE or a similar clause. + public bool HasLockingClause { get; init; } + + /// True for EXPLAIN with ANALYZE, which runs the statement it explains. + public bool ExecutesInner { get; init; } + + /// True for COPY ... PROGRAM, which runs a command on the server. + public bool RunsProgram { get; init; } + + /// + /// The statement wrapped by this one: the query of EXPLAIN, PREPARE, DECLARE CURSOR, CREATE TABLE AS + /// or COPY (query). Null otherwise. + /// + public StatementInfo? Inner { get; init; } + + /// + /// True if the statement, as written, changes no data and no database objects and takes no explicit + /// locks. It errs on the side of false: CALL, DO and EXECUTE count as not read-only because what they + /// run is not visible. A function called from a SELECT can still write, which the parse tree cannot show. + /// + public bool IsReadOnly => Kind switch + { + StatementKind.Select => !HasSelectInto && !HasDataModifyingCte && !HasLockingClause, + StatementKind.Explain => !ExecutesInner || Inner is { IsReadOnly: true }, + StatementKind.Copy => Node is CopyStmt { IsFrom: false } && !RunsProgram && Inner is null or { IsReadOnly: true }, + StatementKind.Cursor => Inner is null or { IsReadOnly: true }, + StatementKind.Prepared => Node is DeallocateStmt || Inner is { IsReadOnly: true }, + StatementKind.Transaction or StatementKind.Session => true, + _ => false + }; +} + +/// +/// Classifies statements by reading their parse tree. +/// +public static class StatementClassifier +{ + /// Classifies every statement of a parse tree, in order. + public static IReadOnlyList Classify(this ParseResult tree) + { + ArgumentNullException.ThrowIfNull(tree); + return tree.Stmts.Select(Classify).ToList(); + } + + /// Classifies one statement. + public static StatementInfo Classify(this RawStmt statement) + { + ArgumentNullException.ThrowIfNull(statement); + return Classify(statement.Stmt?.Unwrap()); + } + + private static StatementInfo Classify(IMessage? node) + { + if (node is null) + return new StatementInfo { Node = new Node(), Kind = StatementKind.Other }; + + var inner = node switch + { + ExplainStmt explain => explain.Query, + PrepareStmt prepare => prepare.Query, + DeclareCursorStmt cursor => cursor.Query, + CreateTableAsStmt createAs => createAs.Query, + CopyStmt copy => copy.Query, + _ => null + }; + + var hasDataModifyingCte = false; + var hasLockingClause = false; + node.Walk(visit => + { + // The wrapped statement is classified on its own, as Inner. + if (inner is not null && ReferenceEquals(visit.Node, inner.Unwrap())) + return WalkAction.SkipChildren; + + switch (visit.Node) + { + case CommonTableExpr cte when cte.Ctequery?.Unwrap() is InsertStmt or UpdateStmt or DeleteStmt or MergeStmt: + hasDataModifyingCte = true; + break; + case SelectStmt { LockingClause.Count: > 0 }: + hasLockingClause = true; + break; + } + + return WalkAction.Continue; + }); + + return new StatementInfo + { + Node = node, + Kind = KindOf(node), + HasDataModifyingCte = hasDataModifyingCte, + HasLockingClause = hasLockingClause || node is SelectStmt { LockingClause.Count: > 0 }, + HasSelectInto = node is SelectStmt { IntoClause: not null }, + ExecutesInner = node is ExplainStmt explainStmt && explainStmt.Options.Any(IsEnabledAnalyze), + RunsProgram = node is CopyStmt { IsProgram: true }, + Inner = inner?.Unwrap() is { } innerNode ? Classify(innerNode) : null + }; + } + + // EXPLAIN (ANALYZE), (ANALYZE true) and (ANALYZE on) run the statement; (ANALYZE false) does not. + private static bool IsEnabledAnalyze(Node option) + { + if (option.DefElem is not { Defname: "analyze" } analyze) + return false; + + return analyze.Arg?.Unwrap() switch + { + null => true, + Boolean boolean => boolean.Boolval, + Integer integer => integer.Ival != 0, + String text => text.Sval is "true" or "on" or "1", + _ => true + }; + } + + private static StatementKind KindOf(IMessage node) + { + return node switch + { + SelectStmt => StatementKind.Select, + InsertStmt => StatementKind.Insert, + UpdateStmt => StatementKind.Update, + DeleteStmt => StatementKind.Delete, + MergeStmt => StatementKind.Merge, + TruncateStmt => StatementKind.Truncate, + CopyStmt => StatementKind.Copy, + TransactionStmt or ConstraintsSetStmt => StatementKind.Transaction, + VariableSetStmt or VariableShowStmt or DiscardStmt or ListenStmt or UnlistenStmt or NotifyStmt + or LoadStmt => StatementKind.Session, + DeclareCursorStmt or FetchStmt or ClosePortalStmt => StatementKind.Cursor, + PrepareStmt or ExecuteStmt or DeallocateStmt => StatementKind.Prepared, + ExplainStmt => StatementKind.Explain, + CallStmt or DoStmt => StatementKind.Call, + VacuumStmt or ClusterStmt or ReindexStmt or CheckPointStmt or RefreshMatViewStmt => StatementKind.Maintenance, + LockStmt => StatementKind.Lock, + // Not statements that can stand at the top level. + RawStmt or ReturnStmt or PLAssignStmt or SetOperationStmt => StatementKind.Other, + // Every remaining statement type creates, alters or drops something, or changes privileges. + _ when node.Descriptor.Name.EndsWith("Stmt", StringComparison.Ordinal) => StatementKind.Ddl, + _ => StatementKind.Other + }; + } +} From e3abf67f0f4adcfdf6731348b8383672d375a945 Mon Sep 17 00:00:00 2001 From: mysticmind Date: Fri, 9 Oct 2026 20:10:22 +0530 Subject: [PATCH 2/2] Add node builders, common rewrites and typed visitors; keep everything under 2.1.0 - Ast builds nodes that match what the parser produces: columns, constants, operators, AND/OR/NOT, IN, calls, casts, tables, select list and ORDER BY items, and Ast.Expression from SQL text - QueryRewrites: AddWhere, SetLimit, CapLimit, SetOffset, ToCount, QualifyTables, RenameSchema, RenameTable, ToDropStatement, EnsureIfNotExists, EnsureOrReplace - NodeVisitor: a handler per node type, several visitors in one walk, each skipping or stopping on its own - EqualsIgnoringLocations now also ignores the start and end positions of parenthesized lists, which made IN lists compare unequal - Version back to 2.1.0, since 2.1.0 was never published; one What's new in 2.1 section in the README --- README.md | 71 ++++- src/PackageSmokeTest/Program.cs | 9 + src/PgSqlParser.Tests/AstTests.cs | 118 ++++++++ src/PgSqlParser.Tests/NodeToolsTests.cs | 2 +- .../RewriteAndVisitorTests.cs | 251 ++++++++++++++++ src/PgSqlParser.Tests/TreeEditingTests.cs | 2 + src/PgSqlParser/Ast.cs | 234 +++++++++++++++ src/PgSqlParser/NodeVisitor.cs | 114 +++++++ src/PgSqlParser/ParseTreeExtensions.cs | 9 +- src/PgSqlParser/PgSqlParser.csproj | 2 +- src/PgSqlParser/QueryRewrites.cs | 277 ++++++++++++++++++ 11 files changed, 1073 insertions(+), 16 deletions(-) create mode 100644 src/PgSqlParser.Tests/AstTests.cs create mode 100644 src/PgSqlParser.Tests/RewriteAndVisitorTests.cs create mode 100644 src/PgSqlParser/Ast.cs create mode 100644 src/PgSqlParser/NodeVisitor.cs create mode 100644 src/PgSqlParser/QueryRewrites.cs diff --git a/README.md b/README.md index 71ac31f..f661eb1 100644 --- a/README.md +++ b/README.md @@ -261,6 +261,48 @@ Console.WriteLine(a.EqualsIgnoringLocations(b)); // Output: True ``` +#### Building nodes and common changes + +`Ast` builds nodes without writing out their fields. Each method gives the same node the parser would for the matching SQL, and `Ast.Expression` parses a piece of SQL when that is easier: + +```csharp +var condition = Ast.And( + Ast.Eq(Ast.Column("tenant_id"), Ast.Param(1)), + Ast.IsNull(Ast.Column("deleted_at"))); +// The same as Ast.Expression("tenant_id = $1 AND deleted_at IS NULL") +``` + +The usual changes to a statement are ready made. They work on the tree, so they hold however the SQL was written: + +```csharp +var tree = Parser.Parse("SELECT id, name FROM customers WHERE active ORDER BY name").GetValueOrThrow(); + +tree.AddWhere(condition); // combined with the existing WHERE using AND +tree.CapLimit(100); // adds LIMIT 100, or lowers a larger one +tree.QualifyTables("app"); // gives unqualified tables a schema, leaving CTE names alone + +Console.WriteLine(tree.Deparse().GetValueOrThrow()); +// Output: SELECT id, name FROM app.customers WHERE active AND tenant_id = $1 AND deleted_at IS NULL ORDER BY name LIMIT 100 + +Console.WriteLine(tree.ToCount().Deparse().GetValueOrThrow()); +// Output: SELECT count(*) FROM (SELECT id, name FROM app.customers WHERE active AND tenant_id = $1 AND deleted_at IS NULL ORDER BY name LIMIT 100) q +``` + +Also available: `SetLimit`, `SetOffset`, `RenameSchema`, `RenameTable`, and for DDL `ToDropStatement` (the DROP that undoes a CREATE), `EnsureIfNotExists` and `EnsureOrReplace`. + +#### Typed visitors + +A `NodeVisitor` holds a handler per node type. Several visitors can share one walk, and each skips or stops on its own: + +```csharp +var tables = new List(); +var functions = new List(); + +tree.Walk( + new NodeVisitor().On((table, visit) => tables.Add(table.Relname)), + new NodeVisitor().On((call, visit) => functions.Add(call.Funcname[^1].String.Sval))); +``` + A few helpers go with these: `Parser.ParseExpression` and `Parser.ParseTypeName` parse a single expression or type name on its own, `node.AsNode()` wraps a node for a property or list that takes any node, and `PgIdentifier.Quote` quotes a name the way PostgreSQL's `quote_ident` does. `Walk` and `Descendants` visit parents before their children and siblings in field order, which is not always the order of the query text. `GetLocation()` returns a node's location if it has one; see [Offsets and non-ASCII text](#offsets-and-non-ascii-text) for its unit. Only statements record a length, so `GetText` is available for statements and not for other nodes. @@ -665,29 +707,36 @@ Deeply nested queries, such as a long chain of operators without parentheses (`a Protobuf's own recursive operations on a parse tree, such as `ToString()`, `Clone()` and `Equals()`, run on your thread. On a very deeply nested tree they can still overflow a small stack. -## What's new in 2.2 +## What's new in 2.1 + +All additions; nothing from 2.0 changes. -All additions; nothing from earlier versions changes. +**Understanding a statement** - **Classify statements**: `Parser.Classify` gives each statement's kind, whether it is read-only, and facts such as a data-modifying CTE, `SELECT INTO`, a locking clause or `EXPLAIN ANALYZE`. - **Resolved references**: `GetReferences` returns the tables, functions and columns of a query as nodes, with table roles (read, write, DDL), CTE detection and columns matched to their table. -- **Format in one call**: `Parser.Format` pretty prints a query or a script and keeps its comments. -- **Readable errors**: `Error.Format` and `Error.GetLineAndColumn`. +- **Parameters**: `ParameterRefs` finds each `$n`, where it is, and the type it is cast to. -## What's new in 2.1 - -All additions; nothing from 2.0 changes. +**Working with the tree** - **Deparse a single node**: `node.Deparse()` turns one clause, expression or table reference back into SQL. - **More control when walking**: skip a subtree or stop, and each visit knows its ancestors, the property it came from and its top-level statement. -- **Change trees**: `Rewrite` keeps, replaces or removes nodes in place. +- **Typed visitors**: `NodeVisitor` with a handler per node type, and several visitors in one walk. - **Compare trees**: `EqualsIgnoringLocations` treats queries that differ only in formatting as equal. + +**Changing a statement** + +- **Build nodes**: `Ast.Column`, `Ast.Const`, `Ast.Eq`, `Ast.And`, `Ast.Call`, `Ast.Table`, `Ast.Expression` and more, each giving the same node the parser would. +- **Rewrite**: `Rewrite` keeps, replaces or removes nodes in place. +- **Common changes ready made**: `AddWhere`, `SetLimit`, `CapLimit`, `SetOffset`, `ToCount`, `QualifyTables`, `RenameSchema`, `RenameTable`, `ToDropStatement`, `EnsureIfNotExists`, `EnsureOrReplace`. + +**Working with text** + +- **Format in one call**: `Parser.Format` pretty prints a query or a script and keeps its comments. +- **Readable errors**: `Error.Format` and `Error.GetLineAndColumn`. - **Parse fragments**: `ParseExpression` and `ParseTypeName`. -- **Parameters**: `ParameterRefs` finds each `$n`, where it is, and the type it is cast to. - **Identifier quoting**: `PgIdentifier.Quote`, following PostgreSQL's `quote_ident`. -See [Navigating the parse tree](#navigating-the-parse-tree) and [ParameterRefs](#parameterrefs). - ## What's new in 2.0 - **PostgreSQL 18**: built on libpg_query 18.1.0, which uses the PostgreSQL 18.6 parser and includes its security fix for `pg_query_normalize`. diff --git a/src/PackageSmokeTest/Program.cs b/src/PackageSmokeTest/Program.cs index 021b448..fde4ea5 100644 --- a/src/PackageSmokeTest/Program.cs +++ b/src/PackageSmokeTest/Program.cs @@ -72,6 +72,15 @@ void Check(string name, bool ok) && refs.Columns.All(c => c.Table is not null)); Check("Format", Parser.Format("select a from t -- note\n where x = 1").Value == "SELECT a\nFROM t\nWHERE\n -- note\n x = 1"); Check("Error.Format", Parser.Parse("SELECT FROM WHERE").Error!.Format("SELECT FROM WHERE").EndsWith(" ^")); +var rewritten = Parser.Parse("SELECT id FROM customers WHERE active").Value!; +rewritten.AddWhere(Ast.And(Ast.Eq(Ast.Column("tenant_id"), Ast.Param(1)), Ast.IsNull(Ast.Column("deleted_at")))); +rewritten.CapLimit(100); +rewritten.QualifyTables("app"); +Check("Ast and rewrites", rewritten.Deparse().Value == "SELECT id FROM app.customers WHERE active AND tenant_id = $1 AND deleted_at IS NULL LIMIT 100"); +Check("ToDropStatement", Parser.Parse("CREATE TABLE s.t (a int)").Value!.ToDropStatement(ifExists: true).Deparse().Value == "DROP TABLE IF EXISTS s.t"); +var visitedTables = 0; +rewritten.Walk(new NodeVisitor().On((_, _) => visitedTables++), new NodeVisitor().On((_, _) => WalkAction.Stop)); +Check("NodeVisitor", visitedTables == 1); Check("ParameterRefs", Parser.ParameterRefs("SELECT $1, '$2', $3").Value?.Select(p => p.Number).SequenceEqual([1, 3]) == true); Check("ParsePlpgsqlFunctions", Parser.ParsePlpgsqlFunctions( "CREATE FUNCTION f() RETURNS int AS $$ BEGIN RETURN 1; END; $$ LANGUAGE plpgsql").Value?[0].Queries().SequenceEqual(["1"]) == true); diff --git a/src/PgSqlParser.Tests/AstTests.cs b/src/PgSqlParser.Tests/AstTests.cs new file mode 100644 index 0000000..4a4eb1f --- /dev/null +++ b/src/PgSqlParser.Tests/AstTests.cs @@ -0,0 +1,118 @@ +using Google.Protobuf; +using Shouldly; +using Xunit; + +namespace PgSqlParser.Tests; + +public class AstTests +{ + public static TheoryData Expressions => new() + { + { "name", Ast.Column("name") }, + { "t.name", Ast.Column("t", "name") }, + { "s.t.name", Ast.Column("s", "t", "name") }, + { "1", Ast.Const(1) }, + { "0", Ast.Const(0) }, + { "5000000000", Ast.Const(5000000000) }, + { "1.5", Ast.Const(1.5m) }, + { "2.25", Ast.Const(2.25) }, + { "'it''s'", Ast.Const("it's") }, + { "'😀 é'", Ast.Const("😀 é") }, + { "true", Ast.Const(true) }, + { "false", Ast.Const(false) }, + { "NULL", Ast.Null() }, + { "$1", Ast.Param(1) }, + { "a = 1", Ast.Eq(Ast.Column("a"), Ast.Const(1)) }, + { "a <> 1", Ast.NotEq(Ast.Column("a"), Ast.Const(1)) }, + { "a < 1", Ast.Lt(Ast.Column("a"), Ast.Const(1)) }, + { "a <= 1", Ast.LtEq(Ast.Column("a"), Ast.Const(1)) }, + { "a > 1", Ast.Gt(Ast.Column("a"), Ast.Const(1)) }, + { "a >= 1", Ast.GtEq(Ast.Column("a"), Ast.Const(1)) }, + { "a || b", Ast.Op("||", Ast.Column("a"), Ast.Column("b")) }, + { "a = 1 AND b = 2", Ast.And(Ast.Eq(Ast.Column("a"), Ast.Const(1)), Ast.Eq(Ast.Column("b"), Ast.Const(2))) }, + { "a OR b OR c", Ast.Or(Ast.Column("a"), Ast.Or(Ast.Column("b"), Ast.Column("c"))) }, + { "a AND b AND c", Ast.And(Ast.And(Ast.Column("a"), Ast.Column("b")), Ast.Column("c")) }, + { "a AND (b OR c)", Ast.And(Ast.Column("a"), Ast.Or(Ast.Column("b"), Ast.Column("c"))) }, + { "NOT a", Ast.Not(Ast.Column("a")) }, + { "a IS NULL", Ast.IsNull(Ast.Column("a")) }, + { "a IS NOT NULL", Ast.IsNotNull(Ast.Column("a")) }, + { "a IN (1, 2, 3)", Ast.In(Ast.Column("a"), Ast.Const(1), Ast.Const(2), Ast.Const(3)) }, + { "a NOT IN ('x')", Ast.NotIn(Ast.Column("a"), Ast.Const("x")) }, + { "lower(name)", Ast.Call("lower", Ast.Column("name")) }, + { "now()", Ast.Call("now") }, + { "pg_catalog.lower(a, 1)", Ast.Call(["pg_catalog", "lower"], Ast.Column("a"), Ast.Const(1)) }, + { "count(*)", Ast.CountStar() }, + { "$1::int", Ast.Cast(Ast.Param(1), "int") }, + { "a::numeric(10, 2)", Ast.Cast(Ast.Column("a"), "numeric(10,2)") }, + { "tenant_id = $1 AND deleted_at IS NULL", Ast.Expression("tenant_id = $1 and deleted_at is null") } + }; + + [Theory] + [MemberData(nameof(Expressions))] + public void ExpressionBuildersMatchTheParser(string sql, IMessage built) + { + // The same tree the parser gives, and it deparses to the same text. + built.AsNode().EqualsIgnoringLocations(Parser.ParseExpression(sql).GetValueOrThrow()).ShouldBeTrue(sql); + built.Deparse().GetValueOrThrow().ShouldBe(sql); + } + + [Fact] + public void TablesTargetsAndOrderBy() + { + var parsed = Parser.Parse("SELECT a, lower(b) AS lb, * FROM public.users u ORDER BY a, b DESC").GetValueOrThrow() + .Stmts[0].Stmt.SelectStmt; + + Ast.Table("public", "users", "u").EqualsIgnoringLocations(parsed.FromClause[0].RangeVar).ShouldBeTrue(); + Ast.Table("users").Deparse().GetValueOrThrow().ShouldBe("users"); + Ast.Table("My Schema", "order").Deparse().GetValueOrThrow().ShouldBe("\"My Schema\".\"order\""); + + Ast.Target(Ast.Column("a")).EqualsIgnoringLocations(parsed.TargetList[0].ResTarget).ShouldBeTrue(); + Ast.Target(Ast.Call("lower", Ast.Column("b")), "lb").EqualsIgnoringLocations(parsed.TargetList[1].ResTarget).ShouldBeTrue(); + Ast.Target(Ast.Star()).EqualsIgnoringLocations(parsed.TargetList[2].ResTarget).ShouldBeTrue(); + Ast.Star("t").Deparse().GetValueOrThrow().ShouldBe("t.*"); + + Ast.OrderBy(Ast.Column("a")).EqualsIgnoringLocations(parsed.SortClause[0].SortBy).ShouldBeTrue(); + Ast.OrderBy(Ast.Column("b"), descending: true).EqualsIgnoringLocations(parsed.SortClause[1].SortBy).ShouldBeTrue(); + } + + [Fact] + public void SelectMatchesTheParser() + { + var built = Ast.Select( + [Ast.Column("id"), Ast.Target(Ast.Call("lower", Ast.Column("name")), "n")], + Ast.Table("public", "users", "u"), + Ast.And(Ast.Eq(Ast.Column("u", "tenant"), Ast.Param(1)), Ast.IsNull(Ast.Column("deleted_at")))); + + const string sql = "SELECT id, lower(name) AS n FROM public.users u WHERE u.tenant = $1 AND deleted_at IS NULL"; + built.Deparse().GetValueOrThrow().ShouldBe(sql); + built.EqualsIgnoringLocations(Parser.Parse(sql).GetValueOrThrow().Stmts[0].Stmt.SelectStmt).ShouldBeTrue(); + + Ast.Select([Ast.Const(1)]).Deparse().GetValueOrThrow().ShouldBe("SELECT 1"); + } + + [Fact] + public void BuiltNodesCanGoIntoAParsedTree() + { + var tree = Parser.Parse("SELECT id FROM users").GetValueOrThrow(); + var select = tree.Stmts[0].Stmt.SelectStmt; + + select.WhereClause = Ast.Eq(Ast.Column("tenant_id"), Ast.Param(1)).AsNode(); + select.TargetList.Add(Ast.Target(Ast.CountStar(), "n").AsNode()); + select.SortClause.Add(Ast.OrderBy(Ast.Column("id"), descending: true).AsNode()); + + tree.Deparse().GetValueOrThrow().ShouldBe("SELECT id, count(*) AS n FROM users WHERE tenant_id = $1 ORDER BY id DESC"); + } + + [Fact] + public void RejectsInvalidArguments() + { + Should.Throw(() => Ast.Column()); + Should.Throw(() => Ast.And(Ast.Column("a"))); + Should.Throw(() => Ast.In(Ast.Column("a"))); + Should.Throw(() => Ast.Call([])); + Should.Throw(() => Ast.Const((string)null!)); + Should.Throw(() => Ast.Eq(null!, Ast.Const(1))); + Should.Throw(() => Ast.Expression("a, b")); + Should.Throw(() => Ast.Cast(Ast.Column("a"), "not a type!")); + } +} diff --git a/src/PgSqlParser.Tests/NodeToolsTests.cs b/src/PgSqlParser.Tests/NodeToolsTests.cs index 54395ad..8aedfb0 100644 --- a/src/PgSqlParser.Tests/NodeToolsTests.cs +++ b/src/PgSqlParser.Tests/NodeToolsTests.cs @@ -162,7 +162,7 @@ public void WalkWithAVisitorMatchesTheEnumerableOrder() }); seen.ShouldBe(tree.Walk().ToList()); - Should.Throw(() => tree.Walk(null!)); + Should.Throw(() => tree.Walk((Func)null!)); Should.Throw(() => ((IMessage)null!).Walk(_ => WalkAction.Continue)); } diff --git a/src/PgSqlParser.Tests/RewriteAndVisitorTests.cs b/src/PgSqlParser.Tests/RewriteAndVisitorTests.cs new file mode 100644 index 0000000..c81c77f --- /dev/null +++ b/src/PgSqlParser.Tests/RewriteAndVisitorTests.cs @@ -0,0 +1,251 @@ +using Google.Protobuf; +using Shouldly; +using Xunit; + +namespace PgSqlParser.Tests; + +public class RewriteAndVisitorTests +{ + private static ParseResult Parse(string query) => Parser.Parse(query).GetValueOrThrow(); + + private static string Sql(IMessage node) => node.Deparse().GetValueOrThrow(); + + [Theory] + [InlineData("SELECT * FROM t", "SELECT * FROM t WHERE tenant_id = $1")] + [InlineData("SELECT * FROM t WHERE a = 1", "SELECT * FROM t WHERE a = 1 AND tenant_id = $1")] + [InlineData("SELECT * FROM t WHERE a = 1 AND b = 2", "SELECT * FROM t WHERE a = 1 AND b = 2 AND tenant_id = $1")] + [InlineData("SELECT * FROM t WHERE a = 1 OR b = 2", "SELECT * FROM t WHERE (a = 1 OR b = 2) AND tenant_id = $1")] + [InlineData("UPDATE t SET a = 1", "UPDATE t SET a = 1 WHERE tenant_id = $1")] + [InlineData("DELETE FROM t WHERE a = 1", "DELETE FROM t WHERE a = 1 AND tenant_id = $1")] + [InlineData("SELECT * FROM t ORDER BY a LIMIT 5", "SELECT * FROM t WHERE tenant_id = $1 ORDER BY a LIMIT 5")] + public void AddWhere(string before, string after) + { + var tree = Parse(before); + + tree.AddWhere(Ast.Eq(Ast.Column("tenant_id"), Ast.Param(1))); + + Sql(tree).ShouldBe(after); + // The result is a tree the parser itself would build. + tree.EqualsIgnoringLocations(Parse(after)).ShouldBeTrue(); + } + + [Fact] + public void AddWhereWorksOnNodesAndRejectsWhatHasNoWhere() + { + var tree = Parse("SELECT * FROM t WHERE id IN (SELECT id FROM u)"); + var inner = tree.Descendants().Last(); + + inner.AddWhere(Ast.Expression("active")); + tree.Stmts[0].AddWhere(Ast.Expression("x = 1")); + + Sql(tree).ShouldBe("SELECT * FROM t WHERE id IN (SELECT id FROM u WHERE active) AND x = 1"); + + Should.Throw(() => Parse("SELECT 1 UNION SELECT 2").AddWhere(Ast.Expression("true"))) + .Message.ShouldBe("A WHERE condition cannot be added to a UNION, INTERSECT or EXCEPT."); + Should.Throw(() => Parse("INSERT INTO t VALUES (1)").AddWhere(Ast.Expression("true"))) + .Message.ShouldBe("A WHERE condition cannot be added to a InsertStmt."); + Should.Throw(() => Parse("VALUES (1)").AddWhere(Ast.Expression("true"))); + Should.Throw(() => Parse("SELECT 1; SELECT 2").AddWhere(Ast.Expression("true"))); + } + + [Theory] + [InlineData("SELECT * FROM t", "SELECT * FROM t LIMIT 100")] + [InlineData("SELECT * FROM t LIMIT 10", "SELECT * FROM t LIMIT 10")] + [InlineData("SELECT * FROM t LIMIT 100", "SELECT * FROM t LIMIT 100")] + [InlineData("SELECT * FROM t LIMIT 5000", "SELECT * FROM t LIMIT 100")] + [InlineData("SELECT * FROM t LIMIT ALL", "SELECT * FROM t LIMIT 100")] + [InlineData("SELECT * FROM t LIMIT $1", "SELECT * FROM t LIMIT LEAST($1, 100)")] + [InlineData("SELECT * FROM t ORDER BY a LIMIT 500 OFFSET 20", "SELECT * FROM t ORDER BY a LIMIT 100 OFFSET 20")] + [InlineData("SELECT 1 UNION SELECT 2", "SELECT 1 UNION SELECT 2 LIMIT 100")] + public void CapLimit(string before, string after) + { + var tree = Parse(before); + + tree.CapLimit(100); + + Sql(tree).ShouldBe(after); + Parse(Sql(tree)).EqualsIgnoringLocations(Parse(after)).ShouldBeTrue(); + } + + [Fact] + public void SetLimitOffsetAndCount() + { + var tree = Parse("SELECT id FROM t WHERE a = 1 ORDER BY id LIMIT 3"); + + tree.SetLimit(50); + tree.SetOffset(100); + Sql(tree).ShouldBe("SELECT id FROM t WHERE a = 1 ORDER BY id LIMIT 50 OFFSET 100"); + + var count = tree.ToCount(); + Sql(count).ShouldBe("SELECT count(*) FROM (SELECT id FROM t WHERE a = 1 ORDER BY id LIMIT 50 OFFSET 100) q"); + Sql(Parse("SELECT 1").ToCount("rows")).ShouldBe("SELECT count(*) FROM (SELECT 1) rows"); + // The original is not changed by ToCount. + Sql(tree).ShouldBe("SELECT id FROM t WHERE a = 1 ORDER BY id LIMIT 50 OFFSET 100"); + + Should.Throw(() => Parse("DELETE FROM t").SetLimit(1)) + .Message.ShouldBe("A LIMIT applies to a SELECT, not to a DeleteStmt."); + } + + [Fact] + public void QualifyAndRenameTables() + { + var tree = Parse("WITH r AS (SELECT * FROM orders) SELECT * FROM customers c JOIN r ON true JOIN other.things t ON true"); + + tree.QualifyTables("app").ShouldBe(2); + Sql(tree).ShouldBe("WITH r AS (SELECT * FROM app.orders) SELECT * FROM app.customers c JOIN r ON true JOIN other.things t ON true"); + + tree.RenameSchema("app", "tenant_1").ShouldBe(2); + tree.RenameTable("things", "stuff", schema: "other").ShouldBe(1); + tree.RenameTable("r", "renamed").ShouldBe(0, "r is a CTE here, not a table"); + Sql(tree).ShouldBe("WITH r AS (SELECT * FROM tenant_1.orders) SELECT * FROM tenant_1.customers c JOIN r ON true JOIN other.stuff t ON true"); + + var dml = Parse("UPDATE accounts SET a = 1; INSERT INTO accounts SELECT * FROM staging"); + dml.RenameTable("accounts", "accounts_v2").ShouldBe(2); + Sql(dml).ShouldBe("UPDATE accounts_v2 SET a = 1; INSERT INTO accounts_v2 SELECT * FROM staging"); + } + + [Theory] + [InlineData("CREATE TABLE s.t (a int)", "DROP TABLE s.t")] + [InlineData("CREATE TABLE \"My Table\" (a int)", "DROP TABLE \"My Table\"")] + [InlineData("CREATE UNIQUE INDEX idx ON s.t (a)", "DROP INDEX s.idx")] + [InlineData("CREATE VIEW v AS SELECT 1", "DROP VIEW v")] + [InlineData("CREATE MATERIALIZED VIEW mv AS SELECT 1", "DROP MATERIALIZED VIEW mv")] + [InlineData("CREATE TABLE copy AS SELECT 1", "DROP TABLE copy")] + [InlineData("CREATE SEQUENCE s.seq", "DROP SEQUENCE s.seq")] + [InlineData("CREATE SCHEMA app", "DROP SCHEMA app")] + [InlineData("CREATE EXTENSION hstore", "DROP EXTENSION hstore")] + [InlineData("CREATE TYPE mood AS ENUM ('a', 'b')", "DROP TYPE mood")] + [InlineData("CREATE TYPE pair AS (a int, b int)", "DROP TYPE pair")] + [InlineData("CREATE DOMAIN positive AS int CHECK (VALUE > 0)", "DROP DOMAIN positive")] + [InlineData("CREATE FUNCTION s.f(a int, b text, OUT c int) RETURNS int AS 'select 1' LANGUAGE sql", "DROP FUNCTION s.f(int, text)")] + [InlineData("CREATE PROCEDURE p(VARIADIC a int[]) AS 'select 1' LANGUAGE sql", "DROP PROCEDURE p(VARIADIC int[])")] + [InlineData("CREATE TRIGGER trg BEFORE INSERT ON s.t FOR EACH ROW EXECUTE FUNCTION f()", "DROP TRIGGER trg ON s.t")] + [InlineData("CREATE POLICY p ON t USING (true)", "DROP POLICY p ON t")] + public void ToDropStatement(string create, string drop) + { + var tree = Parse(create); + + Sql(tree.ToDropStatement()).ShouldBe(drop); + // IF EXISTS goes after the object kind, which is two words for a materialized view. + var kindEnd = drop.StartsWith("DROP MATERIALIZED VIEW ", StringComparison.Ordinal) ? 23 : drop.IndexOf(' ', 5) + 1; + Sql(tree.ToDropStatement(ifExists: true, cascade: true)).ShouldBe(drop.Insert(kindEnd, "IF EXISTS ") + " CASCADE"); + } + + [Fact] + public void ToDropStatementRejectsWhatItCannotUndo() + { + Should.Throw(() => Parse("SELECT 1").ToDropStatement()) + .Message.ShouldBe("There is no DROP statement for a SelectStmt."); + // An index without a name gets one from PostgreSQL, which the statement does not say. + Should.Throw(() => Parse("CREATE INDEX ON t (a)").ToDropStatement()); + } + + [Theory] + [InlineData("CREATE TABLE t (a int)", "CREATE TABLE IF NOT EXISTS t (a int)")] + [InlineData("CREATE INDEX i ON t (a)", "CREATE INDEX IF NOT EXISTS i ON t USING btree (a)")] + [InlineData("CREATE SEQUENCE s", "CREATE SEQUENCE IF NOT EXISTS s")] + [InlineData("CREATE SCHEMA app", "CREATE SCHEMA IF NOT EXISTS app")] + [InlineData("CREATE EXTENSION hstore", "CREATE EXTENSION IF NOT EXISTS hstore")] + public void EnsureIfNotExists(string before, string after) + { + var tree = Parse(before); + + tree.EnsureIfNotExists().ShouldBeTrue(); + + Sql(tree).ShouldBe(after); + } + + [Fact] + public void EnsureOrReplaceAndUnsupportedStatements() + { + var view = Parse("CREATE VIEW v AS SELECT 1"); + view.EnsureOrReplace().ShouldBeTrue(); + Sql(view).ShouldBe("CREATE OR REPLACE VIEW v AS SELECT 1"); + + var function = Parse("CREATE FUNCTION f() RETURNS int AS 'select 1' LANGUAGE sql"); + function.EnsureOrReplace().ShouldBeTrue(); + Sql(function).ShouldStartWith("CREATE OR REPLACE FUNCTION f()"); + + Parse("CREATE TABLE t (a int)").EnsureOrReplace().ShouldBeFalse(); + Parse("CREATE VIEW v AS SELECT 1").EnsureIfNotExists().ShouldBeFalse(); + Parse("CREATE INDEX ON t (a)").EnsureIfNotExists().ShouldBeFalse("IF NOT EXISTS needs an index name"); + } + + [Fact] + public void TypedVisitorDispatchesByNodeType() + { + var tree = Parse("SELECT lower(a) FROM t1 JOIN t2 ON true WHERE b IN (SELECT c FROM t3)"); + var tables = new List(); + var functions = new List(); + + tree.Walk(new NodeVisitor() + .On((table, _) => tables.Add(table.Relname)) + .On((call, visit) => functions.Add($"{call.Funcname[0].String.Sval} in {visit.FieldName}"))); + + tables.ShouldBe(["t1", "t2", "t3"]); + functions.ShouldBe(["lower in Val"]); + } + + [Fact] + public void TypedVisitorCanSkipAndStop() + { + var tree = Parse("SELECT a FROM t1 WHERE b IN (SELECT c FROM t2) AND d IN (SELECT e FROM t3)"); + + var outside = new List(); + tree.Walk(new NodeVisitor() + .On((_, _) => WalkAction.SkipChildren) + .On((table, _) => outside.Add(table.Relname))); + outside.ShouldBe(["t1"]); + + var firstTwo = new List(); + tree.Walk(new NodeVisitor().On((table, _) => + { + firstTwo.Add(table.Relname); + return firstTwo.Count == 2 ? WalkAction.Stop : WalkAction.Continue; + })); + firstTwo.ShouldBe(["t1", "t2"]); + } + + [Fact] + public void SeveralVisitorsShareOneWalkIndependently() + { + var tree = Parse("SELECT a FROM t1 WHERE b IN (SELECT c FROM t2) AND d IN (SELECT e FROM t3)"); + var all = new List(); + var outside = new List(); + var first = new List(); + var subqueries = 0; + + tree.Walk( + new NodeVisitor().On((table, _) => all.Add(table.Relname)), + new NodeVisitor() + .On((_, _) => WalkAction.SkipChildren) + .On((table, _) => outside.Add(table.Relname)), + new NodeVisitor().On((table, _) => + { + first.Add(table.Relname); + return WalkAction.Stop; + }), + new NodeVisitor().On((_, _) => subqueries++)); + + // One visitor skipping or stopping does not affect the others. + all.ShouldBe(["t1", "t2", "t3"]); + outside.ShouldBe(["t1"]); + first.ShouldBe(["t1"]); + subqueries.ShouldBe(2); + } + + [Fact] + public void SeveralHandlersForOneTypeAndNoVisitors() + { + var tree = Parse("SELECT a FROM t1, t2"); + var log = new List(); + + tree.Walk(new NodeVisitor() + .On((table, _) => log.Add($"first {table.Relname}")) + .On((table, _) => log.Add($"second {table.Relname}"))); + + log.ShouldBe(["first t1", "second t1", "first t2", "second t2"]); + tree.Walk(Array.Empty()); + Should.Throw(() => new NodeVisitor().On((Action)null!)); + } +} diff --git a/src/PgSqlParser.Tests/TreeEditingTests.cs b/src/PgSqlParser.Tests/TreeEditingTests.cs index 7a701ae..a72f71a 100644 --- a/src/PgSqlParser.Tests/TreeEditingTests.cs +++ b/src/PgSqlParser.Tests/TreeEditingTests.cs @@ -16,6 +16,8 @@ public class TreeEditingTests [InlineData("SELECT * FROM t WHERE a <> 1", "SELECT * FROM t WHERE a != 1")] [InlineData("CREATE TABLE t (a varchar(10), b int)", "CREATE TABLE t (a character varying(10), b integer)")] [InlineData("SELECT 1; SELECT 2", "SELECT 1;\n\n SELECT 2;")] + [InlineData("SELECT 1 WHERE a IN (1, 2)", "SELECT 1 WHERE a IN ( 1 ,2 )")] + [InlineData("SELECT ARRAY[1, 2], ROW(1, 2)", "SELECT ARRAY[ 1,2 ], ROW( 1,2 )")] public void EqualsIgnoringLocationsForTheSameSql(string left, string right) { var a = Parse(left); diff --git a/src/PgSqlParser/Ast.cs b/src/PgSqlParser/Ast.cs new file mode 100644 index 0000000..d3cf256 --- /dev/null +++ b/src/PgSqlParser/Ast.cs @@ -0,0 +1,234 @@ +using System.Globalization; +using Google.Protobuf; + +namespace PgSqlParser; + +/// +/// Builds parse tree nodes. Each method gives the same node the parser would for the matching SQL, so +/// the result can be put into a tree, compared with parsed SQL and deparsed. +/// Arguments that take a node accept any node type or a wrapper. +/// +public static class Ast +{ + /// Parses an expression into a node: Ast.Expression("a > 1 AND b IS NULL"). + /// The text is not a single expression. + public static Node Expression(string sql) => Parser.ParseExpression(sql).GetValueOrThrow(); + + /// A column reference: Ast.Column("name") or Ast.Column("t", "name"). + public static ColumnRef Column(params string[] parts) + { + ArgumentNullException.ThrowIfNull(parts); + if (parts.Length == 0) + throw new ArgumentException("A column reference needs at least a name.", nameof(parts)); + + var column = new ColumnRef(); + column.Fields.Add(parts.Select(Name)); + return column; + } + + /// *, or table.* when a qualifier is given. + public static ColumnRef Star(params string[] qualifier) + { + ArgumentNullException.ThrowIfNull(qualifier); + + var column = new ColumnRef(); + column.Fields.Add(qualifier.Select(Name)); + column.Fields.Add(new Node { AStar = new A_Star() }); + return column; + } + + public static A_Const Const(int value) => new() { Ival = new Integer { Ival = value } }; + + /// An integer constant. Values outside the 32-bit range are held the way the parser holds them. + public static A_Const Const(long value) + { + return value is >= int.MinValue and <= int.MaxValue + ? Const((int)value) + : new A_Const { Fval = new Float { Fval = value.ToString(CultureInfo.InvariantCulture) } }; + } + + public static A_Const Const(decimal value) => new() { Fval = new Float { Fval = value.ToString(CultureInfo.InvariantCulture) } }; + + public static A_Const Const(double value) => new() { Fval = new Float { Fval = value.ToString("R", CultureInfo.InvariantCulture) } }; + + /// A string constant. The value is taken as it is; quoting happens when the tree is deparsed. + public static A_Const Const(string value) + { + ArgumentNullException.ThrowIfNull(value); + return new A_Const { Sval = new String { Sval = value } }; + } + + public static A_Const Const(bool value) => new() { Boolval = new Boolean { Boolval = value } }; + + public static A_Const Null() => new() { Isnull = true }; + + /// A parameter reference: Ast.Param(1) is $1. + public static ParamRef Param(int number) => new() { Number = number }; + + /// A binary operator: Ast.Op("+", left, right). + public static A_Expr Op(string op, IMessage left, IMessage right) + { + ArgumentNullException.ThrowIfNull(op); + + var expression = new A_Expr { Kind = A_Expr_Kind.AexprOp, Lexpr = Wrap(left), Rexpr = Wrap(right) }; + expression.Name.Add(Name(op)); + return expression; + } + + public static A_Expr Eq(IMessage left, IMessage right) => Op("=", left, right); + public static A_Expr NotEq(IMessage left, IMessage right) => Op("<>", left, right); + public static A_Expr Lt(IMessage left, IMessage right) => Op("<", left, right); + public static A_Expr LtEq(IMessage left, IMessage right) => Op("<=", left, right); + public static A_Expr Gt(IMessage left, IMessage right) => Op(">", left, right); + public static A_Expr GtEq(IMessage left, IMessage right) => Op(">=", left, right); + + /// a AND b AND .... An argument that is itself an AND is merged in, as the parser does. + public static BoolExpr And(params IMessage[] conditions) => Bool(BoolExprType.AndExpr, conditions); + + /// a OR b OR .... An argument that is itself an OR is merged in, as the parser does. + public static BoolExpr Or(params IMessage[] conditions) => Bool(BoolExprType.OrExpr, conditions); + + public static BoolExpr Not(IMessage condition) + { + var expression = new BoolExpr { Boolop = BoolExprType.NotExpr }; + expression.Args.Add(Wrap(condition)); + return expression; + } + + public static NullTest IsNull(IMessage expression) => + new() { Arg = Wrap(expression), Nulltesttype = NullTestType.IsNull }; + + public static NullTest IsNotNull(IMessage expression) => + new() { Arg = Wrap(expression), Nulltesttype = NullTestType.IsNotNull }; + + /// expression IN (value, ...). + public static A_Expr In(IMessage expression, params IMessage[] values) => InList("=", expression, values); + + /// expression NOT IN (value, ...). + public static A_Expr NotIn(IMessage expression, params IMessage[] values) => InList("<>", expression, values); + + /// A function call: Ast.Call("lower", Ast.Column("name")). + public static FuncCall Call(string name, params IMessage[] arguments) => Call([name], arguments); + + /// A function call with a qualified name: Ast.Call(["pg_catalog", "now"]). + public static FuncCall Call(string[] name, params IMessage[] arguments) + { + ArgumentNullException.ThrowIfNull(name); + ArgumentNullException.ThrowIfNull(arguments); + if (name.Length == 0) + throw new ArgumentException("A function call needs a name.", nameof(name)); + + var call = new FuncCall { Funcformat = CoercionForm.CoerceExplicitCall }; + call.Funcname.Add(name.Select(Name)); + call.Args.Add(arguments.Select(Wrap)); + return call; + } + + /// count(*). + public static FuncCall CountStar() + { + var call = Call("count"); + call.AggStar = true; + return call; + } + + /// A cast: Ast.Cast(Ast.Param(1), "int") is $1::int. + /// is not a type name. + public static TypeCast Cast(IMessage expression, string typeName) => + new() { Arg = Wrap(expression), TypeName = Parser.ParseTypeName(typeName).GetValueOrThrow() }; + + /// A table reference: Ast.Table("users"). + public static RangeVar Table(string name) => Table(null, name); + + /// A table reference with an optional schema and alias: Ast.Table("public", "users", "u"). + public static RangeVar Table(string? schema, string name, string? alias = null) + { + ArgumentNullException.ThrowIfNull(name); + + var table = new RangeVar { Relname = name, Schemaname = schema ?? string.Empty, Inh = true, Relpersistence = "p" }; + if (alias is not null) + table.Alias = new Alias { Aliasname = alias }; + + return table; + } + + /// An item of a select list, with an optional alias: expression AS alias. + public static ResTarget Target(IMessage expression, string? alias = null) => + new() { Val = Wrap(expression), Name = alias ?? string.Empty }; + + /// An item of an ORDER BY clause. + public static SortBy OrderBy(IMessage expression, bool descending = false) => new() + { + Node = Wrap(expression), + SortbyDir = descending ? SortByDir.SortbyDesc : SortByDir.SortbyDefault, + SortbyNulls = SortByNulls.Default + }; + + /// + /// A SELECT statement. Items of that are not already select list items + /// are wrapped as such. + /// + public static SelectStmt Select(IEnumerable targets, IMessage? from = null, IMessage? where = null) + { + ArgumentNullException.ThrowIfNull(targets); + + var select = new SelectStmt { LimitOption = LimitOption.Default, Op = SetOperation.SetopNone }; + foreach (var target in targets) + select.TargetList.Add(Unwrapped(target) is ResTarget item ? item.AsNode() : Target(target).AsNode()); + + if (from is not null) + select.FromClause.Add(Wrap(from)); + + if (where is not null) + select.WhereClause = Wrap(where); + + return select; + } + + private static Node Name(string value) + { + ArgumentNullException.ThrowIfNull(value); + return new Node { String = new String { Sval = value } }; + } + + private static Node Wrap(IMessage node) + { + ArgumentNullException.ThrowIfNull(node); + return node.AsNode(); + } + + private static IMessage? Unwrapped(IMessage node) => node is Node wrapper ? wrapper.Unwrap() : node; + + private static BoolExpr Bool(BoolExprType kind, IMessage[] conditions) + { + ArgumentNullException.ThrowIfNull(conditions); + if (conditions.Length < 2) + throw new ArgumentException("AND and OR need at least two conditions.", nameof(conditions)); + + var expression = new BoolExpr { Boolop = kind }; + foreach (var condition in conditions) + { + // a AND (b AND c) is one flat list of three in the parser's tree. + if (Unwrapped(condition) is BoolExpr nested && nested.Boolop == kind) + expression.Args.Add(nested.Args); + else + expression.Args.Add(Wrap(condition)); + } + + return expression; + } + + private static A_Expr InList(string op, IMessage expression, IMessage[] values) + { + ArgumentNullException.ThrowIfNull(values); + if (values.Length == 0) + throw new ArgumentException("IN needs at least one value.", nameof(values)); + + var list = new List(); + list.Items.Add(values.Select(Wrap)); + + var result = new A_Expr { Kind = A_Expr_Kind.AexprIn, Lexpr = Wrap(expression), Rexpr = new Node { List = list } }; + result.Name.Add(Name(op)); + return result; + } +} diff --git a/src/PgSqlParser/NodeVisitor.cs b/src/PgSqlParser/NodeVisitor.cs new file mode 100644 index 0000000..cdde804 --- /dev/null +++ b/src/PgSqlParser/NodeVisitor.cs @@ -0,0 +1,114 @@ +using Google.Protobuf; + +namespace PgSqlParser; + +/// +/// A set of handlers, each for one node type, to run over a parse tree with +/// . Several visitors can share one walk. +/// +public sealed class NodeVisitor +{ + private readonly Dictionary>> _handlers = []; + + /// Calls for every node of type . + public NodeVisitor On(Action handler) where T : class, IMessage + { + ArgumentNullException.ThrowIfNull(handler); + + return On((node, visit) => + { + handler(node, visit); + return WalkAction.Continue; + }); + } + + /// + /// Calls for every node of type . What it returns + /// applies to this visitor only: it can skip the node's children or stop, while other visitors in + /// the same walk carry on. + /// + public NodeVisitor On(Func handler) where T : class, IMessage + { + ArgumentNullException.ThrowIfNull(handler); + + if (!_handlers.TryGetValue(typeof(T), out var handlers)) + _handlers[typeof(T)] = handlers = []; + + handlers.Add((node, visit) => handler((T)node, visit)); + return this; + } + + // Runs this visitor's handlers for the node. With several handlers the strongest action wins. + internal WalkAction Visit(NodeVisit visit) + { + if (!_handlers.TryGetValue(visit.Node.GetType(), out var handlers)) + return WalkAction.Continue; + + var result = WalkAction.Continue; + foreach (var handler in handlers) + { + var action = handler(visit.Node, visit); + if (action > result) + result = action; + } + + return result; + } +} + +/// +/// Walks a parse tree with typed visitors. +/// +public static class NodeVisitorExtensions +{ + /// + /// Walks the nodes below once, in the order of + /// , and runs every visitor's handlers on each node. + /// A visitor that skips a subtree or stops does so for itself only. + /// + public static void Walk(this IMessage root, params NodeVisitor[] visitors) + { + ArgumentNullException.ThrowIfNull(root); + ArgumentNullException.ThrowIfNull(visitors); + + // Per visitor: the depth of the node whose children it is skipping, or -1; and whether it stopped. + var skipDepth = new int[visitors.Length]; + Array.Fill(skipDepth, -1); + var stopped = new bool[visitors.Length]; + var running = visitors.Length; + + root.Walk(visit => + { + var anyoneGoesDeeper = false; + for (var i = 0; i < visitors.Length; i++) + { + if (stopped[i]) + continue; + + // Still below the skipped node. + if (skipDepth[i] >= 0 && visit.Depth > skipDepth[i]) + continue; + + skipDepth[i] = -1; + switch (visitors[i].Visit(visit)) + { + case WalkAction.Stop: + stopped[i] = true; + running--; + break; + case WalkAction.SkipChildren: + skipDepth[i] = visit.Depth; + break; + default: + anyoneGoesDeeper = true; + break; + } + } + + if (running == 0) + return WalkAction.Stop; + + return anyoneGoesDeeper ? WalkAction.Continue : WalkAction.SkipChildren; + }); + } +} diff --git a/src/PgSqlParser/ParseTreeExtensions.cs b/src/PgSqlParser/ParseTreeExtensions.cs index d5d1e26..044e514 100644 --- a/src/PgSqlParser/ParseTreeExtensions.cs +++ b/src/PgSqlParser/ParseTreeExtensions.cs @@ -444,15 +444,18 @@ public static bool EqualsIgnoringLocations(this IMessage? left, IMessage? right) return true; } - // Positions in the query text: "location", "stmt_location", "arg_location" and so on, and the - // statement length that goes with stmt_location. + // Positions in the query text: "location", "stmt_location", "arg_location" and so on, the + // statement length that goes with stmt_location, and the start and end of a parenthesized list + // ("list_start", "rexpr_list_end"). private static bool IsLocationField(FieldDescriptor field) { return field.FieldType == FieldType.Int32 && !field.IsRepeated && (field.Name == LocationFieldName || field.Name.EndsWith("_location", StringComparison.Ordinal) - || field.Name == "stmt_len"); + || field.Name == "stmt_len" + || field.Name.EndsWith("list_start", StringComparison.Ordinal) + || field.Name.EndsWith("list_end", StringComparison.Ordinal)); } // An explicit stack instead of recursion, so a deeply nested tree cannot overflow the call stack. diff --git a/src/PgSqlParser/PgSqlParser.csproj b/src/PgSqlParser/PgSqlParser.csproj index 232d774..f285d3a 100644 --- a/src/PgSqlParser/PgSqlParser.csproj +++ b/src/PgSqlParser/PgSqlParser.csproj @@ -2,7 +2,7 @@ .NET wrapper on libpg_query to access PostgreSQL parser outside of the server environment - 2.2.0 + 2.1.0 Babu Annamalai net8.0;net9.0;net10.0 enable diff --git a/src/PgSqlParser/QueryRewrites.cs b/src/PgSqlParser/QueryRewrites.cs new file mode 100644 index 0000000..355c90c --- /dev/null +++ b/src/PgSqlParser/QueryRewrites.cs @@ -0,0 +1,277 @@ +using Google.Protobuf; + +namespace PgSqlParser; + +/// +/// Common changes to statements, made on the parse tree so they hold for any way the SQL was written. +/// They change the tree in place; deparse it to get the SQL. +/// +public static class QueryRewrites +{ + /// + /// Adds a condition to the WHERE clause of a SELECT, UPDATE or DELETE, combining it with AND if + /// there already is one. + /// + /// + /// The statement is of another kind, or is a UNION, INTERSECT or EXCEPT, which has no WHERE of its own. + /// + public static void AddWhere(this IMessage statement, IMessage condition) + { + ArgumentNullException.ThrowIfNull(statement); + ArgumentNullException.ThrowIfNull(condition); + + switch (Statement(statement)) + { + case SelectStmt { Op: SetOperation.SetopNone } select when select.ValuesLists.Count == 0: + select.WhereClause = Combine(select.WhereClause, condition); + break; + case UpdateStmt update: + update.WhereClause = Combine(update.WhereClause, condition); + break; + case DeleteStmt delete: + delete.WhereClause = Combine(delete.WhereClause, condition); + break; + case var other: + throw new NotSupportedException($"A WHERE condition cannot be added to {Describe(other)}."); + } + } + + /// Sets the LIMIT of a SELECT, replacing any it has. + public static void SetLimit(this IMessage statement, long limit) + { + var select = Select(statement, "A LIMIT"); + select.LimitCount = Ast.Const(limit).AsNode(); + select.LimitOption = LimitOption.Count; + } + + /// + /// Makes sure a SELECT returns at most rows: adds a LIMIT if it has none, + /// lowers a larger constant LIMIT, and wraps any other LIMIT expression in LEAST(..., max). + /// + public static void CapLimit(this IMessage statement, long max) + { + var select = Select(statement, "A LIMIT"); + switch (select.LimitCount?.Unwrap()) + { + // No LIMIT, or LIMIT ALL / LIMIT NULL. + case null or A_Const { Isnull: true }: + select.SetLimit(max); + break; + case A_Const { Ival: { } integer }: + if (integer.Ival > max) + select.SetLimit(max); + break; + case var expression: + var least = new MinMaxExpr { Op = MinMaxOp.IsLeast }; + least.Args.Add(expression.AsNode()); + least.Args.Add(Ast.Const(max).AsNode()); + select.LimitCount = least.AsNode(); + break; + } + } + + /// Sets the OFFSET of a SELECT, replacing any it has. + public static void SetOffset(this IMessage statement, long offset) + { + Select(statement, "An OFFSET").LimitOffset = Ast.Const(offset).AsNode(); + } + + /// + /// Returns a new statement that counts the rows of a SELECT: SELECT count(*) FROM (...) AS alias. + /// The original is used as it is, including its LIMIT, and is not modified. + /// + public static SelectStmt ToCount(this IMessage statement, string alias = "q") + { + var select = Select(statement, "A count"); + var subquery = new RangeSubselect { Subquery = select.AsNode(), Alias = new Alias { Aliasname = alias } }; + return Ast.Select([Ast.CountStar()], subquery); + } + + /// + /// Gives every table reference that has no schema the schema . Names that + /// refer to a CTE are left alone. Returns the number of references changed. + /// + public static int QualifyTables(this IMessage root, string schema) + { + ArgumentNullException.ThrowIfNull(schema); + + var changed = 0; + foreach (var table in root.GetReferences().Tables) + { + if (table.IsCte || table.Schema is not null) + continue; + + table.Node.Schemaname = schema; + changed++; + } + + return changed; + } + + /// + /// Changes every table reference in schema to schema . + /// Returns the number of references changed. Only table references are changed, not functions or types. + /// + public static int RenameSchema(this IMessage root, string from, string to) + { + ArgumentNullException.ThrowIfNull(from); + ArgumentNullException.ThrowIfNull(to); + + var changed = 0; + foreach (var table in root.Descendants()) + { + if (table.Schemaname != from) + continue; + + table.Schemaname = to; + changed++; + } + + return changed; + } + + /// + /// Changes every reference to table (with no schema, or in + /// when given) to . CTE names are left alone. Returns the number of references changed. + /// + public static int RenameTable(this IMessage root, string from, string to, string? schema = null) + { + ArgumentNullException.ThrowIfNull(from); + ArgumentNullException.ThrowIfNull(to); + + var changed = 0; + foreach (var table in root.GetReferences().Tables) + { + if (table.IsCte || table.Name != from || table.Schema != schema) + continue; + + table.Node.Relname = to; + changed++; + } + + return changed; + } + + /// + /// Adds IF NOT EXISTS to a CREATE statement that supports it (table, index, sequence, schema, + /// extension, CREATE TABLE AS, materialized view). Returns false for other statements. + /// + public static bool EnsureIfNotExists(this IMessage statement) + { + switch (Statement(statement)) + { + case CreateStmt create: create.IfNotExists = true; return true; + case IndexStmt { Idxname.Length: > 0 } index: index.IfNotExists = true; return true; + case CreateSeqStmt sequence: sequence.IfNotExists = true; return true; + case CreateSchemaStmt schema: schema.IfNotExists = true; return true; + case CreateExtensionStmt extension: extension.IfNotExists = true; return true; + case CreateTableAsStmt createAs: createAs.IfNotExists = true; return true; + default: return false; + } + } + + /// + /// Adds OR REPLACE to a CREATE statement that supports it (view, function, procedure, trigger, rule). + /// Returns false for other statements. + /// + public static bool EnsureOrReplace(this IMessage statement) + { + switch (Statement(statement)) + { + case ViewStmt view: view.Replace = true; return true; + case CreateFunctionStmt function: function.Replace = true; return true; + case CreateTrigStmt trigger: trigger.Replace = true; return true; + case RuleStmt rule: rule.Replace = true; return true; + default: return false; + } + } + + /// + /// Returns the DROP statement that undoes a CREATE statement, for tables, indexes, views, + /// materialized views, sequences, schemas, types, domains, functions, procedures, triggers, + /// policies and extensions. + /// + /// There is no DROP for this statement here. + public static IMessage ToDropStatement(this IMessage statement, bool ifExists = false, bool cascade = false) + { + var (kind, name) = Statement(statement) switch + { + CreateStmt create => ("TABLE", Name(create.Relation)), + IndexStmt { Idxname.Length: > 0 } index => + ("INDEX", PgIdentifier.Quote([.. Parts(index.Relation.Schemaname), index.Idxname])), + ViewStmt view => ("VIEW", Name(view.View)), + CreateTableAsStmt { Into.Rel: { } relation } createAs => + (createAs.Objtype == ObjectType.ObjectMatview ? "MATERIALIZED VIEW" : "TABLE", Name(relation)), + CreateSeqStmt sequence => ("SEQUENCE", Name(sequence.Sequence)), + CreateSchemaStmt { Schemaname.Length: > 0 } schema => ("SCHEMA", PgIdentifier.Quote(schema.Schemaname)), + CreateExtensionStmt extension => ("EXTENSION", PgIdentifier.Quote(extension.Extname)), + CreateEnumStmt type => ("TYPE", Name(type.TypeName)), + CreateRangeStmt type => ("TYPE", Name(type.TypeName)), + CompositeTypeStmt type => ("TYPE", Name(type.Typevar)), + CreateDomainStmt domain => ("DOMAIN", Name(domain.Domainname)), + CreateFunctionStmt function => + (function.IsProcedure ? "PROCEDURE" : "FUNCTION", $"{Name(function.Funcname)}({Signature(function)})"), + CreateTrigStmt trigger => ("TRIGGER", $"{PgIdentifier.Quote(trigger.Trigname)} ON {Name(trigger.Relation)}"), + CreatePolicyStmt policy => ("POLICY", $"{PgIdentifier.Quote(policy.PolicyName)} ON {Name(policy.Table)}"), + var other => throw new NotSupportedException($"There is no DROP statement for {Describe(other)}.") + }; + + // Written as SQL and parsed, so the tree is exactly what PostgreSQL's parser builds for it. + var sql = $"DROP {kind} {(ifExists ? "IF EXISTS " : "")}{name}{(cascade ? " CASCADE" : "")}"; + return Parser.Parse(sql).GetValueOrThrow().Stmts[0].Stmt.Unwrap()!; + } + + private static IMessage? Statement(IMessage statement) + { + ArgumentNullException.ThrowIfNull(statement); + + return statement switch + { + ParseResult { Stmts.Count: 1 } tree => tree.Stmts[0].Stmt?.Unwrap(), + ParseResult => throw new NotSupportedException("This needs a single statement, and the parse tree has several or none."), + RawStmt raw => raw.Stmt?.Unwrap(), + Node node => node.Unwrap(), + _ => statement + }; + } + + private static SelectStmt Select(IMessage statement, string what) + { + return Statement(statement) as SelectStmt + ?? throw new NotSupportedException($"{what} applies to a SELECT, not to {Describe(Statement(statement))}."); + } + + private static string Describe(IMessage? node) + { + return node switch + { + null => "an empty statement", + SelectStmt { Op: not SetOperation.SetopNone } => "a UNION, INTERSECT or EXCEPT", + SelectStmt { ValuesLists.Count: > 0 } => "a VALUES list", + _ => $"a {node.Descriptor.Name}" + }; + } + + private static Node Combine(Node? existing, IMessage condition) + { + return existing?.Unwrap() is { } current ? Ast.And(current, condition).AsNode() : condition.AsNode(); + } + + private static string[] Parts(string schema) => schema.Length > 0 ? [schema] : []; + + private static string Name(RangeVar relation) => PgIdentifier.Quote([.. Parts(relation.Schemaname), relation.Relname]); + + private static string Name(IEnumerable names) => PgIdentifier.Quote(names.Select(part => part.String.Sval).ToArray()); + + // The argument types that identify a function: everything except OUT and TABLE parameters. + private static string Signature(CreateFunctionStmt function) + { + var types = function.Parameters + .Select(parameter => parameter.FunctionParameter) + .Where(parameter => parameter.Mode is not (FunctionParameterMode.FuncParamOut or FunctionParameterMode.FuncParamTable)) + .Select(parameter => (parameter.Mode == FunctionParameterMode.FuncParamVariadic ? "VARIADIC " : "") + + parameter.ArgType.Deparse().GetValueOrThrow()); + + return string.Join(", ", types); + } +}