Merge Update Keyspace/Table name in prepared Query statement (#1714)

Author: martin-sucha

AUTHORS

@@ -140,3 +140,4 @@ Lauro Ramos Venancio <[email protected]>
 Dmitry Kropachev <[email protected]>
 Oliver Boyle <[email protected]>
 Jackson Fleming <[email protected]>
+Sylwia Szunejko <[email protected]>

CHANGELOG.md

@@ -10,6 +10,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added the InstaclustrPasswordAuthenticator to the list of default approved authenticators.
 - Added the `com.scylladb.auth.SaslauthdAuthenticator` and `com.scylladb.auth.TransitionalAuthenticator`
   to the list of default approved authenticators.
+- Added transferring Keyspace and Table names to the Query from the prepared response and updating
+  information about that every time this information is received
 ### Changed
 
 ### Fixed

cass1batch_test.go

@@ -53,7 +53,7 @@ func TestShouldPrepareFunction(t *testing.T) {
 	}
 
 	for _, test := range shouldPrepareTests {
-		q := &Query{stmt: test.Stmt}
+		q := &Query{stmt: test.Stmt, routingInfo: &queryRoutingInfo{}}
 		if got := q.shouldPrepare(); got != test.Result {
 			t.Fatalf("%q: got %v, expected %v\n", test.Stmt, got, test.Result)
 		}

conn.go

@@ -1378,6 +1378,12 @@ func (c *Conn) executeQuery(ctx context.Context, qry *Query) *Iter {
 			params:        params,
 			customPayload: qry.customPayload,
 		}
+
+		// Set "keyspace" and "table" property in the query if it is present in preparedMetadata
+		qry.routingInfo.mu.Lock()
+		qry.routingInfo.keyspace = info.request.keyspace
+		qry.routingInfo.table = info.request.table
+		qry.routingInfo.mu.Unlock()
 	} else {
 		frame = &writeQueryFrame{
 			statement:     qry.stmt,

frame.go

@@ -918,6 +918,10 @@ type preparedMetadata struct {
 
 	// proto v4+
 	pkeyColumns []int
+
+	keyspace string
+
+	table string
 }
 
 func (r preparedMetadata) String() string {
@@ -952,26 +956,25 @@ func (f *framer) parsePreparedMetadata() preparedMetadata {
 		return meta
 	}
 
-	var keyspace, table string
 	globalSpec := meta.flags&flagGlobalTableSpec == flagGlobalTableSpec
 	if globalSpec {
-		keyspace = f.readString()
-		table = f.readString()
+		meta.keyspace = f.readString()
+		meta.table = f.readString()
 	}
 
 	var cols []ColumnInfo
 	if meta.colCount < 1000 {
 		// preallocate columninfo to avoid excess copying
 		cols = make([]ColumnInfo, meta.colCount)
 		for i := 0; i < meta.colCount; i++ {
-			f.readCol(&cols[i], &meta.resultMetadata, globalSpec, keyspace, table)
+			f.readCol(&cols[i], &meta.resultMetadata, globalSpec, meta.keyspace, meta.table)
 		}
 	} else {
 		// use append, huge number of columns usually indicates a corrupt frame or
 		// just a huge row.
 		for i := 0; i < meta.colCount; i++ {
 			var col ColumnInfo
-			f.readCol(&col, &meta.resultMetadata, globalSpec, keyspace, table)
+			f.readCol(&col, &meta.resultMetadata, globalSpec, meta.keyspace, meta.table)
 			cols = append(cols, col)
 		}
 	}

keyspace_table_test.go

@@ -0,0 +1,81 @@
+//go:build all || integration
+// +build all integration
+
+package gocql
+
+import (
+	"context"
+	"fmt"
+	"testing"
+)
+
+// Keyspace_table checks if Query.Keyspace() is updated based on prepared statement
+func TestKeyspaceTable(t *testing.T) {
+	cluster := createCluster()
+
+	fallback := RoundRobinHostPolicy()
+	cluster.PoolConfig.HostSelectionPolicy = TokenAwareHostPolicy(fallback)
+
+	session, err := cluster.CreateSession()
+	if err != nil {
+		t.Fatal("createSession:", err)
+	}
+
+	cluster.Keyspace = "wrong_keyspace"
+
+	keyspace := "test1"
+	table := "table1"
+
+	err = createTable(session, `DROP KEYSPACE IF EXISTS `+keyspace)
+	if err != nil {
+		t.Fatal("unable to drop keyspace:", err)
+	}
+
+	err = createTable(session, fmt.Sprintf(`CREATE KEYSPACE %s
+	WITH replication = {
+		'class' : 'SimpleStrategy',
+		'replication_factor' : 1
+	}`, keyspace))
+
+	if err != nil {
+		t.Fatal("unable to create keyspace:", err)
+	}
+
+	if err := session.control.awaitSchemaAgreement(); err != nil {
+		t.Fatal(err)
+	}
+
+	err = createTable(session, fmt.Sprintf(`CREATE TABLE %s.%s (pk int, ck int, v int, PRIMARY KEY (pk, ck));
+	`, keyspace, table))
+
+	if err != nil {
+		t.Fatal("unable to create table:", err)
+	}
+
+	if err := session.control.awaitSchemaAgreement(); err != nil {
+		t.Fatal(err)
+	}
+
+	ctx := context.Background()
+
+	// insert a row
+	if err := session.Query(`INSERT INTO test1.table1(pk, ck, v) VALUES (?, ?, ?)`,
+		1, 2, 3).WithContext(ctx).Consistency(One).Exec(); err != nil {
+		t.Fatal(err)
+	}
+
+	var pk int
+
+	/* Search for a specific set of records whose 'pk' column matches
+	 * the value of inserted row. */
+	qry := session.Query(`SELECT pk FROM test1.table1 WHERE pk = ? LIMIT 1`,
+		1).WithContext(ctx).Consistency(One)
+	if err := qry.Scan(&pk); err != nil {
+		t.Fatal(err)
+	}
+
+	// cluster.Keyspace was set to "wrong_keyspace", but during prepering statement
+	// Keyspace in Query should be changed to "test" and Table should be changed to table1
+	assertEqual(t, "qry.Keyspace()", "test1", qry.Keyspace())
+	assertEqual(t, "qry.Table()", "table1", qry.Table())
+}

policies_test.go

@@ -54,7 +54,7 @@ func TestHostPolicy_TokenAware_SimpleStrategy(t *testing.T) {
 		return nil, errors.New("not initalized")
 	}
 
-	query := &Query{}
+	query := &Query{routingInfo: &queryRoutingInfo{}}
 	query.getKeyspace = func() string { return keyspace }
 
 	iter := policy.Pick(nil)
@@ -201,7 +201,7 @@ func TestHostPolicy_TokenAware_NilHostInfo(t *testing.T) {
 	}
 	policy.SetPartitioner("OrderedPartitioner")
 
-	query := &Query{}
+	query := &Query{routingInfo: &queryRoutingInfo{}}
 	query.getKeyspace = func() string { return "myKeyspace" }
 	query.RoutingKey([]byte("20"))
 
@@ -259,7 +259,7 @@ func TestCOWList_Add(t *testing.T) {
 
 // TestSimpleRetryPolicy makes sure that we only allow 1 + numRetries attempts
 func TestSimpleRetryPolicy(t *testing.T) {
-	q := &Query{}
+	q := &Query{routingInfo: &queryRoutingInfo{}}
 
 	// this should allow a total of 3 tries.
 	rt := &SimpleRetryPolicy{NumRetries: 2}
@@ -317,7 +317,7 @@ func TestExponentialBackoffPolicy(t *testing.T) {
 
 func TestDowngradingConsistencyRetryPolicy(t *testing.T) {
 
-	q := &Query{cons: LocalQuorum}
+	q := &Query{cons: LocalQuorum, routingInfo: &queryRoutingInfo{}}
 
 	rewt0 := &RequestErrWriteTimeout{
 		Received:  0,
@@ -478,7 +478,7 @@ func TestHostPolicy_TokenAware(t *testing.T) {
 		return nil, errors.New("not initialized")
 	}
 
-	query := &Query{}
+	query := &Query{routingInfo: &queryRoutingInfo{}}
 	query.getKeyspace = func() string { return keyspace }
 
 	iter := policy.Pick(nil)
@@ -580,7 +580,7 @@ func TestHostPolicy_TokenAware_NetworkStrategy(t *testing.T) {
 		return nil, errors.New("not initialized")
 	}
 
-	query := &Query{}
+	query := &Query{routingInfo: &queryRoutingInfo{}}
 	query.getKeyspace = func() string { return keyspace }
 
 	iter := policy.Pick(nil)
@@ -707,7 +707,7 @@ func TestHostPolicy_TokenAware_RackAware(t *testing.T) {
 	policyWithFallbackInternal.getKeyspaceName = policyInternal.getKeyspaceName
 	policyWithFallbackInternal.getKeyspaceMetadata = policyInternal.getKeyspaceMetadata
 
-	query := &Query{}
+	query := &Query{routingInfo: &queryRoutingInfo{}}
 	query.getKeyspace = func() string { return keyspace }
 
 	iter := policy.Pick(nil)

query_executor.go

@@ -15,6 +15,7 @@ type ExecutableQuery interface {
 	speculativeExecutionPolicy() SpeculativeExecutionPolicy
 	GetRoutingKey() ([]byte, error)
 	Keyspace() string
+	Table() string
 	IsIdempotent() bool
 
 	withContext(context.Context) ExecutableQuery