feat(psql version): add version management for MERGE statement with automatic RETURNING support

Author: Roman A. Grigorovich

dialect/psql/merge_test.go

@@ -1,8 +1,10 @@
 package psql_test
 
 import (
+	"context"
 	"testing"
 
+	"github.com/stephenafamo/bob"
 	"github.com/stephenafamo/bob/dialect/psql"
 	"github.com/stephenafamo/bob/dialect/psql/mm"
 	"github.com/stephenafamo/bob/dialect/psql/sm"
@@ -505,3 +507,149 @@ func TestMerge(t *testing.T) {
 
 	testutils.RunTests(t, examples, formatter)
 }
+
+func TestMergeWithVersion(t *testing.T) {
+	t.Run("version 17+ adds RETURNING automatically with mm.Returning", func(t *testing.T) {
+		ctx := context.Background()
+		ctx = psql.SetVersion(ctx, 17)
+
+		q := psql.Merge(
+			mm.Into("target"),
+			mm.Using("source").As("s").On(
+				psql.Quote("s", "id").EQ(psql.Quote("target", "id")),
+			),
+			mm.WhenMatched(
+				mm.ThenUpdate(
+					mm.SetCol("name").ToExpr(psql.Quote("s", "name")),
+				),
+			),
+			mm.Returning("*"),
+		)
+
+		sql, args, err := bob.Build(ctx, q)
+		if err != nil {
+			t.Fatalf("error: %v", err)
+		}
+
+		expectedSQL := `MERGE INTO target USING source AS "s" ON "s"."id" = "target"."id" WHEN MATCHED THEN UPDATE SET "name" = "s"."name" RETURNING *`
+		diff, err := testutils.QueryDiff(expectedSQL, sql, formatter)
+		if err != nil {
+			t.Fatalf("error: %v", err)
+		}
+		if diff != "" {
+			t.Errorf("SQL mismatch:\n%s\nGot: %s", diff, sql)
+		}
+		if len(args) != 0 {
+			t.Errorf("expected no args, got %v", args)
+		}
+	})
+
+	t.Run("version 16 does not affect MERGE with explicit RETURNING", func(t *testing.T) {
+		ctx := context.Background()
+		ctx = psql.SetVersion(ctx, 16)
+
+		q := psql.Merge(
+			mm.Into("target"),
+			mm.Using("source").As("s").On(
+				psql.Quote("s", "id").EQ(psql.Quote("target", "id")),
+			),
+			mm.WhenMatched(
+				mm.ThenUpdate(
+					mm.SetCol("name").ToExpr(psql.Quote("s", "name")),
+				),
+			),
+			mm.Returning("*"),
+		)
+
+		sql, args, err := bob.Build(ctx, q)
+		if err != nil {
+			t.Fatalf("error: %v", err)
+		}
+
+		// RETURNING should still be present because it was explicitly added
+		expectedSQL := `MERGE INTO target USING source AS "s" ON "s"."id" = "target"."id" WHEN MATCHED THEN UPDATE SET "name" = "s"."name" RETURNING *`
+		diff, err := testutils.QueryDiff(expectedSQL, sql, formatter)
+		if err != nil {
+			t.Fatalf("error: %v", err)
+		}
+		if diff != "" {
+			t.Errorf("SQL mismatch:\n%s\nGot: %s", diff, sql)
+		}
+		if len(args) != 0 {
+			t.Errorf("expected no args, got %v", args)
+		}
+	})
+
+	t.Run("no version set - MERGE without RETURNING", func(t *testing.T) {
+		ctx := context.Background()
+		// No version set
+
+		q := psql.Merge(
+			mm.Into("target"),
+			mm.Using("source").As("s").On(
+				psql.Quote("s", "id").EQ(psql.Quote("target", "id")),
+			),
+			mm.WhenMatched(
+				mm.ThenUpdate(
+					mm.SetCol("name").ToExpr(psql.Quote("s", "name")),
+				),
+			),
+		)
+
+		sql, args, err := bob.Build(ctx, q)
+		if err != nil {
+			t.Fatalf("error: %v", err)
+		}
+
+		// No RETURNING because no version set
+		expectedSQL := `MERGE INTO target USING source AS "s" ON "s"."id" = "target"."id" WHEN MATCHED THEN UPDATE SET "name" = "s"."name"`
+		diff, err := testutils.QueryDiff(expectedSQL, sql, formatter)
+		if err != nil {
+			t.Fatalf("error: %v", err)
+		}
+		if diff != "" {
+			t.Errorf("SQL mismatch:\n%s\nGot: %s", diff, sql)
+		}
+		if len(args) != 0 {
+			t.Errorf("expected no args, got %v", args)
+		}
+	})
+
+	t.Run("version 17+ with WhenNotMatchedBySource (PG17 feature)", func(t *testing.T) {
+		ctx := context.Background()
+		ctx = psql.SetVersion(ctx, 17)
+
+		q := psql.Merge(
+			mm.Into("target"),
+			mm.Using("source").As("s").On(
+				psql.Quote("s", "id").EQ(psql.Quote("target", "id")),
+			),
+			mm.WhenMatched(
+				mm.ThenUpdate(
+					mm.SetCol("name").ToExpr(psql.Quote("s", "name")),
+				),
+			),
+			mm.WhenNotMatchedBySource(
+				mm.ThenDelete(),
+			),
+			mm.Returning(psql.Quote("target", "id")),
+		)
+
+		sql, args, err := bob.Build(ctx, q)
+		if err != nil {
+			t.Fatalf("error: %v", err)
+		}
+
+		expectedSQL := `MERGE INTO target USING source AS "s" ON "s"."id" = "target"."id" WHEN MATCHED THEN UPDATE SET "name" = "s"."name" WHEN NOT MATCHED BY SOURCE THEN DELETE RETURNING "target"."id"`
+		diff, err := testutils.QueryDiff(expectedSQL, sql, formatter)
+		if err != nil {
+			t.Fatalf("error: %v", err)
+		}
+		if diff != "" {
+			t.Errorf("SQL mismatch:\n%s\nGot: %s", diff, sql)
+		}
+		if len(args) != 0 {
+			t.Errorf("expected no args, got %v", args)
+		}
+	})
+}

dialect/psql/table.go

@@ -146,8 +146,9 @@ func (t *Table[T, Tslice, Tset, C]) Delete(queryMods ...bob.Mod[*dialect.DeleteQ
 
 // Starts a Merge query for this table
 // The caller must provide USING and WHEN clauses via queryMods
-// Note: RETURNING clause is NOT added automatically because it requires PostgreSQL 17+
-// Use mm.Returning() explicitly if needed and you are on PostgreSQL 17+
+// RETURNING clause is automatically added if version >= 17 is set in context.
+// Use psql.SetVersion(ctx, 17) to enable automatic RETURNING for MERGE.
+// For older versions, use mm.Returning() explicitly if needed.
 func (t *Table[T, Tslice, Tset, C]) Merge(queryMods ...bob.Mod[*dialect.MergeQuery]) *ormMergeQuery[T, Tslice] {
 	q := &ormMergeQuery[T, Tslice]{
 		ExecQuery: orm.ExecQuery[*dialect.MergeQuery]{
@@ -157,6 +158,16 @@ func (t *Table[T, Tslice, Tset, C]) Merge(queryMods ...bob.Mod[*dialect.MergeQue
 		Scanner: t.scanner,
 	}
 
+	q.Expression.AppendContextualModFunc(
+		func(ctx context.Context, q *dialect.MergeQuery) (context.Context, error) {
+			// RETURNING in MERGE requires version 17+
+			if VersionAtLeast(ctx, 17) && !q.HasReturning() {
+				q.AppendReturning(t.Columns)
+			}
+			return ctx, nil
+		},
+	)
+
 	q.Apply(queryMods...)
 
 	return q

dialect/psql/version.go

@@ -0,0 +1,32 @@
+package psql
+
+import "context"
+
+// VersionKey is a context key for storing the database version.
+// Use SetVersion to set the version in context.
+type VersionKey struct{}
+
+// SetVersion sets the major version (e.g., 15, 16, 17) in the context.
+// This is used to enable version-specific features like MERGE with RETURNING (version 17+).
+//
+// Example:
+//
+//	ctx := psql.SetVersion(ctx, 17)
+func SetVersion(ctx context.Context, version int) context.Context {
+	return context.WithValue(ctx, VersionKey{}, version)
+}
+
+// GetVersion returns the major version from the context.
+// Returns 0 if the version is not set.
+func GetVersion(ctx context.Context) int {
+	if v, ok := ctx.Value(VersionKey{}).(int); ok {
+		return v
+	}
+	return 0
+}
+
+// VersionAtLeast checks if the version in context is at least the given version.
+// Returns false if version is not set in context.
+func VersionAtLeast(ctx context.Context, minVersion int) bool {
+	return GetVersion(ctx) >= minVersion
+}