feat: add task queue (#9321)

# Which Problems Are Solved

To integrate river as a task queue we need to ensure the migrations of
river are executed.

# How the Problems Are Solved

- A new schema was added to the Zitadel database called "queue"
- Added a repeatable setup step to Zitadel which executes the
[migrations of
river](https://riverqueue.com/docs/migrations#go-migration-api).

# Additional Changes

- Added more hooks to the databases to properly set the schema for the
task queue

# Additional Context

- Closes https://github.com/zitadel/zitadel/issues/9280

internal/database/cockroach/crdb.go

@@ -97,6 +97,27 @@ func (c *Config) Connect(useAdmin bool) (*sql.DB, *pgxpool.Pool, error) {
 		}
 	}
 
+	if len(connConfig.BeforeAcquire) > 0 {
+		config.BeforeAcquire = func(ctx context.Context, conn *pgx.Conn) bool {
+			for _, f := range connConfig.BeforeAcquire {
+				if err := f(ctx, conn); err != nil {
+					return false
+				}
+			}
+			return true
+		}
+	}
+	if len(connConfig.AfterRelease) > 0 {
+		config.AfterRelease = func(conn *pgx.Conn) bool {
+			for _, f := range connConfig.AfterRelease {
+				if err := f(conn); err != nil {
+					return false
+				}
+			}
+			return true
+		}
+	}
+
 	if connConfig.MaxOpenConns != 0 {
 		config.MaxConns = int32(connConfig.MaxOpenConns)
 	}

internal/database/dialect/connections.go

@@ -18,7 +18,9 @@ var (
 type ConnectionConfig struct {
 	MaxOpenConns,
 	MaxIdleConns uint32
-	AfterConnect []func(ctx context.Context, c *pgx.Conn) error
+	AfterConnect  []func(ctx context.Context, c *pgx.Conn) error
+	BeforeAcquire []func(ctx context.Context, c *pgx.Conn) error
+	AfterRelease  []func(c *pgx.Conn) error
 }
 
 var afterConnectFuncs []func(ctx context.Context, c *pgx.Conn) error
@@ -27,6 +29,18 @@ func RegisterAfterConnect(f func(ctx context.Context, c *pgx.Conn) error) {
 	afterConnectFuncs = append(afterConnectFuncs, f)
 }
 
+var beforeAcquireFuncs []func(ctx context.Context, c *pgx.Conn) error
+
+func RegisterBeforeAcquire(f func(ctx context.Context, c *pgx.Conn) error) {
+	beforeAcquireFuncs = append(beforeAcquireFuncs, f)
+}
+
+var afterReleaseFuncs []func(c *pgx.Conn) error
+
+func RegisterAfterRelease(f func(c *pgx.Conn) error) {
+	afterReleaseFuncs = append(afterReleaseFuncs, f)
+}
+
 func RegisterDefaultPgTypeVariants[T any](m *pgtype.Map, name, arrayName string) {
 	// T
 	var value T
@@ -58,8 +72,10 @@ func RegisterDefaultPgTypeVariants[T any](m *pgtype.Map, name, arrayName string)
 // The pusherRatio and spoolerRatio must be between 0 and 1.
 func NewConnectionConfig(openConns, idleConns uint32) *ConnectionConfig {
 	return &ConnectionConfig{
-		MaxOpenConns: openConns,
-		MaxIdleConns: idleConns,
-		AfterConnect: afterConnectFuncs,
+		MaxOpenConns:  openConns,
+		MaxIdleConns:  idleConns,
+		AfterConnect:  afterConnectFuncs,
+		BeforeAcquire: beforeAcquireFuncs,
+		AfterRelease:  afterReleaseFuncs,
 	}
 }

internal/database/postgres/pg.go

@@ -81,13 +81,37 @@ func (c *Config) Connect(useAdmin bool) (*sql.DB, *pgxpool.Pool, error) {
 		return nil, nil, err
 	}
 
-	config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
-		for _, f := range connConfig.AfterConnect {
-			if err := f(ctx, conn); err != nil {
-				return err
+	if len(connConfig.AfterConnect) > 0 {
+		config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
+			for _, f := range connConfig.AfterConnect {
+				if err := f(ctx, conn); err != nil {
+					return err
+				}
 			}
+			return nil
+		}
+	}
+
+	if len(connConfig.BeforeAcquire) > 0 {
+		config.BeforeAcquire = func(ctx context.Context, conn *pgx.Conn) bool {
+			for _, f := range connConfig.BeforeAcquire {
+				if err := f(ctx, conn); err != nil {
+					return false
+				}
+			}
+			return true
+		}
+	}
+
+	if len(connConfig.AfterRelease) > 0 {
+		config.AfterRelease = func(conn *pgx.Conn) bool {
+			for _, f := range connConfig.AfterRelease {
+				if err := f(conn); err != nil {
+					return false
+				}
+			}
+			return true
 		}
-		return nil
 	}
 
 	if connConfig.MaxOpenConns != 0 {

internal/queue/queue.go

@@ -0,0 +1,75 @@
+package queue
+
+import (
+	"context"
+	"sync"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/riverqueue/river/riverdriver"
+	"github.com/riverqueue/river/riverdriver/riverpgxv5"
+	"github.com/riverqueue/river/rivermigrate"
+
+	"github.com/zitadel/zitadel/internal/database"
+	"github.com/zitadel/zitadel/internal/database/dialect"
+)
+
+const (
+	schema          = "queue"
+	applicationName = "zitadel_queue"
+)
+
+var conns = &sync.Map{}
+
+type queueKey struct{}
+
+func WithQueue(parent context.Context) context.Context {
+	return context.WithValue(parent, queueKey{}, struct{}{})
+}
+
+func init() {
+	dialect.RegisterBeforeAcquire(func(ctx context.Context, c *pgx.Conn) error {
+		if _, ok := ctx.Value(queueKey{}).(struct{}); !ok {
+			return nil
+		}
+		_, err := c.Exec(ctx, "SET search_path TO "+schema+"; SET application_name TO "+applicationName)
+		if err != nil {
+			return err
+		}
+		conns.Store(c, struct{}{})
+		return nil
+	})
+	dialect.RegisterAfterRelease(func(c *pgx.Conn) error {
+		_, ok := conns.LoadAndDelete(c)
+		if !ok {
+			return nil
+		}
+		_, err := c.Exec(context.Background(), "SET search_path TO DEFAULT; SET application_name TO "+dialect.DefaultAppName)
+		return err
+	})
+}
+
+// Queue abstracts the underlying queuing library
+// For more information see github.com/riverqueue/river
+// TODO(adlerhurst): maybe it makes more sense to split the effective queue from the migrator.
+type Queue struct {
+	driver riverdriver.Driver[pgx.Tx]
+}
+
+func New(client *database.DB) *Queue {
+	return &Queue{driver: riverpgxv5.New(client.Pool)}
+}
+
+func (q *Queue) ExecuteMigrations(ctx context.Context) error {
+	_, err := q.driver.GetExecutor().Exec(ctx, "CREATE SCHEMA IF NOT EXISTS "+schema)
+	if err != nil {
+		return err
+	}
+
+	migrator, err := rivermigrate.New(q.driver, nil)
+	if err != nil {
+		return err
+	}
+	ctx = WithQueue(ctx)
+	_, err = migrator.Migrate(ctx, rivermigrate.DirectionUp, nil)
+	return err
+}