From 9c663b1ba299fb945c99acaa0542ee0b9fecd586 Mon Sep 17 00:00:00 2001
From: Steven <stevenlgtm@gmail.com>
Date: Wed, 10 Jan 2024 00:03:47 +0800
Subject: [PATCH] fix: merge mysql dsn with params

---
 store/db/mysql/mysql.go | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/store/db/mysql/mysql.go b/store/db/mysql/mysql.go
index 55984ef2..2c8da727 100644
--- a/store/db/mysql/mysql.go
+++ b/store/db/mysql/mysql.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"net/url"
 
 	"github.com/go-sql-driver/mysql"
 	"github.com/pkg/errors"
@@ -23,9 +24,13 @@ func NewDB(profile *profile.Profile) (store.Driver, error) {
 	// Open MySQL connection with parameter.
 	// multiStatements=true is required for migration.
 	// See more in: https://github.com/go-sql-driver/mysql#multistatements
-	dsn := fmt.Sprintf("%s?multiStatements=true", profile.DSN)
+	dsn, err := mergeDSNWithParams(profile.DSN, map[string]string{
+		"multiStatements": "true",
+	})
+	if err != nil {
+		return nil, err
+	}
 
-	var err error
 	driver := DB{profile: profile}
 	driver.config, err = mysql.ParseDSN(dsn)
 	if err != nil {
@@ -104,3 +109,17 @@ func (d *DB) GetCurrentDBSize(ctx context.Context) (int64, error) {
 func (d *DB) Close() error {
 	return d.db.Close()
 }
+
+func mergeDSNWithParams(baseDSN string, params map[string]string) (string, error) {
+	parsedDSN, err := url.Parse(baseDSN)
+	if err != nil {
+		return "", errors.Wrapf(err, "failed to parse DSN: %s", baseDSN)
+	}
+
+	existingParams := parsedDSN.Query()
+	for key, value := range params {
+		existingParams.Add(key, value)
+	}
+	parsedDSN.RawQuery = existingParams.Encode()
+	return parsedDSN.String(), nil
+}