Drop ClickHouse query condition cache on every replica after restore

cmd/clickhouse/check_and_finalize.go

@@ -1,14 +1,18 @@
 package clickhouse
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"time"
 
+	"github.com/ClickHouse/clickhouse-go/v2/lib/proto"
 	"github.com/spf13/cobra"
 	"github.com/stackvista/stackstate-backup-cli/cmd/cmdutils"
 	"github.com/stackvista/stackstate-backup-cli/internal/app"
 	ch "github.com/stackvista/stackstate-backup-cli/internal/clients/clickhouse"
 	"github.com/stackvista/stackstate-backup-cli/internal/foundation/config"
+	"github.com/stackvista/stackstate-backup-cli/internal/foundation/logger"
 	"github.com/stackvista/stackstate-backup-cli/internal/orchestration/portforward"
 	"github.com/stackvista/stackstate-backup-cli/internal/orchestration/restore"
 	"github.com/stackvista/stackstate-backup-cli/internal/orchestration/scale"
@@ -135,7 +139,7 @@ func waitAndFinalize(appCtx *app.Context, chClient ch.Interface, operationID str
 // finalizeRestore finalizes the restore by executing SQL, scaling up, and releasing lock
 func finalizeRestore(appCtx *app.Context) error {
 	if err := executePostRestoreSQL(appCtx); err != nil {
-		appCtx.Logger.Warningf("Post-restore SQL failed: %v", err)
+		warnPostRestoreSQLFailed(appCtx.Logger, err)
 	}
 
 	appCtx.Logger.Println()
@@ -186,14 +190,59 @@ func executePostRestoreSQL(appCtx *app.Context) error {
 		_ = closeConn()
 	}()
 
-	// Execute post-restore SQL command
-	query := "SYSTEM DROP MARK CACHE"
-
-	appCtx.Logger.Debugf("Executing SQL: %s", query)
-	if err := conn.Exec(appCtx.Context, query); err != nil {
-		return fmt.Errorf("failed to execute SQL: %w", err)
+	if err := runPostRestoreSQL(appCtx.Context, conn, appCtx.Logger); err != nil {
+		return err
 	}
 
 	appCtx.Logger.Successf("Post-restore SQL executed successfully")
 	return nil
 }
+
+type sqlExecer interface {
+	Exec(ctx context.Context, query string, args ...any) error
+}
+
+type postRestoreStatement struct {
+	query string
+	// ClickHouse before 25.3 has no query condition cache and rejects the statement as a syntax error.
+	skipIfUnsupported bool
+}
+
+const (
+	clusterName     = "default"
+	syntaxErrorCode = 62
+)
+
+// The service port-forward reaches a single pod, so caches are dropped through the chart's cluster.
+// The query condition cache keeps pre-restore results because restored parts reuse table UUIDs and part names.
+var postRestoreStatements = []postRestoreStatement{
+	{query: "SYSTEM DROP MARK CACHE ON CLUSTER " + clusterName},
+	{query: "SYSTEM DROP QUERY CONDITION CACHE ON CLUSTER " + clusterName, skipIfUnsupported: true},
+}
+
+// Every statement is attempted: an unreachable replica can fail one drop while the others still apply.
+func runPostRestoreSQL(ctx context.Context, conn sqlExecer, log *logger.Logger) error {
+	var errs []error
+	for _, stmt := range postRestoreStatements {
+		log.Debugf("Executing SQL: %s", stmt.query)
+		err := conn.Exec(ctx, stmt.query)
+		var exception *proto.Exception
+		if stmt.skipIfUnsupported && errors.As(err, &exception) && exception.Code == syntaxErrorCode {
+			log.Infof("Skipping SQL not supported by this ClickHouse version: %s", stmt.query)
+			continue
+		}
+		if err != nil {
+			errs = append(errs, fmt.Errorf("failed to execute SQL %q: %w", stmt.query, err))
+		}
+	}
+	return errors.Join(errs...)
+}
+
+func warnPostRestoreSQLFailed(log *logger.Logger, err error) {
+	log.Warningf("Post-restore SQL failed: %v", err)
+	log.Warningf("Stale ClickHouse caches can make services read pre-restore data; traces schema migrations may re-run and crash-loop.")
+	log.Warningf("Run these statements manually against ClickHouse once all replicas are up:")
+	for _, stmt := range postRestoreStatements {
+		log.Warningf("  %s", stmt.query)
+	}
+}

cmd/clickhouse/check_and_finalize_test.go

@@ -0,0 +1,96 @@
+package clickhouse
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"testing"
+
+	"github.com/ClickHouse/clickhouse-go/v2/lib/proto"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/stackvista/stackstate-backup-cli/internal/foundation/logger"
+)
+
+type recordingExecer struct {
+	errs    map[string]error
+	queries []string
+}
+
+func (r *recordingExecer) Exec(_ context.Context, query string, _ ...any) error {
+	r.queries = append(r.queries, query)
+	return r.errs[query]
+}
+
+const (
+	dropMarkCache           = "SYSTEM DROP MARK CACHE ON CLUSTER default"
+	dropQueryConditionCache = "SYSTEM DROP QUERY CONDITION CACHE ON CLUSTER default"
+)
+
+func TestRunPostRestoreSQL(t *testing.T) {
+	syntaxError := &proto.Exception{Code: syntaxErrorCode, Message: "Syntax error"}
+	clusterMissing := &proto.Exception{Code: 701, Message: "Requested cluster 'default' not found"}
+	replicaTimeout := &proto.Exception{Code: 159, Message: "Distributed DDL task is not finished on 1 of 3 hosts"}
+
+	tests := []struct {
+		name        string
+		errs        map[string]error
+		wantQueries []string
+		wantErrs    []string
+	}{
+		{
+			name:        "drops mark and query condition caches on every replica",
+			wantQueries: []string{dropMarkCache, dropQueryConditionCache},
+		},
+		{
+			name:        "skips query condition cache on ClickHouse without it",
+			errs:        map[string]error{dropQueryConditionCache: fmt.Errorf("wrapped: %w", syntaxError)},
+			wantQueries: []string{dropMarkCache, dropQueryConditionCache},
+		},
+		{
+			name:        "does not skip a syntax error from a required statement",
+			errs:        map[string]error{dropMarkCache: syntaxError},
+			wantQueries: []string{dropMarkCache, dropQueryConditionCache},
+			wantErrs:    []string{dropMarkCache},
+		},
+		{
+			name:        "fails on other query condition cache errors",
+			errs:        map[string]error{dropQueryConditionCache: clusterMissing},
+			wantQueries: []string{dropMarkCache, dropQueryConditionCache},
+			wantErrs:    []string{"cluster 'default' not found"},
+		},
+		{
+			name:        "still drops query condition cache when mark cache drop times out",
+			errs:        map[string]error{dropMarkCache: replicaTimeout},
+			wantQueries: []string{dropMarkCache, dropQueryConditionCache},
+			wantErrs:    []string{dropMarkCache, "Distributed DDL task"},
+		},
+		{
+			name: "reports every failed statement",
+			errs: map[string]error{
+				dropMarkCache:           errors.New("connection reset"),
+				dropQueryConditionCache: errors.New("broken pipe"),
+			},
+			wantQueries: []string{dropMarkCache, dropQueryConditionCache},
+			wantErrs:    []string{"connection reset", "broken pipe"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			conn := &recordingExecer{errs: tt.errs}
+
+			err := runPostRestoreSQL(context.Background(), conn, logger.New(true, false))
+
+			assert.Equal(t, tt.wantQueries, conn.queries)
+			if len(tt.wantErrs) == 0 {
+				require.NoError(t, err)
+				return
+			}
+			for _, want := range tt.wantErrs {
+				assert.ErrorContains(t, err, want)
+			}
+		})
+	}
+}