👍 add update parser (#12)

okazaki-kk · web-flow · commit e33c4ca215a0 · 2023-07-09T13:53:56.000+09:00
* add test

* add parser
diff --git a/internal/parser/ast/ast.go b/internal/parser/ast/ast.go
@@ -30,6 +30,17 @@ type SelectStatement struct {
 	Offset  *OffsetStatement
 }
 
+type UpdateStatement struct {
+	Table string
+	Set   []SetStatement
+	Where *WhereStatement
+}
+
+type SetStatement struct {
+	Column string
+	Value  Expression
+}
+
 // ResultStatement node represents a returning expression in a SELECT statement.
 type ResultStatement struct {
 	Expr Expression
@@ -85,6 +96,8 @@ func (s *LimitStatement) statementNode()          {}
 func (s *OffsetStatement) statementNode()         {}
 func (s *InsertStatement) statementNode()         {}
 func (s *CreateDatabaseStatement) statementNode() {}
+func (s *UpdateStatement) statementNode()         {}
+func (s *SetStatement) statementNode()            {}
 
 // IdentExpr node represents an identifier.
 type IdentExpr struct {
diff --git a/internal/parser/parser.go b/internal/parser/parser.go
@@ -39,6 +39,8 @@ func (p *Parser) parseStatement() (ast.Statement, error) {
 		return p.parseSelectStatement()
 	case token.INSERT:
 		return p.parseInsertStatement()
+	case token.UPDATE:
+		return p.parseUpdateStatement()
 	case token.CREATE:
 		p.nextToken()
 		return p.parseCreateStatement()
@@ -123,6 +125,27 @@ func (p *Parser) parseInsertStatement() (ast.Statement, error) {
 	return &insert, nil
 }
 
+func (p *Parser) parseUpdateStatement() (ast.Statement, error) {
+	p.nextToken()
+
+	table, err := p.parseIdent()
+	if err != nil {
+		return nil, err
+	}
+
+	set, err := p.parseSetStatement()
+	if err != nil {
+		return nil, err
+	}
+
+	where, err := p.parseWhereStatement()
+	if err != nil {
+		return nil, err
+	}
+
+	return &ast.UpdateStatement{Table: table.Name, Set: set, Where: where}, nil
+}
+
 func (p *Parser) parseCreateStatement() (ast.Statement, error) {
 	switch p.token.Type {
 	case token.TABLE:
@@ -289,6 +312,46 @@ func (p *Parser) parseFromStatement() (*ast.FromStatement, error) {
 	return &from, nil
 }
 
+func (p *Parser) parseSetStatement() ([]ast.SetStatement, error) {
+	if err := p.expect(token.SET); err != nil {
+		return nil, err
+	}
+
+	columns := make([]ast.SetStatement, 0)
+
+	for {
+		column, err := p.parseIdent()
+		if err != nil {
+			return nil, err
+		}
+
+		if err = p.expect(token.EQ); err != nil {
+			return nil, err
+		}
+
+		value, err := p.parsePrimaryExpr()
+		if err != nil {
+			return nil, err
+		}
+
+		columns = append(columns, ast.SetStatement{
+			Column: column.Name,
+			Value:  value,
+		})
+
+		if p.peekToken.Type == token.EOF || p.peekToken.Type == token.WHERE {
+			p.nextToken()
+			break
+		}
+
+		if err = p.expect(token.COMMA); err != nil {
+			return nil, err
+		}
+	}
+
+	return columns, nil
+}
+
 func (p *Parser) parseWhereStatement() (*ast.WhereStatement, error) {
 	if p.token.Type != token.WHERE {
 		return nil, nil
diff --git a/internal/parser/parser_test.go b/internal/parser/parser_test.go
@@ -489,3 +489,67 @@ func TestParser_CreateTable(t *testing.T) {
 		})
 	}
 }
+
+func TestParser_Update(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		input string
+		stmt  ast.Statement
+	}{
+		{
+			input: "UPDATE customers SET name = 'vlad', salary = 10*100 WHERE id = 1",
+			stmt: &ast.UpdateStatement{
+				Table: "customers",
+				Set: []ast.SetStatement{
+					{
+						Column: "name",
+						Value: &ast.ScalarExpr{
+							Type:    token.TEXT,
+							Literal: "vlad",
+						},
+					},
+					{
+						Column: "salary",
+						Value: &ast.ConditionExpr{
+							Left: &ast.ScalarExpr{
+								Type:    token.INT,
+								Literal: "10",
+							},
+							Operator: token.ASTERISK,
+							Right: &ast.ScalarExpr{
+								Type:    token.INT,
+								Literal: "100",
+							},
+						},
+					},
+				},
+				Where: &ast.WhereStatement{
+					Expr: &ast.ConditionExpr{
+						Left: &ast.IdentExpr{
+							Name: "id",
+						},
+						Operator: token.EQ,
+						Right: &ast.ScalarExpr{
+							Type:    token.INT,
+							Literal: "1",
+						},
+					},
+				},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+
+		t.Run(test.input, func(t *testing.T) {
+			t.Parallel()
+
+			p := New(lexer.New(test.input))
+			stmts, err := p.Parse()
+			assert.NoError(t, err)
+			assert.Equal(t, test.stmt, stmts)
+		})
+	}
+}
