# RestoreOverride: Component Image Tags and TiKV Additional IOPS Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** Extend `RestoreOverride` config to support per-component image tag overrides and TiKV DATA disk additional IOPS, applying them when building `RestoreDedicatedClusterRequest`.

**Architecture:** Add new fields to `RestoreOverride` struct, create helper functions to convert config map to `ClusterConfiguration` proto, and integrate overrides in `buildCreateClusterReqFor3LayerArch()`.

**Tech Stack:** Go, protobuf, viper config, testify

## Global Constraints

- Follow existing `RestoreOverride` pattern: fields are optional, zero values mean "no override"
- Component image tag keys must be valid: "tidb", "tikv", "pd", "tiflash", "tiproxy"
- Use backup's tenant/project/cluster IDs for feature gate matching (consistent with existing BRImageTag pattern)
- New fields must be backward compatible - existing configs continue to work unchanged

---

### Task 1: Add New Fields to RestoreOverride Struct

**Files:**
- Modify: `internal/config/config.go:843-858`

**Interfaces:**
- Produces: `RestoreOverride.ComponentImageTags map[string]string` and `RestoreOverride.TiKVDataAdditionalIops int32`

- [ ] **Step 1: Read the current RestoreOverride struct**

Read `internal/config/config.go` lines 843-858 to see the current structure.

- [ ] **Step 2: Add new fields to RestoreOverride**

Add the following fields after `BRImageTag`:

```go
type RestoreOverride struct {
	Gate FeatureGate `mapstructure:"gate"`

	// CompactionAcceleratedPiTR controls whether a backup compaction(compactbackup) is performed before PiTR.
	CompactionAcceleratedPiTR bool `mapstructure:"compactionAcceleratedPiTR" default:"false"`

	// AllowConcurrentPiTR bypasses the PiTR serialization constraint.
	AllowConcurrentPiTR bool `mapstructure:"allowConcurrentPiTR" default:"false"`

	BRImageTag string `mapstructure:"brImageTag"` // Override BR image tag, e.g., "v8.5.3"

	// ComponentImageTags overrides per-component image tags for the restored cluster.
	// Valid keys: "tidb", "tikv", "pd", "tiflash", "tiproxy"
	// Example: {"tidb": "v8.5.0-custom", "tikv": "v8.5.0-custom"}
	ComponentImageTags map[string]string `mapstructure:"componentImageTags"`

	// TiKVDataAdditionalIops adds additional IOPS to TiKV DATA disk.
	// Final IOPS = SKU base IOPS + tikvDataAdditionalIops
	TiKVDataAdditionalIops int32 `mapstructure:"tikvDataAdditionalIops"`
}
```

- [ ] **Step 3: Verify the changes compile**

Run: `cd /home/nil/pingcap/workbench/repos/dataflow-service && go build ./internal/config/...`

Expected: No errors

- [ ] **Step 4: Commit**

```bash
git add internal/config/config.go
git commit -m "feat(config): add ComponentImageTags and TiKVDataAdditionalIops to RestoreOverride"
```

---

### Task 2: Add Helper Function to Build ClusterConfiguration

**Files:**
- Create: `internal/config/restore_override_helpers.go`
- Test: `internal/config/restore_override_helpers_test.go`

**Interfaces:**
- Produces: `BuildClusterConfiguration(componentImageTags map[string]string) *globalv1.ClusterConfiguration`

- [ ] **Step 1: Write the failing test**

Create `internal/config/restore_override_helpers_test.go`:

