Repository: sql-formatter-org/sql-formatter
Pull request #260 (2022-06-28 13:56:16)

8.0: Detect array subscript operator in Parser
That is, detecting syntax like `someArray[5]`.


Patch:
diff --git a/src/core/ExpressionFormatter.ts b/src/core/ExpressionFormatter.ts
index e52761de..81dcf668 100644
--- a/src/core/ExpressionFormatter.ts
+++ b/src/core/ExpressionFormatter.ts
@@ -7,6 +7,7 @@ import Params from './Params';
 import { isReserved, type Token, TokenType } from './token';
 import {
   AllColumnsAsterisk,
+  ArraySubscript,
   AstNode,
   BetweenPredicate,
   FunctionCall,
@@ -46,6 +47,9 @@ export default class ExpressionFormatter {
         case 'function_call':
           this.formatFunctionCall(node);
           break;
+        case 'array_subscript':
+          this.formatArraySubscript(node);
+          break;
         case 'parenthesis':
           this.formatParenthesis(node);
           break;
@@ -71,6 +75,11 @@ export default class ExpressionFormatter {
     this.formatParenthesis(node.parenthesis);
   }
 
+  private formatArraySubscript(node: ArraySubscript) {
+    this.query.add(this.show(node.arrayToken));
+    this.formatParenthesis(node.parenthesis);
+  }
+
   private formatParenthesis(node: Parenthesis) {
     const inline = this.inlineBlock.isInlineBlock(node);
 
diff --git a/src/core/InlineBlock.ts b/src/core/InlineBlock.ts
index 86c9735c..9b95bcab 100644
--- a/src/core/InlineBlock.ts
+++ b/src/core/InlineBlock.ts
@@ -24,6 +24,12 @@ export default class InlineBlock {
 
     for (const node of parenthesis.children) {
       switch (node.type) {
+        case 'function_call':
+          length += node.nameToken.value.length + this.inlineWidth(node.parenthesis);
+          break;
+        case 'array_subscript':
+          length += node.arrayToken.value.length + this.inlineWidth(node.parenthesis);
+          break;
         case 'parenthesis':
           length += this.inlineWidth(node);
           break;
@@ -32,6 +38,9 @@ export default class InlineBlock {
           break;
         case 'limit_clause':
           return Infinity;
+        case 'all_columns_asterisk':
+          length += 1;
+          break;
         case 'token':
           length += node.token.value.length;
           if (this.isForbiddenToken(node.token)) {
diff --git a/src/core/Parser.ts b/src/core/Parser.ts
index 409cee00..75546a29 100644
--- a/src/core/Parser.ts
+++ b/src/core/Parser.ts
@@ -1,6 +1,7 @@
 /* eslint-disable no-cond-assign */
 import {
   AllColumnsAsterisk,
+  ArraySubscript,
   AstNode,
   BetweenPredicate,
   FunctionCall,
@@ -50,6 +51,7 @@ export default class Parser {
   private expression(): AstNode {
     return (
       this.functionCall() ||
+      this.arraySubscript() ||
       this.parenthesis() ||
       this.betweenPredicate() ||
       this.limitClause() ||
@@ -73,6 +75,20 @@ export default class Parser {
     return undefined;
   }
 
+  private arraySubscript(): ArraySubscript | undefined {
+    if (
+      (this.look().type === TokenType.RESERVED_KEYWORD || this.look().type === TokenType.IDENT) &&
+      this.look(1).value === '['
+    ) {
+      return {
+        type: 'array_subscript',
+        arrayToken: this.next(),
+        parenthesis: this.parenthesis() as Parenthesis,
+      };
+    }
+    return undefined;
+  }
+
   private parenthesis(): Parenthesis | undefined {
     if (this.look().type === TokenType.OPEN_PAREN) {
       const children: AstNode[] = [];
diff --git a/src/core/ast.ts b/src/core/ast.ts
index 021b0a1b..6d060e87 100644
--- a/src/core/ast.ts
+++ b/src/core/ast.ts
@@ -17,6 +17,13 @@ export type FunctionCall = {
   parenthesis: Parenthesis;
 };
 
+// <ident>[<expr>]
+export type ArraySubscript = {
+  type: 'array_subscript';
+  arrayToken: Token;
+  parenthesis: Parenthesis;
+};
+
 export type Parenthesis = {
   type: 'parenthesis';
   children: AstNode[];
@@ -49,6 +56,7 @@ export type AllColumnsAsterisk = {
 
 export type AstNode =
   | FunctionCall
+  | ArraySubscript
   | Parenthesis
   | BetweenPredicate
   | LimitClause
diff --git a/src/core/token.ts b/src/core/token.ts
index e75b5cfa..d7a14298 100644
--- a/src/core/token.ts
+++ b/src/core/token.ts
@@ -46,6 +46,7 @@ export const testToken =
 export const isToken = {
   AS: testToken({ value: 'AS', type: TokenType.RESERVED_KEYWORD }),
   AND: testToken({ value: 'AND', type: TokenType.RESERVED_LOGICAL_OPERATOR }),
+  ARRAY: testToken({ value: 'ARRAY', type: TokenType.RESERVED_KEYWORD }),
   BETWEEN: testToken({ value: 'BETWEEN', type: TokenType.RESERVED_KEYWORD }),
   CASE: testToken({ value: 'CASE', type: TokenType.RESERVED_CASE_START }),
   CAST: testToken({ value: 'CAST', type: TokenType.RESERVED_KEYWORD }),
@@ -55,6 +56,7 @@ export const isToken = {
   LIMIT: testToken({ value: 'LIMIT', type: TokenType.RESERVED_COMMAND }),
   SELECT: testToken({ value: 'SELECT', type: TokenType.RESERVED_COMMAND }),
   SET: testToken({ value: 'SET', type: TokenType.RESERVED_COMMAND }),
+  STRUCT: testToken({ value: 'STRUCT', type: TokenType.RESERVED_KEYWORD }),
   TABLE: testToken({ value: 'TABLE', type: TokenType.RESERVED_KEYWORD }),
   WINDOW: testToken({ value: 'WINDOW', type: TokenType.RESERVED_COMMAND }),
   WITH: testToken({ value: 'WITH', type: TokenType.RESERVED_COMMAND }),
diff --git a/src/languages/bigquery.formatter.ts b/src/languages/bigquery.formatter.ts
index a6c76ecc..ca57b324 100644
--- a/src/languages/bigquery.formatter.ts
+++ b/src/languages/bigquery.formatter.ts
@@ -1,6 +1,6 @@
 import Formatter from 'src/core/Formatter';
 import Tokenizer from 'src/core/Tokenizer';
-import { EOF_TOKEN, TokenType, type Token } from 'src/core/token';
+import { EOF_TOKEN, isToken, TokenType, type Token } from 'src/core/token';
 import { dedupe } from 'src/utils';
 
 /**
@@ -884,7 +884,7 @@ function combineParameterizedTypes(tokens: Token[]) {
     const token = tokens[i];
     const nextToken = tokens[i + 1] || EOF_TOKEN;
 
-    if ((token.value === 'ARRAY' || token.value === 'STRUCT') && nextToken.value === '<') {
+    if ((isToken.ARRAY(token) || isToken.STRUCT(token)) && nextToken.value === '<') {
       const endIndex = findClosingAngleBracketIndex(tokens, i + 1);
       const typeDefTokens = tokens.slice(i, endIndex + 1);
       processed.push({


Test patch:
diff --git a/test/bigquery.test.ts b/test/bigquery.test.ts
index d896d101..669dbd6f 100644
--- a/test/bigquery.test.ts
+++ b/test/bigquery.test.ts
@@ -107,14 +107,17 @@ describe('BigQueryFormatter', () => {
     `);
   });
 
-  it('supports parametric ARRAY and STRUCT', () => {
-    const result = format('SELECT STRUCT<ARRAY<INT64>>([]), ARRAY<FLOAT>[1] FROM tbl');
-    expect(result).toBe(dedent`
+  it('supports parametric ARRAY', () => {
+    expect(format('SELECT ARRAY<FLOAT>[1]')).toBe(dedent`
       SELECT
-        STRUCT<ARRAY<INT64>>([]),
         ARRAY<FLOAT>[1]
-      FROM
-        tbl
+    `);
+  });
+
+  it('supports parametric STRUCT', () => {
+    expect(format('SELECT STRUCT<ARRAY<INT64>>([])')).toBe(dedent`
+      SELECT
+        STRUCT<ARRAY<INT64>>([])
     `);
   });
 
diff --git a/test/unit/Parser.test.ts b/test/unit/Parser.test.ts
index 9be3338e..e1f9589f 100644
--- a/test/unit/Parser.test.ts
+++ b/test/unit/Parser.test.ts
@@ -9,6 +9,8 @@ describe('Parser', () => {
       reservedBinaryCommands: ['UNION', 'JOIN'],
       reservedJoinConditions: ['ON', 'USING'],
       reservedKeywords: ['BETWEEN', 'LIKE', 'SQRT'],
+      openParens: ['(', '['],
+      closeParens: [')', ']'],
       stringTypes: ["''"],
       identTypes: ['""'],
     }).tokenize(sql);
@@ -109,6 +111,68 @@ describe('Parser', () => {
     `);
   });
 
+  it('parses array subscript', () => {
+    expect(parse('SELECT my_array[OFFSET(5)]')).toMatchInlineSnapshot(`
+      Array [
+        Object {
+          "children": Array [
+            Object {
+              "token": Object {
+                "text": "SELECT",
+                "type": "RESERVED_COMMAND",
+                "value": "SELECT",
+                "whitespaceBefore": "",
+              },
+              "type": "token",
+            },
+            Object {
+              "arrayToken": Object {
+                "text": "my_array",
+                "type": "IDENT",
+                "value": "my_array",
+                "whitespaceBefore": " ",
+              },
+              "parenthesis": Object {
+                "children": Array [
+                  Object {
+                    "nameToken": Object {
+                      "text": "OFFSET",
+                      "type": "IDENT",
+                      "value": "OFFSET",
+                      "whitespaceBefore": "",
+                    },
+                    "parenthesis": Object {
+                      "children": Array [
+                        Object {
+                          "token": Object {
+                            "text": "5",
+                            "type": "NUMBER",
+                            "value": "5",
+                            "whitespaceBefore": "",
+                          },
+                          "type": "token",
+                        },
+                      ],
+                      "closeParen": ")",
+                      "openParen": "(",
+                      "type": "parenthesis",
+                    },
+                    "type": "function_call",
+                  },
+                ],
+                "closeParen": "]",
+                "openParen": "[",
+                "type": "parenthesis",
+              },
+              "type": "array_subscript",
+            },
+          ],
+          "type": "statement",
+        },
+      ]
+    `);
+  });
+
   it('parses parenthesized expressions', () => {
     expect(parse('SELECT (birth_year - (CURRENT_DATE + 1))')).toMatchInlineSnapshot(`
       Array [