package main

import (
	"context"
	"database/sql"
	"flag"
	"fmt"
	"log"
	"strings"
	"time"

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

func main() {
	host := flag.String("host", "127.0.0.1", "the host to check for database readiness")
	port := flag.String("port", "4000", "the port to check for database readiness")
	tlsEnabled := flag.Bool("tlsEnabled", true, "whether TLS is enabled on tidb clusters")
	timeoutSeconds := flag.Int("timeoutSeconds", 3, "the timeout in seconds to wait for the database to be ready")
	flag.Parse()

	connectionString := fmt.Sprintf("test_readiness:@tcp(%s:%s)/?timeout=2s", *host, *port)
	if *tlsEnabled {
		connectionString = fmt.Sprintf("test_readiness:@tcp(%s:%s)/?timeout=2s&tls=skip-verify", *host, *port)
	}

	db, err := sql.Open("mysql", connectionString)
	if err != nil {
		log.Fatal("failed to connect to the database: ", err)
	}
	defer func() {
		_ = db.Close()
	}()

	if err := waitReady(db, time.Duration(*timeoutSeconds)*time.Second); err != nil {
		log.Fatal(err)
	}
}

func waitReady(db *sql.DB, timeout time.Duration) error {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	if err := db.PingContext(ctx); err != nil {
		if strings.Contains(err.Error(), "Access denied") {
			// wrong username or password
			return nil
		}
		// other errors may be: connection refused.
	} else {
		// happens to be a user and doesn't have a password.
		return nil
	}

	return fmt.Errorf("the database is not ready")
}