Fix filter columns transformer for oldKeys

commit_hash:7d727c7b872d9219bcb4fdc47773dfb8b538b1d9

pkg/transformer/registry/filter/filter_columns_transformer.go

@@ -180,10 +180,38 @@ func (f *FilterColumnsTransformer) trimChangeItem(original abstract.ChangeItem,
 		transformed.ColumnValues[newIndex] = original.ColumnValues[origIndex]
 	}
 
-	transformed.OldKeys = copyOldKeys(&original.OldKeys)
+	transformed.OldKeys = copyAndTrimOldKeys(&original.OldKeys, filteredSchema.ColNames)
 	return transformed, nil
 }
 
+func copyAndTrimOldKeys(oldKeys *abstract.OldKeysType, filteredColumns set.Set[string]) abstract.OldKeysType {
+	newLen := 0
+	for _, keyName := range oldKeys.KeyNames {
+		if filteredColumns.Contains(keyName) {
+			newLen++
+		}
+	}
+	if newLen == len(oldKeys.KeyNames) {
+		return copyOldKeys(oldKeys)
+	}
+	trimmedOldKeys := abstract.OldKeysType{
+		KeyNames:  make([]string, 0, newLen),
+		KeyTypes:  make([]string, 0, newLen),
+		KeyValues: make([]interface{}, 0, newLen),
+	}
+	for i, keyName := range oldKeys.KeyNames {
+		if !filteredColumns.Contains(keyName) {
+			continue
+		}
+		trimmedOldKeys.KeyNames = append(trimmedOldKeys.KeyNames, oldKeys.KeyNames[i])
+		trimmedOldKeys.KeyValues = append(trimmedOldKeys.KeyValues, oldKeys.KeyValues[i])
+		if len(oldKeys.KeyTypes) == len(oldKeys.KeyNames) {
+			trimmedOldKeys.KeyTypes = append(trimmedOldKeys.KeyTypes, oldKeys.KeyTypes[i])
+		}
+	}
+	return trimmedOldKeys
+}
+
 func (f *FilterColumnsTransformer) Suitable(table abstract.TableID, schema *abstract.TableSchema) bool {
 	return MatchAnyTableNameVariant(f.Tables, table) && f.validSchema(schema.Columns())
 }

pkg/transformer/registry/filter/filter_columns_transformer_test.go

@@ -6,6 +6,7 @@ import (
 	"github.com/doublecloud/transfer/internal/logger"
 	"github.com/doublecloud/transfer/pkg/abstract"
 	"github.com/doublecloud/transfer/pkg/util"
+	"github.com/doublecloud/transfer/pkg/util/set"
 	"github.com/stretchr/testify/require"
 )
 
@@ -156,3 +157,72 @@ func BenchmarkColumnsFilter(b *testing.B) {
 		b.ReportAllocs()
 	})
 }
+
+func TestCopyAndTrimOldKeys(t *testing.T) {
+	tests := []struct {
+		name             string
+		oldKeys          *abstract.OldKeysType
+		filteredColumns  *set.Set[string]
+		expectedKeyNames []string
+		expectedKeyTypes []string
+		expectedValues   []interface{}
+	}{
+		{
+			name: "Empty input",
+			oldKeys: &abstract.OldKeysType{
+				KeyNames:  []string{},
+				KeyTypes:  []string{},
+				KeyValues: []interface{}{},
+			},
+			filteredColumns:  set.New[string](),
+			expectedKeyNames: []string{},
+			expectedKeyTypes: []string{},
+			expectedValues:   []interface{}{},
+		},
+		{
+			name: "All keys filtered out",
+			oldKeys: &abstract.OldKeysType{
+				KeyNames:  []string{"key1", "key2"},
+				KeyTypes:  []string{"string", "int"},
+				KeyValues: []interface{}{"value1", 42},
+			},
+			filteredColumns:  set.New[string](),
+			expectedKeyNames: []string{},
+			expectedKeyTypes: []string{},
+			expectedValues:   []interface{}{},
+		},
+		{
+			name: "Some keys filtered",
+			oldKeys: &abstract.OldKeysType{
+				KeyNames:  []string{"key1", "key2", "key3"},
+				KeyTypes:  []string{"string", "int", "bool"},
+				KeyValues: []interface{}{"value1", 42, true},
+			},
+			filteredColumns:  set.New("key1", "key3"),
+			expectedKeyNames: []string{"key1", "key3"},
+			expectedKeyTypes: []string{"string", "bool"},
+			expectedValues:   []interface{}{"value1", true},
+		},
+		{
+			name: "All keys kept",
+			oldKeys: &abstract.OldKeysType{
+				KeyNames:  []string{"key1", "key2"},
+				KeyTypes:  []string{"string", "int"},
+				KeyValues: []interface{}{"value1", 42},
+			},
+			filteredColumns:  set.New("key1", "key2"),
+			expectedKeyNames: []string{"key1", "key2"},
+			expectedKeyTypes: []string{"string", "int"},
+			expectedValues:   []interface{}{"value1", 42},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := copyAndTrimOldKeys(tt.oldKeys, *tt.filteredColumns)
+			require.Equal(t, tt.expectedKeyNames, result.KeyNames)
+			require.Equal(t, tt.expectedKeyTypes, result.KeyTypes)
+			require.Equal(t, tt.expectedValues, result.KeyValues)
+		})
+	}
+}