mirror of
https://github.com/zitadel/zitadel.git
synced 2025-03-01 03:07:22 +00:00
415bc32ed6
# 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
76 lines
1.8 KiB
Go
76 lines
1.8 KiB
Go
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
|
|
}
|