Treat a reserved word after AS as an alias

src/lexer/disambiguateTokens.ts

@@ -12,13 +12,17 @@ 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 after AS to IDENTIFIER when that word cannot start
+ * a clause there, leaving `CREATE TABLE t AS SELECT ...` alone.
+ *
  * 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]`.
  */
 export function disambiguateTokens(tokens: Token[]): Token[] {
   return tokens
     .map(propertyNameKeywordToIdent)
+    .map(keywordAliasAfterAs)
     .map(funcNameToIdent)
     .map(dataTypeToParameterizedDataType)
     .map(identToArrayIdent)
@@ -39,6 +43,46 @@ 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.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);

test/behavesLikeSqlFormatter.ts

@@ -277,4 +277,15 @@ export default function behavesLikeSqlFormatter(format: FormatFn) {
         tbl;
     `);
   });
+
+  // Issue #801
+  it('supports reserved word as alias after AS', () => {
+    const result = format('SELECT id AS set FROM tbl AS set;');
+    expect(result).toBe(dedent`
+      SELECT
+        id AS set
+      FROM
+        tbl AS set;
+    `);
+  });
 }

test/features/update.ts

@@ -39,6 +39,16 @@ export default function supportsUpdate(format: FormatFn, { whereCurrentOf }: Upd
     `);
   });
 
+  // Issue #801
+  it('keeps SET as a clause after a reserved word alias', () => {
+    const result = format('UPDATE tbl AS set SET x = 1;');
+    expect(result).toBe(dedent`
+      UPDATE tbl AS set
+      SET
+        x = 1;
+    `);
+  });
+
   if (whereCurrentOf) {
     it('formats UPDATE statement with cursor position', () => {
       const result = format("UPDATE Customers SET Name='John' WHERE CURRENT OF my_cursor;");