chore/enterpriseportal: propagate context to migration database context (#63455)

Follow-up to #63448 - we now get Redis spans, but not the database
operations that happen throughout a migration. Maybe this will do the
thing?

## Test plan

n/a
This commit is contained in:
Robert Lin 2024-06-24 18:18:15 -07:00 committed by GitHub
parent e9a31eacc4
commit ffddca87ea
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 16 additions and 9 deletions

View File

@ -79,7 +79,7 @@ func maybeMigrate(ctx context.Context, logger log.Logger, contract runtime.Contr
fmt.Sprintf("%s:auto-migrate", dbName),
15*time.Second,
func() error {
ctx := context.WithoutCancel(ctx)
ctx := context.WithoutCancel(ctx) // do not interrupt once we start
span.AddEvent("lock.acquired")
versionKey := fmt.Sprintf("%s:db_version", dbName)
@ -98,7 +98,8 @@ func maybeMigrate(ctx context.Context, logger log.Logger, contract runtime.Contr
// Create a session that ignore debug logging.
sess := conn.Session(&gorm.Session{
Logger: gormlogger.Default.LogMode(gormlogger.Warn),
Context: ctx,
Logger: gormlogger.Default.LogMode(gormlogger.Warn),
})
// Auto-migrate database table definitions.
for _, table := range tables.All() {

View File

@ -6,12 +6,13 @@ import (
"fmt"
"time"
"github.com/openfga/openfga/assets"
openfga_assets "github.com/openfga/openfga/assets"
"github.com/pressly/goose/v3"
"github.com/redis/go-redis/v9"
"github.com/sourcegraph/log"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/codes"
"go.opentelemetry.io/otel/trace"
"gorm.io/driver/postgres"
"gorm.io/gorm"
gormlogger "gorm.io/gorm/logger"
@ -31,7 +32,10 @@ type metadata struct {
// migrateAndReconcile migrates the "msp-iam" database schema (when needed) and
// reconciles the framework metadata.
func migrateAndReconcile(ctx context.Context, logger log.Logger, sqlDB *sql.DB, redisClient *redis.Client) (_ *metadata, err error) {
ctx, span := iamTracer.Start(ctx, "iam.migrateAndReconcile")
ctx, span := iamTracer.Start(ctx, "iam.migrateAndReconcile",
trace.WithAttributes(
attribute.String("database", databaseName),
))
defer func() {
if err != nil {
span.RecordError(err)
@ -53,10 +57,9 @@ func migrateAndReconcile(ctx context.Context, logger log.Logger, sqlDB *sql.DB,
return nil, errors.Wrap(err, "open connection")
}
goose.SetBaseFS(assets.EmbedMigrations)
goose.SetBaseFS(openfga_assets.EmbedMigrations)
goose.SetLogger(&gooseLoggerShim{Logger: logger})
currentVersion, err := goose.GetDBVersion(sqlDB)
currentVersion, err := goose.GetDBVersionContext(ctx, sqlDB)
if err != nil {
return nil, errors.Wrap(err, "get DB version")
}
@ -70,11 +73,13 @@ func migrateAndReconcile(ctx context.Context, logger log.Logger, sqlDB *sql.DB,
fmt.Sprintf("%s:auto-migrate", databaseName),
15*time.Second,
func() error {
ctx := context.WithoutCancel(ctx) // do not interrupt once we start
span.AddEvent("lock.acquired")
// Create a session that ignore debug logging.
sess := conn.Session(&gorm.Session{
Logger: gormlogger.Default.LogMode(gormlogger.Warn),
Context: ctx,
Logger: gormlogger.Default.LogMode(gormlogger.Warn),
})
// Auto-migrate database table definitions.
for _, table := range []any{&metadata{}} {
@ -86,10 +91,11 @@ func migrateAndReconcile(ctx context.Context, logger log.Logger, sqlDB *sql.DB,
}
// Migrate OpenFGA's database schema.
span.AddEvent("automigrate.openfga")
err = goose.UpContext(
ctx,
sqlDB,
assets.PostgresMigrationDir,
openfga_assets.PostgresMigrationDir,
)
if err != nil {
return errors.Wrap(err, "run OpenFGA migrations")