chore: improve create table query

uptrace · Mar 8, 2022 · 9f0d6ab · 9f0d6ab
1 parent 1c1b01f
commit 9f0d6ab
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 2 deletions.
diff --git a/ch/query_table_create.go b/ch/query_table_create.go
@@ -13,6 +13,7 @@ type CreateTableQuery struct {
 
 	ifNotExists bool
 	engine      chschema.QueryWithArgs
+	ttl         chschema.QueryWithArgs
 	partition   chschema.QueryWithArgs
 }
 
@@ -62,6 +63,26 @@ func (q *CreateTableQuery) IfNotExists() *CreateTableQuery {
 	return q
 }
 
+func (q *CreateTableQuery) Engine(query string, args ...any) *CreateTableQuery {
+	q.engine = chschema.SafeQuery(query, args)
+	return q
+}
+
+func (q *CreateTableQuery) TTL(query string, args ...any) *CreateTableQuery {
+	q.ttl = chschema.SafeQuery(query, args)
+	return q
+}
+
+func (q *CreateTableQuery) Partition(query string, args ...any) *CreateTableQuery {
+	q.partition = chschema.SafeQuery(query, args)
+	return q
+}
+
+func (q *CreateTableQuery) Setting(query string, args ...any) *CreateTableQuery {
+	q.settings = append(q.settings, chschema.SafeQuery(query, args))
+	return q
+}
+
 //------------------------------------------------------------------------------
 
 func (q *CreateTableQuery) Operation() string {
@@ -132,6 +153,11 @@ func (q *CreateTableQuery) AppendQuery(fmter chschema.Formatter, b []byte) (_ []
 		b = append(b, "MergeTree()"...)
 	}
 
+	b, err = q.appendPartition(fmter, b)
+	if err != nil {
+		return nil, err
+	}
+
 	b = append(b, " ORDER BY "...)
 	if len(q.table.PKs) > 0 {
 		b = append(b, '(')
@@ -146,7 +172,15 @@ func (q *CreateTableQuery) AppendQuery(fmter chschema.Formatter, b []byte) (_ []
 		b = append(b, "tuple()"...)
 	}
 
-	b, err = q.appendPartition(fmter, b)
+	if !q.ttl.IsZero() {
+		b = append(b, " TTL "...)
+		b, err = q.ttl.AppendQuery(fmter, b)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	b, err = q.appendSettings(fmter, b)
 	if err != nil {
 		return nil, err
 	}
@@ -161,7 +195,7 @@ func (q *CreateTableQuery) appendPartition(fmter chschema.Formatter, b []byte) (
 
 	b = append(b, " PARTITION BY "...)
 	if !q.partition.IsZero() {
-		return q.engine.AppendQuery(fmter, b)
+		return q.partition.AppendQuery(fmter, b)
 	}
 	return append(b, q.table.CHPartition...), nil
 }

diff --git a/ch/query_test.go b/ch/query_test.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/bradleyjkemp/cupaloy"
 	"github.com/uptrace/go-clickhouse/ch"
@@ -61,6 +62,19 @@ func TestQuery(t *testing.T) {
 				Model((*Model)(nil)).
 				Sample("?", 1000)
 		},
+		func(db *ch.DB) chschema.QueryAppender {
+			type Model struct {
+				ch.CHModel `ch:"table:spans,partition:toYYYYMM(time)"`
+
+				ID   uint64
+				Text string    `ch:",lc"` // low cardinality column
+				Time time.Time `ch:",pk"` // ClickHouse primary key for order by
+			}
+			return db.NewCreateTable().Model((*Model)(nil)).
+				TTL("time + INTERVAL 30 DAY DELETE").
+				Partition("toDate(time)").
+				Setting("ttl_only_drop_parts = 1")
+		},
 	}
 
 	db := chDB()

diff --git a/ch/testdata/snapshots/TestQuery-10 b/ch/testdata/snapshots/TestQuery-10
@@ -0,0 +1 @@
+CREATE TABLE "spans" (id UInt64, text LowCardinality(String), time DateTime) Engine = MergeTree() PARTITION BY toDate(time) ORDER BY (time) TTL time + INTERVAL 30 DAY DELETE SETTINGS ttl_only_drop_parts = 1