fix: drop stale __rill_temp_* ClickHouse tables left behind by interrupted model runs

runtime/drivers/clickhouse/clickhouse.go

@@ -16,6 +16,7 @@ import (
 	"github.com/mitchellh/mapstructure"
 	"github.com/rilldata/rill/runtime/drivers"
 	"github.com/rilldata/rill/runtime/pkg/activity"
+	"github.com/rilldata/rill/runtime/pkg/ctxsync"
 	"github.com/rilldata/rill/runtime/pkg/priorityqueue"
 	"github.com/rilldata/rill/runtime/storage"
 	"go.opentelemetry.io/otel/attribute"
@@ -383,12 +384,17 @@ func (d driver) Open(ctx context.Context, connectorName, instanceID string, conf
 		cancel:          cancel,
 		metaSem:         semaphore.NewWeighted(1),
 		olapSem:         priorityqueue.NewSemaphore(olapSemSize),
+		tempTableMu:     ctxsync.NewRWMutex(),
 		opts:            opts,
 		embed:           embed,
 	}
 
 	c.used()
 	go c.periodicallyEmitStats()
+	if conf.Mode == modeReadWrite {
+		// Only readwrite connections create temporary tables, and dropping them needs DDL privileges.
+		go c.periodicallyDropStaleTempTables()
+	}
 
 	return c, nil
 }
@@ -433,6 +439,10 @@ type Connection struct {
 	metaSem *semaphore.Weighted
 	olapSem *priorityqueue.Semaphore
 
+	// tempTableMu serializes the creation of temporary tables (read lock) with the sweep of stale ones (write lock).
+	// Temporary tables have deterministic names, so without it an execution could recreate a stale table after the sweep listed it and before it dropped it.
+	tempTableMu ctxsync.RWMutex
+
 	// options used to open clickhouse connections
 	opts *clickhouse.Options
 	// embed is embedded clickhouse server for local run
@@ -699,6 +709,32 @@ func (c *Connection) periodicallyEmitStats() {
 	}
 }
 
