Skip to content
This repository was archived by the owner on Sep 30, 2024. It is now read-only.

Commit ffddca8

Browse files
authored
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
1 parent e9a31ea commit ffddca8

File tree

2 files changed

+16
-9
lines changed

2 files changed

+16
-9
lines changed

cmd/enterprise-portal/internal/database/migrate.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ func maybeMigrate(ctx context.Context, logger log.Logger, contract runtime.Contr
7979
fmt.Sprintf("%s:auto-migrate", dbName),
8080
15*time.Second,
8181
func() error {
82-
ctx := context.WithoutCancel(ctx)
82+
ctx := context.WithoutCancel(ctx) // do not interrupt once we start
8383
span.AddEvent("lock.acquired")
8484

8585
versionKey := fmt.Sprintf("%s:db_version", dbName)
@@ -98,7 +98,8 @@ func maybeMigrate(ctx context.Context, logger log.Logger, contract runtime.Contr
9898

9999
// Create a session that ignore debug logging.
100100
sess := conn.Session(&gorm.Session{
101-
Logger: gormlogger.Default.LogMode(gormlogger.Warn),
101+
Context: ctx,
102+
Logger: gormlogger.Default.LogMode(gormlogger.Warn),
102103
})
103104
// Auto-migrate database table definitions.
104105
for _, table := range tables.All() {

lib/managedservicesplatform/iam/database.go

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,13 @@ import (
66
"fmt"
77
"time"
88

9-
"github.com/openfga/openfga/assets"
9+
openfga_assets "github.com/openfga/openfga/assets"
1010
"github.com/pressly/goose/v3"
1111
"github.com/redis/go-redis/v9"
1212
"github.com/sourcegraph/log"
1313
"go.opentelemetry.io/otel/attribute"
1414
"go.opentelemetry.io/otel/codes"
15+
"go.opentelemetry.io/otel/trace"
1516
"gorm.io/driver/postgres"
1617
"gorm.io/gorm"
1718
gormlogger "gorm.io/gorm/logger"
@@ -31,7 +32,10 @@ type metadata struct {
3132
// migrateAndReconcile migrates the "msp-iam" database schema (when needed) and
3233
// reconciles the framework metadata.
3334
func migrateAndReconcile(ctx context.Context, logger log.Logger, sqlDB *sql.DB, redisClient *redis.Client) (_ *metadata, err error) {
34-
ctx, span := iamTracer.Start(ctx, "iam.migrateAndReconcile")
35+
ctx, span := iamTracer.Start(ctx, "iam.migrateAndReconcile",
36+
trace.WithAttributes(
37+
attribute.String("database", databaseName),
38+
))
3539
defer func() {
3640
if err != nil {
3741
span.RecordError(err)
@@ -53,10 +57,9 @@ func migrateAndReconcile(ctx context.Context, logger log.Logger, sqlDB *sql.DB,
5357
return nil, errors.Wrap(err, "open connection")
5458
}
5559

56-
goose.SetBaseFS(assets.EmbedMigrations)
60+
goose.SetBaseFS(openfga_assets.EmbedMigrations)
5761
goose.SetLogger(&gooseLoggerShim{Logger: logger})
58-
59-
currentVersion, err := goose.GetDBVersion(sqlDB)
62+
currentVersion, err := goose.GetDBVersionContext(ctx, sqlDB)
6063
if err != nil {
6164
return nil, errors.Wrap(err, "get DB version")
6265
}
@@ -70,11 +73,13 @@ func migrateAndReconcile(ctx context.Context, logger log.Logger, sqlDB *sql.DB,
7073
fmt.Sprintf("%s:auto-migrate", databaseName),
7174
15*time.Second,
7275
func() error {
76+
ctx := context.WithoutCancel(ctx) // do not interrupt once we start
7377
span.AddEvent("lock.acquired")
7478

7579
// Create a session that ignore debug logging.
7680
sess := conn.Session(&gorm.Session{
77-
Logger: gormlogger.Default.LogMode(gormlogger.Warn),
81+
Context: ctx,
82+
Logger: gormlogger.Default.LogMode(gormlogger.Warn),
7883
})
7984
// Auto-migrate database table definitions.
8085
for _, table := range []any{&metadata{}} {
@@ -86,10 +91,11 @@ func migrateAndReconcile(ctx context.Context, logger log.Logger, sqlDB *sql.DB,
8691
}
8792

8893
// Migrate OpenFGA's database schema.
94+
span.AddEvent("automigrate.openfga")
8995
err = goose.UpContext(
9096
ctx,
9197
sqlDB,
92-
assets.PostgresMigrationDir,
98+
openfga_assets.PostgresMigrationDir,
9399
)
94100
if err != nil {
95101
return errors.Wrap(err, "run OpenFGA migrations")

0 commit comments

Comments
 (0)