fix: truncate table with option in pg

Author: taozhi8833998

pegjs/postgresql.pegjs

@@ -1335,28 +1335,46 @@ drop_stmt
       };
     }
 
+truncate_table_name
+  = t:table_name __ s:STAR? {
+    // => table_name & { suffix?: string }
+    tableList.add(`truncate::${t.db}::${t.table}`)
+    if (s) t.suffix = s
+    return t
+  }
+truncate_table_name_list
+  = head:truncate_table_name tail:(__ COMMA __ truncate_table_name)* {
+    // => truncate_table_name[]
+      return createList(head, tail)
+    }
 truncate_stmt
   = a:KW_TRUNCATE  __
     kw:KW_TABLE? __
-    t:table_ref_list {
+    on: 'ONLY'i? __
+    t:truncate_table_name_list __
+    id: (('RESTART'i / 'CONTINUE'i) __ 'IDENTITY'i)? __
+    op:('CASCADE'i / 'RESTRICT'i)? {
       /*
       export interface truncate_stmt_node {
         type: 'trucate';
         keyword: 'table';
+        prefix?: string;
         name: table_ref_list;
+        suffix: string[];
       }
       => AstStatement<truncate_stmt_node>
       */
-      if(t) t.forEach(tt => tableList.add(`${a}::${tt.db}::${tt.table}`));
       return {
         tableList: Array.from(tableList),
         columnList: columnListTableAlias(columnList),
         ast: {
           type: a.toLowerCase(),
           keyword: kw && kw.toLowerCase() || 'table',
-          name: t
+          prefix: on,
+          name: t,
+          suffix: [id && [id[0], id[2]].join(' '), op].filter(v => v).map(v => ({ type: 'origin', value: v }))
         }
-      };
+      }
     }
 
 use_stmt

src/command.js

