Merge pull request #331 from scylladb/dk/add_columns_remove

fix(columns): add remove method
scylladb · Jun 1, 2023 · 8794c6a · 8794c6a
2 parents 713e574 + 51b92e8
commit 8794c6a
Show file tree

Hide file tree

Showing 3 changed files with 78 additions and 1 deletion.
diff --git a/pkg/generators/statement_generator.go b/pkg/generators/statement_generator.go
@@ -622,7 +622,7 @@ func genDropColumnStmt(t *testschema.Table, keyspace string, colNum int) (*typed
 	return &typedef.Stmts{
 		List: stmts,
 		PostStmtHook: func() {
-			t.Columns = append(t.Columns[:colNum], t.Columns[colNum+1:]...)
+			t.Columns = t.Columns.Remove(colNum)
 			t.ResetQueryCache()
 		},
 	}, nil

diff --git a/pkg/testschema/columns.go b/pkg/testschema/columns.go
@@ -84,6 +84,13 @@ func (c Columns) Names() []string {
 	return names
 }
 
+func (c Columns) Remove(idx int) Columns {
+	out := c
+	copy(out[idx:], out[idx+1:])
+	out[len(out)-1] = nil
+	return out[:len(c)-1]
+}
+
 func (c Columns) ToJSONMap(values map[string]interface{}, r *rand.Rand, p *typedef.PartitionRangeConfig) map[string]interface{} {
 	for _, k := range c {
 		switch t := k.Type.(type) {

diff --git a/pkg/testschema/columns_test.go b/pkg/testschema/columns_test.go
@@ -16,6 +16,7 @@ package testschema_test
 
 import (
 	"encoding/json"
+	"strings"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -139,3 +140,72 @@ func TestMarshalUnmarshal(t *testing.T) {
 		t.Errorf("schema not the same after marshal/unmarshal, diff=%s", diff)
 	}
 }
+
+func TestPrimitives(t *testing.T) {
+	sc := &typedef.SchemaConfig{
+		MaxPartitionKeys:  3,
+		MinPartitionKeys:  2,
+		MaxClusteringKeys: 3,
+		MinClusteringKeys: 2,
+		MaxColumns:        3,
+		MinColumns:        2,
+		MaxTupleParts:     2,
+		MaxUDTParts:       2,
+	}
+
+	cols := testschema.Columns{
+		&testschema.ColumnDef{
+			Name: "pk_mv_0",
+			Type: generators.GenListType(sc),
+		},
+		&testschema.ColumnDef{
+			Name: "pk_mv_1",
+			Type: generators.GenTupleType(sc),
+		},
+		&testschema.ColumnDef{
+			Name: "ct_1",
+			Type: &coltypes.CounterType{},
+		},
+	}
+	if cols.Len() != 3 {
+		t.Errorf("%d != %d", cols.Len(), 3)
+	}
+	colNames := strings.Join(cols.Names(), ",")
+	if colNames != "pk_mv_0,pk_mv_1,ct_1" {
+		t.Errorf("%s != %s", colNames, "pk_mv_0,pk_mv_1,ct_1")
+	}
+	if cols.NonCounters().Len() != 2 {
+		t.Errorf("%d != %d", cols.NonCounters().Len(), 2)
+	}
+	colNames = strings.Join(cols.NonCounters().Names(), ",")
+	if colNames != "pk_mv_0,pk_mv_1" {
+		t.Errorf("%s != %s", colNames, "pk_mv_0,pk_mv_1")
+	}
+
+	cols = cols.Remove(2)
+	if cols.Len() != 2 {
+		t.Errorf("%d != %d", cols.Len(), 2)
+	}
+	colNames = strings.Join(cols.Names(), ",")
+	if colNames != "pk_mv_0,pk_mv_1" {
+		t.Errorf("%s != %s", colNames, "pk_mv_0,pk_mv_1")
+	}
+
+	cols = cols.Remove(0)
+	if cols.Len() != 1 {
+		t.Errorf("%d != %d", cols.Len(), 1)
+	}
+	colNames = strings.Join(cols.Names(), ",")
+	if colNames != "pk_mv_1" {
+		t.Errorf("%s != %s", colNames, "pk_mv_1")
+	}
+
+	cols = cols.Remove(0)
+	if cols.Len() != 0 {
+		t.Errorf("%d != %d", cols.Len(), 0)
+	}
+	colNames = strings.Join(cols.Names(), ",")
+	if colNames != "" {
+		t.Errorf("%s != %s", colNames, "")
+	}
+}