Repository: gocql/gocql
Pull request #1714 (2023-07-11 07:19:19)

Update Keyspace/Table name in prepared Query statement
Previously  `gocql.TokenAwareHostPolicy` always used keyspace explicitly specified in `.Keyspace` in `ClusterConfig` regardless of the keyspace in the Query. This would mean that the token awareness feature would not work properly if the user did not specify `.Keyspace` in `ClusterConfig` or did a query to a different keyspace.

This PR introduces transferring Keyspace and Table names to the Query from the PREPARED response  and updating information about that every time this information is received. After this change, the token awareness feature can work properly without user having to specify a keyspace in `ClusterConfig`.

I created automated integration test ( `keyspace_table_test.go`) on 3 nodes cluster. 

Fixes: #1621 


avelanarius: cc @martin-sucha 

Patch:
diff --git a/AUTHORS b/AUTHORS
index 5f2d361..3275707 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -140,3 +140,4 @@ Lauro Ramos Venancio <lauro.venancio@incognia.com>
 Dmitry Kropachev <dmitry.kropachev@gmail.com>
 Oliver Boyle <pleasedontspamme4321+gocql@gmail.com>
 Jackson Fleming <jackson.fleming@instaclustr.com>
+Sylwia Szunejko <sylwia.szunejko@scylladb.com>
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 94bc1de..4674c48 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/conn.go b/conn.go
index 06a969e..d780bcd 100644
--- a/conn.go
+++ b/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,
diff --git a/frame.go b/frame.go
index fff55a4..44be787 100644
--- a/frame.go
+++ b/frame.go
@@ -918,6 +918,10 @@ type preparedMetadata struct {
 
 	// proto v4+
 	pkeyColumns []int
+
+	keyspace string
+
+	table string
 }
 
 func (r preparedMetadata) String() string {
@@ -952,11 +956,10 @@ 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
@@ -964,14 +967,14 @@ func (f *framer) parsePreparedMetadata() preparedMetadata {
 		// 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)
 		}
 	}