@@ -13,7 +13,7 @@ function callToSQL(stmt) {
 }
 
 function commonCmdToSQL(stmt) {
-  const { type, keyword, name, prefix } = stmt
+  const { type, keyword, name, prefix, suffix } = stmt
   const clauses = [toUpper(type), toUpper(keyword), toUpper(prefix)]
   switch (keyword) {
     case 'table':
@@ -36,6 +36,7 @@ function commonCmdToSQL(stmt) {
     default:
       break
   }
+  if (suffix) clauses.push(suffix.map(exprToSQL).filter(hasVal).join(' '))
   return clauses.filter(hasVal).join(' ')
 }
 

src/tables.js

@@ -74,7 +74,7 @@ function tableTumbleToSQL(tumble) {
 
 function tableToSQL(tableInfo) {
   if (toUpper(tableInfo.type) === 'UNNEST') return unnestToSQL(tableInfo)
-  const { table, db, as, expr, operator, prefix: prefixStr, schema, server, tablesample, table_hint } = tableInfo
+  const { table, db, as, expr, operator, prefix: prefixStr, schema, server, suffix, tablesample, table_hint } = tableInfo
   const serverName = identifierToSql(server)
   const database = identifierToSql(db)
   const schemaStr = identifierToSql(schema)
@@ -97,7 +97,7 @@ function tableToSQL(tableInfo) {
         tableName = exprToSQL(expr)
     }
   }
-  tableName = [toUpper(prefixStr), tableName].filter(hasVal).join(' ')
+  tableName = [toUpper(prefixStr), tableName, toUpper(suffix)].filter(hasVal).join(' ')
   let str = [serverName, database, schemaStr, tableName].filter(hasVal).join('.')
   if (tableInfo.parentheses) str = `(${str})`
   const result = [str]

test/postgres.spec.js

@@ -1422,6 +1422,20 @@ describe('Postgres', () => {
         'SELECT * FROM "schema"."func"("start_time", "end_time")'
       ]
     },
+    {
+      title: 'truncate table',
+      sql: [
+        'TRUNCATE TABLE employee RESTART IDENTITY cascade',
+        'TRUNCATE TABLE "employee" RESTART IDENTITY CASCADE'
+      ]
+    },
+    {
+      title: 'truncate all table',
+      sql: [
+        'TRUNCATE TABLE ONLY employee * CONTINUE IDENTITY RESTRICT',
+        'TRUNCATE TABLE ONLY "employee" * CONTINUE IDENTITY RESTRICT'
+      ]
+    },
   ]
   function neatlyNestTestedSQL(sqlList){
     sqlList.forEach(sqlInfo => {

test/redshift.spec.js

@@ -15,6 +15,21 @@ describe('redshift', () => {
     Last_VALUE(hermes_health_score IGNORE NULLS) OVER ( partition by sf_account_id, month_id ) as hermes_health_score_monthly,  row_number() OVER ( PARTITION BY sf_account_id, date_trunc('week', date_)
     ORDER BY date_ desc ) AS o_key_week  FROM dwh.dwh_health_score_hermes  WHERE date_trunc('month', date_) >= '2023-01-01'  Qualify o_key_week = 1`
     expect(getParsedSql(sql)).to.be.equal(`SELECT sf_account_id, date_trunc('week', date_) AS "week_id", date_trunc('month', date_) AS "month_id", Last_VALUE(hermes_health_score IGNORE NULLS) OVER (PARTITION BY sf_account_id, week_id) AS "hermes_health_score", Last_VALUE(hermes_health_score IGNORE NULLS) OVER (PARTITION BY sf_account_id, month_id) AS "hermes_health_score_monthly", row_number() OVER (PARTITION BY sf_account_id, date_trunc('week', date_) ORDER BY date_ DESC) AS "o_key_week" FROM "dwh"."dwh_health_score_hermes" WHERE date_trunc('month', date_) >= '2023-01-01' QUALIFY o_key_week = 1`)
+    sql = `with pv as (
+      select
+        action_date,
+        visitor_id_v,
+        visit_country_name,
+        referer_channel_group,
+        email,
+        sgid,
+        mp."brand/non-brand" as is_brand
+      from
+        dwh_fact_pageviews pv
+        left join ppc_keywords_mapping mp using (campaign_keyword)
+      )
+      select is_brand as "b/nb" from pv where mp."brand/non-brand" = 'brand'`
+      expect(getParsedSql(sql)).to.be.equal(`WITH "pv" AS (SELECT action_date, visitor_id_v, visit_country_name, referer_channel_group, email, sgid, "mp"."brand/non-brand" AS "is_brand" FROM "dwh_fact_pageviews" AS "pv" LEFT JOIN "ppc_keywords_mapping" AS "mp" USING ("campaign_keyword")) SELECT is_brand AS "b/nb" FROM "pv" WHERE "mp"."brand/non-brand" = 'brand'`)
   })
 
 })

test/transactsql.spec.js

@@ -150,6 +150,11 @@ describe('transactsql', () => {
     expect(getParsedSql(sql)).to.equal('SELECT [source].[dbo].[movie].[id] FROM [source].[dbo].[movie]')
     sql = 'SELECT * FROM source.dbo.movie WHERE source.dbo.movie.genre_id = 1'
     expect(getParsedSql(sql)).to.equal('SELECT * FROM [source].[dbo].[movie] WHERE [source].[dbo].[movie].[genre_id] = 1')
+    sql = 'SELECT TOP 1000 [production].[categories].[category_name], COUNT([production].[products].[product_id]) AS [product_count]\n' +
+            'FROM [production].[products]\n' +
+            'INNER JOIN [production].[categories] ON [production].[products].[category_id] = [production].[categories].[category_id]\n' +
+            'GROUP BY [production].[categories].[category_name]'
+    expect(getParsedSql(sql)).to.be.equal("SELECT TOP 1000 [production].[categories].[category_name], COUNT([production].[products].[product_id]) AS [product_count] FROM [production].[products] INNER JOIN [production].[categories] ON [production].[products].[category_id] = [production].[categories].[category_id] GROUP BY [production].[categories].[category_name]")
   })
 
   it('should support with clause', () => {