From 22e0d30655e5ef9eb55ef4b4e682d80777b29cd2 Mon Sep 17 00:00:00 2001
From: Eric Fritz <eric@eric-fritz.com>
Date: Mon, 20 May 2024 21:46:52 -0500
Subject: [PATCH] WIP.

---
 internal/syntax/lexing/lexer.go       |  1 +
 internal/syntax/parsing/statements.go | 61 +++++++++++++++++++++------
 internal/syntax/parsing/types.go      |  3 +-
 internal/syntax/tokens/token_type.go  |  1 +
 4 files changed, 53 insertions(+), 13 deletions(-)

diff --git a/internal/syntax/lexing/lexer.go b/internal/syntax/lexing/lexer.go
index 6430238..f040af5 100644
--- a/internal/syntax/lexing/lexer.go
+++ b/internal/syntax/lexing/lexer.go
@@ -69,6 +69,7 @@ var keywordSet = map[string]tokens.TokenType{
 	"using":      tokens.TokenTypeUsing,
 	"values":     tokens.TokenTypeValues,
 	"where":      tokens.TokenTypeWhere,
+	"with":       tokens.TokenTypeWith,
 }
 
 var punctuationMap = map[rune]tokens.TokenType{
diff --git a/internal/syntax/parsing/statements.go b/internal/syntax/parsing/statements.go
index bdcaf44..c46cc5e 100644
--- a/internal/syntax/parsing/statements.go
+++ b/internal/syntax/parsing/statements.go
@@ -26,7 +26,7 @@ func (p *parser) initStatementParsers() {
 
 // statement := ddlStatement | ( [ `EXPLAIN` ] explainableStatement )
 // ddlStatement := ( `CREATE` createTail ) | ( `ALTER` alterTail )
-// explainableStatement := ( `SELECT` selectTail ) | ( `INSERT` insertTail ) | ( `UPDATE` updateTail ) | ( `DELETE` deleteTail )
+// explainableStatement := [ `WITH` name `AS` `(` selectInsertUpdateOrDelete `)` [, ...] ] selectInsertUpdateOrDelete
 func (p *parser) parseStatement() (Query, error) {
 	for tokenType, parser := range p.ddlParsers {
 		token := p.current()
@@ -35,26 +35,63 @@ func (p *parser) parseStatement() (Query, error) {
 		}
 	}
 
-	isExplain := false
-	if p.advanceIf(isType(tokens.TokenTypeExplain)) {
-		isExplain = true
+	isExplain := p.advanceIf(isType(tokens.TokenTypeExplain))
+
+	type namedNode struct {
+		name string
+		node queries.Node
 	}
+	var namedNodes []namedNode
+	if p.advanceIf(isType(tokens.TokenTypeWith)) {
+		for {
+			name, err := p.parseIdent()
+			if err != nil {
+				return nil, err
+			}
 
-	for tokenType, parser := range p.explainableParsers {
-		token := p.current()
-		if p.advanceIf(isType(tokenType)) {
-			node, err := parser(token)
+			if _, err := p.mustAdvance(isType(tokens.TokenTypeAs)); err != nil {
+				return nil, err
+			}
+
+			node, err := parseParenthesized(p, func() (queries.Node, error) {
+				return p.parseSelectInsertUpdateOrDelete()
+			})
 			if err != nil {
 				return nil, err
 			}
 
-			if isExplain {
-				node = explain.NewExplain(node)
+			namedNodes = append(namedNodes, namedNode{name, node})
+
+			if !p.advanceIf(isType(tokens.TokenTypeComma)) {
+				break
 			}
+		}
+	}
 
-			return queries.NewQuery(node), nil
+	node, err := p.parseSelectInsertUpdateOrDelete()
+	if err != nil {
+		return nil, err
+	}
+
+	if len(namedNodes) > 0 {
+		fmt.Printf("> %#v\n", namedNodes)
+	}
+
+	if isExplain {
+		node = explain.NewExplain(node)
+	}
+
+	return queries.NewQuery(node), nil
+}
+
+// selectInsertUpdateOrDelete := ( `SELECT` selectTail ) | ( `INSERT` insertTail ) | ( `UPDATE` updateTail ) | ( `DELETE` deleteTail )
+func (p *parser) parseSelectInsertUpdateOrDelete() (queries.Node, error) {
+	for tokenType, parser := range p.explainableParsers {
+		token := p.current()
+		if p.advanceIf(isType(tokenType)) {
+			return parser(token)
 		}
 	}
 
-	return nil, fmt.Errorf("expected start of statement (near %s)", p.current().Text)
+	return nil, fmt.Errorf("expected start of select, insert, update, or delete statement (near %s)", p.current().Text)
 }
diff --git a/internal/syntax/parsing/types.go b/internal/syntax/parsing/types.go
index f2d83d8..ef157db 100644
--- a/internal/syntax/parsing/types.go
+++ b/internal/syntax/parsing/types.go
@@ -5,6 +5,7 @@ import (
 	"strings"
 
 	"github.com/efritz/gostgres/internal/shared"
+	"github.com/efritz/gostgres/internal/syntax/tokens"
 )
 
 // basicType := ident
@@ -38,7 +39,7 @@ func (p *parser) parseBasicType() (shared.Type, error) {
 		typ = shared.TypeBool
 		// TODO - use multi-phrase keyword(s)
 	case "timestamp":
-		if !p.advanceIf(isIdent("with"), isIdent("time"), isIdent("zone")) {
+		if !p.advanceIf(isType(tokens.TokenTypeWith), isIdent("time"), isIdent("zone")) {
 			return shared.TypeUnknown, fmt.Errorf("unknown type %q", "timestamp")
 		}
 		typ = shared.TypeTimestampTz
diff --git a/internal/syntax/tokens/token_type.go b/internal/syntax/tokens/token_type.go
index 31ec83c..025ae71 100644
--- a/internal/syntax/tokens/token_type.go
+++ b/internal/syntax/tokens/token_type.go
@@ -67,6 +67,7 @@ const (
 	TokenTypeUsing
 	TokenTypeValues
 	TokenTypeWhere
+	TokenTypeWith
 
 	//
 	// Single-character operators