feat: wrap ALTER-nodes into transaction statements

agneum · agneum · commit f142964dcbbc · 2022-02-15T18:51:58.000+07:00
diff --git a/engine/internal/schema/diff/diff.go b/engine/internal/schema/diff/diff.go
@@ -50,7 +50,6 @@ func detectNodeType(node *pg_query.Node) []*pg_query.Node {
 		DropStmt(stmt)
 
 	case *pg_query.Node_AlterTableStmt:
-		fmt.Println("Alter Type")
 		return AlterStmt(node)
 
 	case *pg_query.Node_SelectStmt:
@@ -88,10 +87,6 @@ func AlterStmt(node *pg_query.Node) []*pg_query.Node {
 	for _, cmd := range initialCommands {
 		switch v := cmd.Node.(type) {
 		case *pg_query.Node_AlterTableCmd:
-			fmt.Printf("%#v\n", v)
-			fmt.Printf("%#v\n", v.AlterTableCmd.Def.Node)
-			fmt.Println(v.AlterTableCmd.Subtype.Enum())
-
 			switch v.AlterTableCmd.Subtype {
 			case pg_query.AlterTableType_AT_AddColumn:
 				def := v.AlterTableCmd.Def.GetColumnDef()
@@ -106,13 +101,13 @@ func AlterStmt(node *pg_query.Node) []*pg_query.Node {
 				if index, ok := constraintsMap[pg_query.ConstrType_CONSTR_DEFAULT]; ok {
 					def.Constraints = make([]*pg_query.Node, 0)
 
-					alterStmts = append(alterStmts, node)
+					alterStmts = append(alterStmts, wrapTransaction([]*pg_query.Node{node})...)
 
 					defaultDefinitionTemp := fmt.Sprintf(`alter table %s alter column %s set default %v;`,
 						alterTableStmt.GetRelation().GetRelname(), def.Colname,
 						constraints[index].GetConstraint().GetRawExpr().GetAConst().GetVal().GetInteger().GetIval())
 
-					alterStmts = append(alterStmts, generateNodes(defaultDefinitionTemp)...)
+					alterStmts = append(alterStmts, wrapTransaction(generateNodes(defaultDefinitionTemp))...)
 
 					// TODO: Update rows
 
@@ -125,12 +120,12 @@ func AlterStmt(node *pg_query.Node) []*pg_query.Node {
 				constraint := v.AlterTableCmd.Def.GetConstraint()
 				constraint.SkipValidation = true
 
-				alterStmts = append(alterStmts, node)
+				alterStmts = append(alterStmts, wrapTransaction([]*pg_query.Node{node})...)
 
-				validationTemp := fmt.Sprintf(`begin; alter table %s validate constraint %s; commit;`,
+				validationTemp := fmt.Sprintf(`alter table %s validate constraint %s;`,
 					alterTableStmt.GetRelation().GetRelname(), constraint.GetConname())
 
-				alterStmts = append(alterStmts, generateNodes(validationTemp)...)
+				alterStmts = append(alterStmts, wrapTransaction(generateNodes(validationTemp))...)
 
 			default:
 				alterStmts = append(alterStmts, node)
@@ -160,3 +155,14 @@ func generateNodes(nodeTemplate string) []*pg_query.Node {
 
 	return nodes
 }
+
+// wrapTransaction wraps nodes into transaction statements.
+func wrapTransaction(nodes []*pg_query.Node) []*pg_query.Node {
+	begin := makeBeginTransactionStmt()
+	commit := makeCommitTransactionStmt()
+
+	nodes = append([]*pg_query.Node{begin}, nodes...)
+	nodes = append(nodes, commit)
+
+	return nodes
+}
diff --git a/engine/internal/schema/diff/diff_test.go b/engine/internal/schema/diff/diff_test.go
@@ -26,14 +26,14 @@ const expectedTpl = `CREATE UNIQUE INDEX CONCURRENTLY title_idx ON films USING b
 
 DROP INDEX CONCURRENTLY title_idx;
 
-ALTER TABLE distributors ADD CONSTRAINT zipchk CHECK (char_length(zipcode) = 5) NOT VALID;
+BEGIN; ALTER TABLE distributors ADD CONSTRAINT zipchk CHECK (char_length(zipcode) = 5) NOT VALID; COMMIT;
 BEGIN; ALTER TABLE distributors VALIDATE CONSTRAINT zipchk; COMMIT;
 
-ALTER TABLE distributors ADD CONSTRAINT distfk FOREIGN KEY (address) REFERENCES addresses (address) NOT VALID;
+BEGIN; ALTER TABLE distributors ADD CONSTRAINT distfk FOREIGN KEY (address) REFERENCES addresses (address) NOT VALID; COMMIT;
 BEGIN; ALTER TABLE distributors VALIDATE CONSTRAINT distfk; COMMIT;
 
-ALTER TABLE pgbench_accounts ADD COLUMN test int;
-ALTER TABLE pgbench_accounts ALTER COLUMN test SET DEFAULT 0`
+BEGIN; ALTER TABLE pgbench_accounts ADD COLUMN test int; COMMIT;
+BEGIN; ALTER TABLE pgbench_accounts ALTER COLUMN test SET DEFAULT 0; COMMIT`
 
 var space = regexp.MustCompile(`\s+`)
 
diff --git a/engine/internal/schema/diff/makefuncs.go b/engine/internal/schema/diff/makefuncs.go
@@ -0,0 +1,25 @@
+package diff
+
+import (
+	pg_query "github.com/pganalyze/pg_query_go/v2"
+)
+
+func makeBeginTransactionStmt() *pg_query.Node {
+	return &pg_query.Node{
+		Node: &pg_query.Node_TransactionStmt{
+			TransactionStmt: &pg_query.TransactionStmt{
+				Kind: pg_query.TransactionStmtKind_TRANS_STMT_BEGIN,
+			},
+		},
+	}
+}
+
+func makeCommitTransactionStmt() *pg_query.Node {
+	return &pg_query.Node{
+		Node: &pg_query.Node_TransactionStmt{
+			TransactionStmt: &pg_query.TransactionStmt{
+				Kind: pg_query.TransactionStmtKind_TRANS_STMT_COMMIT,
+			},
+		},
+	}
+}
