app/common/database.go

package common

import (
	"context"
	"database/sql"
	"fmt"
	"log/slog"

	"github.com/XSAM/otelsql"
	"github.com/google/wire"
	"go.opentelemetry.io/otel/metric"
	"go.opentelemetry.io/otel/trace"

	"github.com/openmeterio/openmeter/app/config"
	"github.com/openmeterio/openmeter/openmeter/ent/db"
	"github.com/openmeterio/openmeter/pkg/framework/entutils/entdriver"
	"github.com/openmeterio/openmeter/pkg/framework/pgdriver"
	"github.com/openmeterio/openmeter/tools/migrate"
)

var Database = wire.NewSet(
	wire.Struct(new(Migrator), "*"),

	NewPostgresDriver,
	NewDB,
	NewEntPostgresDriver,
	NewEntClient,
)

// Migrator executes database migrations.
type Migrator struct {
	Config config.PostgresConfig
	Client *db.Client
	Logger *slog.Logger
}

func (m Migrator) Migrate(ctx context.Context) error {
	if !m.Config.AutoMigrate.Enabled() {
		m.Logger.Debug("auto migration is disabled")

		return nil
	}

	m.Logger.Info("running migrations", slog.String("strategy", string(m.Config.AutoMigrate)))

	migrator, err := migrate.New(migrate.MigrateOptions{
		ConnectionString: m.Config.AsURL(),
		Migrations:       migrate.OMMigrationsConfig,
		Logger:           m.Logger,
	})
	if err != nil {
		return fmt.Errorf("failed to create migrator: %w", err)
	}
	defer migrator.CloseOrLogError()

	switch m.Config.AutoMigrate {
	case config.AutoMigrateMigration:
		if err := migrator.Up(); err != nil {
			return fmt.Errorf("failed to migrate db: %w", err)
		}
	case config.AutoMigrateMigrationJob:
		if err := migrator.WaitForMigrationJob(); err != nil {
			return fmt.Errorf("failed to wait for migration job: %w", err)
		}
	}

	m.Logger.Info("database initialized")

	return nil
}

// AdoptLegacyEnt is the explicit upgrade-job entrypoint for databases previously managed by Ent.
func (m Migrator) AdoptLegacyEnt(ctx context.Context) error {
	driver, err := pgdriver.NewPostgresDriver(ctx, m.Config.AsURL())
	if err != nil {
		return fmt.Errorf("open database for legacy Ent adoption: %w", err)
	}
	defer driver.Close()

	return migrate.AdoptLegacyEnt(ctx, driver.DB(), m.Config.AsURL(), m.Logger)
}

func NewPostgresDriver(
	ctx context.Context,
	conf config.PostgresConfig,
	meterProvider metric.MeterProvider,
	meter metric.Meter,
	tracerProvider trace.TracerProvider,
	logger *slog.Logger,
) (*pgdriver.Driver, func(), error) {
	driver, err := pgdriver.NewPostgresDriver(
		ctx,
		conf.AsURL(),
		pgdriver.WithMetricMeter(meter),
		pgdriver.WithTracerProvider(tracerProvider),
		pgdriver.WithMeterProvider(meterProvider),
		pgdriver.WithSpanOptions(otelsql.SpanOptions{
			OmitConnPrepare:      true,
			OmitRows:             true,
			OmitConnectorConnect: true,
		}),
	)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to initialize postgres driver: %w", err)
	}

	return driver, func() {
		err := driver.Close()
		if err != nil {
			logger.Error("failed to close postgres driver", "error", err)
		}
	}, nil
}

// TODO: add closer function?
func NewDB(driver *pgdriver.Driver) *sql.DB {
	return driver.DB()
}

func NewEntPostgresDriver(db *sql.DB, logger *slog.Logger) (*entdriver.EntPostgresDriver, func()) {
	driver := entdriver.NewEntPostgresDriver(db)

	return driver, func() {
		err := driver.Close()
		if err != nil {
			logger.Error("failed to close ent driver", "error", err)
		}
	}
}

// TODO: add closer function?
func NewEntClient(driver *entdriver.EntPostgresDriver) *db.Client {
	return driver.Client()
}