diff --git a/query_executor.go b/query_executor.go
index 58976a7..d5b53b0 100644
--- a/query_executor.go
+++ b/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
diff --git a/session.go b/session.go
index d5ff9ec..fe903f4 100644
--- a/session.go
+++ b/session.go
@@ -87,7 +87,7 @@ type Session struct {
 
 var queryPool = &sync.Pool{
 	New: func() interface{} {
-		return &Query{refCount: 1}
+		return &Query{routingInfo: &queryRoutingInfo{}, refCount: 1}
 	},
 }
 
@@ -630,6 +630,9 @@ func (s *Session) routingKeyInfo(ctx context.Context, stmt string) (*routingKeyI
 		return nil, nil
 	}
 
+	table := info.request.table
+	keyspace := info.request.keyspace
+
 	if len(info.request.pkeyColumns) > 0 {
 		// proto v4 dont need to calculate primary key columns
 		types := make([]TypeInfo, len(info.request.pkeyColumns))
@@ -638,17 +641,16 @@ func (s *Session) routingKeyInfo(ctx context.Context, stmt string) (*routingKeyI
 		}
 
 		routingKeyInfo := &routingKeyInfo{
-			indexes: info.request.pkeyColumns,
-			types:   types,
+			indexes:  info.request.pkeyColumns,
+			types:    types,
+			keyspace: keyspace,
+			table:    table,
 		}
 
 		inflight.value = routingKeyInfo
 		return routingKeyInfo, nil
 	}
 
-	// get the table metadata
-	table := info.request.columns[0].Table
-
 	var keyspaceMetadata *KeyspaceMetadata
 	keyspaceMetadata, inflight.err = s.KeyspaceMetadata(info.request.columns[0].Keyspace)
 	if inflight.err != nil {
@@ -672,8 +674,10 @@ func (s *Session) routingKeyInfo(ctx context.Context, stmt string) (*routingKeyI
 
 	size := len(partitionKey)
 	routingKeyInfo := &routingKeyInfo{
-		indexes: make([]int, size),
-		types:   make([]TypeInfo, size),
+		indexes:  make([]int, size),
+		types:    make([]TypeInfo, size),
+		keyspace: keyspace,
+		table:    table,
 	}
 
 	for keyIndex, keyColumn := range partitionKey {
@@ -909,6 +913,18 @@ type Query struct {
 	// used by control conn queries to prevent triggering a write to systems
 	// tables in AWS MCS see
 	skipPrepare bool
+
+	// routingInfo is a pointer because Query can be copied and copyable struct can't hold a mutex.
+	routingInfo *queryRoutingInfo
+}
+
+type queryRoutingInfo struct {
+	// mu protects contents of queryRoutingInfo.
+	mu sync.RWMutex
+
+	keyspace string
+
+	table string
 }
 
 func (q *Query) defaultsFromSession() {
@@ -1104,6 +1120,10 @@ func (q *Query) Keyspace() string {
 	if q.getKeyspace != nil {
 		return q.getKeyspace()
 	}
+	if q.routingInfo.keyspace != "" {
+		return q.routingInfo.keyspace
+	}
+
 	if q.session == nil {
 		return ""
 	}
@@ -1112,6 +1132,11 @@ func (q *Query) Keyspace() string {
 	return q.session.cfg.Keyspace
 }
 
+// Table returns name of the table the query will be executed against.
+func (q *Query) Table() string {
+	return q.routingInfo.table
+}
+
 // GetRoutingKey gets the routing key to use for routing this query. If
 // a routing key has not been explicitly set, then the routing key will
 // be constructed if possible using the keyspace's schema and the query
@@ -1134,6 +1159,12 @@ func (q *Query) GetRoutingKey() ([]byte, error) {
 		return nil, err
 	}
 
+	if routingKeyInfo != nil {
+		q.routingInfo.mu.Lock()
+		q.routingInfo.keyspace = routingKeyInfo.keyspace
+		q.routingInfo.table = routingKeyInfo.table
+		q.routingInfo.mu.Unlock()
+	}
 	return createRoutingKey(routingKeyInfo, q.values)
 }
 
@@ -1349,7 +1380,7 @@ func (q *Query) Release() {
 
 // reset zeroes out all fields of a query so that it can be safely pooled.
 func (q *Query) reset() {
-	*q = Query{refCount: 1}
+	*q = Query{routingInfo: &queryRoutingInfo{}, refCount: 1}
 }
 
 func (q *Query) incRefCount() {
@@ -1691,6 +1722,9 @@ type Batch struct {
 	cancelBatch           func()
 	keyspace              string
 	metrics               *queryMetrics
+
+	// routingInfo is a pointer because Query can be copied and copyable struct can't hold a mutex.
+	routingInfo *queryRoutingInfo
 }
 
 // NewBatch creates a new batch operation without defaults from the cluster
@@ -1698,9 +1732,10 @@ type Batch struct {
 // Deprecated: use session.NewBatch instead
 func NewBatch(typ BatchType) *Batch {
 	return &Batch{
-		Type:    typ,
-		metrics: &queryMetrics{m: make(map[string]*hostMetrics)},
-		spec:    &NonSpeculativeExecution{},
+		Type:        typ,
+		metrics:     &queryMetrics{m: make(map[string]*hostMetrics)},
+		spec:        &NonSpeculativeExecution{},
+		routingInfo: &queryRoutingInfo{},
 	}
 }
 
@@ -1719,6 +1754,7 @@ func (s *Session) NewBatch(typ BatchType) *Batch {
 		keyspace:         s.cfg.Keyspace,
 		metrics:          &queryMetrics{m: make(map[string]*hostMetrics)},
 		spec:             &NonSpeculativeExecution{},
+		routingInfo:      &queryRoutingInfo{},
 	}
 
 	s.mu.RUnlock()
@@ -1743,6 +1779,11 @@ func (b *Batch) Keyspace() string {
 	return b.keyspace
 }
 
+// Batch has no reasonable eqivalent of Query.Table().
+func (b *Batch) Table() string {
+	return b.routingInfo.table
+}
+
 // Attempts returns the number of attempts made to execute the batch.
 func (b *Batch) Attempts() int {
 	return b.metrics.attempts()
@@ -2014,8 +2055,10 @@ type routingKeyInfoLRU struct {
 }
 
 type routingKeyInfo struct {
-	indexes []int
-	types   []TypeInfo
+	indexes  []int
+	types    []TypeInfo
+	keyspace string
+	table    string
 }
 
 func (r *routingKeyInfo) String() string {


Test patch:
diff --git a/cass1batch_test.go b/cass1batch_test.go
index 43d73ae..f8a796b 100644
--- a/cass1batch_test.go
+++ b/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)
 		}
diff --git a/keyspace_table_test.go b/keyspace_table_test.go
new file mode 100644
index 0000000..2040d0a
--- /dev/null
+++ b/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())
+}
diff --git a/policies_test.go b/policies_test.go
index 19c38ed..170e9e7 100644
--- a/policies_test.go
+++ b/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)
diff --git a/session_test.go b/session_test.go
index a7c8b80..cbe3f67 100644
--- a/session_test.go
+++ b/session_test.go
@@ -100,7 +100,7 @@ func (f funcQueryObserver) ObserveQuery(ctx context.Context, o ObservedQuery) {
 }
 
 func TestQueryBasicAPI(t *testing.T) {
-	qry := &Query{}
+	qry := &Query{routingInfo: &queryRoutingInfo{}}
 
 	// Initiate host
 	ip := "127.0.0.1"
@@ -164,7 +164,7 @@ func TestQueryBasicAPI(t *testing.T) {
 func TestQueryShouldPrepare(t *testing.T) {
 	toPrepare := []string{"select * ", "INSERT INTO", "update table", "delete from", "begin batch"}
 	cantPrepare := []string{"create table", "USE table", "LIST keyspaces", "alter table", "drop table", "grant user", "revoke user"}
-	q := &Query{}
+	q := &Query{routingInfo: &queryRoutingInfo{}}
 
 	for i := 0; i < len(toPrepare); i++ {
 		q.stmt = toPrepare[i]