Error out if table/column is missing (#681)

agedemenli · web-flow · commit 9a3eaed7a5ca · 2025-02-13T15:17:51.000+03:00
Graceful error for the segfault scenario described in #641
diff --git a/pkg/migrations/op_add_column.go b/pkg/migrations/op_add_column.go
@@ -19,6 +19,9 @@ var _ Operation = (*OpAddColumn)(nil)
 
 func (o *OpAddColumn) Start(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, s *schema.Schema) (*schema.Table, error) {
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return nil, TableDoesNotExistError{Name: o.Table}
+	}
 
 	if err := addColumn(ctx, conn, *o, table, tr); err != nil {
 		return nil, fmt.Errorf("failed to start add column operation: %w", err)
@@ -120,6 +123,9 @@ func (o *OpAddColumn) Complete(ctx context.Context, conn db.DB, tr SQLTransforme
 	// Add the column to the in-memory schema so that Complete steps in subsequent
 	// operations can see the new column.
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return TableDoesNotExistError{Name: o.Table}
+	}
 	table.AddColumn(o.Column.Name, &schema.Column{
 		Name: o.Column.Name,
 	})
@@ -129,7 +135,13 @@ func (o *OpAddColumn) Complete(ctx context.Context, conn db.DB, tr SQLTransforme
 
 func (o *OpAddColumn) Rollback(ctx context.Context, conn db.DB, tr SQLTransformer, s *schema.Schema) error {
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return TableDoesNotExistError{Name: o.Table}
+	}
 	column := table.GetColumn(o.Column.Name)
+	if column == nil {
+		return ColumnDoesNotExistError{Table: o.Table, Name: o.Column.Name}
+	}
 
 	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
 		pq.QuoteIdentifier(table.Name),
diff --git a/pkg/migrations/op_alter_column.go b/pkg/migrations/op_alter_column.go
@@ -17,7 +17,13 @@ var _ Operation = (*OpAlterColumn)(nil)
 
 func (o *OpAlterColumn) Start(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, s *schema.Schema) (*schema.Table, error) {
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return nil, TableDoesNotExistError{Name: o.Table}
+	}
 	column := table.GetColumn(o.Column)
+	if column == nil {
+		return nil, ColumnDoesNotExistError{Table: o.Table, Name: o.Column}
+	}
 	ops := o.subOperations()
 
 	// Duplicate the column on the underlying table.
@@ -114,7 +120,13 @@ func (o *OpAlterColumn) Complete(ctx context.Context, conn db.DB, tr SQLTransfor
 
 	// Rename the new column to the old column name
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return TableDoesNotExistError{Name: o.Table}
+	}
 	column := table.GetColumn(o.Column)
+	if column == nil {
+		return ColumnDoesNotExistError{Table: o.Table, Name: o.Column}
+	}
 	if err := RenameDuplicatedColumn(ctx, conn, table, column); err != nil {
 		return err
 	}
@@ -124,7 +136,13 @@ func (o *OpAlterColumn) Complete(ctx context.Context, conn db.DB, tr SQLTransfor
 
 func (o *OpAlterColumn) Rollback(ctx context.Context, conn db.DB, tr SQLTransformer, s *schema.Schema) error {
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return TableDoesNotExistError{Name: o.Table}
+	}
 	column := table.GetColumn(o.Column)
+	if column == nil {
+		return ColumnDoesNotExistError{Table: o.Table, Name: o.Column}
+	}
 
 	// Perform any operation specific rollback steps
 	ops := o.subOperations()
diff --git a/pkg/migrations/op_change_type.go b/pkg/migrations/op_change_type.go
@@ -21,6 +21,9 @@ var _ Operation = (*OpChangeType)(nil)
 
 func (o *OpChangeType) Start(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, s *schema.Schema) (*schema.Table, error) {
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return nil, TableDoesNotExistError{Name: o.Table}
+	}
 
 	return table, nil
 }
diff --git a/pkg/migrations/op_create_constraint.go b/pkg/migrations/op_create_constraint.go
@@ -17,9 +17,16 @@ var _ Operation = (*OpCreateConstraint)(nil)
 
 func (o *OpCreateConstraint) Start(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, s *schema.Schema) (*schema.Table, error) {
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return nil, TableDoesNotExistError{Name: o.Table}
+	}
+
 	columns := make([]*schema.Column, len(o.Columns))
 	for i, colName := range o.Columns {
 		columns[i] = table.GetColumn(colName)
+		if columns[i] == nil {
+			return nil, ColumnDoesNotExistError{Table: o.Table, Name: colName}
+		}
 	}
 
 	// Duplicate each column using its final name after migration completion
@@ -137,8 +144,14 @@ func (o *OpCreateConstraint) Complete(ctx context.Context, conn db.DB, tr SQLTra
 
 	// rename new columns to old name
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return TableDoesNotExistError{Name: o.Table}
+	}
 	for _, col := range o.Columns {
 		column := table.GetColumn(col)
+		if column == nil {
+			return ColumnDoesNotExistError{Table: o.Table, Name: col}
+		}
 		if err := RenameDuplicatedColumn(ctx, conn, table, column); err != nil {
 			return err
 		}
@@ -149,6 +162,9 @@ func (o *OpCreateConstraint) Complete(ctx context.Context, conn db.DB, tr SQLTra
 
 func (o *OpCreateConstraint) Rollback(ctx context.Context, conn db.DB, tr SQLTransformer, s *schema.Schema) error {
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return TableDoesNotExistError{Name: o.Table}
+	}
 
 	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s %s",
 		pq.QuoteIdentifier(table.Name),
diff --git a/pkg/migrations/op_create_index.go b/pkg/migrations/op_create_index.go
@@ -17,6 +17,9 @@ var _ Operation = (*OpCreateIndex)(nil)
 
 func (o *OpCreateIndex) Start(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, s *schema.Schema) (*schema.Table, error) {
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return nil, TableDoesNotExistError{Name: o.Table}
+	}
 
 	// create index concurrently
 	stmtFmt := "CREATE INDEX CONCURRENTLY %s ON %s"
diff --git a/pkg/migrations/op_drop_column.go b/pkg/migrations/op_drop_column.go
@@ -30,6 +30,15 @@ func (o *OpDropColumn) Start(ctx context.Context, conn db.DB, latestSchema strin
 		}
 	}
 
+	table := s.GetTable(o.Table)
+	if table == nil {
+		return nil, TableDoesNotExistError{Name: o.Table}
+	}
+	column := table.GetColumn(o.Column)
+	if column == nil {
+		return nil, ColumnDoesNotExistError{Table: o.Table, Name: o.Column}
+	}
+
 	s.GetTable(o.Table).RemoveColumn(o.Column)
 	return nil, nil
 }
diff --git a/pkg/migrations/op_drop_constraint.go b/pkg/migrations/op_drop_constraint.go
@@ -16,10 +16,16 @@ var _ Operation = (*OpDropConstraint)(nil)
 
 func (o *OpDropConstraint) Start(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, s *schema.Schema) (*schema.Table, error) {
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return nil, TableDoesNotExistError{Name: o.Table}
+	}
 
 	// By this point Validate() should have run which ensures the constraint exists and that we only have
 	// one column associated with it.
 	column := table.GetColumn(table.GetConstraintColumns(o.Name)[0])
+	if column == nil {
+		return nil, ColumnDoesNotExistError{Table: o.Table, Name: table.GetConstraintColumns(o.Name)[0]}
+	}
 
 	// Create a copy of the column on the underlying table.
 	d := NewColumnDuplicator(conn, table, column).WithoutConstraint(o.Name)
diff --git a/pkg/migrations/op_drop_multicolumn_constraint.go b/pkg/migrations/op_drop_multicolumn_constraint.go
@@ -16,12 +16,18 @@ var _ Operation = (*OpDropMultiColumnConstraint)(nil)
 
 func (o *OpDropMultiColumnConstraint) Start(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, s *schema.Schema) (*schema.Table, error) {
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return nil, TableDoesNotExistError{Name: o.Table}
+	}
 
 	// Get all columns covered by the constraint to be dropped
 	constraintColumns := table.GetConstraintColumns(o.Name)
 	columns := make([]*schema.Column, len(constraintColumns))
 	for i, c := range constraintColumns {
 		columns[i] = table.GetColumn(c)
+		if columns[i] == nil {
+			return nil, ColumnDoesNotExistError{Table: o.Table, Name: c}
+		}
 	}
 
 	// Duplicate each of the columns covered by the constraint to be dropped
diff --git a/pkg/migrations/op_drop_not_null.go b/pkg/migrations/op_drop_not_null.go
@@ -21,6 +21,9 @@ var _ Operation = (*OpDropNotNull)(nil)
 
 func (o *OpDropNotNull) Start(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, s *schema.Schema) (*schema.Table, error) {
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return nil, TableDoesNotExistError{Name: o.Table}
+	}
 
 	return table, nil
 }
diff --git a/pkg/migrations/op_drop_table.go b/pkg/migrations/op_drop_table.go
@@ -15,6 +15,9 @@ var _ Operation = (*OpDropTable)(nil)
 
 func (o *OpDropTable) Start(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, s *schema.Schema) (*schema.Table, error) {
 	table := s.GetTable(o.Name)
+	if table == nil {
+		return nil, TableDoesNotExistError{Name: o.Name}
+	}
 
 	// Soft-delete the table in order that a create table operation in the same
 	// migration can create a table with the same name
diff --git a/pkg/migrations/op_rename_column.go b/pkg/migrations/op_rename_column.go
@@ -15,6 +15,13 @@ var _ Operation = (*OpRenameColumn)(nil)
 func (o *OpRenameColumn) Start(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, s *schema.Schema) (*schema.Table, error) {
 	// Rename the table in the in-memory schema.
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return nil, TableDoesNotExistError{Name: o.Table}
+	}
+	column := table.GetColumn(o.From)
+	if column == nil {
+		return nil, ColumnDoesNotExistError{Table: o.Table, Name: o.From}
+	}
 	table.RenameColumn(o.From, o.To)
 
 	return nil, nil
diff --git a/pkg/migrations/op_set_check.go b/pkg/migrations/op_set_check.go
@@ -25,6 +25,9 @@ var _ Operation = (*OpSetCheckConstraint)(nil)
 
 func (o *OpSetCheckConstraint) Start(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, s *schema.Schema) (*schema.Table, error) {
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return nil, TableDoesNotExistError{Name: o.Table}
+	}
 
 	// Add the check constraint to the new column as NOT VALID.
 	if err := o.addCheckConstraint(ctx, conn, s); err != nil {
diff --git a/pkg/migrations/op_set_comment.go b/pkg/migrations/op_set_comment.go
@@ -22,6 +22,9 @@ var _ Operation = (*OpSetComment)(nil)
 
 func (o *OpSetComment) Start(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, s *schema.Schema) (*schema.Table, error) {
 	tbl := s.GetTable(o.Table)
+	if tbl == nil {
+		return nil, TableDoesNotExistError{Name: o.Table}
+	}
 
 	return tbl, addCommentToColumn(ctx, conn, o.Table, TemporaryName(o.Column), o.Comment)
 }
diff --git a/pkg/migrations/op_set_default.go b/pkg/migrations/op_set_default.go
@@ -24,7 +24,13 @@ var _ Operation = (*OpSetDefault)(nil)
 
 func (o *OpSetDefault) Start(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, s *schema.Schema) (*schema.Table, error) {
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return nil, TableDoesNotExistError{Name: o.Table}
+	}
 	column := table.GetColumn(o.Column)
+	if column == nil {
+		return nil, ColumnDoesNotExistError{Table: o.Table, Name: o.Column}
+	}
 
 	var err error
 	if o.Default == nil {
diff --git a/pkg/migrations/op_set_fk.go b/pkg/migrations/op_set_fk.go
@@ -83,9 +83,22 @@ func (o *OpSetForeignKey) Validate(ctx context.Context, s *schema.Schema) error
 
 func (o *OpSetForeignKey) addForeignKeyConstraint(ctx context.Context, conn db.DB, s *schema.Schema) error {
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return TableDoesNotExistError{Name: o.Table}
+	}
 	column := table.GetColumn(o.Column)
+	if column == nil {
+		return ColumnDoesNotExistError{Table: o.Table, Name: o.Column}
+	}
 	referencedTable := s.GetTable(o.References.Table)
+	if referencedTable == nil {
+		return TableDoesNotExistError{Name: o.References.Table}
+	}
+
 	referencedColumn := referencedTable.GetColumn(o.References.Column)
+	if referencedColumn == nil {
+		return ColumnDoesNotExistError{Table: o.References.Table, Name: o.References.Column}
+	}
 
 	sql := fmt.Sprintf("ALTER TABLE %s ADD ", pq.QuoteIdentifier(table.Name))
 	writer := &ConstraintSQLWriter{
diff --git a/pkg/migrations/op_set_notnull.go b/pkg/migrations/op_set_notnull.go
@@ -22,7 +22,13 @@ var _ Operation = (*OpSetNotNull)(nil)
 
 func (o *OpSetNotNull) Start(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, s *schema.Schema) (*schema.Table, error) {
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return nil, TableDoesNotExistError{Name: o.Table}
+	}
 	column := table.GetColumn(o.Column)
+	if column == nil {
+		return nil, ColumnDoesNotExistError{Table: o.Table, Name: o.Column}
+	}
 
 	// Add an unchecked NOT NULL constraint to the new column.
 	if err := addNotNullConstraint(ctx, conn, table.Name, o.Column, column.Name); err != nil {
diff --git a/pkg/migrations/op_set_unique.go b/pkg/migrations/op_set_unique.go
@@ -23,7 +23,13 @@ var _ Operation = (*OpSetUnique)(nil)
 
 func (o *OpSetUnique) Start(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, s *schema.Schema) (*schema.Table, error) {
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return nil, TableDoesNotExistError{Name: o.Table}
+	}
 	column := table.GetColumn(o.Column)
+	if column == nil {
+		return nil, ColumnDoesNotExistError{Table: o.Table, Name: o.Column}
+	}
 
 	return table, createUniqueIndexConcurrently(ctx, conn, s.Name, o.Name, table.Name, []string{column.Name})
 }

Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,9 @@ var _ Operation = (*OpChangeType)(nil)`
`21`	`21`
`22`	`22`	`func (o OpChangeType) Start(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, s schema.Schema) (*schema.Table, error) {`
`23`	`23`	`table := s.GetTable(o.Table)`
	`24`	`+ if table == nil {`
	`25`	`+ return nil, TableDoesNotExistError{Name: o.Table}`
	`26`	`+ }`
`24`	`27`
`25`	`28`	`return table, nil`
`26`	`29`	`}`
Original file line number	Diff line number	Diff line change
`@@ -30,6 +30,15 @@ func (o *OpDropColumn) Start(ctx context.Context, conn db.DB, latestSchema strin`
`30`	`30`	`}`
`31`	`31`	`}`
`32`	`32`
	`33`	`+ table := s.GetTable(o.Table)`
	`34`	`+ if table == nil {`
	`35`	`+ return nil, TableDoesNotExistError{Name: o.Table}`
	`36`	`+ }`
	`37`	`+ column := table.GetColumn(o.Column)`
	`38`	`+ if column == nil {`
	`39`	`+ return nil, ColumnDoesNotExistError{Table: o.Table, Name: o.Column}`
	`40`	`+ }`
	`41`	`+`
`33`	`42`	`s.GetTable(o.Table).RemoveColumn(o.Column)`
`34`	`43`	`return nil, nil`
`35`	`44`	`}`