From faaf5d3708ca1d69e32b0e6468ec6b15a4d75ac6 Mon Sep 17 00:00:00 2001
From: Vladimir Mihailenco <vladimir.webdev@gmail.com>
Date: Thu, 7 Sep 2023 11:03:37 +0300
Subject: [PATCH] fix: escape cluster name

---
 ch/chschema/formatter.go |  4 ++++
 ch/chschema/sqlfmt.go    |  2 +-
 ch/query_table_create.go |  4 ++--
 ch/query_table_drop.go   |  4 ++--
 ch/query_view_create.go  | 10 +++++-----
 ch/query_view_drop.go    | 14 +++++++-------
 chmigrate/migrator.go    |  2 +-
 7 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/ch/chschema/formatter.go b/ch/chschema/formatter.go
index 3fcd868..1aa39cf 100644
--- a/ch/chschema/formatter.go
+++ b/ch/chschema/formatter.go
@@ -31,6 +31,10 @@ func (f Formatter) AppendIdent(b []byte, ident string) []byte {
 	return AppendIdent(b, ident)
 }
 
+func (f Formatter) AppendFQN(b []byte, ident string) []byte {
+	return AppendFQN(b, ident)
+}
+
 func (f Formatter) WithArg(arg NamedArgAppender) Formatter {
 	return Formatter{
 		args: f.args.WithArg(arg),
diff --git a/ch/chschema/sqlfmt.go b/ch/chschema/sqlfmt.go
index 874c596..230bd30 100644
--- a/ch/chschema/sqlfmt.go
+++ b/ch/chschema/sqlfmt.go
@@ -34,7 +34,7 @@ type FQN string
 var _ QueryAppender = (*FQN)(nil)
 
 func (s FQN) AppendQuery(fmter Formatter, b []byte) ([]byte, error) {
-	return fmter.AppendIdent(b, string(s)), nil
+	return fmter.AppendFQN(b, string(s)), nil
 }
 
 func AppendFQN(b []byte, field string) []byte {
diff --git a/ch/query_table_create.go b/ch/query_table_create.go
index 0757e06..3f2d14c 100644
--- a/ch/query_table_create.go
+++ b/ch/query_table_create.go
@@ -69,8 +69,8 @@ func (q *CreateTableQuery) IfNotExists() *CreateTableQuery {
 	return q
 }
 
-func (q *CreateTableQuery) OnCluster(query string, args ...any) *CreateTableQuery {
-	q.onCluster = chschema.SafeQuery(query, args)
+func (q *CreateTableQuery) OnCluster(cluster string) *CreateTableQuery {
+	q.onCluster = chschema.UnsafeIdent(cluster)
 	return q
 }
 
diff --git a/ch/query_table_drop.go b/ch/query_table_drop.go
index f9feca7..aee9139 100644
--- a/ch/query_table_drop.go
+++ b/ch/query_table_drop.go
@@ -57,8 +57,8 @@ func (q *DropTableQuery) IfExists() *DropTableQuery {
 	return q
 }
 
-func (q *DropTableQuery) OnCluster(query string, args ...any) *DropTableQuery {
-	q.onCluster = chschema.SafeQuery(query, args)
+func (q *DropTableQuery) OnCluster(cluster string) *DropTableQuery {
+	q.onCluster = chschema.UnsafeIdent(cluster)
 	return q
 }
 
diff --git a/ch/query_view_create.go b/ch/query_view_create.go
index 14052c4..a333247 100644
--- a/ch/query_view_create.go
+++ b/ch/query_view_create.go
@@ -14,7 +14,7 @@ type CreateViewQuery struct {
 	materialized bool
 	ifNotExists  bool
 	view         chschema.QueryWithArgs
-	cluster      chschema.QueryWithArgs
+	onCluster    chschema.QueryWithArgs
 	to           chschema.QueryWithArgs
 	where        whereQuery
 	group        []chschema.QueryWithArgs
@@ -53,12 +53,12 @@ func (q *CreateViewQuery) ViewExpr(query string, args ...any) *CreateViewQuery {
 }
 
 func (q *CreateViewQuery) OnCluster(cluster string) *CreateViewQuery {
-	q.cluster = chschema.UnsafeIdent(cluster)
+	q.onCluster = chschema.UnsafeIdent(cluster)
 	return q
 }
 
 func (q *CreateViewQuery) OnClusterExpr(query string, args ...any) *CreateViewQuery {
-	q.cluster = chschema.SafeQuery(query, args)
+	q.onCluster = chschema.SafeQuery(query, args)
 	return q
 }
 
@@ -197,9 +197,9 @@ func (q *CreateViewQuery) AppendQuery(fmter chschema.Formatter, b []byte) (_ []b
 		return nil, err
 	}
 
-	if !q.cluster.IsEmpty() {
+	if !q.onCluster.IsEmpty() {
 		b = append(b, " ON CLUSTER "...)
-		b, err = q.cluster.AppendQuery(fmter, b)
+		b, err = q.onCluster.AppendQuery(fmter, b)
 		if err != nil {
 			return nil, err
 		}
diff --git a/ch/query_view_drop.go b/ch/query_view_drop.go
index 900118d..8a3ccad 100644
--- a/ch/query_view_drop.go
+++ b/ch/query_view_drop.go
@@ -11,9 +11,9 @@ import (
 type DropViewQuery struct {
 	baseQuery
 
-	ifExists bool
-	view     chschema.QueryWithArgs
-	cluster  chschema.QueryWithArgs
+	ifExists  bool
+	view      chschema.QueryWithArgs
+	onCluster chschema.QueryWithArgs
 }
 
 var _ Query = (*DropViewQuery)(nil)
@@ -54,12 +54,12 @@ func (q *DropViewQuery) ViewExpr(query string, args ...any) *DropViewQuery {
 }
 
 func (q *DropViewQuery) OnCluster(cluster string) *DropViewQuery {
-	q.cluster = chschema.UnsafeIdent(cluster)
+	q.onCluster = chschema.UnsafeIdent(cluster)
 	return q
 }
 
 func (q *DropViewQuery) OnClusterExpr(query string, args ...any) *DropViewQuery {
-	q.cluster = chschema.SafeQuery(query, args)
+	q.onCluster = chschema.SafeQuery(query, args)
 	return q
 }
 
@@ -84,9 +84,9 @@ func (q *DropViewQuery) AppendQuery(fmter chschema.Formatter, b []byte) (_ []byt
 		return nil, err
 	}
 
-	if !q.cluster.IsEmpty() {
+	if !q.onCluster.IsEmpty() {
 		b = append(b, " ON CLUSTER "...)
-		b, err = q.cluster.AppendQuery(fmter, b)
+		b, err = q.onCluster.AppendQuery(fmter, b)
 		if err != nil {
 			return nil, err
 		}
diff --git a/chmigrate/migrator.go b/chmigrate/migrator.go
index 6b62015..659f7c7 100644
--- a/chmigrate/migrator.go
+++ b/chmigrate/migrator.go
@@ -40,7 +40,7 @@ func WithOnCluster(cluster string) MigratorOption {
 }
 
 // WithMarkAppliedOnSuccess sets the migrator to only mark migrations as applied/unapplied
-// when their up/down is successful
+// when their up/down is successful.
 func WithMarkAppliedOnSuccess(enabled bool) MigratorOption {
 	return func(m *Migrator) {
 		m.markAppliedOnSuccess = enabled