support constraints

Author: goccy

sqlparser/ast.go

@@ -513,7 +513,8 @@ type DDL struct {
 
 type CreateTable struct {
 	*DDL
-	Columns []*ColumnDef
+	Columns     []*ColumnDef
+	Constraints []*Constraint
 }
 
 // DDL strings.
@@ -757,6 +758,28 @@ func (node Nextval) WalkSubtree(visit Visit) error {
 	return Walk(visit, node.Expr)
 }
 
+type ConstraintType int
+
+const (
+	ConstraintNoConstraint ConstraintType = iota
+	ConstraintPrimaryKey
+	ConstraintKey
+	ConstraintIndex
+	ConstraintUniq
+	ConstraintUniqKey
+	ConstraintUniqIndex
+	ConstraintForeignKey
+	ConstraintFulltext
+)
+
+// Constraint is constraint for table definition.
+type Constraint struct {
+	Type ConstraintType
+	Name string
+	// Used for PRIMARY KEY, UNIQUE, ......
+	Keys []ColIdent
+}
+
 type ColumnDef struct {
 	Name string
 	Type string

sqlparser/type_converter.go

@@ -5,17 +5,30 @@ import (
 )
 
 func convertFromCreateTableStmt(stmt *ast.CreateTableStmt, ddl *DDL) Statement {
-	var columns []*ColumnDef
+	columns := []*ColumnDef{}
 	for _, col := range stmt.Cols {
 		columns = append(columns, &ColumnDef{
 			Name:  col.Name.Name.String(),
 			Type:  col.Tp.String(),
 			Elems: col.Tp.Elems,
 		})
 	}
+	constraints := []*Constraint{}
+	for _, constraint := range stmt.Constraints {
+		keys := []ColIdent{}
+		for _, key := range constraint.Keys {
+			keys = append(keys, NewColIdent(key.Column.Name.String()))
+		}
+		constraints = append(constraints, &Constraint{
+			Type: ConstraintType(constraint.Tp),
+			Name: constraint.Name,
+			Keys: keys,
+		})
+	}
 	return &CreateTable{
-		DDL:     ddl,
-		Columns: columns,
+		DDL:         ddl,
+		Columns:     columns,
+		Constraints: constraints,
 	}
 }