Add support for UNIQUE table constraint

chaisql · Apr 13, 2021 · 7d3e769 · 7d3e769
1 parent 96b18eb
commit 7d3e769
Show file tree

Hide file tree

Showing 3 changed files with 93 additions and 40 deletions.
diff --git a/sql/parser/create.go b/sql/parser/create.go
@@ -2,7 +2,6 @@ package parser
 
 import (
 	"github.com/genjidb/genji/database"
-	"github.com/genjidb/genji/document"
 	"github.com/genjidb/genji/expr"
 	"github.com/genjidb/genji/query"
 	"github.com/genjidb/genji/sql/scanner"
@@ -95,36 +94,20 @@ func (p *Parser) parseConstraints(stmt *query.CreateTableStmt) error {
 	for {
 		// we start by checking if it is a table constraint,
 		// as it's easier to determine
-		tc, err := p.parseTableConstraint()
+		ok, err := p.parseTableConstraint(stmt)
 		if err != nil {
 			return err
 		}
 
 		// no table constraint found
-		if tc == nil && parsingTableConstraints {
+		if !ok && parsingTableConstraints {
 			tok, pos, lit := p.ScanIgnoreWhitespace()
 			return newParseError(scanner.Tokstr(tok, lit), []string{"CONSTRAINT", ")"}, pos)
 		}
 
 		// only PRIMARY KEY(path) is currently supported.
-		if tc != nil {
+		if ok {
 			parsingTableConstraints = true
-
-			if pk := stmt.Info.GetPrimaryKey(); pk != nil {
-				return stringutil.Errorf("table %q has more than one primary key", stmt.TableName)
-			}
-			fc := stmt.Info.FieldConstraints.Get(tc.primaryKey)
-			if fc == nil {
-				err = stmt.Info.FieldConstraints.Add(&database.FieldConstraint{
-					Path:         tc.primaryKey,
-					IsPrimaryKey: true,
-				})
-				if err != nil {
-					return err
-				}
-			} else {
-				fc.IsPrimaryKey = true
-			}
 		}
 
 		// if set to false, we are still parsing field definitions
@@ -225,37 +208,81 @@ func (p *Parser) parseFieldConstraint(fc *database.FieldConstraint) error {
 	}
 }
 
-func (p *Parser) parseTableConstraint() (*tableConstraint, error) {
-	var tc tableConstraint
+func (p *Parser) parseTableConstraint(stmt *query.CreateTableStmt) (bool, error) {
 	var err error
 
 	tok, _, _ := p.ScanIgnoreWhitespace()
 	switch tok {
 	case scanner.PRIMARY:
-		// Parse "KEY"
-		if tok, pos, lit := p.ScanIgnoreWhitespace(); tok != scanner.KEY {
-			return nil, newParseError(scanner.Tokstr(tok, lit), []string{"KEY"}, pos)
+		// Parse "KEY ("
+		err = p.parseTokens(scanner.KEY, scanner.LPAREN)
+		if err != nil {
+			return false, err
+		}
+
+		primaryKeyPath, err := p.parsePath()
+		if err != nil {
+			return false, err
+		}
+
+		// Parse ")"
+		err = p.parseTokens(scanner.RPAREN)
+		if err != nil {
+			return false, err
+		}
+
+		if pk := stmt.Info.GetPrimaryKey(); pk != nil {
+			return false, stringutil.Errorf("table %q has more than one primary key", stmt.TableName)
+		}
+		fc := stmt.Info.FieldConstraints.Get(primaryKeyPath)
+		if fc == nil {
+			err = stmt.Info.FieldConstraints.Add(&database.FieldConstraint{
+				Path:         primaryKeyPath,
+				IsPrimaryKey: true,
+			})
+			if err != nil {
+				return false, err
+			}
+		} else {
+			fc.IsPrimaryKey = true
 		}
 
+		return true, nil
+	case scanner.UNIQUE:
 		// Parse "("
-		if tok, pos, lit := p.ScanIgnoreWhitespace(); tok != scanner.LPAREN {
-			return nil, newParseError(scanner.Tokstr(tok, lit), []string{"("}, pos)
+		err = p.parseTokens(scanner.LPAREN)
+		if err != nil {
+			return false, err
 		}
 
-		tc.primaryKey, err = p.parsePath()
+		uniquePath, err := p.parsePath()
 		if err != nil {
-			return nil, err
+			return false, err
 		}
 
 		// Parse ")"
-		if tok, pos, lit := p.ScanIgnoreWhitespace(); tok != scanner.RPAREN {
-			return nil, newParseError(scanner.Tokstr(tok, lit), []string{")"}, pos)
+		err = p.parseTokens(scanner.RPAREN)
+		if err != nil {
+			return false, err
 		}
 
-		return &tc, nil
+		fc := stmt.Info.FieldConstraints.Get(uniquePath)
+		if fc == nil {
+			err = stmt.Info.FieldConstraints.Add(&database.FieldConstraint{
+				Path:     uniquePath,
+				IsUnique: true,
+			})
+			if err != nil {
+				return false, err
+			}
+		} else {
+			fc.IsUnique = true
+		}
+
+		return true, nil
 	default:
 		p.Unscan()
-		return nil, nil
+		return false, nil
 	}
 }
 
@@ -312,7 +339,3 @@ func (p *Parser) parseCreateIndexStatement(unique bool) (query.CreateIndexStmt,
 
 	return stmt, nil
 }
-
-type tableConstraint struct {
-	primaryKey document.Path
-}
diff --git a/sql/parser/create_test.go b/sql/parser/create_test.go
@@ -133,6 +133,36 @@ func TestParserCreateTable(t *testing.T) {
 		{"With table constraints / field constraint after table constraint", "CREATE TABLE test(PRIMARY KEY (bar), foo INTEGER)", nil, true},
 		{"With table constraints / duplicate pk", "CREATE TABLE test(foo INTEGER PRIMARY KEY, PRIMARY KEY (bar))", nil, true},
 		{"With table constraints / duplicate pk on same path", "CREATE TABLE test(foo INTEGER PRIMARY KEY, PRIMARY KEY (foo))", nil, true},
+		{"With table constraints / UNIQUE on defined field", "CREATE TABLE test(foo INTEGER, bar NOT NULL, UNIQUE (foo))",
+			query.CreateTableStmt{
+				TableName: "test",
+				Info: database.TableInfo{
+					FieldConstraints: []*database.FieldConstraint{
+						{Path: document.Path(parsePath(t, "foo")), Type: document.IntegerValue, IsUnique: true},
+						{Path: document.Path(parsePath(t, "bar")), IsNotNull: true},
+					},
+				},
+			}, false},
+		{"With table constraints / UNIQUE on undefined field", "CREATE TABLE test(foo INTEGER, UNIQUE (bar))",
+			query.CreateTableStmt{
+				TableName: "test",
+				Info: database.TableInfo{
+					FieldConstraints: []*database.FieldConstraint{
+						{Path: document.Path(parsePath(t, "foo")), Type: document.IntegerValue},
+						{Path: document.Path(parsePath(t, "bar")), IsUnique: true},
+					},
+				},
+			}, false},
+		{"With table constraints / UNIQUE twice", "CREATE TABLE test(foo INTEGER UNIQUE, UNIQUE (foo))",
+			query.CreateTableStmt{
+				TableName: "test",
+				Info: database.TableInfo{
+					FieldConstraints: []*database.FieldConstraint{
+						{Path: document.Path(parsePath(t, "foo")), Type: document.IntegerValue, IsUnique: true},
+					},
+				},
+			}, false},
+		{"With table constraints / duplicate pk on same path", "CREATE TABLE test(foo INTEGER PRIMARY KEY, PRIMARY KEY (foo))", nil, true},
 		{"With multiple primary keys", "CREATE TABLE test(foo PRIMARY KEY, bar PRIMARY KEY)",
 			query.CreateTableStmt{}, true},
 		{"With all supported fixed size data types",

diff --git a/sql/parser/parser.go b/sql/parser/parser.go
@@ -210,9 +210,9 @@ func (p *Parser) Unscan() {
 	p.s.Unscan()
 }
 
-// parseConsecutiveTokens parses all the given tokens one after the other.
+// parseTokens parses all the given tokens one after the other.
 // It returns an error if one of the token is missing.
-func (p *Parser) parseConsecutiveTokens(tokens ...scanner.Token) error {
+func (p *Parser) parseTokens(tokens ...scanner.Token) error {
 	for _, t := range tokens {
 		if tok, pos, lit := p.ScanIgnoreWhitespace(); tok != t {
 			return newParseError(scanner.Tokstr(tok, lit), []string{t.String()}, pos)
@@ -236,7 +236,7 @@ func (p *Parser) parseOptional(tokens ...scanner.Token) (bool, error) {
 		return true, nil
 	}
 
-	err := p.parseConsecutiveTokens(tokens[1:]...)
+	err := p.parseTokens(tokens[1:]...)
 	return err == nil, err
 }