Fixed bug with generated values not being properly quoted

server/functions/framework/quick_function.go

@@ -136,7 +136,17 @@ func (q *QuickFunction1) Children() []sql.Expression {
 
 // WithChildren implements the interface sql.Expression.
 func (q *QuickFunction1) WithChildren(children ...sql.Expression) (sql.Expression, error) {
-	return nil, errors.Errorf("cannot change the children for `%T`", q)
+	if len(children) != 1 {
+		return nil, sql.ErrInvalidChildrenNumber.New(q, len(children), 1)
+	}
+
+	if children[0].Type().Equals(q.Argument.Type()) {
+		nq := *q
+		nq.Argument = children[0]
+		return &nq, nil
+	}
+
+	return nil, errors.Errorf("cannot change the types of children for `%T`", q)
 }
 
 // specificFuncImpl implements the interface sql.Expression.
@@ -244,7 +254,18 @@ func (q *QuickFunction2) Children() []sql.Expression {
 
 // WithChildren implements the interface sql.Expression.
 func (q *QuickFunction2) WithChildren(children ...sql.Expression) (sql.Expression, error) {
-	return nil, errors.Errorf("cannot change the children for `%T`", q)
+	if len(children) != 2 {
+		return nil, sql.ErrInvalidChildrenNumber.New(q, len(children), 2)
+	}
+
+	if children[0].Type().Equals(q.Arguments[0].Type()) &&
+		children[1].Type().Equals(q.Arguments[1].Type()) {
+		nq := *q
+		nq.Arguments = [2]sql.Expression{children[0], children[1]}
+		return &nq, nil
+	}
+
+	return nil, errors.Errorf("cannot change the types of children for `%T`", q)
 }
 
 // specificFuncImpl implements the interface sql.Expression.
@@ -352,7 +373,19 @@ func (q *QuickFunction3) Children() []sql.Expression {
 
 // WithChildren implements the interface sql.Expression.
 func (q *QuickFunction3) WithChildren(children ...sql.Expression) (sql.Expression, error) {
-	return nil, errors.Errorf("cannot change the children for `%T`", q)
+	if len(children) != 3 {
+		return nil, sql.ErrInvalidChildrenNumber.New(q, len(children), 3)
+	}
+
+	if children[0].Type().Equals(q.Arguments[0].Type()) &&
+		children[1].Type().Equals(q.Arguments[1].Type()) &&
+		children[2].Type().Equals(q.Arguments[2].Type()) {
+		nq := *q
+		nq.Arguments = [3]sql.Expression{children[0], children[1], children[2]}
+		return &nq, nil
+	}
+
+	return nil, errors.Errorf("cannot change the types of children for `%T`", q)
 }
 
 // specificFuncImpl implements the interface sql.Expression.

server/node/create_table.go

@@ -30,6 +30,7 @@ type CreateTable struct {
 
 var _ sql.ExecSourceRel = (*CreateTable)(nil)
 var _ sql.SchemaTarget = (*CreateTable)(nil)
+var _ sql.Expressioner = (*CreateTable)(nil)
 
 // NewCreateTable returns a new *CreateTable.
 func NewCreateTable(createTable *plan.CreateTable, sequences []*CreateSequence) *CreateTable {
@@ -117,3 +118,20 @@ func (c CreateTable) WithTargetSchema(schema sql.Schema) (sql.Node, error) {
 
 	return &c, nil
 }
+
+// Expressions implements the sql.Expressioner interface.
+func (c *CreateTable) Expressions() []sql.Expression {
+	return c.gmsCreateTable.Expressions()
+}
+
+// WithExpressions implements the sql.Expressioner interface.
+func (c *CreateTable) WithExpressions(expression ...sql.Expression) (sql.Node, error) {
+	nc := *c
+	n, err := nc.gmsCreateTable.WithExpressions(expression...)
+	if err != nil {
+		return nil, err
+	}
+
+	nc.gmsCreateTable = n.(*plan.CreateTable)
+	return &nc, nil
+}

testing/go/alter_table_test.go

@@ -115,6 +115,55 @@ func TestAlterTable(t *testing.T) {
 				},
 			},
 		},
+		{
+			Name: "Add Check Constraint with IN tuple",
+			SetUpScript: []string{
+				"create table t1 (pk int primary key, c1 int);",
+				"insert into t1 values (1,1);",
+			},
+			Assertions: []ScriptTestAssertion{
+				{
+					// Add a check constraint that is already violated by the existing data
+					Query:       "ALTER TABLE t1 ADD CONSTRAINT constraint1 CHECK (c1 in (100));",
+					ExpectedErr: "violated",
+				},
+				{
+					// Add a check constraint
+					Query:    "ALTER TABLE t1 ADD CONSTRAINT constraint1 CHECK (c1 in (1,2));",
+					Expected: []sql.Row{},
+				},
+				{
+					Query:    "INSERT INTO t1 VALUES (2, 2);",
+					Expected: []sql.Row{},
+				},
+				{
+					Query:       "INSERT INTO t1 VALUES (3, 101);",
+					ExpectedErr: "violated",
+				},
+			},
+		},
+		{
+			Name: "Add Check Constraint and another constraint in same statement",
+			SetUpScript: []string{
+				"create table t1 (pk int, c1 int);",
+				"insert into t1 values (1,1);",
+			},
+			Assertions: []ScriptTestAssertion{
+				{
+					// Add a check constraint
+					Query:    " ALTER TABLE t1 ADD CONSTRAINT check_a CHECK (c1 IN (1)), ALTER c1 SET NOT NULL;",
+					Expected: []sql.Row{},
+				},
+				{
+					Query:       "INSERT INTO t1 VALUES (2, 2);",
+					ExpectedErr: "violated",
+				},
+				{
+					Query:       "INSERT INTO t1 VALUES (1, NULL);",
+					ExpectedErr: "non-nullable",
+				},
+			},
+		},
 		{
 			Name: "Drop Constraint",
 			SetUpScript: []string{

testing/go/create_table_test.go

@@ -106,6 +106,27 @@ func TestCreateTable(t *testing.T) {
 				},
 			},
 		},
+		{
+			Name: "check constraint with a function",
+			Assertions: []ScriptTestAssertion{
+				{
+					Query:    "CREATE TABLE mytbl (a text CHECK (length(a) > 2) PRIMARY KEY, b text);",
+					Expected: []sql.Row{},
+				},
+				{
+					Query:    "insert into mytbl values ('abc', 'def');",
+					Expected: []sql.Row{},
+				},
+				{
+					Query:       "insert into mytbl values ('de', 'abc');",
+					ExpectedErr: `Check constraint`,
+				},
+				{
+					Query:    "select * from mytbl;",
+					Expected: []sql.Row{{"abc", "def"}},
+				},
+			},
+		},
 		{
 			Skip: true, // TODO: vitess does not support multiple check constraint on a single column
 			Name: "Create table with multiple check constraints on a single column",
@@ -234,8 +255,7 @@ func TestCreateTable(t *testing.T) {
 			},
 		},
 		{
-			Name: "generated column with reference to another column that needs quote",
-			Skip: true, // generated columns are not being properly quoted, this is a bug in GMS
+			Name: "generated column with space in column name",
 			SetUpScript: []string{
 				`create table t1 (
     			a varchar(10) primary key,