Merge pull request #342 from scylladb/dk/move_gen_stmt_funcs_into_jobs

move all generate stmt functions into jobs package

pkg/generators/schema_check_stmt_test.go → pkg/jobs/gen_check_stmt_test.go

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 //nolint:thelper
-package generators
+package jobs
 
 import (
 	"path"

pkg/generators/suite_const_test.go → pkg/jobs/gen_const_test.go

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package generators
+package jobs
 
 import (
 	"flag"

pkg/jobs/gen_ddl_stmt.go

@@ -0,0 +1,136 @@
+// Copyright 2019 ScyllaDB
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package jobs
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/pkg/errors"
+	"golang.org/x/exp/rand"
+
+	"github.com/scylladb/gemini/pkg/builders"
+	"github.com/scylladb/gemini/pkg/coltypes"
+	"github.com/scylladb/gemini/pkg/generators"
+	"github.com/scylladb/gemini/pkg/testschema"
+	"github.com/scylladb/gemini/pkg/typedef"
+)
+
+func GenDDLStmt(s *testschema.Schema, t *testschema.Table, r *rand.Rand, _ *typedef.PartitionRangeConfig, sc *typedef.SchemaConfig) (*typedef.Stmts, error) {
+	maxVariant := 1
+	if len(t.Columns) > 0 {
+		maxVariant = 2
+	}
+	switch n := r.Intn(maxVariant + 2); n {
+	// case 0: // Alter column not supported in Cassandra from 3.0.11
+	//	return t.alterColumn(s.Keyspace.Name)
+	case 2:
+		colNum := r.Intn(len(t.Columns))
+		return genDropColumnStmt(t, s.Keyspace.Name, colNum)
+	default:
+		column := testschema.ColumnDef{Name: generators.GenColumnName("col", len(t.Columns)+1), Type: generators.GenColumnType(len(t.Columns)+1, sc)}
+		return genAddColumnStmt(t, s.Keyspace.Name, &column)
+	}
+}
+
+func appendValue(columnType typedef.Type, r *rand.Rand, p *typedef.PartitionRangeConfig, values []interface{}) []interface{} {
+	return append(values, columnType.GenValue(r, p)...)
+}
+
+func genAddColumnStmt(t *testschema.Table, keyspace string, column *testschema.ColumnDef) (*typedef.Stmts, error) {
+	var stmts []*typedef.Stmt
+	if c, ok := column.Type.(*coltypes.UDTType); ok {
+		createType := "CREATE TYPE IF NOT EXISTS %s.%s (%s);"
+		var typs []string
+		for name, typ := range c.Types {
+			typs = append(typs, name+" "+typ.CQLDef())
+		}
+		stmt := fmt.Sprintf(createType, keyspace, c.TypeName, strings.Join(typs, ","))
+		stmts = append(stmts, &typedef.Stmt{
+			StmtCache: &typedef.StmtCache{
+				Query: &builders.AlterTableBuilder{
+					Stmt: stmt,
+				},
+			},
+		})
+	}
+	stmt := "ALTER TABLE " + keyspace + "." + t.Name + " ADD " + column.Name + " " + column.Type.CQLDef()
+	stmts = append(stmts, &typedef.Stmt{
+		StmtCache: &typedef.StmtCache{
+			Query: &builders.AlterTableBuilder{
+				Stmt: stmt,
+			},
+		},
+	})
+	return &typedef.Stmts{
+		List: stmts,
+		PostStmtHook: func() {
+			t.Columns = append(t.Columns, column)
+			t.ResetQueryCache()
+		},
+	}, nil
+}
+
+//nolint:unused
+func alterColumn(t *testschema.Table, keyspace string) ([]*typedef.Stmt, func(), error) {
+	var stmts []*typedef.Stmt
+	idx := rand.Intn(len(t.Columns))
+	column := t.Columns[idx]
+	oldType, isSimpleType := column.Type.(coltypes.SimpleType)
+	if !isSimpleType {
+		return nil, func() {}, errors.Errorf("complex type=%s cannot be altered", column.Name)
+	}
+	compatTypes := coltypes.CompatibleColumnTypes[oldType]
+	if len(compatTypes) == 0 {
+		return nil, func() {}, errors.Errorf("simple type=%s has no compatible coltypes so it cannot be altered", column.Name)
+	}
+	newType := compatTypes.Random()
+	newColumn := testschema.ColumnDef{Name: column.Name, Type: newType}
+	stmt := "ALTER TABLE " + keyspace + "." + t.Name + " ALTER " + column.Name + " TYPE " + column.Type.CQLDef()
+	stmts = append(stmts, &typedef.Stmt{
+		StmtCache: &typedef.StmtCache{
+			Query: &builders.AlterTableBuilder{
+				Stmt: stmt,
+			},
+			QueryType: typedef.AlterColumnStatementType,
+		},
+	})
+	return stmts, func() {
+		t.Columns[idx] = &newColumn
+		t.ResetQueryCache()
+	}, nil
+}
+
+func genDropColumnStmt(t *testschema.Table, keyspace string, colNum int) (*typedef.Stmts, error) {
+	var stmts []*typedef.Stmt
+
+	column := t.Columns[colNum]
+	stmt := "ALTER TABLE " + keyspace + "." + t.Name + " DROP " + column.Name
+	stmts = append(stmts, &typedef.Stmt{
+		StmtCache: &typedef.StmtCache{
+			Query: &builders.AlterTableBuilder{
+				Stmt: stmt,
+			},
+			QueryType: typedef.DropColumnStatementType,
+		},
+	})
+	return &typedef.Stmts{
+		List: stmts,
+		PostStmtHook: func() {
+			t.Columns = t.Columns.Remove(colNum)
+			t.ResetQueryCache()
+		},
+	}, nil
+}

pkg/generators/schema_ddl_stmt_test.go → pkg/jobs/gen_ddl_stmt_test.go

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 //nolint:thelper
-package generators
+package jobs
 
 import (
 	"path"

pkg/jobs/gen_mutate_stmt.go

@@ -0,0 +1,199 @@
+// Copyright 2019 ScyllaDB
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package jobs
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"golang.org/x/exp/rand"
+
+	"github.com/scylladb/gocqlx/v2/qb"
+
+	"github.com/scylladb/gemini/pkg/coltypes"
+	"github.com/scylladb/gemini/pkg/generators"
+	"github.com/scylladb/gemini/pkg/testschema"
+	"github.com/scylladb/gemini/pkg/typedef"
+)
+
+func GenMutateStmt(s *testschema.Schema, t *testschema.Table, g generators.GeneratorInterface, r *rand.Rand, p *typedef.PartitionRangeConfig, deletes bool) (*typedef.Stmt, error) {
+	t.RLock()
+	defer t.RUnlock()
+
+	valuesWithToken := g.Get()
+	if valuesWithToken == nil {
+		return nil, nil
+	}
+	useLWT := false
+	if p.UseLWT && r.Uint32()%10 == 0 {
+		useLWT = true
+	}
+
+	if !deletes {
+		return genInsertOrUpdateStmt(s, t, valuesWithToken, r, p, useLWT)
+	}
+	switch n := rand.Intn(1000); n {
+	case 10, 100:
+		return genDeleteRows(s, t, valuesWithToken, r, p)
+	default:
+		switch rand.Intn(2) {
+		case 0:
+			if t.KnownIssues[typedef.KnownIssuesJSONWithTuples] {
+				return genInsertOrUpdateStmt(s, t, valuesWithToken, r, p, useLWT)
+			}
+			return genInsertJSONStmt(s, t, valuesWithToken, r, p)
+		default:
+			return genInsertOrUpdateStmt(s, t, valuesWithToken, r, p, useLWT)
+		}
+	}
+}
+
+func genInsertOrUpdateStmt(
+	s *testschema.Schema,
+	t *testschema.Table,
+	valuesWithToken *typedef.ValueWithToken,
+	r *rand.Rand,
+	p *typedef.PartitionRangeConfig,
+	useLWT bool,
+) (*typedef.Stmt, error) {
+	if t.IsCounterTable() {
+		return genUpdateStmt(s, t, valuesWithToken, r, p)
+	}
+	return genInsertStmt(s, t, valuesWithToken, r, p, useLWT)
+}
+
+func genUpdateStmt(_ *testschema.Schema, t *testschema.Table, valuesWithToken *typedef.ValueWithToken, r *rand.Rand, p *typedef.PartitionRangeConfig) (*typedef.Stmt, error) {
+	stmtCache := t.GetQueryCache(typedef.CacheUpdate)
+	nonCounters := t.Columns.NonCounters()
+	values := make(typedef.Values, 0, t.PartitionKeys.LenValues()+t.ClusteringKeys.LenValues()+nonCounters.LenValues())
+	for _, cdef := range nonCounters {
+		values = appendValue(cdef.Type, r, p, values)
+	}
+	values = values.CopyFrom(valuesWithToken.Value)
+	for _, ck := range t.ClusteringKeys {
+		values = appendValue(ck.Type, r, p, values)
+	}
+	return &typedef.Stmt{
+		StmtCache:       stmtCache,
+		ValuesWithToken: valuesWithToken,
+		Values:          values,
+	}, nil
+}
+
+func genInsertStmt(
+	_ *testschema.Schema,
+	t *testschema.Table,
+	valuesWithToken *typedef.ValueWithToken,
+	r *rand.Rand,
+	p *typedef.PartitionRangeConfig,
+	useLWT bool,
+) (*typedef.Stmt, error) {
+	values := make(typedef.Values, 0, t.PartitionKeys.LenValues()+t.ClusteringKeys.LenValues()+t.Columns.LenValues())
+	values = values.CopyFrom(valuesWithToken.Value)
+	for _, ck := range t.ClusteringKeys {
+		values = append(values, ck.Type.GenValue(r, p)...)
+	}
+	for _, col := range t.Columns {
+		values = append(values, col.Type.GenValue(r, p)...)
+	}
+	cacheType := typedef.CacheInsert
+	if useLWT {
+		cacheType = typedef.CacheInsertIfNotExists
+	}
+	stmtCache := t.GetQueryCache(cacheType)
+	return &typedef.Stmt{
+		StmtCache:       stmtCache,
+		ValuesWithToken: valuesWithToken,
+		Values:          values,
+	}, nil
+}
+
+func genInsertJSONStmt(
+	s *testschema.Schema,
+	table *testschema.Table,
+	valuesWithToken *typedef.ValueWithToken,
+	r *rand.Rand,
+	p *typedef.PartitionRangeConfig,
+) (*typedef.Stmt, error) {
+	var v string
+	var ok bool
+	if table.IsCounterTable() {
+		return nil, nil
+	}
+	vs := valuesWithToken.Value.Copy()
+	values := make(map[string]interface{})
+	for i, pk := range table.PartitionKeys {
+		switch t := pk.Type.(type) {
+		case coltypes.SimpleType:
+			if t != coltypes.TYPE_BLOB {
+				values[pk.Name] = vs[i]
+				continue
+			}
+			v, ok = vs[i].(string)
+			if ok {
+				values[pk.Name] = "0x" + v
+			}
+		case *coltypes.TupleType:
+			tupVals := make([]interface{}, len(t.Types))
+			for j := 0; j < len(t.Types); j++ {
+				if t.Types[j] == coltypes.TYPE_BLOB {
+					v, ok = vs[i+j].(string)
+					if ok {
+						v = "0x" + v
+					}
+					vs[i+j] = v
+				}
+				tupVals[i] = vs[i+j]
+				i++
+			}
+			values[pk.Name] = tupVals
+		default:
+			panic(fmt.Sprintf("unknown type: %s", t.Name()))
+		}
+	}
+	values = table.ClusteringKeys.ToJSONMap(values, r, p)
+	values = table.Columns.ToJSONMap(values, r, p)
+
+	jsonString, err := json.Marshal(values)
+	if err != nil {
+		return nil, err
+	}
+
+	builder := qb.Insert(s.Keyspace.Name + "." + table.Name).Json()
+	return &typedef.Stmt{
+		StmtCache: &typedef.StmtCache{
+			Query:     builder,
+			Types:     []typedef.Type{coltypes.TYPE_TEXT},
+			QueryType: typedef.InsertStatement,
+		},
+		ValuesWithToken: valuesWithToken,
+		Values:          []interface{}{string(jsonString)},
+	}, nil
+}
+
+func genDeleteRows(_ *testschema.Schema, t *testschema.Table, valuesWithToken *typedef.ValueWithToken, r *rand.Rand, p *typedef.PartitionRangeConfig) (*typedef.Stmt, error) {
+	stmtCache := t.GetQueryCache(typedef.CacheDelete)
+	values := valuesWithToken.Value.Copy()
+	if len(t.ClusteringKeys) > 0 {
+		ck := t.ClusteringKeys[0]
+		values = appendValue(ck.Type, r, p, values)
+		values = appendValue(ck.Type, r, p, values)
+	}
+	return &typedef.Stmt{
+		StmtCache:       stmtCache,
+		ValuesWithToken: valuesWithToken,
+		Values:          values,
+	}, nil
+}

pkg/generators/schema_mutation_stmt_test.go → pkg/jobs/gen_mutate_stmt_test.go

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 //nolint:thelper
-package generators
+package jobs
 
 import (
 	"path"

pkg/generators/suite_utils_test.go → pkg/jobs/gen_utils_test.go

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package generators
+package jobs
 
 import (
 	"encoding/json"
@@ -28,6 +28,7 @@ import (
 
 	"github.com/scylladb/gemini/pkg/builders"
 	"github.com/scylladb/gemini/pkg/coltypes"
+	"github.com/scylladb/gemini/pkg/generators"
 	"github.com/scylladb/gemini/pkg/replication"
 	"github.com/scylladb/gemini/pkg/routingkey"
 	"github.com/scylladb/gemini/pkg/tableopts"
@@ -454,7 +455,7 @@ func getFromOptions(t testInterface, table *testschema.Table, option, optionsNum
 		case "addSt":
 			funcOpts.addType = testschema.ColumnDef{
 				Type: createColumnSimpleType(t, optionsNum),
-				Name: GenColumnName("col", len(table.Columns)+1),
+				Name: generators.GenColumnName("col", len(table.Columns)+1),
 			}
 		}