refactor(migrate): use initialized pgx driver

cmd/balance-worker/main.go

@@ -239,7 +239,7 @@ func main() {
 
 	entClient := entPostgresDriver.Client()
 
-	if err := startup.DB(ctx, conf.Postgres, entClient); err != nil {
+	if err := startup.DB(ctx, conf.Postgres, entClient, postgresDriver.DB()); err != nil {
 		logger.Error("failed to initialize database", "error", err)
 		os.Exit(1)
 	}

cmd/notification-service/main.go

@@ -242,7 +242,7 @@ func main() {
 	entClient := entPostgresDriver.Client()
 
 	// Run database schema creation
-	if err := startup.DB(ctx, conf.Postgres, entClient); err != nil {
+	if err := startup.DB(ctx, conf.Postgres, entClient, postgresDriver.DB()); err != nil {
 		logger.Error("failed to initialize database", "error", err)
 		os.Exit(1)
 	}

cmd/server/main.go

@@ -351,7 +351,7 @@ func main() {
 
 	entClient := entPostgresDriver.Client()
 
-	if err := startup.DB(ctx, conf.Postgres, entClient); err != nil {
+	if err := startup.DB(ctx, conf.Postgres, entClient, postgresDriver.DB()); err != nil {
 		logger.Error("failed to initialize database", "error", err)
 		os.Exit(1)
 	}

openmeter/entitlement/metered/utils_test.go

@@ -77,8 +77,12 @@ func setupConnector(t *testing.T) (meteredentitlement.Connector, *dependencies)
 	m.Lock()
 	defer m.Unlock()
 	// migrate db
-	if err := migrate.Up(testdb.URL); err != nil {
-		t.Fatalf("failed to migrate db: %s", err.Error())
+	if m, err := migrate.Default(testdb.SQLDriver); err == nil {
+		if err := m.Up(); err != nil {
+			t.Fatalf("failed to migrate db: %s", err.Error())
+		}
+	} else {
+		t.Fatalf("failed to create migrate instance: %s", err.Error())
 	}
 
 	mockPublisher := eventbus.NewMock(t)

openmeter/productcatalog/adapter/feature_test.go

@@ -220,8 +220,12 @@ func TestCreateFeature(t *testing.T) {
 			dbClient := testdb.EntDriver.Client()
 			defer dbClient.Close()
 
-			if err := migrate.Up(testdb.URL); err != nil {
-				t.Fatalf("failed to migrate db: %s", err.Error())
+			if m, err := migrate.Default(testdb.SQLDriver); err == nil {
+				if err := m.Up(); err != nil {
+					t.Fatalf("failed to migrate db: %s", err.Error())
+				}
+			} else {
+				t.Fatalf("failed to create migrate instance: %s", err.Error())
 			}
 
 			dbConnector := adapter.NewPostgresFeatureRepo(dbClient, testutils.NewLogger(t))
@@ -239,8 +243,12 @@ func TestCreateFeature(t *testing.T) {
 		dbClient := testdb.EntDriver.Client()
 		defer dbClient.Close()
 
-		if err := migrate.Up(testdb.URL); err != nil {
-			t.Fatalf("failed to migrate db: %s", err.Error())
+		if m, err := migrate.Default(testdb.SQLDriver); err == nil {
+			if err := m.Up(); err != nil {
+				t.Fatalf("failed to migrate db: %s", err.Error())
+			}
+		} else {
+			t.Fatalf("failed to create migrate instance: %s", err.Error())
 		}
 
 		dbConnector := adapter.NewPostgresFeatureRepo(dbClient, testutils.NewLogger(t))

openmeter/registry/startup/db.go

@@ -2,26 +2,31 @@ package startup
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 
 	"github.com/openmeterio/openmeter/config"
 	"github.com/openmeterio/openmeter/openmeter/ent/db"
 	"github.com/openmeterio/openmeter/tools/migrate"
 )
 
-func DB(ctx context.Context, cfg config.PostgresConfig, db *db.Client) error {
+func DB(ctx context.Context, cfg config.PostgresConfig, client *db.Client, db *sql.DB) error {
 	if !cfg.AutoMigrate.Enabled() {
 		return nil
 	}
 
 	switch cfg.AutoMigrate {
 	case config.AutoMigrateEnt:
-		if err := db.Schema.Create(ctx); err != nil {
+		if err := client.Schema.Create(ctx); err != nil {
 			return fmt.Errorf("failed to migrate db: %w", err)
 		}
 	case config.AutoMigrateMigration:
-		if err := migrate.Up(cfg.URL); err != nil {
-			return fmt.Errorf("failed to migrate db: %w", err)
+		if m, err := migrate.Default(db); err == nil {
+			if err := m.Up(); err != nil {
+				return fmt.Errorf("failed to migrate db: %w", err)
+			}
+		} else {
+			return fmt.Errorf("failed to create migrate instance: %w", err)
 		}
 	}
 

openmeter/testutils/pg_driver.go

@@ -91,6 +91,7 @@ func InitPostgresDB(t *testing.T) *TestDB {
 
 	return &TestDB{
 		PGDriver:  postgresDriver,
+		SQLDriver: postgresDriver.DB(),
 		EntDriver: entDriver,
 		URL:       dbConf.URL(),
 	}

test/entitlement/regression/framework_test.go

@@ -62,8 +62,14 @@ func setupDependencies(t *testing.T) Dependencies {
 	driver := testutils.InitPostgresDB(t)
 	// init db
 	dbClient := db.NewClient(db.Driver(driver.EntDriver.Driver()))
-	if err := migrate.Up(driver.URL); err != nil {
-		t.Fatalf("failed to migrate db: %s", err.Error())
+
+	// Migrate
+	if m, err := migrate.Default(driver.SQLDriver); err == nil {
+		if err := m.Up(); err != nil {
+			t.Fatalf("failed to migrate db: %s", err.Error())
+		}
+	} else {
+		t.Fatalf("failed to create migrate instance: %s", err.Error())
 	}
 
 	// Init product catalog

tools/migrate/migrate.go

@@ -2,12 +2,12 @@
 package migrate
 
 import (
+	"database/sql"
 	"embed"
 	"io/fs"
-	"net/url"
 
 	"github.com/golang-migrate/migrate/v4"
-	_ "github.com/golang-migrate/migrate/v4/database/postgres"
+	"github.com/golang-migrate/migrate/v4/database/postgres"
 	"github.com/golang-migrate/migrate/v4/source/iofs"
 )
 
@@ -20,21 +20,48 @@ type Migrate = migrate.Migrate
 //go:embed migrations
 var OMMigrations embed.FS
 
+type Options struct {
+	DB       *sql.DB
+	FS       fs.FS
+	FSPath   string
+	PGConfig *postgres.Config
+}
+
 // NewMigrate creates a new migrate instance.
-func NewMigrate(conn string, fs fs.FS, fsPath string) (*Migrate, error) {
-	d, err := iofs.New(fs, fsPath)
+func NewMigrate(opt Options) (*Migrate, error) {
+	d, err := iofs.New(opt.FS, opt.FSPath)
 	if err != nil {
 		return nil, err
 	}
-	return migrate.NewWithSourceInstance("iofs", d, conn)
-}
 
-func Up(conn string) error {
-	conn, err := SetMigrationTableName(conn, MigrationsTable)
+	driver, err := postgres.WithInstance(opt.DB, opt.PGConfig)
 	if err != nil {
-		return err
+		return nil, err
 	}
-	m, err := NewMigrate(conn, OMMigrations, "migrations")
+
+	return migrate.NewWithInstance("iofs", d, "postgres", driver)
+}
+
+func Default(db *sql.DB) (*Migrate, error) {
+	return NewMigrate(Options{
+		DB:     db,
+		FS:     OMMigrations,
+		FSPath: "migrations",
+		PGConfig: &postgres.Config{
+			MigrationsTable: MigrationsTable,
+		},
+	})
+}
+
+func Up(db *sql.DB) error {
+	m, err := NewMigrate(Options{
+		DB:     db,
+		FS:     OMMigrations,
+		FSPath: "migrations",
+		PGConfig: &postgres.Config{
+			MigrationsTable: MigrationsTable,
+		},
+	})
 	if err != nil {
 		return err
 	}
@@ -46,16 +73,3 @@ func Up(conn string) error {
 	}
 	return nil
 }
-
-func SetMigrationTableName(conn, tableName string) (string, error) {
-	parsedURL, err := url.Parse(conn)
-	if err != nil {
-		return "", err
-	}
-
-	values := parsedURL.Query()
-	values.Set("x-migrations-table", tableName)
-	parsedURL.RawQuery = values.Encode()
-
-	return parsedURL.String(), nil
-}

tools/migrate/migrate_test.go

@@ -12,7 +12,7 @@ func TestUpDownUp(t *testing.T) {
 	testDB := testutils.InitPostgresDB(t)
 	defer testDB.PGDriver.Close()
 
-	migrator, err := migrate.NewMigrate(testDB.URL, migrate.OMMigrations, "migrations")
+	migrator, err := migrate.Default(testDB.SQLDriver)
 	if err != nil {
 		t.Fatal(err)
 	}