Chore/support clickhouse cluster

.gitignore

@@ -8,3 +8,4 @@
 /vendor
 dist
 node_modules
+.vscode

README.md

@@ -253,6 +253,33 @@ DATABASE_URL="sqlite:/tmp/database.sqlite3"
 DATABASE_URL="clickhouse://username:password@127.0.0.1:9000/database_name"
 ```
 
+To work with ClickHouse cluster, there are 4 connection query parameters that can be supplied:
+
+- `on_cluster` - Indicataion to use cluster statements and replicated migration table. (default: `false`) If this parameter is not supplied, other cluster related query parameters are ignored.
+```sh
+DATABASE_URL="clickhouse://username:password@127.0.0.1:9000/database_name?on_cluster"
+
+DATABASE_URL="clickhouse://username:password@127.0.0.1:9000/database_name?on_cluster=true"
+```
+
+- `cluster_macro` (Optional) - Macro value to be used for ON CLUSTER statements and for the replciated migration table engine zookeeper path. (default: `{cluster}`)
+
+```sh
+DATABASE_URL="clickhouse://username:password@127.0.0.1:9000/database_name?on_cluster&cluster_macro={my_cluster}"
+```
+
+- `replica_macro` (Optional) - Macro value to be used for the replica name in the replciated migration table engine. (default: `{replica}`)
+
+```sh
+DATABASE_URL="clickhouse://username:password@127.0.0.1:9000/database_name?on_cluster&replica_macro={my_replica}"
+```
+
+- `zoo_path` (Optional) - The path to the table migration in ClickHouse/Zoo Keeper. (default: `/clickhouse/tables/<cluster_macro>/{table}`)
+
+```sh
+DATABASE_URL="clickhouse://username:password@127.0.0.1:9000/database_name?on_cluster&zoo_path=/zk/path/tables"
+```
+
 [See other supported connection options](https://github.com/ClickHouse/clickhouse-go#dsn).
 
 ### Creating Migrations

docker-compose.yml

@@ -10,8 +10,12 @@ services:
       - mysql
       - postgres
       - clickhouse
+      - clickhouse-cluster-01
+      - clickhouse-cluster-02
     environment:
       CLICKHOUSE_TEST_URL: clickhouse://clickhouse:9000/dbmate_test
+      CLICKHOUSE_CLUSTER_01_TEST_URL: clickhouse://ch-cluster-01:9000/dbmate_test
+      CLICKHOUSE_CLUSTER_02_TEST_URL: clickhouse://ch-cluster-02:9000/dbmate_test
       MYSQL_TEST_URL: mysql://root:root@mysql/dbmate_test
       POSTGRES_TEST_URL: postgres://postgres:postgres@postgres/dbmate_test?sslmode=disable
       SQLITE_TEST_URL: sqlite3:/tmp/dbmate_test.sqlite3
@@ -35,3 +39,27 @@ services:
 
   clickhouse:
     image: clickhouse/clickhouse-server:22.8
+
+  zookeeper:
+    image: zookeeper
+    hostname: zookeeper
+
+  clickhouse-cluster-01:
+    image: clickhouse/clickhouse-server:22.8
+    hostname: ch-cluster-01
+    environment:
+      - CLICKHOUSE_CONFIG=/etc/clickhouse-server/config.xml
+    depends_on:
+      - zookeeper
+    volumes:
+      - ./pkg/driver/clickhouse/testdata/cluster_config/ch-cluster-01:/etc/clickhouse-server
+
+  clickhouse-cluster-02:
+    image: clickhouse/clickhouse-server:22.8
+    hostname: ch-cluster-02
+    environment:
+      - CLICKHOUSE_CONFIG=/etc/clickhouse-server/config.xml
+    depends_on:
+      - zookeeper
+    volumes:
+      - ./pkg/driver/clickhouse/testdata/cluster_config/ch-cluster-02:/etc/clickhouse-server

pkg/driver/clickhouse/clickhouse.go

@@ -25,6 +25,7 @@ type Driver struct {
 	migrationsTableName string
 	databaseURL         *url.URL
 	log                 io.Writer
+	clusterParameters   *ClusterParameters
 }
 
 // NewDriver initializes the driver
@@ -33,6 +34,7 @@ func NewDriver(config dbmate.DriverConfig) dbmate.Driver {
 		migrationsTableName: config.MigrationsTableName,
 		databaseURL:         config.DatabaseURL,
 		log:                 config.Log,
+		clusterParameters:   ExtractClusterParametersFromURL(config.DatabaseURL),
 	}
 }
 
@@ -74,6 +76,8 @@ func connectionString(initialURL *url.URL) string {
 
 	u.RawQuery = query.Encode()
 
+	u = ClearClusterParametersFromURL(u)
+
 	return u.String()
 }
 
@@ -95,6 +99,15 @@ func (drv *Driver) openClickHouseDB() (*sql.DB, error) {
 	return sql.Open("clickhouse", clickhouseURL.String())
 }
 
+func (drv *Driver) onClusterClause() string {
+	clusterClause := ""
+	if drv.clusterParameters.OnCluster {
+		escapedClusterMacro := drv.escapeString(drv.clusterParameters.ClusterMacro)
+		clusterClause = fmt.Sprintf(" ON CLUSTER '%s'", escapedClusterMacro)
+	}
+	return clusterClause
+}
+
 func (drv *Driver) databaseName() string {
 	name := strings.TrimLeft(dbutil.MustParseURL(connectionString(drv.databaseURL)).Path, "/")
 	if name == "" {
@@ -115,6 +128,12 @@ func (drv *Driver) quoteIdentifier(str string) string {
 	return fmt.Sprintf(`"%s"`, str)
 }
 
+func (drv *Driver) escapeString(str string) string {
+	quoteEscaper := strings.NewReplacer(`'`, `\'`, `\`, `\\`)
+	str = quoteEscaper.Replace(str)
+	return str
+}
+
 // CreateDatabase creates the specified database
 func (drv *Driver) CreateDatabase() error {
 	name := drv.databaseName()
@@ -126,7 +145,9 @@ func (drv *Driver) CreateDatabase() error {
 	}
 	defer dbutil.MustClose(db)
 
-	_, err = db.Exec("create database " + drv.quoteIdentifier(name))
+	q := fmt.Sprintf("CREATE DATABASE %s%s", drv.quoteIdentifier(name), drv.onClusterClause())
+
+	_, err = db.Exec(q)
 
 	return err
 }
