Adding support for setting/dropping column attributes: type and not n…

…ull constraint (needed for Postgres' more modular ALTER TABLE syntax)
dolthub · Oct 2, 2024 · 1c84e91 · 1c84e91
1 parent 07dfd6f
commit 1c84e91
Showing 1 changed file with 56 additions and 0 deletions.
diff --git a/sql/planbuilder/ddl.go b/sql/planbuilder/ddl.go
@@ -526,6 +526,14 @@ func (b *Builder) buildAlterTableClause(inScope *scope, ddl *ast.DDL) []*scope {
 			outScopes = append(outScopes, b.buildAlterCollationSpec(tableScope, ddl, rt))
 		}
 
+		if ddl.NotNullSpec != nil {
+			outScopes = append(outScopes, b.buildAlterNotNull(tableScope, ddl, rt))
+		}
+
+		if ddl.ColumnTypeSpec != nil {
+			outScopes = append(outScopes, b.buildAlterChangeColumnType(tableScope, ddl, rt))
+		}
+
 		for _, s := range outScopes {
 			if ts, ok := s.node.(sql.SchemaTarget); ok {
 				s.node = b.modifySchemaTarget(s, ts, rt.Schema())
@@ -924,6 +932,54 @@ func (b *Builder) buildAlterAutoIncrement(inScope *scope, ddl *ast.DDL, table *p
 	return
 }
 
+func (b *Builder) buildAlterNotNull(inScope *scope, ddl *ast.DDL, table *plan.ResolvedTable) (outScope *scope) {
+	outScope = inScope
+	spec := ddl.NotNullSpec
+	for _, c := range table.Schema() {
+		if strings.EqualFold(c.Name, spec.Column.String()) {
+			switch strings.ToLower(spec.Action) {
+			case ast.SetStr:
+				// Set NOT NULL constraint
+				c.Nullable = false
+			case ast.DropStr:
+				// Drop NOT NULL constraint
+				c.Nullable = true
+			default:
+				err := sql.ErrUnsupportedFeature.New(ast.String(ddl))
+				b.handleErr(err)
+			}
+
+			modifyColumn := plan.NewModifyColumnResolved(table, c.Name, *c, nil)
+			outScope.node = b.modifySchemaTarget(inScope, modifyColumn, table.Schema())
+			return
+		}
+	}
+	err := sql.ErrTableColumnNotFound.New(table.Name(), spec.Column.String())
+	b.handleErr(err)
+	return
+}
+
+func (b *Builder) buildAlterChangeColumnType(inScope *scope, ddl *ast.DDL, table *plan.ResolvedTable) (outScope *scope) {
+	outScope = inScope
+	spec := ddl.ColumnTypeSpec
+	for _, c := range table.Schema() {
+		if strings.EqualFold(c.Name, spec.Column.String()) {
+			typ, err := types.ColumnTypeToType(&spec.Type)
+			if err != nil {
+				b.handleErr(err)
+				return
+			}
+			c.Type = typ
+			modifyColumn := plan.NewModifyColumnResolved(table, c.Name, *c, nil)
+			outScope.node = b.modifySchemaTarget(inScope, modifyColumn, table.Schema())
+			return
+		}
+	}
+	err := sql.ErrTableColumnNotFound.New(table.Name(), spec.Column.String())
+	b.handleErr(err)
+	return
+}
+
 func (b *Builder) buildAlterDefault(inScope *scope, ddl *ast.DDL, table *plan.ResolvedTable) (outScope *scope) {
 	outScope = inScope
 	switch strings.ToLower(ddl.DefaultSpec.Action) {