From 2aba292171a6f6ff7d89ff9f755aa2acb3cb5b38 Mon Sep 17 00:00:00 2001
From: Christian De la Hoz <christian.hoz@anixe.pl>
Date: Wed, 1 Sep 2021 19:15:20 +0200
Subject: [PATCH] with cte can define columns

---
 corpus/with_statement.txt | 50 +++++++++++++++++++++++++++++++++++++++
 grammar.js                |  2 +-
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/corpus/with_statement.txt b/corpus/with_statement.txt
index 1786eca..9f92869 100644
--- a/corpus/with_statement.txt
+++ b/corpus/with_statement.txt
@@ -23,6 +23,33 @@ with w as (select * from foo) select * from w;
         (from_table
           (identifier))))))
 
+================================================================================
+with column names
+================================================================================
+with w(foo, bar) as (select * from foo) select * from w;
+--------------------------------------------------------------------------------
+
+(source_file
+  (select_statement
+    (with_query
+      (with_query_item
+        (identifier)
+        (identifier)
+        (identifier)
+        (select_statement
+          (select_item
+            (star))
+          (select_from
+            (from_item
+              (from_table
+                (identifier)))))))
+    (select_item
+      (star))
+    (select_from
+      (from_item
+        (from_table
+          (identifier))))))
+
 ================================================================================
 materialized
 ================================================================================
@@ -110,6 +137,29 @@ with new as (insert into foo values(1)) select 1;
     (select_item
       (number))))
 
+================================================================================
+with update
+================================================================================
+with new as (update foo set bar = 1 returning *) select 1;
+--------------------------------------------------------------------------------
+
+(source_file
+  (select_statement
+    (with_query
+      (with_query_item
+        (identifier)
+        (update_statement
+          (identifier)
+          (update_set
+            (identifier)
+            (update_value
+              (number)))
+          (returning
+            (select_item
+              (star))))))
+    (select_item
+      (number))))
+
 ================================================================================
 many
 ================================================================================
diff --git a/grammar.js b/grammar.js
index f53209f..012cd93 100644
--- a/grammar.js
+++ b/grammar.js
@@ -506,7 +506,7 @@ module.exports = grammar({
     with_query: $ => seq(kw("with"), commaSep1($.with_query_item)),
     with_query_item: $ => seq(
       $.identifier,
-      // TODO(chrde): columns
+      optional($._list_of_identifiers),
       kw("as"),
       optional(choice(kw("materialized"), seq(kw("not"), kw("materialized")))),
       "(", $._with_query_statement, ")"