Improve diff output

cmd/generate.go

@@ -584,6 +584,7 @@ func executeTargetSQL(ctx context.Context, config *internal.Config, wd string, t
 // for the given target and migration connections. This feature is not yet
 // available in pg-schema-diff, but planned.
 //
+// nolint:godox
 // TODO: This function should probably be moved to the internal package.
 func generateMissingPermissionStatements(
 	ctx context.Context,

internal/diff.go

@@ -16,19 +16,21 @@ import (
 
 const diffMaxOpenConns = 100
 
-func Diff(ctx context.Context, postgresConn, targetConn, migrateConn *pgx.Conn) (string, error) {
-	fromDB := stdlib.OpenDB(*targetConn.Config())
+// Diff generates a SQL script to migrate the schema from the 'from' database to match the 'to' database.
+// nolint:gocognit,cyclop
+func Diff(ctx context.Context, postgresConn, fromConn, toConn *pgx.Conn) (string, error) {
+	fromDB := stdlib.OpenDB(*fromConn.Config())
 	fromDB.SetMaxOpenConns(diffMaxOpenConns)
 	defer fromDB.Close()
-	err := fromDB.Ping()
+	err := fromDB.PingContext(ctx)
 	if err != nil {
 		return "", fmt.Errorf("failed to open 'from database' connection: %w", err)
 	}
 
-	toDB := stdlib.OpenDB(*migrateConn.Config())
+	toDB := stdlib.OpenDB(*toConn.Config())
 	toDB.SetMaxOpenConns(diffMaxOpenConns)
 	defer toDB.Close()
-	err = toDB.Ping()
+	err = toDB.PingContext(ctx)
 	if err != nil {
 		return "", fmt.Errorf("failed to open 'to database' connection: %w", err)
 	}
@@ -37,8 +39,9 @@ func Diff(ctx context.Context, postgresConn, targetConn, migrateConn *pgx.Conn)
 	defer func() {
 		// call close funcs in inverted order (last first)
 		for i := len(deferredCloseFuncs) - 1; i >= 0; i-- {
-			if err := deferredCloseFuncs[i](); err != nil {
-				log.Printf("failed to close resource: %v", err)
+			errClose := deferredCloseFuncs[i]()
+			if errClose != nil {
+				log.Printf("failed to close resource: %v", errClose)
 			}
 		}
 	}()
@@ -47,10 +50,11 @@ func Diff(ctx context.Context, postgresConn, targetConn, migrateConn *pgx.Conn)
 		config.Database = dbName
 		tempDB := stdlib.OpenDB(*config)
 		tempDB.SetMaxOpenConns(diffMaxOpenConns)
-		err := tempDB.Ping()
-		if err != nil {
+		errPing := tempDB.PingContext(ctx)
+		if errPing != nil {
 			tempDB.Close()
-			return nil, fmt.Errorf("failed to connect to temp database: %w", err)
+
+			return nil, fmt.Errorf("failed to connect to temp database: %w", errPing)
 		}
 		deferredCloseFuncs = append(deferredCloseFuncs, tempDB.Close)
 
@@ -64,59 +68,50 @@ func Diff(ctx context.Context, postgresConn, targetConn, migrateConn *pgx.Conn)
 		diff.DBSchemaSource(fromDB),
 		diff.DBSchemaSource(toDB),
 		diff.WithTempDbFactory(tempFactory), // Required to validate the generated diff statements.
-		diff.WithNoConcurrentIndexOps(),     // Concurrent index creation is not available in transactions that are used by go-migrate.
+		diff.WithNoConcurrentIndexOps(),     // Concurrent index creation is not available in transactions.
 	)
 	if err != nil {
 		return "", fmt.Errorf("failed to generate diff plan: %w", err)
 	}
 
 	// Ignore the proposed deletion of the schema_migrations table.
 	plan.Statements = slices.DeleteFunc(plan.Statements, func(statement diff.Statement) bool {
-		if statement.DDL == `DROP TABLE "public"."schema_migrations"` {
-			return true
-		}
-		return false
+		return statement.DDL == `DROP TABLE "public"."schema_migrations"`
 	})
 
-	return planToSql(plan), nil
-}
-
-// planToSql converts the plan to one large runnable SQL script.
-//
-// This function is copied verbatim from the pg-schema-diff cmd (package main)
-// https://github.com/stripe/pg-schema-diff/blob/v1.0.2/cmd/pg-schema-diff/plan_cmd.go#L607C1-L628C2
-// TODO: We could open an issue to ask them to expose this function, although it
-// could be beneficial for trek to write our own custom sql output format.
-// Licensed MIT by pg-schema-diff
-func planToSql(plan diff.Plan) string {
 	sb := strings.Builder{}
+	var lastStatementTimeout int64
+	var lastLockTimeout int64
 	for i, stmt := range plan.Statements {
-		sb.WriteString("/*\n")
-		sb.WriteString(fmt.Sprintf("Statement %d\n", i))
+		statementTimeout := stmt.Timeout.Milliseconds()
+		lockTimeout := stmt.LockTimeout.Milliseconds()
+		if lastStatementTimeout != statementTimeout || lastLockTimeout != lockTimeout {
+			if lastStatementTimeout != statementTimeout {
+				lastStatementTimeout = statementTimeout
+				sb.WriteString(fmt.Sprintf("SET SESSION statement_timeout = %d;\n", statementTimeout))
+			}
+			if lastLockTimeout != lockTimeout {
+				lastLockTimeout = lockTimeout
+				sb.WriteString(fmt.Sprintf("SET SESSION lock_timeout = %d;\n", lockTimeout))
+			}
+			sb.WriteString("\n")
+		}
 		if len(stmt.Hazards) > 0 {
+			sb.WriteString("/* Hazards:\n")
 			for _, hazard := range stmt.Hazards {
-				sb.WriteString(fmt.Sprintf("  - %s\n", hazardToPrettyS(hazard)))
+				if hazard.Message != "" {
+					sb.WriteString(fmt.Sprintf(" - %s: %s\n", hazard.Type, hazard.Message))
+				} else {
+					sb.WriteString(fmt.Sprintf(" - %s\n", hazard.Type))
+				}
 			}
+			sb.WriteString("*/\n")
 		}
-		sb.WriteString("*/\n")
-		sb.WriteString(fmt.Sprintf("SET SESSION statement_timeout = %d;\n", stmt.Timeout.Milliseconds()))
-		sb.WriteString(fmt.Sprintf("SET SESSION lock_timeout = %d;\n", stmt.LockTimeout.Milliseconds()))
-		sb.WriteString(fmt.Sprintf("%s;", stmt.DDL))
+		sb.WriteString(fmt.Sprintf("%s;\n", stmt.DDL))
 		if i < len(plan.Statements)-1 {
-			sb.WriteString("\n\n")
+			sb.WriteString("\n")
 		}
 	}
-	return sb.String()
-}
 
-// hazardToPrettyS converts a migration hazard to a pretty string.
-//
-// Copied verbatim from the pg-schema-diff cmd, together with planToSql.
-// Licensed MIT by pg-schema-diff
-func hazardToPrettyS(hazard diff.MigrationHazard) string {
-	if len(hazard.Message) > 0 {
-		return fmt.Sprintf("%s: %s", hazard.Type, hazard.Message)
-	} else {
-		return hazard.Type
-	}
+	return sb.String(), nil
 }

mise.toml

@@ -0,0 +1,2 @@
+[tools]
+golangci-lint = "2.7.2"

tests/output/migrations/002_schemas.up.sql

@@ -1,13 +1,7 @@
-/*
-Statement 0
-*/
 SET SESSION statement_timeout = 3000;
 SET SESSION lock_timeout = 3000;
+
 CREATE SCHEMA "factory";
 
-/*
-Statement 1
-*/
-SET SESSION statement_timeout = 3000;
-SET SESSION lock_timeout = 3000;
 CREATE SCHEMA "warehouse";
+

tests/output/migrations/003_tables.up.sql

@@ -1,21 +1,15 @@
-/*
-Statement 0
-*/
 SET SESSION statement_timeout = 3000;
 SET SESSION lock_timeout = 3000;
+
 CREATE TABLE "factory"."machines" (
 	"name" text COLLATE "pg_catalog"."default" NOT NULL,
 	"toys_produced" bigint NOT NULL
 );
 
-/*
-Statement 1
-*/
-SET SESSION statement_timeout = 3000;
-SET SESSION lock_timeout = 3000;
 CREATE TABLE "warehouse"."storage_locations" (
 	"shelf" bigint NOT NULL,
 	"total_capacity" bigint NOT NULL,
 	"used_capacity" bigint NOT NULL,
 	"current_toy_type" text COLLATE "pg_catalog"."default" NOT NULL
 );
+

tests/output/migrations/004_sequences.up.sql

@@ -1,69 +1,41 @@
-/*
-Statement 0
-  - HAS_UNTRACKABLE_DEPENDENCIES: This sequence has no owner, so it cannot be tracked. It may be in use by a table or function.
-*/
 SET SESSION statement_timeout = 3000;
 SET SESSION lock_timeout = 3000;
+
+/* Hazards:
+ - HAS_UNTRACKABLE_DEPENDENCIES: This sequence has no owner, so it cannot be tracked. It may be in use by a table or function.
+*/
 CREATE SEQUENCE "factory"."seq_machines_id"
 	AS bigint
 	INCREMENT BY 1
 	MINVALUE 0 MAXVALUE 2147483647
 	START WITH 1 CACHE 1 NO CYCLE
 ;
 
-/*
-Statement 1
-  - HAS_UNTRACKABLE_DEPENDENCIES: This sequence has no owner, so it cannot be tracked. It may be in use by a table or function.
+/* Hazards:
+ - HAS_UNTRACKABLE_DEPENDENCIES: This sequence has no owner, so it cannot be tracked. It may be in use by a table or function.
 */
-SET SESSION statement_timeout = 3000;
-SET SESSION lock_timeout = 3000;
 CREATE SEQUENCE "warehouse"."seq_storage_locations_id"
 	AS bigint
 	INCREMENT BY 1
 	MINVALUE 0 MAXVALUE 2147483647
 	START WITH 1 CACHE 1 NO CYCLE
 ;
 
-/*
-Statement 2
-*/
-SET SESSION statement_timeout = 3000;
-SET SESSION lock_timeout = 3000;
 ALTER TABLE "factory"."machines" ADD COLUMN "id" bigint DEFAULT nextval('factory.seq_machines_id'::regclass) NOT NULL;
 
-/*
-Statement 3
-  - ACQUIRES_SHARE_LOCK: Non-concurrent index creates will lock out writes to the table during the duration of the index build.
+/* Hazards:
+ - ACQUIRES_SHARE_LOCK: Non-concurrent index creates will lock out writes to the table during the duration of the index build.
 */
-SET SESSION statement_timeout = 3000;
-SET SESSION lock_timeout = 3000;
 CREATE UNIQUE INDEX machines_pk ON factory.machines USING btree (id);
 
-/*
-Statement 4
-*/
-SET SESSION statement_timeout = 3000;
-SET SESSION lock_timeout = 3000;
 ALTER TABLE "factory"."machines" ADD CONSTRAINT "machines_pk" PRIMARY KEY USING INDEX "machines_pk";
 
-/*
-Statement 5
-*/
-SET SESSION statement_timeout = 3000;
-SET SESSION lock_timeout = 3000;
 ALTER TABLE "warehouse"."storage_locations" ADD COLUMN "id" bigint DEFAULT nextval('warehouse.seq_storage_locations_id'::regclass) NOT NULL;
 
-/*
-Statement 6
-  - ACQUIRES_SHARE_LOCK: Non-concurrent index creates will lock out writes to the table during the duration of the index build.
+/* Hazards:
+ - ACQUIRES_SHARE_LOCK: Non-concurrent index creates will lock out writes to the table during the duration of the index build.
 */
-SET SESSION statement_timeout = 3000;
-SET SESSION lock_timeout = 3000;
 CREATE UNIQUE INDEX storage_locations_pk ON warehouse.storage_locations USING btree (id);
 
-/*
-Statement 7
-*/
-SET SESSION statement_timeout = 3000;
-SET SESSION lock_timeout = 3000;
 ALTER TABLE "warehouse"."storage_locations" ADD CONSTRAINT "storage_locations_pk" PRIMARY KEY USING INDEX "storage_locations_pk";
+

tests/output/migrations/005_checks.up.sql

@@ -1,13 +1,7 @@
-/*
-Statement 0
-*/
 SET SESSION statement_timeout = 3000;
 SET SESSION lock_timeout = 3000;
+
 ALTER TABLE "warehouse"."storage_locations" ADD CONSTRAINT "ck_capacity" CHECK((total_capacity >= used_capacity)) NOT VALID;
 
-/*
-Statement 1
-*/
-SET SESSION statement_timeout = 3000;
-SET SESSION lock_timeout = 3000;
 ALTER TABLE "warehouse"."storage_locations" VALIDATE CONSTRAINT "ck_capacity";
+

tests/output/migrations/006_triggers.up.sql

@@ -1,9 +1,9 @@
-/*
-Statement 0
-  - HAS_UNTRACKABLE_DEPENDENCIES: Dependencies, i.e. other functions used in the function body, of non-sql functions cannot be tracked. As a result, we cannot guarantee that function dependencies are ordered properly relative to this statement. For adds, this means you need to ensure that all functions this function depends on are created/altered before this statement.
-*/
 SET SESSION statement_timeout = 3000;
 SET SESSION lock_timeout = 3000;
+
+/* Hazards:
+ - HAS_UNTRACKABLE_DEPENDENCIES: Dependencies, i.e. other functions used in the function body, of non-sql functions cannot be tracked. As a result, we cannot guarantee that function dependencies are ordered properly relative to this statement. For adds, this means you need to ensure that all functions this function depends on are created/altered before this statement.
+*/
 CREATE OR REPLACE FUNCTION factory.tr_machines_toys_produced_increase()
  RETURNS trigger
  LANGUAGE plpgsql
@@ -17,9 +17,5 @@ END;
 $function$
 ;
 
-/*
-Statement 1
-*/
-SET SESSION statement_timeout = 3000;
-SET SESSION lock_timeout = 3000;
 CREATE TRIGGER toys_produced_increase BEFORE UPDATE OF toys_produced ON factory.machines FOR EACH ROW EXECUTE FUNCTION factory.tr_machines_toys_produced_increase();
+