From b0c78bfaffd500338381494dc5ba648969176a2f Mon Sep 17 00:00:00 2001
From: Vasanth-96 <ramavath.naik@itilite.com>
Date: Wed, 7 May 2025 00:40:59 +0530
Subject: [PATCH] Add support for MariaDB connection in configuration and URI
 generation

---
 go/cmd/gh-ost/main.go       |  1 +
 go/mysql/connection.go      | 11 ++++++++++-
 go/mysql/connection_test.go | 14 ++++++++++++++
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/go/cmd/gh-ost/main.go b/go/cmd/gh-ost/main.go
index a1670cdd4..5c7a3a1f9 100644
--- a/go/cmd/gh-ost/main.go
+++ b/go/cmd/gh-ost/main.go
@@ -52,6 +52,7 @@ func main() {
 	flag.StringVar(&migrationContext.AssumeMasterHostname, "assume-master-host", "", "(optional) explicitly tell gh-ost the identity of the master. Format: some.host.com[:port] This is useful in master-master setups where you wish to pick an explicit master, or in a tungsten-replicator where gh-ost is unable to determine the master")
 	flag.IntVar(&migrationContext.InspectorConnectionConfig.Key.Port, "port", 3306, "MySQL port (preferably a replica, not the master)")
 	flag.Float64Var(&migrationContext.InspectorConnectionConfig.Timeout, "mysql-timeout", 0.0, "Connect, read and write timeout for MySQL")
+	flag.BoolVar(&migrationContext.InspectorConnectionConfig.IsMariaDB, "mariadb", false, "Set to 'true' if connecting to MariaDB instance")
 	flag.StringVar(&migrationContext.CliUser, "user", "", "MySQL user")
 	flag.StringVar(&migrationContext.CliPassword, "password", "", "MySQL password")
 	flag.StringVar(&migrationContext.CliMasterUser, "master-user", "", "MySQL user on master, if different from that on replica. Requires --assume-master-host")
diff --git a/go/mysql/connection.go b/go/mysql/connection.go
index f728fc7fe..65bfa3498 100644
--- a/go/mysql/connection.go
+++ b/go/mysql/connection.go
@@ -31,6 +31,7 @@ type ConnectionConfig struct {
 	Timeout              float64
 	TransactionIsolation string
 	Charset              string
+	IsMariaDB            bool
 }
 
 func NewConnectionConfig() *ConnectionConfig {
@@ -51,6 +52,7 @@ func (this *ConnectionConfig) DuplicateCredentials(key InstanceKey) *ConnectionC
 		Timeout:              this.Timeout,
 		TransactionIsolation: this.TransactionIsolation,
 		Charset:              this.Charset,
+		IsMariaDB:            this.IsMariaDB,
 	}
 
 	if this.tlsConfig != nil {
@@ -152,12 +154,19 @@ func (this *ConnectionConfig) GetDBUri(databaseName string) string {
 		this.Charset = "utf8mb4,utf8,latin1"
 	}
 
+	var isolationParam string
+	if this.IsMariaDB {
+		isolationParam = fmt.Sprintf("tx_isolation=%q", this.TransactionIsolation)
+	} else {
+		isolationParam = fmt.Sprintf("transaction_isolation=%q", this.TransactionIsolation)
+	}
+
 	connectionParams := []string{
 		"autocommit=true",
 		"interpolateParams=true",
 		fmt.Sprintf("charset=%s", this.Charset),
 		fmt.Sprintf("tls=%s", tlsOption),
-		fmt.Sprintf("transaction_isolation=%q", this.TransactionIsolation),
+		isolationParam,
 		fmt.Sprintf("timeout=%fs", this.Timeout),
 		fmt.Sprintf("readTimeout=%fs", this.Timeout),
 		fmt.Sprintf("writeTimeout=%fs", this.Timeout),
diff --git a/go/mysql/connection_test.go b/go/mysql/connection_test.go
index 59761b047..79371bdfb 100644
--- a/go/mysql/connection_test.go
+++ b/go/mysql/connection_test.go
@@ -113,3 +113,17 @@ func TestGetDBTLSConfigKey(t *testing.T) {
 	configKey := GetDBTLSConfigKey("myhost")
 	require.Equal(t, "ghost-myhost", configKey)
 }
+
+func TestGetDBUriMariaDB(t *testing.T) {
+	c := NewConnectionConfig()
+	c.Key = InstanceKey{Hostname: "mariadb-host", Port: 3307}
+	c.User = "maria"
+	c.Password = "dba"
+	c.Timeout = 2.0
+	c.TransactionIsolation = transactionIsolation
+	c.Charset = "utf8mb4"
+	c.IsMariaDB = true
+
+	uri := c.GetDBUri("analytics")
+	require.Equal(t, `maria:dba@tcp(mariadb-host:3307)/analytics?autocommit=true&interpolateParams=true&charset=utf8mb4&tls=false&tx_isolation="REPEATABLE-READ"&timeout=2.000000s&readTimeout=2.000000s&writeTimeout=2.000000s`, uri)
+}