query-orchestration/internal/database/migrations.go

package database

import (
	"context"
	"database/sql"
	"fmt"
	"log"
	"os"
	"path"

	"github.com/golang-migrate/migrate/v4"
	"github.com/golang-migrate/migrate/v4/database/postgres"
	_ "github.com/golang-migrate/migrate/v4/source/file"
	"github.com/jackc/pgx/v5"
)

func createDB(ctx context.Context) {
	driver := "postgres"
	opts := ""
	conf := mustGetConfig()
	if conf.disableSSL {
		opts += "sslmode=disable"
	}
	connStr := fmt.Sprintf("%s://%s:%s@%s:%s?%s", driver, conf.user, conf.password, conf.host, conf.port, opts)

	conn, err := pgx.Connect(ctx, connStr)
	if err != nil {
		log.Panicf("Error opening database: %v", err)
	}

	err = conn.Ping(ctx)
	if err != nil {
		log.Panicf("Error pinging database: %v", err)
	}

	rs, err := conn.Query(ctx, fmt.Sprintf("SELECT 'CREATE DATABASE %s' WHERE NOT EXISTS (SELECT FROM pg_database WHERE datname = '%s')", conf.name, conf.name))
	if err != nil {
		log.Panicf("Error creating database: %v", err)
	}

	if rs.Next() {
		log.Printf("Database created: %s", conf.name)
	} else {
		log.Printf("Database already exists: %s", conf.name)
	}
}

type MigrationConfig struct {
	BasePath         string
	ConnectionString string
}

func RunMigrations(ctx context.Context, config *MigrationConfig) {
	// createDB(ctx) - Add back only if necessary

	connStr := config.ConnectionString
	if connStr == "" {
		connStr = mustGetConnectionString()
	}

	basePath := config.BasePath
	if basePath == "" {
		basePath = os.Getenv("PWD")
	}

	driver := "postgres"

	db, err := sql.Open(driver, connStr)
	if err != nil {
		log.Panicf("failed to open database: %v", err)
	}
	defer db.Close()

	if err = db.Ping(); err != nil {
		log.Panic("failed to connect to database:", err)
	}

	conn, err := postgres.WithInstance(db, &postgres.Config{})
	if err != nil {
		log.Panic("failed to create migrations: ", err)
	}

	migPath := "file://" + path.Join(basePath, "database/migrations")
	m, err := migrate.NewWithDatabaseInstance(
		migPath,
		driver,
		conn,
	)
	if err != nil {
		log.Panicf("failed to create migrate instance: %v", err)
	}

	err = m.Up()
	if err == migrate.ErrNoChange {
		log.Println("No changes required")
	} else if err != nil {
		log.Panicf("failed to apply migrations: %v", err)
	} else {
		log.Println("Migrations applied successfully!")
	}
}