fix(generator): reimplement genMultiplePartitionClusteringRangeQueryMv

scylladb · Jun 7, 2023 · c5078d1 · c5078d1
1 parent eb92080
commit c5078d1
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 38 deletions.
diff --git a/pkg/generators/statement_generator.go b/pkg/generators/statement_generator.go
@@ -371,8 +371,7 @@ func genClusteringRangeQueryMv(
 	values := vs.Value.Copy()
 	mv := t.MaterializedViews[mvNum]
 	if mv.HaveNonPrimaryKey() {
-		var mvValues []interface{}
-		mvValues = append(mvValues, mv.NonPrimaryKey.Type.GenValue(r, p)...)
+		mvValues := append([]interface{}{}, mv.NonPrimaryKey.Type.GenValue(r, p)...)
 		values = append(mvValues, values...)
 	}
 	builder := qb.Select(s.Keyspace.Name + "." + mv.Name)
@@ -474,44 +473,50 @@ func genMultiplePartitionClusteringRangeQueryMv(
 ) *typedef.Stmt {
 	t.RLock()
 	defer t.RUnlock()
+
+	mv := t.MaterializedViews[mvNum]
+	clusteringKeys := mv.ClusteringKeys
+	pkValues := mv.PartitionKeysLenValues()
+	valuesCount := pkValues*numQueryPKs + clusteringKeys[:maxClusteringRels].LenValues() + clusteringKeys[maxClusteringRels].Type.LenValue()*2
+	mvKey := mv.NonPrimaryKey
+
 	var (
-		values []interface{}
-		typs   []typedef.Type
+		mvKeyLen int
+		baseID   int
 	)
-	mv := t.MaterializedViews[mvNum]
+	if mvKey != nil {
+		mvKeyLen = mvKey.Type.LenValue()
+		baseID = 1
+		valuesCount += mv.PartitionKeys.LenValues() * numQueryPKs
+	}
+	values := make(typedef.Values, pkValues*numQueryPKs, valuesCount)
+	typs := make(typedef.Types, pkValues*numQueryPKs, valuesCount)
 	builder := qb.Select(s.Keyspace.Name + "." + mv.Name)
-	switch mv.HaveNonPrimaryKey() {
-	case true:
-		for i, pk := range mv.PartitionKeys {
-			builder = builder.Where(qb.InTuple(pk.Name, numQueryPKs))
-			for j := 0; j < numQueryPKs; j++ {
-				vs := g.GetOld()
-				if vs == nil {
-					return nil
-				}
-				if i == 0 {
-					values = appendValue(pk.Type, r, p, values)
-					typs = append(typs, pk.Type)
-				} else {
-					values = append(values, vs.Value[i-1])
-					typs = append(typs, pk.Type)
-				}
-			}
+
+	for _, pk := range mv.PartitionKeys {
+		builder = builder.Where(qb.InTuple(pk.Name, numQueryPKs))
+	}
+
+	if mvKey != nil {
+		// Fill values for Materialized view primary key
+		for j := 0; j < numQueryPKs; j++ {
+			typs[j] = mvKey.Type
+			copy(values[j*mvKeyLen:], mvKey.Type.GenValue(r, p))
 		}
-	case false:
-		for i, pk := range mv.PartitionKeys {
-			builder = builder.Where(qb.InTuple(pk.Name, numQueryPKs))
-			for j := 0; j < numQueryPKs; j++ {
-				vs := g.GetOld()
-				if vs == nil {
-					return nil
-				}
-				values = append(values, vs.Value[i])
-				typs = append(typs, pk.Type)
-			}
+	}
+
+	for j := 0; j < numQueryPKs; j++ {
+		vs := g.GetOld()
+		if vs == nil {
+			return nil
+		}
+		for id := range vs.Value {
+			idx := (baseID+id)*numQueryPKs + j
+			typs[idx] = mv.PartitionKeys[baseID+id].Type
+			values[idx] = vs.Value[id]
 		}
 	}
-	clusteringKeys := mv.ClusteringKeys
+
 	if len(clusteringKeys) > 0 {
 		for i := 0; i < maxClusteringRels; i++ {
 			ck := clusteringKeys[i]

diff --git a/pkg/testschema/schema.go b/pkg/testschema/schema.go
@@ -19,10 +19,11 @@ import (
 )
 
 type MaterializedView struct {
-	NonPrimaryKey  *ColumnDef
-	Name           string  `json:"name"`
-	PartitionKeys  Columns `json:"partition_keys"`
-	ClusteringKeys Columns `json:"clustering_keys"`
+	NonPrimaryKey          *ColumnDef
+	Name                   string  `json:"name"`
+	PartitionKeys          Columns `json:"partition_keys"`
+	ClusteringKeys         Columns `json:"clustering_keys"`
+	partitionKeysLenValues int
 }
 
 type Schema struct {
@@ -33,3 +34,10 @@ type Schema struct {
 func (m *MaterializedView) HaveNonPrimaryKey() bool {
 	return m.NonPrimaryKey != nil
 }
+
+func (m *MaterializedView) PartitionKeysLenValues() int {
+	if m.partitionKeysLenValues == 0 && m.PartitionKeys != nil {
+		m.partitionKeysLenValues = m.PartitionKeys.LenValues()
+	}
+	return m.partitionKeysLenValues
+}