Merge pull request #2524 from taozhi8833998/feat-func-call-array-pg

feat: support func call in array index in pg

Author: taozhi8833998

ast/postgresql.ts

@@ -939,7 +939,7 @@ export type column_list_items = column_list_item[];
 
 export type column_clause = 'ALL' | '*' | column_list_item[] | column_list_items;
 
-export type array_index = { brackets: boolean, number: number };
+export type array_index = { brackets: boolean, index: literal_numeric | literal_string | func_call };
 
 export type array_index_list = array_index[];
 

pegjs/postgresql.pegjs

@@ -3425,8 +3425,8 @@ column_clause
   / column_list_items
 
 array_index
-  = LBRAKE __ n:(literal_numeric / literal_string) __ RBRAKE {
-    // => { brackets: boolean, number: number }
+  = LBRAKE __ n:(literal_numeric / literal_string / func_call) __ RBRAKE {
+    // => { brackets: boolean, index: literal_numeric | literal_string | func_call }
     return {
       brackets: true,
       index: n

src/column.js

@@ -29,7 +29,7 @@ function arrayIndexToSQL(arrayIndexList) {
   if (!arrayIndexList || arrayIndexList.length === 0) return ''
   const result = []
   for (const arrayIndex of arrayIndexList) {
-    let arrayIndexStr = arrayIndex.brackets ? `[${literalToSQL(arrayIndex.index)}]` : `${arrayIndex.notation}${literalToSQL(arrayIndex.index)}`
+    let arrayIndexStr = arrayIndex.brackets ? `[${exprToSQL(arrayIndex.index)}]` : `${arrayIndex.notation}${exprToSQL(arrayIndex.index)}`
     if (arrayIndex.property) arrayIndexStr = `${arrayIndexStr}.${literalToSQL(arrayIndex.property)}`
     result.push(arrayIndexStr)
   }

test/postgres.spec.js

@@ -1893,6 +1893,26 @@ describe('Postgres', () => {
         'CREATE TABLE "sal_emp" (name TEXT, pay_by_quarter INTEGER[], schedule TEXT[][])'
       ]
     },
+    {
+      title: 'bracket array index with func call',
+      sql: [
+        `SELECT
+          p.description AS description,
+          p.invoice_date AS invoice_date
+        FROM anonymized_schema.payments p
+        JOIN anonymized_schema.entities e
+          ON p.entity_id = e.entity_id
+        LEFT JOIN anonymized_schema.entities pe
+          ON p.parent_entity_id = pe.entity_id
+        LEFT JOIN anonymized_schema.entities ppe
+          ON pe.parent_ids[ array_upper(pe.parent_ids, 1) ] = ppe.entity_id
+        WHERE e.entity_type IN ('A', 'B')
+          AND p.invoice_date >= '2025-08-01'
+          AND p.invoice_date <  '2025-09-01'
+          AND CAST(p.amount AS NUMERIC) > 0`,
+        `SELECT "p".description AS "description", "p".invoice_date AS "invoice_date" FROM "anonymized_schema"."payments" AS "p" INNER JOIN "anonymized_schema"."entities" AS "e" ON "p".entity_id = "e".entity_id LEFT JOIN "anonymized_schema"."entities" AS "pe" ON "p".parent_entity_id = "pe".entity_id LEFT JOIN "anonymized_schema"."entities" AS "ppe" ON "pe".parent_ids[array_upper("pe".parent_ids, 1)] = "ppe".entity_id WHERE "e".entity_type IN ('A', 'B') AND "p".invoice_date >= '2025-08-01' AND "p".invoice_date < '2025-09-01' AND CAST("p".amount AS NUMERIC) > 0`
+      ]
+    },
   ]
   function neatlyNestTestedSQL(sqlList){
     sqlList.forEach(sqlInfo => {