docker.go

package main

import (
	"context"
	"database/sql"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/jackc/pgx/v5"
)

type ContainerInfo struct {
	ContainerID string
	HostPort    string
	ConnString  string
}

func EnsureDatabaseContainer(dbType string) (*ContainerInfo, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
	defer cancel()

	if err := exec.CommandContext(ctx, "docker", "version").Run(); err != nil {
		return nil, fmt.Errorf("docker is not running or not installed: %w", err)
	}

	var (
		imageName    string
		internalPort string
		envArgs      []string
	)

	switch dbType {
	case "PostgreSQL":
		imageName = "postgres:16-alpine"
		internalPort = "5432"
		envArgs = []string{
			"-e", "POSTGRES_USER=postgres",
			"-e", "POSTGRES_PASSWORD=postgres",
			"-e", "POSTGRES_DB=mockdb",
		}
	case "MySQL":
		imageName = "mysql:8.0"
		internalPort = "3306"
		envArgs = []string{
			"-e", "MYSQL_ROOT_PASSWORD=postgres",
			"-e", "MYSQL_DATABASE=mockdb",
		}
	default:
		return nil, fmt.Errorf("unsupported database type: %s", dbType)
	}

	pullCmd := exec.CommandContext(ctx, "docker", "pull", imageName)
	if out, err := pullCmd.CombinedOutput(); err != nil {
		return nil, fmt.Errorf("failed to pull image %s: %w\noutput: %s", imageName, err, string(out))
	}

	hostPort, err := getFreePort()
	if err != nil {
		return nil, fmt.Errorf("failed to allocate host port: %w", err)
	}

	portMapping := fmt.Sprintf("127.0.0.1:%s:%s/tcp", hostPort, internalPort)
	args := []string{
		"run", "-d", "--rm",
		"-p", portMapping,
	}
	args = append(args, envArgs...)
	args = append(args, imageName)

	runCmd := exec.CommandContext(ctx, "docker", args...)
	out, err := runCmd.Output()
	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			return nil, fmt.Errorf("failed to start container: %w\nstderr: %s", err, string(exitErr.Stderr))
		}
		return nil, fmt.Errorf("failed to start container: %w", err)
	}

	containerID := strings.TrimSpace(string(out))
	if containerID == "" {
		return nil, fmt.Errorf("docker run returned empty container ID")
	}

	var connStr string
	switch dbType {
	case "PostgreSQL":
		connStr = fmt.Sprintf("postgres://postgres:postgres@127.0.0.1:%s/mockdb?sslmode=disable", hostPort)
	case "MySQL":
		connStr = fmt.Sprintf("root:postgres@tcp(127.0.0.1:%s)/mockdb", hostPort)
	}

	if err := waitForDatabase(ctx, dbType, connStr); err != nil {
		_ = exec.CommandContext(context.Background(), "docker", "stop", "-t", "5", containerID).Run()
		return nil, fmt.Errorf("database health check failed: %w", err)
	}

	return &ContainerInfo{
		ContainerID: containerID,
		HostPort:    hostPort,
		ConnString:  connStr,
	}, nil
}

func setupSQLite() (string, error) {
	homeDir, err := os.UserHomeDir()
	if err != nil {
		return "", err
	}
	return filepath.Join(homeDir, ".data-mocker.sqlite"), nil
}

func getFreePort() (string, error) {
	addr, err := net.ResolveTCPAddr("tcp", "127.0.0.1:0")
	if err != nil {
		return "", err
	}
	l, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return "", err
	}
	defer l.Close()
	return fmt.Sprintf("%d", l.Addr().(*net.TCPAddr).Port), nil
}

func waitForDatabase(ctx context.Context, dbType, connStr string) error {
	timeout := time.After(30 * time.Second)
	ticker := time.NewTicker(500 * time.Millisecond)
	defer ticker.Stop()

	for {
		select {
		case <-timeout:
			return fmt.Errorf("timed out waiting for database to be ready")
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
			var ready bool
			var err error

			switch dbType {
			case "PostgreSQL":
				ready, err = pingPostgres(connStr)
			case "MySQL":
				ready, err = pingMySQL(connStr)
			}

			if ready {
				return nil
			}
			if err != nil {
				continue
			}
		}
	}
}

func pingPostgres(connStr string) (bool, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	conn, err := pgx.Connect(ctx, connStr)
	if err != nil {
		return false, err
	}
	defer conn.Close(ctx)

	var one int
	err = conn.QueryRow(ctx, "SELECT 1").Scan(&one)
	return err == nil && one == 1, err
}

func pingMySQL(connStr string) (bool, error) {
	db, err := sql.Open("mysql", connStr)
	if err != nil {
		return false, err
	}
	defer db.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	if err := db.PingContext(ctx); err != nil {
		return false, err
	}
	return true, nil
}