From bcd3a426da9603f21dfac49eb9f8e4e43ebda79a Mon Sep 17 00:00:00 2001
From: Christopher Pruijsen <christopher.pruijsen@gmail.com>
Date: Fri, 18 Sep 2026 15:15:24 +0100
Subject: [PATCH] Treat a reserved word after AS as an alias

Some dialects allow reserved words as aliases, as in
`SELECT id AS set FROM tbl`. The word is tokenized as a RESERVED_* token
and the parser then reads it as the start of a clause, so the formatter
breaks the line and reindents.

Directly after AS a reserved word can only be an alias name, so convert
it to IDENTIFIER there. Only the token types that cannot legitimately
follow AS are converted, leaving `CREATE TABLE t AS SELECT ...` and
`PREPARE foo AS UPDATE ...` working.

The bare alias form, as in `FROM pg_settings set`, needs statement
position tracking and is not covered here.

Refs #801
---
 src/lexer/disambiguateTokens.ts | 45 +++++++++++++++++++++++++++++++++
 test/behavesLikeSqlFormatter.ts | 11 ++++++++
 test/postgresql.test.ts         | 26 +++++++++++++++++++
 3 files changed, 82 insertions(+)

diff --git a/src/lexer/disambiguateTokens.ts b/src/lexer/disambiguateTokens.ts
index f97a72c705..db096c322b 100644
--- a/src/lexer/disambiguateTokens.ts
+++ b/src/lexer/disambiguateTokens.ts
@@ -12,6 +12,9 @@ import { isReserved, Token, TokenType } from './token.js';
  * When IDENTIFIER or RESERVED_DATA_TYPE token is followed by "["
  * converts it to ARRAY_IDENTIFIER or ARRAY_KEYWORD accordingly.
  *
+ * Converts a reserved word directly after AS to IDENTIFIER, as it can only
+ * be an alias name there.
+ *
  * This is needed to avoid ambiguity in parser which expects function names
  * to always be followed by open-paren, and to distinguish between
  * array accessor `foo[1]` and array literal `[1, 2, 3]`.
@@ -19,6 +22,7 @@ import { isReserved, Token, TokenType } from './token.js';
 export function disambiguateTokens(tokens: Token[]): Token[] {
   return tokens
     .map(propertyNameKeywordToIdent)
+    .map(keywordAliasAfterAs)
     .map(funcNameToIdent)
     .map(dataTypeToParameterizedDataType)
     .map(identToArrayIdent)
@@ -39,6 +43,47 @@ const propertyNameKeywordToIdent = (token: Token, i: number, tokens: Token[]): T
   return token;
 };
 
+/**
+ * Some dialects allow reserved words as aliases, as in `SELECT id AS set FROM tbl`.
+ * Such a word is tokenized as a RESERVED_* token, which the parser then treats as
+ * the start of a clause. Directly after AS it can only be an alias name, so we
+ * convert it to IDENTIFIER.
+ *
+ * Only the token types that cannot legitimately follow AS are converted, leaving
+ * `CREATE TABLE t AS SELECT ...` and `PREPARE foo AS UPDATE ...` working.
+ */
+const keywordAliasAfterAs = (token: Token, i: number, tokens: Token[]): Token => {
+  if (canBeAliasAfterAs(token)) {
+    const prevToken = prevNonCommentToken(tokens, i);
+    if (prevToken && isAsKeyword(prevToken)) {
+      return { ...token, type: TokenType.IDENTIFIER, text: token.raw };
+    }
+  }
+  return token;
+};
+
+const isAsKeyword = (token: Token): boolean =>
+  (token.type === TokenType.RESERVED_KEYWORD || token.type === TokenType.RESERVED_KEYWORD_PHRASE) &&
+  token.text === 'AS';
+
+const canBeAliasAfterAs = (token: Token): boolean =>
+  token.type === TokenType.RESERVED_SET_OPERATION ||
+  token.type === TokenType.RESERVED_JOIN ||
+  token.type === TokenType.LIMIT ||
+  token.type === TokenType.BETWEEN ||
+  token.type === TokenType.CASE ||
+  token.type === TokenType.END ||
+  token.type === TokenType.WHEN ||
+  token.type === TokenType.ELSE ||
+  token.type === TokenType.THEN ||
+  token.type === TokenType.AND ||
+  token.type === TokenType.OR ||
+  token.type === TokenType.XOR ||
+  // SET is the clause keyword used as an alias in #801. The other
+  // RESERVED_CLAUSE words can follow AS for real (SELECT, VALUES, WITH,
+  // INSERT, UPDATE, DELETE, EXECUTE, ...), so they stay keywords.
+  (token.type === TokenType.RESERVED_CLAUSE && token.text === 'SET');
+
 const funcNameToIdent = (token: Token, i: number, tokens: Token[]): Token => {
   if (token.type === TokenType.RESERVED_FUNCTION_NAME) {
     const nextToken = nextNonCommentToken(tokens, i);
diff --git a/test/behavesLikeSqlFormatter.ts b/test/behavesLikeSqlFormatter.ts
index 1b9e636ba6..cfde304b19 100644
--- a/test/behavesLikeSqlFormatter.ts
+++ b/test/behavesLikeSqlFormatter.ts
@@ -277,4 +277,15 @@ export default function behavesLikeSqlFormatter(format: FormatFn) {
         tbl;
     `);
   });
+
+  // Issue #801
+  it('supports reserved word as column alias after AS', () => {
+    const result = format('SELECT id AS set FROM tbl;');
+    expect(result).toBe(dedent`
+      SELECT
+        id AS set
+      FROM
+        tbl;
+    `);
+  });
 }
diff --git a/test/postgresql.test.ts b/test/postgresql.test.ts
index e8e99adfea..24c3f1def4 100644
--- a/test/postgresql.test.ts
+++ b/test/postgresql.test.ts
@@ -324,4 +324,30 @@ describe('PostgreSqlFormatter', () => {
       EXECUTE FUNCTION example_function ();
     `);
   });
+
+  // Issue #801
+  it('supports reserved word as table alias after AS', () => {
+    expect(format(`SELECT set.foo FROM settings AS set;`)).toBe(dedent`
+      SELECT
+        set.foo
+      FROM
+        settings AS set;
+    `);
+    expect(format(`SELECT * FROM pg_settings AS set WHERE set.name = $9;`)).toBe(dedent`
+      SELECT
+        *
+      FROM
+        pg_settings AS set
+      WHERE
+        set.name = $9;
+    `);
+  });
+
+  it('keeps SET as a clause after a reserved-word alias', () => {
+    expect(format(`UPDATE tbl AS set SET x = 1;`)).toBe(dedent`
+      UPDATE tbl AS set
+      SET
+        x = 1;
+    `);
+  });
 });