feat: support generator table rows in snowflake

taozhi8833998 · taozhi8833998 · commit cef7d9eac732 · 2024-07-14T21:32:59.000+08:00
diff --git a/pegjs/snowflake.pegjs b/pegjs/snowflake.pegjs
@@ -2436,6 +2436,14 @@ table_join
     };
   }
 
+table_generator_option
+  = t:('ROWCOUNT'i / 'TIMELIMIT'i) __ '=>' __ v:literal_numeric {
+    return {
+      type: t.toLowerCase(),
+      symbol: '=>',
+      value: v
+    }
+  }
 
 //NOTE that, the table assigned to `var` shouldn't write in `table_join`
 table_base
@@ -2453,6 +2461,17 @@ table_base
       ...getLocationObject(),
     };
   }
+  / KW_TABLE __ LPAREN __ 'GENERATOR'i __ LPAREN __ g:(table_generator_option)* __ RPAREN __ RPAREN __ alias:value_alias_clause? {
+    return {
+      expr: {
+        keyword: 'table',
+        type: 'generator',
+        generators: g,
+      },
+      as: alias,
+      ...getLocationObject(),
+    };
+  }
   / l:('LATERAL'i)? __ LPAREN __ stmt:(union_stmt / value_clause) __ RPAREN __ alias:value_alias_clause? {
     // => { prefix?: string; expr: union_stmt | value_clause; as?: alias_clause; }
     if (Array.isArray(stmt)) stmt = { type: 'values', values: stmt }
diff --git a/src/tables.js b/src/tables.js
@@ -3,7 +3,7 @@ import { columnRefToSQL } from './column'
 import { exprToSQL } from './expr'
 import { valuesToSQL } from './insert'
 import { intervalToSQL } from './interval'
-import { commonOptionConnector, hasVal, identifierToSql, literalToSQL, toUpper } from './util'
+import { commonOptionConnector, commonTypeValue, hasVal, identifierToSql, literalToSQL, toUpper } from './util'
 
 function unnestToSQL(unnestExpr) {
   const { type, as, expr, with_offset: withOffset } = unnestExpr
@@ -98,6 +98,12 @@ function temporalTableToSQL(stmt) {
   return [toUpper(keyword), temporalTableOptionToSQL(expr)].filter(hasVal).join(' ')
 }
 
+function generateVirtualTable(stmt) {
+  const { keyword, type, generators } = stmt
+  const generatorSQL = generators.map(generator => commonTypeValue(generator).join(' ')).join(', ')
+  return `${toUpper(keyword)}(${toUpper(type)}(${generatorSQL}))`
+}
+
 function tableToSQL(tableInfo) {
   if (toUpper(tableInfo.type) === 'UNNEST') return unnestToSQL(tableInfo)
   const { table, db, as, expr, operator, prefix: prefixStr, schema, server, suffix, tablesample, temporal_table, table_hint } = tableInfo
@@ -119,6 +125,9 @@ function tableToSQL(tableInfo) {
       case 'tumble':
         tableName = tableTumbleToSQL(expr)
         break
+      case 'generator':
+        tableName = generateVirtualTable(expr)
+        break
       default:
         tableName = exprToSQL(expr)
     }
diff --git a/test/snowflake.spec.js b/test/snowflake.spec.js
@@ -263,6 +263,26 @@ describe('snowflake', () => {
         'SELECT * FROM (((SELECT * FROM "some_table") UNION ALL (SELECT * FROM "some_table")) UNION ALL (SELECT * FROM "some_table"))'
       ]
     },
+    {
+      title: 'generate virtual table rows',
+      sql: [
+        `select *, date(concat(left(date,7),'-01')) as date_start_month, left(date,7) as month from
+          ((select
+            dateadd(day, '-' || seq4(), current_date()) as date
+          from
+            table(generator(rowcount => 1095))
+          where date>='2022-10-01') as dates
+          left join
+          (select id, name, category, start_date, end_date, days, cost, cost/days as cost_per_day
+          from
+          (select id, name, category, start_date, end_date, cost, datediff('day',start_date, end_date)+1 as days
+          from ui_other_costs
+          group by 1,2,3,4,5,6)
+          group by 1,2,3,4,5,6,7) as cost)
+          where dates.date between cost.start_date and cost.end_date`,
+          `SELECT *, date(concat(left("date", 7), '-01')) AS "date_start_month", left("date", 7) AS "month" FROM ((SELECT dateadd("day", '-' || seq4(), CURRENT_DATE()) AS "date" FROM TABLE(GENERATOR(ROWCOUNT => 1095)) WHERE "date" >= '2022-10-01') AS "dates" LEFT JOIN (SELECT "id", "name", "category", "start_date", "end_date", "days", "cost", "cost" / "days" AS "cost_per_day" FROM (SELECT "id", "name", "category", "start_date", "end_date", "cost", datediff('day', "start_date", "end_date") + 1 AS "days" FROM "ui_other_costs" GROUP BY 1, 2, 3, 4, 5, 6) GROUP BY 1, 2, 3, 4, 5, 6, 7) AS "cost") WHERE "dates"."date" BETWEEN "cost"."start_date" AND "cost"."end_date"`
+      ]
+    }
   ]
   SQL_LIST.forEach(sqlInfo => {
     const { title, sql } = sqlInfo
