Convert DROP COLUMN SQL to pgroll operation (#521)

Part of #504

pkg/sql2pgroll/alter_table.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 
 	pgq "github.com/pganalyze/pg_query_go/v6"
+
 	"github.com/xataio/pgroll/pkg/migrations"
 )
 
@@ -35,6 +36,8 @@ func convertAlterTableStmt(stmt *pgq.AlterTableStmt) (migrations.Operations, err
 			op, err = convertAlterTableAlterColumnType(stmt, alterTableCmd)
 		case pgq.AlterTableType_AT_AddConstraint:
 			op, err = convertAlterTableAddConstraint(stmt, alterTableCmd)
+		case pgq.AlterTableType_AT_DropColumn:
+			op, err = convertAlterTableDropColumn(stmt, alterTableCmd)
 		}
 
 		if err != nil {
@@ -155,6 +158,34 @@ func convertAlterTableAddUniqueConstraint(stmt *pgq.AlterTableStmt, constraint *
 	}, nil
 }
 
+// convertAlterTableDropColumn converts SQL statements like:
+//
+// `ALTER TABLE foo DROP COLUMN bar
+//
+// to an OpDropColumn operation.
+func convertAlterTableDropColumn(stmt *pgq.AlterTableStmt, cmd *pgq.AlterTableCmd) (migrations.Operation, error) {
+	if !canConvertDropColumn(cmd) {
+		return nil, nil
+	}
+
+	return &migrations.OpDropColumn{
+		Table:  stmt.GetRelation().GetRelname(),
+		Column: cmd.GetName(),
+		Down:   PlaceHolderSQL,
+	}, nil
+}
+
+// canConvertDropColumn checks whether we can convert the command without losing any information.
+func canConvertDropColumn(cmd *pgq.AlterTableCmd) bool {
+	if cmd.MissingOk {
+		return false
+	}
+	if cmd.Behavior == pgq.DropBehavior_DROP_CASCADE {
+		return false
+	}
+	return true
+}
+
 // canConvertUniqueConstraint checks if the unique constraint `constraint` can
 // be faithfully converted to an OpCreateConstraint operation without losing
 // information.

pkg/sql2pgroll/alter_table_test.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
 	"github.com/xataio/pgroll/pkg/migrations"
 	"github.com/xataio/pgroll/pkg/sql2pgroll"
 	"github.com/xataio/pgroll/pkg/sql2pgroll/expect"
@@ -43,6 +44,14 @@ func TestConvertAlterTableStatements(t *testing.T) {
 			sql:        "ALTER TABLE foo ADD CONSTRAINT bar UNIQUE (a, b)",
 			expectedOp: expect.CreateConstraintOp2,
 		},
+		{
+			sql:        "ALTER TABLE foo DROP COLUMN bar",
+			expectedOp: expect.DropColumnOp1,
+		},
+		{
+			sql:        "ALTER TABLE foo DROP COLUMN bar RESTRICT ",
+			expectedOp: expect.DropColumnOp1,
+		},
 	}
 
 	for _, tc := range tests {
@@ -57,7 +66,7 @@ func TestConvertAlterTableStatements(t *testing.T) {
 	}
 }
 
-func TestUnconvertableAlterTableAddConstraintStatements(t *testing.T) {
+func TestUnconvertableAlterTableStatements(t *testing.T) {
 	t.Parallel()
 
 	tests := []string{
@@ -72,6 +81,10 @@ func TestUnconvertableAlterTableAddConstraintStatements(t *testing.T) {
 		// operations when changing data type.
 		`ALTER TABLE foo ALTER COLUMN a SET DATA TYPE text COLLATE "en_US"`,
 		"ALTER TABLE foo ALTER COLUMN a SET DATA TYPE text USING 'foo'",
+
+		// CASCADE and IF EXISTS clauses are not represented by OpDropColumn
+		"ALTER TABLE foo DROP COLUMN bar CASCADE",
+		"ALTER TABLE foo DROP COLUMN IF EXISTS bar",
 	}
 
 	for _, sql := range tests {

pkg/sql2pgroll/convert.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 
 	pgq "github.com/pganalyze/pg_query_go/v6"
+
 	"github.com/xataio/pgroll/pkg/migrations"
 )
 

pkg/sql2pgroll/create_table.go

@@ -4,6 +4,7 @@ package sql2pgroll
 
 import (
 	pgq "github.com/pganalyze/pg_query_go/v6"
+
 	"github.com/xataio/pgroll/pkg/migrations"
 )
 

pkg/sql2pgroll/create_table_test.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
 	"github.com/xataio/pgroll/pkg/migrations"
 	"github.com/xataio/pgroll/pkg/sql2pgroll"
 	"github.com/xataio/pgroll/pkg/sql2pgroll/expect"

pkg/sql2pgroll/expect/drop_column.go

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package expect
+
+import (
+	"github.com/xataio/pgroll/pkg/migrations"
+	"github.com/xataio/pgroll/pkg/sql2pgroll"
+)
+
+var DropColumnOp1 = &migrations.OpDropColumn{
+	Table:  "foo",
+	Column: "bar",
+	Down:   sql2pgroll.PlaceHolderSQL,
+}

pkg/sql2pgroll/rename.go

@@ -4,6 +4,7 @@ package sql2pgroll
 
 import (
 	pgq "github.com/pganalyze/pg_query_go/v6"
+
 	"github.com/xataio/pgroll/pkg/migrations"
 )
 

pkg/sql2pgroll/rename_test.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
 	"github.com/xataio/pgroll/pkg/migrations"
 	"github.com/xataio/pgroll/pkg/sql2pgroll"
 	"github.com/xataio/pgroll/pkg/sql2pgroll/expect"