diff --git a/runtime/drivers/clickhouse/crud.go b/runtime/drivers/clickhouse/crud.go index 8799113096e..86f21a72eaf 100644 --- a/runtime/drivers/clickhouse/crud.go +++ b/runtime/drivers/clickhouse/crud.go @@ -8,6 +8,7 @@ import ( "strings" "time" + "github.com/google/uuid" "github.com/rilldata/rill/runtime/drivers" "github.com/rilldata/rill/runtime/pkg/graceful" "github.com/rilldata/rill/runtime/pkg/observability" @@ -249,17 +250,26 @@ func (c *Connection) dropTable(ctx context.Context, name string) error { Priority: 100, }) case "DICTIONARY": + // Resolve the source table up front, since the dependency that identifies it disappears with the dictionary. + srcTable, err := c.dictionarySourceTable(ctx, name) + if err != nil { + return err + } // first drop the dictionary - err := c.Exec(ctx, &drivers.Statement{ + err = c.Exec(ctx, &drivers.Statement{ Query: fmt.Sprintf("DROP DICTIONARY IF EXISTS %s %s", safeSQLName(name), onClusterClause), Priority: 100, }) - // then drop the temp table - _ = c.Exec(ctx, &drivers.Statement{ - Query: fmt.Sprintf("DROP TABLE IF EXISTS %s %s", safeSQLName(tempTableForDictionary(name)), onClusterClause), - Priority: 100, - }) - return err + if err != nil { + return err + } + // then drop the table it sourced from, which is now unreferenced + if srcTable != "" { + if dropErr := c.dropTable(ctx, srcTable); dropErr != nil && !errors.Is(dropErr, drivers.ErrNotFound) { + c.logger.Warn("clickhouse: failed to drop dictionary source table", zap.String("name", srcTable), zap.Error(dropErr), observability.ZapCtx(ctx)) + } + } + return nil case "TABLE": // drop the main table // use IF EXISTS so drops succeed in cluster mode even for tables that don't exist on every node, @@ -526,15 +536,39 @@ func (c *Connection) createDictionary(ctx context.Context, name, sql string, out }) } - // create a temp table first - // NOTE :: this can only be dropped when the dictionary is dropped - tempTable := tempTableForDictionary(name) - err := c.createTable(ctx, tempTable, sql, outputProps) + if outputProps.PrimaryKey == "" { + return fmt.Errorf("clickhouse: no primary key specified for dictionary %q", name) + } + + // Note the table the dictionary currently sources from, so it can be dropped once the dictionary is repointed. + oldSrcTable, err := c.dictionarySourceTable(ctx, name) + if err != nil { + return err + } + + // Write the new data to its own table instead of reusing oldSrcTable, which ClickHouse would not let us replace + // while the dictionary depends on it. It also means the dictionary keeps serving until the new data is ready. + srcTable := newDictionarySourceTable(name) + var repointed bool + defer func() { + if repointed { + return + } + ctx, cancel := graceful.WithMinimumDuration(ctx, 15*time.Second) + defer cancel() + + err := c.dropTable(ctx, srcTable) + if err != nil && !errors.Is(err, drivers.ErrNotFound) { + c.logger.Warn("clickhouse: failed to drop dictionary source table", zap.String("name", srcTable), zap.Error(err), observability.ZapCtx(ctx)) + } + }() + + err = c.createTable(ctx, srcTable, sql, outputProps) if err != nil { return err } err = c.Exec(ctx, &drivers.Statement{ - Query: fmt.Sprintf("INSERT INTO %s %s", safeSQLName(tempTable), sql), + Query: fmt.Sprintf("INSERT INTO %s %s", safeSQLName(srcTable), sql), Priority: 100, }) if err != nil { @@ -543,29 +577,38 @@ func (c *Connection) createDictionary(ctx context.Context, name, sql string, out if outputProps.Columns == "" { // infer columns - outputProps.Columns, err = c.columnClause(ctx, tempTable) + outputProps.Columns, err = c.columnClause(ctx, srcTable) if err != nil { return err } } - if outputProps.PrimaryKey == "" { - return fmt.Errorf("clickhouse: no primary key specified for dictionary %q", name) - } - - srcTbl := fmt.Sprintf("CLICKHOUSE(TABLE %s)", drivers.EscapeStringValue(tempTable)) + srcTbl := fmt.Sprintf("CLICKHOUSE(TABLE %s)", drivers.EscapeStringValue(srcTable)) if outputProps.DictionarySourceUser != "" { if outputProps.DictionarySourcePassword == "" { return fmt.Errorf("clickhouse: no password specified for dictionary user") } - srcTbl = fmt.Sprintf("CLICKHOUSE(TABLE %s USER %s PASSWORD %s)", drivers.EscapeStringValue(tempTable), safeSQLString(outputProps.DictionarySourceUser), safeSQLString(outputProps.DictionarySourcePassword)) + srcTbl = fmt.Sprintf("CLICKHOUSE(TABLE %s USER %s PASSWORD %s)", drivers.EscapeStringValue(srcTable), safeSQLString(outputProps.DictionarySourceUser), safeSQLString(outputProps.DictionarySourcePassword)) } // create dictionary - return c.Exec(ctx, &drivers.Statement{ + err = c.Exec(ctx, &drivers.Statement{ Query: fmt.Sprintf(`CREATE OR REPLACE DICTIONARY %s %s %s PRIMARY KEY %s SOURCE(%s) LAYOUT(HASHED()) LIFETIME(0)`, safeSQLName(name), onClusterClause, outputProps.Columns, outputProps.PrimaryKey, srcTbl), Priority: 100, }) + if err != nil { + return err + } + repointed = true + + // The dictionary no longer depends on its previous source table, so it can be dropped. + if oldSrcTable != "" { + err = c.dropTable(ctx, oldSrcTable) + if err != nil && !errors.Is(err, drivers.ErrNotFound) { + c.logger.Warn("clickhouse: failed to drop previous dictionary source table", zap.String("name", oldSrcTable), zap.Error(err), observability.ZapCtx(ctx)) + } + } + return nil } func (c *Connection) columnClause(ctx context.Context, table string) (string, error) { @@ -738,8 +781,51 @@ func isReplicatedEngine(engine string) bool { return strings.Contains(strings.ToLower(engine), "replicated") } -func tempTableForDictionary(name string) string { - return name + "_dict_temp_" +const dictionarySourceTableInfix = "_dict_temp_" + +// newDictionarySourceTable returns a unique name for a table for the dictionary `name` to source from. +// The name is unique per call because ClickHouse forbids dropping or renaming a table that a dictionary depends on, +// so refreshing a dictionary has to write to a new table and repoint the dictionary at it. +func newDictionarySourceTable(name string) string { + // A dictionary is staged by creating it under a temporary name and renaming it into place, but its source table + // is not renamed along with it. Strip the staging prefix so the source table keeps a readable, stable name. + name = strings.TrimPrefix(name, stagingTablePrefix) + return name + dictionarySourceTableInfix + strings.ReplaceAll(uuid.New().String(), "-", "") +} + +// dictionarySourceTable returns the table the dictionary `name` currently sources from. +// It returns an empty string if the dictionary does not exist or does not source from a table created by Rill. +func (c *Connection) dictionarySourceTable(ctx context.Context, name string) (string, error) { + args := []any{c.config.Database, name} + if c.config.Database == "" { + args = []any{nil, name} + } + res, err := c.Query(ctx, &drivers.Statement{ + Query: "SELECT loading_dependencies_table FROM system.tables WHERE database = coalesce(?, currentDatabase()) AND name = ?", + Args: args, + Priority: 100, + }) + if err != nil { + return "", err + } + defer res.Close() + + var deps []string + for res.Next() { + if err := res.Scan(&deps); err != nil { + return "", err + } + } + if err := res.Err(); err != nil { + return "", err + } + + for _, dep := range deps { + if strings.Contains(dep, dictionarySourceTableInfix) { + return dep, nil + } + } + return "", nil } func safeSQLString(name string) string { diff --git a/runtime/drivers/clickhouse/model_executor_self_test.go b/runtime/drivers/clickhouse/model_executor_self_test.go index 409f2c6e59c..b86f3ca5a90 100644 --- a/runtime/drivers/clickhouse/model_executor_self_test.go +++ b/runtime/drivers/clickhouse/model_executor_self_test.go @@ -2,19 +2,58 @@ package clickhouse_test import ( "fmt" + "regexp" + "sort" "strings" "testing" "time" runtimev1 "github.com/rilldata/rill/proto/gen/rill/runtime/v1" "github.com/rilldata/rill/runtime" + "github.com/rilldata/rill/runtime/drivers" + "github.com/rilldata/rill/runtime/drivers/clickhouse/testclickhouse" + "github.com/rilldata/rill/runtime/pkg/activity" + "github.com/rilldata/rill/runtime/storage" "github.com/rilldata/rill/runtime/testruntime" "github.com/stretchr/testify/require" + "go.uber.org/zap" _ "github.com/rilldata/rill/runtime/resolvers" ) -func TestMaterializeType(t *testing.T) { +func TestClickhouseModels(t *testing.T) { + dsn := testclickhouse.Start(t) + + t.Run("MaterializeType", func(t *testing.T) { testMaterializeType(t, dsn) }) + t.Run("PartitionOverwrite", func(t *testing.T) { testPartitionOverwrite(t, dsn) }) + t.Run("StagedPostExecRunsAgainstFinalTable", func(t *testing.T) { testStagedPostExecRunsAgainstFinalTable(t, dsn) }) + t.Run("StagedDictionaryRefresh", func(t *testing.T) { testStagedDictionaryRefresh(t, dsn) }) + t.Run("DictionaryModelRename", func(t *testing.T) { testDictionaryModelRename(t, dsn) }) +} + +// newInstance creates a runtime instance on the shared ClickHouse instance started by TestClickhouseModels. +// Each instance gets its own database so that tests do not observe tables created by the others. +func newInstance(t *testing.T, dsn string, opts testruntime.InstanceOptions) (*runtime.Runtime, string) { + t.Helper() + + database := nonAlphanumeric.ReplaceAllString(t.Name(), "_") + conn, err := drivers.Open("clickhouse", "", "default", map[string]any{"dsn": dsn, "mode": "readwrite"}, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + require.NoError(t, err) + defer conn.Close() + olap, ok := conn.AsOLAP("") + require.True(t, ok) + require.NoError(t, olap.Exec(t.Context(), &drivers.Statement{Query: fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s", database)})) + + opts.Variables = map[string]string{ + "connector.clickhouse.dsn": fmt.Sprintf("%s/%s", dsn, database), + "connector.clickhouse.mode": "readwrite", + } + return testruntime.NewInstanceWithOptions(t, opts) +} + +var nonAlphanumeric = regexp.MustCompile(`[^a-zA-Z0-9]`) + +func testMaterializeType(t *testing.T, dsn string) { truth, falsity := true, false cases := []struct { name string @@ -52,10 +91,7 @@ func TestMaterializeType(t *testing.T) { files[fmt.Sprintf("%s.yaml", c.name)] = data } - rt, id := testruntime.NewInstanceWithOptions(t, testruntime.InstanceOptions{ - TestConnectors: []string{"clickhouse"}, - Files: files, - }) + rt, id := newInstance(t, dsn, testruntime.InstanceOptions{Files: files}) testruntime.ReconcileParserAndWait(t, rt, id) for _, c := range cases { @@ -74,7 +110,7 @@ func TestMaterializeType(t *testing.T) { } } -func TestPartitionOverwrite(t *testing.T) { +func testPartitionOverwrite(t *testing.T, dsn string) { files := map[string]string{ "rill.yaml": "olap_connector: clickhouse", // Model that creates 10 distinct partitions with 10 rows each. @@ -114,10 +150,7 @@ sql: SELECT number as num FROM numbers(10) `, } - rt, id := testruntime.NewInstanceWithOptions(t, testruntime.InstanceOptions{ - TestConnectors: []string{"clickhouse"}, - Files: files, - }) + rt, id := newInstance(t, dsn, testruntime.InstanceOptions{Files: files}) testruntime.ReconcileParserAndWait(t, rt, id) testruntime.RequireReconcileState(t, rt, id, 4, 0, 0) @@ -150,10 +183,9 @@ sql: SELECT number as num FROM numbers(10) }) } -func TestStagedPostExecRunsAgainstFinalTable(t *testing.T) { - rt, id := testruntime.NewInstanceWithOptions(t, testruntime.InstanceOptions{ - TestConnectors: []string{"clickhouse"}, - StageChanges: true, +func testStagedPostExecRunsAgainstFinalTable(t *testing.T, dsn string) { + rt, id := newInstance(t, dsn, testruntime.InstanceOptions{ + StageChanges: true, Files: map[string]string{ "rill.yaml": "olap_connector: clickhouse", "staged_ch.yaml": ` @@ -174,3 +206,114 @@ post_exec: CREATE TABLE staged_ch_marker ENGINE=Memory AS SELECT count() AS c FR Result: []map[string]any{{"c": 2}}, }) } + +func testStagedDictionaryRefresh(t *testing.T, dsn string) { + model := func(sql string) map[string]string { + return map[string]string{"campaign_name_dict.yaml": fmt.Sprintf(` +type: model +materialize: true +sql: %s +output: + type: dictionary + primary_key: id + dictionary_source_user: default + dictionary_source_password: default +`, sql)} + } + + rt, id := newInstance(t, dsn, testruntime.InstanceOptions{ + StageChanges: true, + Files: map[string]string{ + "rill.yaml": "olap_connector: clickhouse", + "campaign_name_dict.yaml": model(`SELECT toUInt64(1) AS id, 'a' AS name`)["campaign_name_dict.yaml"], + }, + }) + testruntime.ReconcileParserAndWait(t, rt, id) + testruntime.RequireReconcileState(t, rt, id, 2, 0, 0) + requireDictionary(t, rt, id, []map[string]any{{"name": "a"}}) + + // Refreshing must repoint the dictionary at a fresh source table without leaving the old one behind. + testruntime.RefreshAndWait(t, rt, id, &runtimev1.ResourceName{Kind: runtime.ResourceKindModel, Name: "campaign_name_dict"}) + testruntime.RequireReconcileState(t, rt, id, 2, 0, 0) + requireDictionary(t, rt, id, []map[string]any{{"name": "a"}}) + + // A refresh that changes both the data and the schema must be picked up too. + testruntime.PutFiles(t, rt, id, model(`SELECT toUInt64(1) AS id, 'b' AS name, 'x' AS extra`)) + testruntime.ReconcileParserAndWait(t, rt, id) + testruntime.RequireReconcileState(t, rt, id, 2, 0, 0) + requireDictionary(t, rt, id, []map[string]any{{"name": "b", "extra": "x"}}) +} + +func testDictionaryModelRename(t *testing.T, dsn string) { + rt, id := newInstance(t, dsn, testruntime.InstanceOptions{ + StageChanges: true, + Files: map[string]string{ + "rill.yaml": "olap_connector: clickhouse", + "dict_a.yaml": ` +type: model +materialize: true +sql: SELECT toUInt64(1) AS id, 'a' AS name +output: + type: dictionary + primary_key: id + dictionary_source_user: default + dictionary_source_password: default +`, + }, + }) + testruntime.ReconcileParserAndWait(t, rt, id) + testruntime.RequireReconcileState(t, rt, id, 2, 0, 0) + + // Renaming the model renames the dictionary but not the table it sources from, so the source table must not be + // identified by the dictionary's name, and the rename must not strand it. + testruntime.RenameFile(t, rt, id, "dict_a.yaml", "dict_b.yaml") + testruntime.ReconcileParserAndWait(t, rt, id) + testruntime.RequireReconcileState(t, rt, id, 2, 0, 0) + + testruntime.RequireResolve(t, rt, id, &testruntime.RequireResolveOptions{ + Resolver: "sql", + Properties: map[string]any{"sql": `SELECT name FROM dict_b`}, + Result: []map[string]any{{"name": "a"}}, + }) + + // Refreshing after the rename must still find and drop the old source table rather than accumulating one. + testruntime.RefreshAndWait(t, rt, id, &runtimev1.ResourceName{Kind: runtime.ResourceKindModel, Name: "dict_b"}) + testruntime.RequireReconcileState(t, rt, id, 2, 0, 0) + testruntime.RequireResolve(t, rt, id, &testruntime.RequireResolveOptions{ + Resolver: "sql", + Properties: map[string]any{"sql": ` + SELECT count() AS c FROM system.tables + WHERE database = currentDatabase() AND position(name, '_dict_temp_') > 0`}, + Result: []map[string]any{{"c": 1}}, + }) +} + +// requireDictionary asserts the contents of the campaign_name_dict dictionary, and that it is backed by exactly one +// source table with no staging leftovers. ClickHouse forbids dropping a table that a dictionary depends on, so a +// refresh that leaves the dictionary sourcing from the wrong table strands that table permanently. +func requireDictionary(t *testing.T, rt *runtime.Runtime, id string, want []map[string]any) { + t.Helper() + + cols := make([]string, 0, len(want[0])) + for col := range want[0] { + cols = append(cols, col) + } + sort.Strings(cols) + + testruntime.RequireResolve(t, rt, id, &testruntime.RequireResolveOptions{ + Resolver: "sql", + Properties: map[string]any{"sql": fmt.Sprintf("SELECT %s FROM campaign_name_dict", strings.Join(cols, ", "))}, + Result: want, + }) + + testruntime.RequireResolve(t, rt, id, &testruntime.RequireResolveOptions{ + Resolver: "sql", + Properties: map[string]any{"sql": ` + SELECT + countIf(startsWith(name, '__rill_tmp_model_')) AS staged, + countIf(position(name, '_dict_temp_') > 0) AS sources + FROM system.tables + WHERE database = currentDatabase()`}, + Result: []map[string]any{{"staged": 0, "sources": 1}}, + }) +} diff --git a/runtime/drivers/clickhouse/model_manager.go b/runtime/drivers/clickhouse/model_manager.go index efe0e34168c..822b3c0b9ab 100644 --- a/runtime/drivers/clickhouse/model_manager.go +++ b/runtime/drivers/clickhouse/model_manager.go @@ -360,8 +360,10 @@ func boolPtr(b bool) *bool { return &b } +const stagingTablePrefix = "__rill_tmp_model_" + // stagingTableName returns a stable temporary table name for a destination table. // By using a stable temporary table name, we can ensure proper garbage collection without managing additional state. func stagingTableNameFor(table string) string { - return "__rill_tmp_model_" + table + return stagingTablePrefix + table } diff --git a/runtime/drivers/clickhouse/olap_test.go b/runtime/drivers/clickhouse/olap_test.go index 409dc139b9c..63ec4b7f648 100644 --- a/runtime/drivers/clickhouse/olap_test.go +++ b/runtime/drivers/clickhouse/olap_test.go @@ -497,6 +497,30 @@ func testDictionary(t *testing.T, c *Connection, olap drivers.OLAPStore) { require.NoError(t, res.Close()) require.NoError(t, c.dropTable(context.Background(), "dict1")) + + // A dictionary can also be pointed at a table the user manages. Dropping the dictionary must leave that table + // alone, since only the source tables Rill creates for a dictionary are safe to garbage collect. + _, err = c.createTableAsSelect(context.Background(), "user_src", "SELECT 1 AS id, 'Venus' AS planet", &ModelOutputProperties{Engine: "MergeTree"}, "", "") + require.NoError(t, err) + _, err = c.createTableAsSelect(context.Background(), "user_dict", "", &ModelOutputProperties{ + Typ: "DICTIONARY", + Columns: "(id UInt64, planet String)", + EngineFull: "PRIMARY KEY id SOURCE(CLICKHOUSE(TABLE 'user_src')) LAYOUT(HASHED()) LIFETIME(0)", + }, "", "") + require.NoError(t, err) + + require.NoError(t, c.dropTable(context.Background(), "user_dict")) + requireExists(t, olap, "user_src") +} + +func requireExists(t *testing.T, olap drivers.OLAPStore, tbl string) { + result, err := olap.Query(context.Background(), &drivers.Statement{Query: "EXISTS " + tbl}) + require.NoError(t, err) + require.True(t, result.Next()) + var exist bool + require.NoError(t, result.Scan(&exist)) + require.True(t, exist) + require.NoError(t, result.Close()) } func testIntervalType(t *testing.T, olap drivers.OLAPStore) {