@@ -142,14 +163,16 @@ func (drv *Driver) DropDatabase() error {
 	}
 	defer dbutil.MustClose(db)
 
-	_, err = db.Exec("drop database if exists " + drv.quoteIdentifier(name))
+	q := fmt.Sprintf("DROP DATABASE IF EXISTS %s%s", drv.quoteIdentifier(name), drv.onClusterClause())
+
+	_, err = db.Exec(q)
 
 	return err
 }
 
 func (drv *Driver) schemaDump(db *sql.DB, buf *bytes.Buffer, databaseName string) error {
 	buf.WriteString("\n--\n-- Database schema\n--\n\n")
-	buf.WriteString("CREATE DATABASE IF NOT EXISTS " + drv.quoteIdentifier(databaseName) + ";\n\n")
+	buf.WriteString(fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s%s;\n\n", drv.quoteIdentifier(databaseName), drv.onClusterClause()))
 
 	tables, err := dbutil.QueryColumn(db, "show tables")
 	if err != nil {
@@ -250,15 +273,22 @@ func (drv *Driver) MigrationsTableExists(db *sql.DB) (bool, error) {
 
 // CreateMigrationsTable creates the schema migrations table
 func (drv *Driver) CreateMigrationsTable(db *sql.DB) error {
+	engineClause := "ReplacingMergeTree(ts)"
+	if drv.clusterParameters.OnCluster {
+		escapedZooPath := drv.escapeString(drv.clusterParameters.ZooPath)
+		escapedReplicaMacro := drv.escapeString(drv.clusterParameters.ReplicaMacro)
+		engineClause = fmt.Sprintf("ReplicatedReplacingMergeTree('%s', '%s', ts)", escapedZooPath, escapedReplicaMacro)
+	}
+
 	_, err := db.Exec(fmt.Sprintf(`
-		create table if not exists %s (
+		create table if not exists %s%s (
 			version String,
 			ts DateTime default now(),
 			applied UInt8 default 1
-		) engine = ReplacingMergeTree(ts)
+		) engine = %s
 		primary key version
 		order by version
-	`, drv.quotedMigrationsTableName()))
+	`, drv.quotedMigrationsTableName(), drv.onClusterClause(), engineClause))
 
 	return err
 }