Repository: pingcap/tidb
File: pkg/lightning/backend/local/local.go (Go)

Code before:
	}
	return ext
}

// verifyImportedStatistics verifies the imported statistics for external engines.
// It checks if OnDuplicateKeyRecord or OnDuplicateKeyRemove is used (which are not yet implemented),
// and verifies that the imported length matches the expected length.
func verifyImportedStatistics(e engineapi.Engine, importedLength int64) error {
	// Check if OnDuplicateKeyRecord or OnDuplicateKeyRemove is used.
	// These options are not yet implemented for local backend, so we skip
	// the statistics verification and return an error to remind future
	// implementers to add support for this check.
	if extEngine, ok := e.(*external.Engine); ok {
		failpoint.Inject("skipOnDuplicateKeyCheck", func(_ failpoint.Value) {
			failpoint.Return(nil)
		})
		onDup := extEngine.GetOnDup()
		if onDup == engineapi.OnDuplicateKeyRecord || onDup == engineapi.OnDuplicateKeyRemove {
			return errors.Errorf("OnDuplicateKeyRecord and OnDuplicateKeyRemove are not yet implemented for local backend. " +
				"When implementing these options, please also implement the statistics verification for them.")
		}

		// Verify the imported statistics after import.
		// For external engine, use the total number of KVs loaded in LoadIngestData
		// (i.e., len(e.memKVsAndBuffers.kvs) across all batches) as the expected count.
		expectedLength := extEngine.GetTotalLoadedKVsCount()
		if importedLength != expectedLength {
			return errors.Errorf("imported length mismatch, expected %d, got %d", expectedLength, importedLength)
		}
	}
	return nil
}

// ImportEngine imports an engine to TiKV.
func (local *Backend) ImportEngine(
	ctx context.Context,
	engineUUID uuid.UUID,
	regionSplitSize, regionSplitKeys int64,
) error {
	kvRegionSplitSize, kvRegionSplitKeys, err := GetRegionSplitSizeKeys(ctx, local.pdCli, local.tls)
	if err == nil {
		if kvRegionSplitSize > regionSplitSize {
			regionSplitSize = kvRegionSplitSize
		}
		if kvRegionSplitKeys > regionSplitKeys {
			regionSplitKeys = kvRegionSplitKeys
		}
	} else {
		tidblogutil.Logger(ctx).Warn("fail to get region split keys and size", zap.Error(err))
	}


Review comment (suggestion):
```suggestion
		expectedKVCount := extEngine.GetTotalLoadedKVsCount()
```

Code after:
	if !ok {
		return nil
	}
	return ext
}

// verifyImportedStatistics verifies the imported statistics for external engines.
// It checks if OnDuplicateKeyRecord or OnDuplicateKeyRemove is used (which are not yet implemented),
// and verifies that the imported KV count matches the expected one.
func verifyImportedStatistics(e engineapi.Engine, importedKVCount int64) error {
	// Check if OnDuplicateKeyRecord or OnDuplicateKeyRemove is used.
	// These options are not yet implemented for local backend, so we skip
	// the statistics verification and return an error to remind future
	// implementers to add support for this check.
	if extEngine, ok := e.(*external.Engine); ok {
		failpoint.Inject("skipOnDuplicateKeyCheck", func(_ failpoint.Value) {
			failpoint.Return(nil)
		})
		onDup := extEngine.GetOnDup()
		if onDup == engineapi.OnDuplicateKeyRecord || onDup == engineapi.OnDuplicateKeyRemove {
			return errors.Errorf("OnDuplicateKeyRecord and OnDuplicateKeyRemove are not yet implemented for local backend. " +
				"When implementing these options, please also implement the statistics verification for them.")
		}

		// Verify the imported statistics after import.
		// For external engine, use the total number of KVs loaded in LoadIngestData
		// (i.e., len(e.memKVsAndBuffers.kvs) across all batches) as the expected count.
		expectedKVCount := extEngine.GetTotalLoadedKVsCount()
		if importedKVCount != expectedKVCount {
			return errors.Errorf("imported length mismatch, expected %d, got %d", expectedKVCount, importedKVCount)
		}
	}
	return nil
}

// ImportEngine imports an engine to TiKV.
func (local *Backend) ImportEngine(
	ctx context.Context,
	engineUUID uuid.UUID,
	regionSplitSize, regionSplitKeys int64,
) error {
	kvRegionSplitSize, kvRegionSplitKeys, err := GetRegionSplitSizeKeys(ctx, local.pdCli, local.tls)
	if err == nil {
		if kvRegionSplitSize > regionSplitSize {
			regionSplitSize = kvRegionSplitSize
		}
		if kvRegionSplitKeys > regionSplitKeys {
			regionSplitKeys = kvRegionSplitKeys
		}
	} else {
		tidblogutil.Logger(ctx).Warn("fail to get region split keys and size", zap.Error(err))