+// periodicallyDropStaleTempTables drops the temporary tables that model executions left behind, e.g. because the runtime restarted while they ran.
+// It sweeps once immediately and then at most hourly until the connection is closed.
+func (c *Connection) periodicallyDropStaleTempTables() {
+	// The ticker is more frequent than the sweeps, so that a sweep that is due runs as soon as a database that can scale to zero is in use again.
+	ticker := time.NewTicker(time.Minute)
+	defer ticker.Stop()
+
+	var lastSweep time.Time
+	for {
+		// Skip while the database hasn't been used recently and may be scaled to zero, so that the sweep does not wake it up.
+		if time.Since(lastSweep) >= time.Hour && (!c.config.CanScaleToZero || time.Since(c.lastUsedOn()) < 2*time.Minute) {
+			lastSweep = time.Now()
+			err := c.dropStaleTempTables(c.ctx, tempTableTTL)
+			if err != nil && !errors.Is(err, c.ctx.Err()) {
+				c.logger.Warn("clickhouse: failed to drop stale temp tables", zap.Error(err))
+			}
+		}
+
+		select {
+		case <-ticker.C:
+		case <-c.ctx.Done():
+			return
+		}
+	}
+}
+
 // accessibleTables returns the list of tables that are accessible to the current user.
 func (c *Connection) accessibleTables(ctx context.Context) (map[tableKey]struct{}, error) {
 	rows, err := c.readDB.QueryxContext(ctx, `

runtime/drivers/clickhouse/crud.go

@@ -121,7 +121,7 @@ func (c *Connection) insertTableAsSelect(ctx context.Context, name, sql string,
 			return nil, err
 		}
 		// create temp table with the same schema using a deterministic name
-		tempName := fmt.Sprintf("__rill_temp_%s_%x", name, md5.Sum([]byte(sql)))
+		tempName := fmt.Sprintf("%s%s_%x", tempTablePrefix, name, md5.Sum([]byte(sql)))
 		// clean up the temp table
 		defer func() {
 			// cleanup using a different ctx to prevent cleanups being impacted by the main ctx cancellation
@@ -137,28 +137,30 @@ func (c *Connection) insertTableAsSelect(ctx context.Context, name, sql string,
 			}
 		}()
 		// create temp table
+		// Hold the read lock while creating it, so that a sweep of stale temp tables that listed an older table of the same name cannot drop it.
+		err = c.tempTableMu.RLock(ctx)
+		if err != nil {
+			return nil, err
+		}
 		if engine == "Distributed" {
 			// create a local table first
 			err = c.Exec(ctx, &drivers.Statement{
 				Query:    fmt.Sprintf("CREATE OR REPLACE TABLE %s %s AS %s", safeSQLName(localTableName(tempName)), onClusterClause, safeSQLName(localTableName(name))),
 				Priority: 1,
 			})
-			if err != nil {
-				return nil, err
-			}
-			// then create the distributed table
-			err = c.createDistributedTable(ctx, tempName, outputProps)
-			if err != nil {
-				return nil, err
+			if err == nil {
+				// then create the distributed table
+				err = c.createDistributedTable(ctx, tempName, outputProps)
 			}
 		} else {
 			err = c.Exec(ctx, &drivers.Statement{
 				Query:    fmt.Sprintf("CREATE OR REPLACE TABLE %s %s AS %s", safeSQLName(tempName), onClusterClause, safeSQLName(name)),
 				Priority: 1,
 			})
-			if err != nil {
-				return nil, err
-			}
+		}
+		c.tempTableMu.RUnlock()
+		if err != nil {
+			return nil, err
 		}
 
 		// insert into temp table
@@ -235,9 +237,83 @@ func (c *Connection) insertTableAsSelect(ctx context.Context, name, sql string,
 	return nil, fmt.Errorf("incremental insert strategy %q not supported", opts.Strategy)
 }
 
+// tempTablePrefix is the prefix of the temporary tables created while writing a model's table, e.g. for column inference and partition inserts.
+const tempTablePrefix = "__rill_temp_"
+
+// tempTableTTL is the age after which a temporary table is assumed to be left behind by an execution that never dropped it.
+// The default model timeout is 3 hours, so this leaves a wide margin for executions that are still running.
+const tempTableTTL = 24 * time.Hour
+
+// dropStaleTempTables drops the temporary tables in the configured database that were created, or last altered, more than olderThan ago.
+// The execution that creates a temporary table also drops it, but that never happens if the runtime restarts in the meantime,
+// and it may fail, e.g. by timing out while a cancelled insert still holds a lock on the table.
+// Failures to drop individual tables are logged, so that the remaining tables are still dropped.
+func (c *Connection) dropStaleTempTables(ctx context.Context, olderThan time.Duration) error {
+	// Hold the lock until the drops are done: a temporary table created before the listing is not stale, and one created after it waits.
+	if err := c.tempTableMu.Lock(ctx); err != nil {
+		return err
+	}
+	defer c.tempTableMu.Unlock()
+
+	names, err := c.listStaleTempTables(ctx, olderThan)
+	if err != nil {
+		return err
+	}
+	for _, name := range names {
+		err := c.dropTable(ctx, name)
+		if errors.Is(err, drivers.ErrNotFound) {
+			// already dropped, e.g. a `_local` table along with its Distributed table
+			continue
+		}
+		if err != nil {
+			if ctx.Err() != nil {
+				return ctx.Err()
+			}
+			c.logger.Warn("clickhouse: failed to drop stale temp table", zap.String("name", name), zap.Error(err), observability.ZapCtx(ctx))
+			continue
+		}
+		c.logger.Info("clickhouse: dropped stale temp table", zap.String("name", name), observability.ZapCtx(ctx))
+	}
+	return nil
+}
+
+// listStaleTempTables returns the names of the temporary tables in the configured database that were created, or last altered, more than olderThan ago.
+// In cluster mode, it includes the `_local` counterparts of the tables, since their names share the prefix.
+func (c *Connection) listStaleTempTables(ctx context.Context, olderThan time.Duration) ([]string, error) {
+	args := []any{c.config.Database, tempTablePrefix, int64(olderThan.Seconds())}
+	if c.config.Database == "" {
+		args[0] = nil
+	}
+	res, err := c.Query(ctx, &drivers.Statement{
+		Query: `SELECT name FROM system.tables
+			WHERE database = coalesce(?, currentDatabase()) AND startsWith(name, ?) AND metadata_modification_time < now() - toIntervalSecond(?)`,
+		Args:     args,
+		Priority: 100,
+	})
+	if err != nil {
+		return nil, err
+	}
+	defer res.Close()
+
+	var names []string
+	for res.Next() {
+		var name string
+		if err := res.Scan(&name); err != nil {
+			return nil, err
+		}
+		names = append(names, name)
+	}
+	return names, res.Err()
+}
+
 func (c *Connection) dropTable(ctx context.Context, name string) error {
 	typ, err := c.entityType(ctx, c.config.Database, name)
 	if err != nil {
+		// In cluster mode, a table is created as a `_local` table first and a Distributed table over it second.
+		// If the second step never ran, e.g. because the execution was cancelled in between, the `_local` table is dropped on its own.
+		if errors.Is(err, drivers.ErrNotFound) && c.config.Cluster != "" && !strings.HasSuffix(name, "_local") {
+			return c.dropTable(ctx, localTableName(name))
+		}
 		return err
 	}
 
@@ -456,7 +532,7 @@ func (c *Connection) createTable(ctx context.Context, name, sql string, outputPr
 			return fmt.Errorf("clickhouse: no columns specified for table %q", name)
 		}
 		// infer columns
-		v := safeSQLName(fmt.Sprintf("__rill_temp_%s_%x", name, md5.Sum([]byte(sql))))
+		v := safeSQLName(fmt.Sprintf("%s%s_%x", tempTablePrefix, name, md5.Sum([]byte(sql))))
 		defer func() {
 			// cleanup using a different ctx to prevent cleanups being impacted by the main ctx cancellation
 			// this is a best effort cleanup and query can still timeout and we don't want to wait forever due to blocked calls
@@ -466,7 +542,13 @@ func (c *Connection) createTable(ctx context.Context, name, sql string, outputPr
 			defer cancel()
 			_ = c.Exec(ctx, &drivers.Statement{Query: fmt.Sprintf("DROP VIEW IF EXISTS %s %s", v, onClusterClause)})
 		}()
-		err := c.Exec(ctx, &drivers.Statement{Query: fmt.Sprintf("CREATE OR REPLACE VIEW %s %s AS %s", v, onClusterClause, sql)})
+		// Hold the read lock while creating the view, so that a sweep of stale temp tables that listed an older view of the same name cannot drop it.
+		err := c.tempTableMu.RLock(ctx)
+		if err != nil {
+			return err
+		}
+		err = c.Exec(ctx, &drivers.Statement{Query: fmt.Sprintf("CREATE OR REPLACE VIEW %s %s AS %s", v, onClusterClause, sql)})
+		c.tempTableMu.RUnlock()
 		if err != nil {
 			return err
 		}

runtime/drivers/clickhouse/olap_test.go

@@ -6,6 +6,7 @@ import (
 	"net/url"
 	"strings"
 	"testing"
+	"time"
 
 	runtimev1 "github.com/rilldata/rill/proto/gen/rill/runtime/v1"
 	"github.com/rilldata/rill/runtime/drivers"
@@ -32,6 +33,7 @@ func TestClickhouseSingle(t *testing.T) {
 
 	t.Run("WithConnection", func(t *testing.T) { testWithConnection(t, olap) })
 	t.Run("DropTable", func(t *testing.T) { testDropTable(t, c, olap) })
+	t.Run("DropStaleTempTables", func(t *testing.T) { testDropStaleTempTables(t, c, olap) })
 	t.Run("RenameView", func(t *testing.T) { testRenameView(t, c, olap) })
 	t.Run("RenameTable", func(t *testing.T) { testRenameTable(t, c, olap) })
 	t.Run("CreateTableAsSelect", func(t *testing.T) { testCreateTableAsSelect(t, c) })
@@ -64,6 +66,7 @@ func TestClickhouseCluster(t *testing.T) {
 
 	t.Run("WithConnection", func(t *testing.T) { testWithConnection(t, olap) })
 	t.Run("DropTable", func(t *testing.T) { testDropTable(t, c, olap) })
+	t.Run("DropStaleTempTables", func(t *testing.T) { testDropStaleTempTables(t, c, olap) })
 	t.Run("RenameView", func(t *testing.T) { testRenameView(t, c, olap) })
 	t.Run("RenameTable", func(t *testing.T) { testRenameTable(t, c, olap) })
 	t.Run("CreateTableAsSelect", func(t *testing.T) { testCreateTableAsSelect(t, c) })
@@ -120,6 +123,51 @@ func testDropTable(t *testing.T, c *Connection, olap drivers.OLAPStore) {
 	require.NoError(t, err)
 	require.NoError(t, c.dropTable(ctx, "drop_view"))
 	notExists(t, olap, "drop_view")
+
+	if c.config.Cluster != "" {
+		// Only the `_local` table exists if the execution was cancelled before it created the Distributed table over it.
+		err = olap.Exec(ctx, &drivers.Statement{Query: fmt.Sprintf("CREATE TABLE drop_local_only_local %s (id Int32) ENGINE = MergeTree ORDER BY tuple()", c.onClusterClause())})
+		require.NoError(t, err)
+		require.NoError(t, c.dropTable(ctx, "drop_local_only"))
+		notExists(t, olap, "drop_local_only_local")
+	}
+}
+
+func testDropStaleTempTables(t *testing.T, c *Connection, olap drivers.OLAPStore) {
+	ctx := context.Background()
+	props := &ModelOutputProperties{Columns: "(id Int32)", Engine: "MergeTree"}
+
+	// A partition insert's table and a column inference view, as an interrupted execution leaves them behind.
+	staleTable := tempTablePrefix + "stale_tbl_" + strings.Repeat("0", 32)
+	staleView := tempTablePrefix + stagingTableNameFor("stale_tbl") + "_" + strings.Repeat("1", 32)
+	require.NoError(t, c.createTable(ctx, staleTable, "", props))
+	_, err := c.createTableAsSelect(ctx, staleView, "SELECT 1 AS id", &ModelOutputProperties{Typ: "VIEW"}, "", "")
+	require.NoError(t, err)
+	// The model's own table must be left alone.
+	require.NoError(t, c.createTable(ctx, "stale_tbl", "", props))
+	// Only the `_local` table exists if the execution was cancelled before it created the Distributed table over it.
+	localOnly := localTableName(tempTablePrefix + "stale_tbl_" + strings.Repeat("2", 32))
+	if c.config.Cluster != "" {
+		err = olap.Exec(ctx, &drivers.Statement{Query: fmt.Sprintf("CREATE TABLE %s %s (id Int32) ENGINE = MergeTree ORDER BY tuple()", safeSQLName(localOnly), c.onClusterClause())})
+		require.NoError(t, err)
+	}
+
+	// Tables younger than the TTL may still be written to by a running execution.
+	require.NoError(t, c.dropStaleTempTables(ctx, tempTableTTL))
+	require.True(t, tableExists(t, olap, staleTable))
+	require.True(t, tableExists(t, olap, staleView))
+
+	// metadata_modification_time has second resolution, so let the tables age past a zero TTL.
+	time.Sleep(time.Second)
+	require.NoError(t, c.dropStaleTempTables(ctx, 0))
+	notExists(t, olap, staleTable)
+	notExists(t, olap, staleView)
+	if c.config.Cluster != "" {
+		notExists(t, olap, localTableName(staleTable))
+		notExists(t, olap, localOnly)
+	}
+	require.True(t, tableExists(t, olap, "stale_tbl"))
+	require.NoError(t, c.dropTable(ctx, "stale_tbl"))
 }
 
 func testRenameView(t *testing.T, c *Connection, olap drivers.OLAPStore) {
@@ -165,15 +213,19 @@ func testRenameTable(t *testing.T, c *Connection, olap drivers.OLAPStore) {
 }
 
 func notExists(t *testing.T, olap drivers.OLAPStore, tbl string) {
+	require.False(t, tableExists(t, olap, tbl), "table %q should not exist", tbl)
+}
+
+func tableExists(t *testing.T, olap drivers.OLAPStore, tbl string) bool {
 	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.False(t, exist)
 	require.NoError(t, result.Close())
+	return exist
 }
 
 func testCreateTableAsSelect(t *testing.T, c *Connection) {