```go
package config

import (
	"testing"

	globalv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/dedicated/v1"
	"github.com/stretchr/testify/assert"
)

func TestBuildClusterConfiguration(t *testing.T) {
	tests := []struct {
		name     string
		tags     map[string]string
		expected *globalv1.ClusterConfiguration
	}{
		{
			name:     "empty map returns nil",
			tags:     map[string]string{},
			expected: nil,
		},
		{
			name: "single component",
			tags: map[string]string{"tidb": "v8.5.0-custom"},
			expected: &globalv1.ClusterConfiguration{
				ComponentConfigurations: []*globalv1.ClusterConfiguration_ComponentConfiguration{
					{
						ComponentType:     globalv1.ComponentType_TIDB,
						OverwriteImageTag: "v8.5.0-custom",
					},
				},
			},
		},
		{
			name: "multiple components",
			tags: map[string]string{
				"tidb": "v8.5.0-custom",
				"tikv": "v8.5.0-custom",
				"pd":   "v8.5.0-custom",
			},
			expected: &globalv1.ClusterConfiguration{
				ComponentConfigurations: []*globalv1.ClusterConfiguration_ComponentConfiguration{
					{
						ComponentType:     globalv1.ComponentType_PD,
						OverwriteImageTag: "v8.5.0-custom",
					},
					{
						ComponentType:     globalv1.ComponentType_TIDB,
						OverwriteImageTag: "v8.5.0-custom",
					},
					{
						ComponentType:     globalv1.ComponentType_TIKV,
						OverwriteImageTag: "v8.5.0-custom",
					},
				},
			},
		},
		{
			name: "invalid key is ignored",
			tags: map[string]string{
				"tidb":    "v8.5.0-custom",
				"invalid": "v8.5.0-custom",
			},
			expected: &globalv1.ClusterConfiguration{
				ComponentConfigurations: []*globalv1.ClusterConfiguration_ComponentConfiguration{
					{
						ComponentType:     globalv1.ComponentType_TIDB,
						OverwriteImageTag: "v8.5.0-custom",
					},
				},
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := BuildClusterConfiguration(tt.tags)
			if tt.expected == nil {
				assert.Nil(t, result)
			} else {
				assert.NotNil(t, result)
				assert.Equal(t, len(tt.expected.ComponentConfigurations), len(result.ComponentConfigurations))
				// Sort by component type for deterministic comparison
				for i, expected := range tt.expected.ComponentConfigurations {
					assert.Equal(t, expected.ComponentType, result.ComponentConfigurations[i].ComponentType)
					assert.Equal(t, expected.OverwriteImageTag, result.ComponentConfigurations[i].OverwriteImageTag)
				}
			}
		})
	}
}
```

- [ ] **Step 2: Run test to verify it fails**

Run: `cd /home/nil/pingcap/workbench/repos/dataflow-service && go test ./internal/config/... -run TestBuildClusterConfiguration -v`

Expected: FAIL with "BuildClusterConfiguration not defined"

- [ ] **Step 3: Write minimal implementation**

Create `internal/config/restore_override_helpers.go`:

```go
package config

import (
	"sort"

	globalv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/dedicated/v1"
)

// componentNameToComponentType maps user-friendly component names to proto ComponentType enum.
var componentNameToComponentType = map[string]globalv1.ComponentType{
	"tidb":    globalv1.ComponentType_TIDB,
	"tikv":    globalv1.ComponentType_TIKV,
	"pd":      globalv1.ComponentType_PD,
	"tiflash": globalv1.ComponentType_TIFLASH,
	"tiproxy": globalv1.ComponentType_TIPROXY,
}

// BuildClusterConfiguration converts a map of component name to image tag into a ClusterConfiguration proto.
// Returns nil if the map is empty. Invalid component keys are silently ignored.
func BuildClusterConfiguration(componentImageTags map[string]string) *globalv1.ClusterConfiguration {
	if len(componentImageTags) == 0 {
		return nil
	}

	var configs []*globalv1.ClusterConfiguration_ComponentConfiguration
	for name, tag := range componentImageTags {
		componentType, ok := componentNameToComponentType[name]
		if !ok {
			continue // skip invalid component names
		}
		configs = append(configs, &globalv1.ClusterConfiguration_ComponentConfiguration{
			ComponentType:     componentType,
			OverwriteImageTag: tag,
		})
	}

	if len(configs) == 0 {
		return nil
	}

	// Sort by component type for deterministic output
	sort.Slice(configs, func(i, j int) bool {
		return configs[i].ComponentType < configs[j].ComponentType
	})

	return &globalv1.ClusterConfiguration{
		ComponentConfigurations: configs,
	}
}
```

- [ ] **Step 4: Run test to verify it passes**

Run: `cd /home/nil/pingcap/workbench/repos/dataflow-service && go test ./internal/config/... -run TestBuildClusterConfiguration -v`

Expected: PASS

- [ ] **Step 5: Commit**

```bash
git add internal/config/restore_override_helpers.go internal/config/restore_override_helpers_test.go
git commit -m "feat(config): add BuildClusterConfiguration helper for component image tags"
```

---

### Task 3: Integrate Overrides into buildCreateClusterReqFor3LayerArch

**Files:**
- Modify: `app/models/restore/impl/restore_base.go:516-576`

**Interfaces:**
- Consumes: `config.Global().BackupConfig.RestoreOverrides.GetMatched()`
- Consumes: `config.BuildClusterConfiguration()`

- [ ] **Step 1: Read the buildCreateClusterReqFor3LayerArch function**

Read `app/models/restore/impl/restore_base.go` lines 516-576 to understand the current implementation.

- [ ] **Step 2: Add override application after building the request**

After the request is built (around line 570, before the return), add:

```go
	// Apply RestoreOverride if matched
	if override, ok := config.Global().BackupConfig.RestoreOverrides.GetMatched(
		uint64(backup.TenantID), uint64(backup.ProjectID), uint64(backup.ClusterID),
		backupCluster.GetCloudProvider(), backupCluster.GetRegion(),
	); ok {
		// Apply component image tags
		if clusterConfig := config.BuildClusterConfiguration(override.ComponentImageTags); clusterConfig != nil {
			createClusterReq.ClusterConfiguration = clusterConfig
		}

		// Apply TiKV additional IOPS
		if override.TiKVDataAdditionalIops > 0 && createClusterReq.TikvComp != nil {
			createClusterReq.TikvComp.AdditionalIops = override.TiKVDataAdditionalIops
		}
	}
```

- [ ] **Step 3: Add import for config package**

Ensure `config` package is imported (it likely already is based on the codebase).

- [ ] **Step 4: Verify the changes compile**

Run: `cd /home/nil/pingcap/workbench/repos/dataflow-service && go build ./app/models/restore/...`

Expected: No errors

- [ ] **Step 5: Commit**

```bash
git add app/models/restore/impl/restore_base.go
git commit -m "feat(restore): apply RestoreOverride for component images and TiKV IOPS"
```

---

### Task 4: Add Unit Tests for Override Integration

**Files:**
- Create: `app/models/restore/impl/restore_base_test.go` (or add to existing test file)

**Interfaces:**
- Tests: `buildCreateClusterReqFor3LayerArch` with various override configurations

- [ ] **Step 1: Check for existing test file**

Check if `app/models/restore/impl/restore_base_test.go` exists. If not, create it.

- [ ] **Step 2: Write test cases**

Add test cases for the override integration:

```go
func TestBuildCreateClusterReqFor3LayerArch_WithOverrides(t *testing.T) {
	// Test cases:
	// 1. No override matched - fields not set
	// 2. Override with ComponentImageTags - ClusterConfiguration populated
	// 3. Override with TiKVDataAdditionalIops - TikvComp.AdditionalIops set
	// 4. Override with both - both fields set
	// 5. Override with zero TiKVDataAdditionalIops - field not set
}
```

- [ ] **Step 3: Run tests**

Run: `cd /home/nil/pingcap/workbench/repos/dataflow-service && go test ./app/models/restore/... -v`

Expected: PASS

- [ ] **Step 4: Commit**

```bash
git add app/models/restore/impl/restore_base_test.go
git commit -m "test(restore): add unit tests for RestoreOverride integration"
```

---

### Task 5: Update Existing RestoreOverride Tests

**Files:**
- Modify: `app/workflow/action/restore/pitr/create/run_compact_backup_test.go`

**Interfaces:**
- Update test structs to include new fields

- [ ] **Step 1: Read existing test file**

Read `app/workflow/action/restore/pitr/create/run_compact_backup_test.go` to understand existing test patterns.

- [ ] **Step 2: Update test structs if needed**

The existing tests use `config.RestoreOverride` directly. Verify they still compile with the new fields added. The new fields have zero values by default, so existing tests should pass unchanged.

- [ ] **Step 3: Run existing tests**

Run: `cd /home/nil/pingcap/workbench/repos/dataflow-service && go test ./app/workflow/action/restore/... -v`

Expected: PASS (existing tests should still pass)

- [ ] **Step 4: Commit if changes were needed**

```bash
git add app/workflow/action/restore/pitr/create/run_compact_backup_test.go
git commit -m "test: update existing RestoreOverride tests for new fields"
```

---

### Task 6: Add Example Configuration Documentation

**Files:**
- Modify: `docs/superpowers/specs/2026-06-28-restore-override-component-images-and-iops-design.md` (already has examples)

**Interfaces:**
- N/A

- [ ] **Step 1: Verify example config in spec**

The spec already contains an example configuration. Verify it's accurate:

```yaml
backup:
  restoreOverrides:
    - gate:
        all: true
      componentImageTags:
        tidb: "v8.5.0-custom"
        tikv: "v8.5.0-custom"
        pd: "v8.5.0-custom"
      tikvDataAdditionalIops: 1000
```

- [ ] **Step 2: Run full test suite**

Run: `cd /home/nil/pingcap/workbench/repos/dataflow-service && go test ./internal/config/... ./app/models/restore/... -v`

Expected: All tests PASS

- [ ] **Step 3: Final commit**

```bash
git add -A
git commit -m "feat: complete RestoreOverride support for component images and TiKV IOPS"
```

---

## Verification Checklist

After completing all tasks, verify:

- [ ] `RestoreOverride` struct has `ComponentImageTags` and `TiKVDataAdditionalIops` fields
- [ ] `BuildClusterConfiguration` helper correctly maps component names to proto types
- [ ] `buildCreateClusterReqFor3LayerArch` applies overrides when matched
- [ ] All existing tests pass
- [ ] New tests cover the override integration
- [ ] Example configuration in spec is accurate
- [ ] Code compiles without errors