From c0e67d2ec49a43437fac539ec9180b54590160c5 Mon Sep 17 00:00:00 2001
From: Vladimir Mihailenco <vladimir.webdev@gmail.com>
Date: Fri, 19 May 2023 09:30:07 +0300
Subject: [PATCH] chore: improve auto-create database

---
 ch/config.go | 7 +++++++
 ch/db.go     | 7 ++++++-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/ch/config.go b/ch/config.go
index 861e25a..76acc20 100644
--- a/ch/config.go
+++ b/ch/config.go
@@ -29,6 +29,7 @@ type Config struct {
 	User     string
 	Password string
 	Database string
+	Cluster  string
 
 	DialTimeout   time.Duration
 	TLSConfig     *tls.Config
@@ -151,6 +152,12 @@ func WithDatabase(database string) Option {
 	}
 }
 
+func WithCluster(cluster string) Option {
+	return func(db *DB) {
+		db.conf.Cluster = cluster
+	}
+}
+
 // WithDialTimeout configures dial timeout for establishing new connections.
 // Default is 5 seconds.
 func WithDialTimeout(timeout time.Duration) Option {
diff --git a/ch/db.go b/ch/db.go
index 8d0df94..4d7a594 100644
--- a/ch/db.go
+++ b/ch/db.go
@@ -130,10 +130,15 @@ func (db *DB) autoCreateDatabase() {
 	conf := db.conf.clone()
 	conf.Database = ""
 
+	query := "CREATE DATABASE IF NOT EXISTS ?"
+	if conf.Cluster != "" {
+		query += " ON CLUSTER ?"
+	}
+
 	tmp := newDB(conf)
 	defer tmp.Close()
 
-	if _, err := tmp.Exec("CREATE DATABASE IF NOT EXISTS ?", Ident(db.conf.Database)); err != nil {
+	if _, err := tmp.Exec(query, Ident(db.conf.Database), Ident(db.conf.Cluster)); err != nil {
 		internal.Logger.Printf("create database %q failed: %s", db.conf.Database, err)
 	}
 }