Skip to content

Commit 6a2ec04

Browse files
k-anshulk-anshul
committed
Remove cluster check for entity type (#9713)
* use view instead of directly querying system.tables in clusterAllReplicas * tests fixes * review * nit * build fix
1 parent a42b0cc commit 6a2ec04

12 files changed

Lines changed: 150 additions & 114 deletions

File tree

‎runtime/drivers/clickhouse/clickhouse.go‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -860,14 +860,14 @@ func (c *Connection) checkBillingTableExists(ctx context.Context, cluster string
860860
return false, fmt.Errorf("failed to check if billing table exists: %w", err)
861861
}
862862
} else {
863+
// Query through view so ClickHouse applies normal SELECT access checks on system.tables.
864+
// Targeting system.tables directly requires SHOW COLUMNS on system.tables starting in ClickHouse 26.3.
863865
err := c.readDB.QueryRowxContext(ctx, fmt.Sprintf(`
864866
SELECT countIf(found) = count() AS exists_everywhere, countIf(found) > 0 AS exists_somewhere
865-
FROM
866-
(
867-
SELECT hostName() AS host, max((database = 'billing') AND (name = 'events')) AS found
868-
FROM clusterAllReplicas('%s', system.tables)
869-
GROUP BY host
870-
)`, cluster)).Scan(&existsEverywhere, &existsSomewhere)
867+
FROM clusterAllReplicas(%s, view(
868+
SELECT max((database = 'billing') AND (name = 'events')) AS found
869+
FROM system.tables
870+
))`, safeSQLString(cluster))).Scan(&existsEverywhere, &existsSomewhere)
871871
if err != nil {
872872
return false, fmt.Errorf("failed to check if billing table exists in cluster %q: %w", cluster, err)
873873
}

‎runtime/drivers/clickhouse/crud.go‎

Lines changed: 32 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,7 @@ type tableWriteMetrics struct {
2121
}
2222

2323
func (c *Connection) createTableAsSelect(ctx context.Context, name, sql string, outputProps *ModelOutputProperties, beforeCreate, afterCreate string) (*tableWriteMetrics, error) {
24-
var onClusterClause string
25-
if c.config.Cluster != "" {
26-
onClusterClause = "ON CLUSTER " + safeSQLName(c.config.Cluster)
27-
}
24+
onClusterClause := c.onClusterClause()
2825

2926
// Execute beforeCreate query
3027
if beforeCreate != "" {
@@ -116,14 +113,7 @@ func (c *Connection) insertTableAsSelect(ctx context.Context, name, sql string,
116113
}
117114

118115
if opts.Strategy == drivers.IncrementalStrategyPartitionOverwrite {
119-
_, onCluster, err := c.entityType(ctx, c.config.Database, name)
120-
if err != nil {
121-
return nil, err
122-
}
123-
onClusterClause := ""
124-
if onCluster {
125-
onClusterClause = "ON CLUSTER " + safeSQLName(c.config.Cluster)
126-
}
116+
onClusterClause := c.onClusterClause()
127117
// Get the engine info of the given table
128118
engine, err := c.getTableEngine(ctx, name)
129119
if err != nil {
@@ -210,13 +200,9 @@ func (c *Connection) insertTableAsSelect(ctx context.Context, name, sql string,
210200
}
211201

212202
if opts.Strategy == drivers.IncrementalStrategyMerge {
213-
_, onCluster, err := c.entityType(ctx, c.config.Database, name)
214-
if err != nil {
215-
return nil, err
216-
}
217203
// get the engine info of the given table - local table for distributed tables
218204
var n string
219-
if onCluster {
205+
if c.config.Cluster != "" {
220206
n = localTableName(name)
221207
} else {
222208
n = name
@@ -249,45 +235,46 @@ func (c *Connection) insertTableAsSelect(ctx context.Context, name, sql string,
249235
}
250236

251237
func (c *Connection) dropTable(ctx context.Context, name string) error {
252-
typ, onCluster, err := c.entityType(ctx, c.config.Database, name)
238+
typ, err := c.entityType(ctx, c.config.Database, name)
253239
if err != nil {
254240
return err
255241
}
256-
var onClusterClause string
257-
if onCluster {
258-
onClusterClause = "ON CLUSTER " + safeSQLName(c.config.Cluster)
259-
}
242+
243+
onClusterClause := c.onClusterClause()
244+
260245
switch typ {
261246
case "VIEW":
262247
return c.Exec(ctx, &drivers.Statement{
263-
Query: fmt.Sprintf("DROP VIEW %s %s", safeSQLName(name), onClusterClause),
248+
Query: fmt.Sprintf("DROP VIEW IF EXISTS %s %s", safeSQLName(name), onClusterClause),
264249
Priority: 100,
265250
})
266251
case "DICTIONARY":
267252
// first drop the dictionary
268253
err := c.Exec(ctx, &drivers.Statement{
269-
Query: fmt.Sprintf("DROP DICTIONARY %s %s", safeSQLName(name), onClusterClause),
254+
Query: fmt.Sprintf("DROP DICTIONARY IF EXISTS %s %s", safeSQLName(name), onClusterClause),
270255
Priority: 100,
271256
})
272257
// then drop the temp table
273258
_ = c.Exec(ctx, &drivers.Statement{
274-
Query: fmt.Sprintf("DROP TABLE %s %s", safeSQLName(tempTableForDictionary(name)), onClusterClause),
259+
Query: fmt.Sprintf("DROP TABLE IF EXISTS %s %s", safeSQLName(tempTableForDictionary(name)), onClusterClause),
275260
Priority: 100,
276261
})
277262
return err
278263
case "TABLE":
279264
// drop the main table
265+
// use IF EXISTS so drops succeed in cluster mode even for tables that don't exist on every node,
266+
// e.g. tables created before the cluster was configured or without a `_local` counterpart
280267
err := c.Exec(ctx, &drivers.Statement{
281-
Query: fmt.Sprintf("DROP TABLE %s %s", safeSQLName(name), onClusterClause),
268+
Query: fmt.Sprintf("DROP TABLE IF EXISTS %s %s", safeSQLName(name), onClusterClause),
282269
Priority: 100,
283270
})
284271
if err != nil {
285272
return err
286273
}
287274
// then drop the local table in case of cluster
288-
if onCluster && !strings.HasSuffix(name, "_local") {
275+
if c.config.Cluster != "" && !strings.HasSuffix(name, "_local") {
289276
return c.Exec(ctx, &drivers.Statement{
290-
Query: fmt.Sprintf("DROP TABLE %s %s", safeSQLName(localTableName(name)), onClusterClause),
277+
Query: fmt.Sprintf("DROP TABLE IF EXISTS %s %s", safeSQLName(localTableName(name)), onClusterClause),
291278
Priority: 100,
292279
})
293280
}
@@ -298,22 +285,19 @@ func (c *Connection) dropTable(ctx context.Context, name string) error {
298285
}
299286

300287
func (c *Connection) renameEntity(ctx context.Context, oldName, newName string) error {
301-
typ, onCluster, err := c.entityType(ctx, c.config.Database, oldName)
288+
typ, err := c.entityType(ctx, c.config.Database, oldName)
302289
if err != nil {
303290
return err
304291
}
305-
var onClusterClause string
306-
if onCluster {
307-
onClusterClause = "ON CLUSTER " + safeSQLName(c.config.Cluster)
308-
}
292+
onClusterClause := c.onClusterClause()
309293

310294
switch typ {
311295
case "VIEW":
312296
return c.renameView(ctx, oldName, newName, onClusterClause)
313297
case "DICTIONARY":
314298
return c.renameTable(ctx, oldName, newName, onClusterClause)
315299
case "TABLE":
316-
if !onCluster {
300+
if c.config.Cluster == "" {
317301
return c.renameTable(ctx, oldName, newName, onClusterClause)
318302
}
319303
// capture the full engine of the old distributed table
@@ -447,10 +431,7 @@ func (c *Connection) renameTable(ctx context.Context, oldName, newName, onCluste
447431
}
448432

449433
func (c *Connection) createTable(ctx context.Context, name, sql string, outputProps *ModelOutputProperties) error {
450-
var onClusterClause string
451-
if c.config.Cluster != "" {
452-
onClusterClause = "ON CLUSTER " + safeSQLName(c.config.Cluster)
453-
}
434+
onClusterClause := c.onClusterClause()
454435
var create strings.Builder
455436
create.WriteString("CREATE OR REPLACE TABLE ")
456437
if c.config.Cluster != "" {
@@ -512,7 +493,7 @@ func (c *Connection) createDistributedTable(ctx context.Context, name string, ou
512493
if c.config.Cluster == "" {
513494
return fmt.Errorf("clickhouse: cannot create distributed table without a cluster")
514495
}
515-
onClusterClause := "ON CLUSTER " + safeSQLName(c.config.Cluster)
496+
onClusterClause := c.onClusterClause()
516497

517498
var distributed strings.Builder
518499
database := "currentDatabase()"
@@ -534,10 +515,7 @@ func (c *Connection) createDistributedTable(ctx context.Context, name string, ou
534515
}
535516

536517
func (c *Connection) createDictionary(ctx context.Context, name, sql string, outputProps *ModelOutputProperties) error {
537-
var onClusterClause string
538-
if c.config.Cluster != "" {
539-
onClusterClause = "ON CLUSTER " + safeSQLName(c.config.Cluster)
540-
}
518+
onClusterClause := c.onClusterClause()
541519
if sql == "" {
542520
if outputProps.Columns == "" {
543521
return fmt.Errorf("clickhouse: no columns specified for dictionary %q", name)
@@ -642,7 +620,7 @@ func (c *Connection) getTableEngine(ctx context.Context, name string) (string, e
642620
return "", err
643621
}
644622
defer res.Close()
645-
if res.Next() {
623+
for res.Next() {
646624
if err := res.Scan(&engine); err != nil {
647625
return "", err
648626
}
@@ -695,9 +673,8 @@ func (c *Connection) getTablePartitions(ctx context.Context, name string) ([]str
695673
}
696674

697675
func (c *Connection) replacePartition(ctx context.Context, src, dest, part string) error {
698-
var onClusterClause string
676+
onClusterClause := c.onClusterClause()
699677
if c.config.Cluster != "" {
700-
onClusterClause = "ON CLUSTER " + safeSQLName(c.config.Cluster)
701678
dest = localTableName(dest)
702679
src = localTableName(src)
703680
}
@@ -718,9 +695,8 @@ func (c *Connection) syncReplica(ctx context.Context, tableName string) error {
718695
if err != nil {
719696
return err
720697
}
721-
onClusterClause := "ON CLUSTER " + safeSQLName(c.config.Cluster)
722698
return c.Exec(ctx, &drivers.Statement{
723-
Query: fmt.Sprintf("SYSTEM SYNC REPLICA %s %s.%s", onClusterClause, safeSQLName(database), safeSQLName(localTableName(tableName))),
699+
Query: fmt.Sprintf("SYSTEM SYNC REPLICA %s %s.%s", c.onClusterClause(), safeSQLName(database), safeSQLName(localTableName(tableName))),
724700
Priority: 1,
725701
})
726702
}
@@ -750,6 +726,14 @@ func (c *Connection) currentDatabase(ctx context.Context) (string, error) {
750726
return database, nil
751727
}
752728

729+
// onClusterClause returns the ON CLUSTER clause for DDL statements, or an empty string if no cluster is configured.
730+
func (c *Connection) onClusterClause() string {
731+
if c.config.Cluster == "" {
732+
return ""
733+
}
734+
return "ON CLUSTER " + safeSQLName(c.config.Cluster)
735+
}
736+
753737
func isReplicatedEngine(engine string) bool {
754738
return strings.Contains(strings.ToLower(engine), "replicated")
755739
}

‎runtime/drivers/clickhouse/information_schema.go‎

Lines changed: 11 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -474,43 +474,30 @@ func scanTables(rows *sqlx.Rows) ([]*drivers.OlapTable, error) {
474474
return res, nil
475475
}
476476

477-
func (c *Connection) entityType(ctx context.Context, db, name string) (typ string, onCluster bool, err error) {
477+
func (c *Connection) entityType(ctx context.Context, db, name string) (typ string, err error) {
478478
conn, release, err := c.acquireMetaConn(ctx)
479479
if err != nil {
480-
return "", false, err
480+
return "", err
481481
}
482482
defer func() { _ = release() }()
483483

484-
var q string
485-
if c.config.Cluster == "" {
486-
q = `SELECT
487-
multiIf(engine IN ('MaterializedView', 'View'), 'VIEW', engine = 'Dictionary', 'DICTIONARY', 'TABLE') AS type,
488-
0 AS is_on_cluster
489-
FROM system.tables AS t
490-
JOIN system.databases AS db ON t.database = db.name
491-
WHERE t.database = coalesce(?, currentDatabase()) AND t.name = ?`
492-
} else {
493-
q = `SELECT
494-
multiIf(engine IN ('MaterializedView', 'View'), 'VIEW', engine = 'Dictionary', 'DICTIONARY', 'TABLE') AS type,
495-
countDistinct(_shard_num) > 1 AS is_on_cluster
496-
FROM clusterAllReplicas(` + safeSQLName(c.config.Cluster) + `, system.tables) AS t
497-
JOIN system.databases AS db ON t.database = db.name
498-
WHERE t.database = coalesce(?, currentDatabase()) AND t.name = ?
499-
GROUP BY engine, t.name`
500-
}
484+
// A local system.tables lookup suffices even when a cluster is configured:
485+
// all DDL in cluster mode runs ON CLUSTER, so the entity always exists on the connected node.
486+
q := `SELECT multiIf(engine IN ('MaterializedView', 'View'), 'VIEW', engine = 'Dictionary', 'DICTIONARY', 'TABLE') AS type
487+
FROM system.tables
488+
WHERE database = coalesce(?, currentDatabase()) AND name = ?`
501489
var args []any
502490
if db == "" {
503491
args = []any{nil, name}
504492
} else {
505493
args = []any{db, name}
506494
}
507-
row := conn.QueryRowxContext(ctx, q, args...)
508-
err = row.Scan(&typ, &onCluster)
495+
err = conn.QueryRowxContext(ctx, q, args...).Scan(&typ)
509496
if err != nil {
510497
if errors.Is(err, sql.ErrNoRows) {
511-
return "", false, drivers.ErrNotFound
498+
return "", drivers.ErrNotFound
512499
}
513-
return "", false, err
500+
return "", err
514501
}
515-
return typ, onCluster, nil
502+
return typ, nil
516503
}

‎runtime/drivers/clickhouse/olap_test.go‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package clickhouse
33
import (
44
"context"
55
"fmt"
6+
"strings"
67
"testing"
78

89
runtimev1 "github.com/rilldata/rill/proto/gen/rill/runtime/v1"
@@ -29,6 +30,7 @@ func TestClickhouseSingle(t *testing.T) {
2930
require.True(t, ok)
3031

3132
t.Run("WithConnection", func(t *testing.T) { testWithConnection(t, olap) })
33+
t.Run("DropTable", func(t *testing.T) { testDropTable(t, c, olap) })
3234
t.Run("RenameView", func(t *testing.T) { testRenameView(t, c, olap) })
3335
t.Run("RenameTable", func(t *testing.T) { testRenameTable(t, c, olap) })
3436
t.Run("CreateTableAsSelect", func(t *testing.T) { testCreateTableAsSelect(t, c) })
@@ -59,6 +61,7 @@ func TestClickhouseCluster(t *testing.T) {
5961
prepareClusterConn(t, olap, cluster)
6062

6163
t.Run("WithConnection", func(t *testing.T) { testWithConnection(t, olap) })
64+
t.Run("DropTable", func(t *testing.T) { testDropTable(t, c, olap) })
6265
t.Run("RenameView", func(t *testing.T) { testRenameView(t, c, olap) })
6366
t.Run("RenameTable", func(t *testing.T) { testRenameTable(t, c, olap) })
6467
t.Run("CreateTableAsSelect", func(t *testing.T) { testCreateTableAsSelect(t, c) })
@@ -69,6 +72,8 @@ func TestClickhouseCluster(t *testing.T) {
6972
t.Run("SyncReplica_NonDefaultDatabase", func(t *testing.T) { testSyncReplicaNonDefaultDatabase(t, olap, dsn, cluster) })
7073
t.Run("TestDictionary", func(t *testing.T) { testDictionary(t, c, olap) })
7174
t.Run("QueryAttributesAsSettings", func(t *testing.T) { testQueryAttributesAsSettings(t, olap) })
75+
t.Run("EntityTypeRestrictedUser", func(t *testing.T) { testEntityTypeRestrictedUser(t, olap, dsn, cluster) })
76+
7277
}
7378

7479
func testWithConnection(t *testing.T, olap drivers.OLAPStore) {
@@ -98,6 +103,23 @@ func testWithConnection(t *testing.T, olap drivers.OLAPStore) {
98103
require.NoError(t, err)
99104
}
100105

106+
func testDropTable(t *testing.T, c *Connection, olap drivers.OLAPStore) {
107+
ctx := context.Background()
108+
109+
_, err := c.createTableAsSelect(ctx, "drop_table", "SELECT 1 AS id", &ModelOutputProperties{Engine: "MergeTree"}, "", "")
110+
require.NoError(t, err)
111+
require.NoError(t, c.dropTable(ctx, "drop_table"))
112+
notExists(t, olap, "drop_table")
113+
if c.config.Cluster != "" {
114+
notExists(t, olap, localTableName("drop_table"))
115+
}
116+
117+
_, err = c.createTableAsSelect(ctx, "drop_view", "SELECT 1 AS id", &ModelOutputProperties{Typ: "VIEW"}, "", "")
118+
require.NoError(t, err)
119+
require.NoError(t, c.dropTable(ctx, "drop_view"))
120+
notExists(t, olap, "drop_view")
121+
}
122+
101123
func testRenameView(t *testing.T, c *Connection, olap drivers.OLAPStore) {
102124
ctx := context.Background()
103125
opts := &ModelOutputProperties{Typ: "VIEW"}
@@ -1027,3 +1049,46 @@ func testQueryAttributesAsSettings(t *testing.T, olap drivers.OLAPStore) {
10271049
require.Contains(t, err.Error(), "neither a builtin setting nor")
10281050
})
10291051
}
1052+
1053+
func testEntityTypeRestrictedUser(t *testing.T, olap drivers.OLAPStore, dsn, cluster string) {
1054+
ctx := context.Background()
1055+
const (
1056+
username = "rill_restricted"
1057+
password = "test_password"
1058+
)
1059+
1060+
require.NoError(t, olap.Exec(ctx, &drivers.Statement{
1061+
Query: fmt.Sprintf("CREATE USER %s ON CLUSTER %s IDENTIFIED WITH sha256_password BY %s", safeSQLName(username), safeSQLName(cluster), safeSQLString(password)),
1062+
}))
1063+
t.Cleanup(func() {
1064+
_ = olap.Exec(context.Background(), &drivers.Statement{Query: fmt.Sprintf("DROP USER IF EXISTS %s ON CLUSTER %s", safeSQLName(username), safeSQLName(cluster))})
1065+
})
1066+
require.NoError(t, olap.Exec(ctx, &drivers.Statement{
1067+
Query: fmt.Sprintf("GRANT ON CLUSTER %s SELECT ON default.* TO %s", safeSQLName(cluster), safeSQLName(username)),
1068+
}))
1069+
require.NoError(t, olap.Exec(ctx, &drivers.Statement{
1070+
Query: fmt.Sprintf("GRANT ON CLUSTER %s CLUSTER, REMOTE ON *.* TO %s", safeSQLName(cluster), safeSQLName(username)),
1071+
}))
1072+
1073+
// create a dedicated table: other subtests rename or drop the shared tables created by prepareClusterConn
1074+
require.NoError(t, olap.Exec(ctx, &drivers.Statement{
1075+
Query: fmt.Sprintf("CREATE OR REPLACE TABLE entity_type_restricted ON CLUSTER %s (x Int32) engine=MergeTree ORDER BY x", safeSQLName(cluster)),
1076+
}))
1077+
t.Cleanup(func() {
1078+
_ = olap.Exec(context.Background(), &drivers.Statement{Query: fmt.Sprintf("DROP TABLE IF EXISTS entity_type_restricted ON CLUSTER %s", safeSQLName(cluster))})
1079+
})
1080+
1081+
restrictedDSN := strings.Replace(dsn, "clickhouse://default@", fmt.Sprintf("clickhouse://%s:%s@", username, password), 1)
1082+
handle, err := drivers.Open("clickhouse", "", "restricted", map[string]any{"dsn": restrictedDSN, "cluster": cluster}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop())
1083+
require.NoError(t, err)
1084+
defer handle.Close()
1085+
1086+
restrictedConnection := handle.(*Connection)
1087+
typ, err := restrictedConnection.entityType(ctx, "default", "entity_type_restricted")
1088+
require.NoError(t, err)
1089+
require.Equal(t, "TABLE", typ)
1090+
1091+
exists, err := restrictedConnection.checkBillingTableExists(ctx, cluster)
1092+
require.NoError(t, err)
1093+
require.False(t, exists)
1094+
}

‎runtime/drivers/clickhouse/testclickhouse/testclickhouse.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ func Start(t TestingT) string {
3333
ctx := context.Background()
3434
clickHouseContainer, err := clickhouse.Run(
3535
ctx,
36-
"clickhouse/clickhouse-server:25.6.12.10",
36+
"clickhouse/clickhouse-server:26.3.1.896",
3737
clickhouse.WithConfigFile(filepath.Join(testdataPath, "clickhouse-config.xml")),
3838
testcontainers.CustomizeRequestOption(func(req *testcontainers.GenericContainerRequest) error {
3939
cf := testcontainers.ContainerFile{

0 commit comments

Comments
 (0)