From 79a376f9f258393adfc0831e4e5e836ef59a3504 Mon Sep 17 00:00:00 2001
From: Adam Shannon <adamkshannon@gmail.com>
Date: Fri, 20 Sep 2024 11:11:42 -0500
Subject: [PATCH] database: add RunMigrationsContext with tracing

---
 database/migrator.go | 28 ++++++++++++++++++++++------
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/database/migrator.go b/database/migrator.go
index d44f34a6..6cbe2596 100644
--- a/database/migrator.go
+++ b/database/migrator.go
@@ -11,18 +11,31 @@ import (
 	"sync"
 	"time"
 
+	"github.com/moov-io/base/log"
+	"github.com/moov-io/base/telemetry"
+
 	"github.com/golang-migrate/migrate/v4"
 	"github.com/golang-migrate/migrate/v4/database"
 	migmysql "github.com/golang-migrate/migrate/v4/database/mysql"
 	"github.com/golang-migrate/migrate/v4/source"
 	"github.com/golang-migrate/migrate/v4/source/iofs"
 
-	"github.com/moov-io/base/log"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/trace"
 )
 
 var migrationMutex sync.Mutex
 
 func RunMigrations(logger log.Logger, config DatabaseConfig, opts ...MigrateOption) error {
+	return RunMigrationsContext(context.Background(), logger, config, opts...)
+}
+
+func RunMigrationsContext(ctx context.Context, logger log.Logger, config DatabaseConfig, opts ...MigrateOption) error {
+	_, span := telemetry.StartSpan(ctx, "run-migrations", trace.WithAttributes(
+		attribute.String("db.database_name", config.DatabaseName),
+	))
+	defer span.End()
+
 	logger.Info().Log("Running Migrations")
 
 	// apply all of our optional arguments
@@ -54,24 +67,26 @@ func RunMigrations(logger log.Logger, config DatabaseConfig, opts ...MigrateOpti
 		m.LockTimeout = *o.timeout
 	}
 
-	currentVersion, dirty, err := m.Version()
+	previousVersion, dirty, err := m.Version()
 	if err != nil {
 		if err != migrate.ErrNilVersion {
 			return logger.Fatal().LogErrorf("Error getting current DB version: %w", err).Err()
 		}
 		// set sane values
-		currentVersion = 0
+		previousVersion = 0
 		dirty = false
 	}
+	span.SetAttributes(attribute.Int64("db.previous_version", int64(previousVersion)))
+
 	err = m.Up()
 	migrationMutex.Unlock()
 
 	switch err {
 	case nil:
 	case migrate.ErrNoChange:
-		logger.Info().Logf("Database already at version %d (dirty: %b)", currentVersion, dirty)
+		logger.Info().Logf("Database already at version %d (dirty: %b)", previousVersion, dirty)
 	default:
-		return logger.Fatal().LogErrorf("Error running migrations (current: %d, dirty: %b): %w", currentVersion, dirty, err).Err()
+		return logger.Fatal().LogErrorf("Error running migrations (current: %d, dirty: %b): %w", previousVersion, dirty, err).Err()
 	}
 
 	newVersion, newDirty, err := m.Version()
@@ -83,8 +98,9 @@ func RunMigrations(logger log.Logger, config DatabaseConfig, opts ...MigrateOpti
 		newVersion = 0
 		newDirty = false
 	}
+	span.SetAttributes(attribute.Int64("db.new_version", int64(newVersion)))
 
-	logger.Info().Logf("Migrations complete: previous: %d (dirty:%v) -> current: %d (dirty:%v)", currentVersion, dirty, newVersion, newDirty)
+	logger.Info().Logf("Migrations complete: previous: %d (dirty:%v) -> new: %d (dirty:%v)", previousVersion, dirty, newVersion, newDirty)
 
 	return nil
 }