Merge upstream changes to commit '0aa180bad'

Author: illia-li

cassandra_test.go

@@ -690,6 +690,92 @@ func TestCAS(t *testing.T) {
 	}
 }
 
+func TestConsistencySerial(t *testing.T) {
+	session := createSession(t)
+	defer session.Close()
+
+	type testStruct struct {
+		name               string
+		id                 int
+		consistency        Consistency
+		expectedPanicValue string
+	}
+
+	testCases := []testStruct{
+		{
+			name:               "Any",
+			consistency:        Any,
+			expectedPanicValue: "Serial consistency can only be SERIAL or LOCAL_SERIAL got ANY",
+		}, {
+			name:               "One",
+			consistency:        One,
+			expectedPanicValue: "Serial consistency can only be SERIAL or LOCAL_SERIAL got ONE",
+		}, {
+			name:               "Two",
+			consistency:        Two,
+			expectedPanicValue: "Serial consistency can only be SERIAL or LOCAL_SERIAL got TWO",
+		}, {
+			name:               "Three",
+			consistency:        Three,
+			expectedPanicValue: "Serial consistency can only be SERIAL or LOCAL_SERIAL got THREE",
+		}, {
+			name:               "Quorum",
+			consistency:        Quorum,
+			expectedPanicValue: "Serial consistency can only be SERIAL or LOCAL_SERIAL got QUORUM",
+		}, {
+			name:               "LocalQuorum",
+			consistency:        LocalQuorum,
+			expectedPanicValue: "Serial consistency can only be SERIAL or LOCAL_SERIAL got LOCAL_QUORUM",
+		}, {
+			name:               "EachQuorum",
+			consistency:        EachQuorum,
+			expectedPanicValue: "Serial consistency can only be SERIAL or LOCAL_SERIAL got EACH_QUORUM",
+		}, {
+			name:               "Serial",
+			id:                 8,
+			consistency:        Serial,
+			expectedPanicValue: "",
+		}, {
+			name:               "LocalSerial",
+			id:                 9,
+			consistency:        LocalSerial,
+			expectedPanicValue: "",
+		}, {
+			name:               "LocalOne",
+			consistency:        LocalOne,
+			expectedPanicValue: "Serial consistency can only be SERIAL or LOCAL_SERIAL got LOCAL_ONE",
+		},
+	}
+
+	err := session.Query("CREATE TABLE IF NOT EXISTS gocql_test.consistency_serial (id int PRIMARY KEY)").Exec()
+	if err != nil {
+		t.Fatalf("can't create consistency_serial table:%v", err)
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			if tc.expectedPanicValue == "" {
+				err = session.Query("INSERT INTO gocql_test.consistency_serial (id) VALUES (?)", tc.id).SerialConsistency(tc.consistency).Exec()
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				var receivedID int
+				err = session.Query("SELECT * FROM gocql_test.consistency_serial WHERE id=?", tc.id).Scan(&receivedID)
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				require.Equal(t, tc.id, receivedID)
+			} else {
+				require.PanicsWithValue(t, tc.expectedPanicValue, func() {
+					session.Query("INSERT INTO gocql_test.consistency_serial (id) VALUES (?)", tc.id).SerialConsistency(tc.consistency)
+				})
+			}
+		})
+	}
+}
+
 func TestDurationType(t *testing.T) {
 	session := createSession(t)
 	defer session.Close()
@@ -2779,7 +2865,6 @@ func TestUnsetColBatch(t *testing.T) {
 	}
 	var id, mInt, count int
 	var mText string
-
 	if err := session.Query("SELECT count(*) FROM gocql_test.batchUnsetInsert;").Scan(&count); err != nil {
 		t.Fatalf("Failed to select with err: %v", err)
 	} else if count != 2 {
@@ -2814,3 +2899,52 @@ func TestQuery_NamedValues(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+// This test ensures that queries are sent to the specified host only
+func TestQuery_SetHostID(t *testing.T) {
+	session := createSession(t)
+	defer session.Close()
+
+	hosts := session.GetHosts()
+
+	const iterations = 5
+	for _, expectedHost := range hosts {
+		for i := 0; i < iterations; i++ {
+			var actualHostID string
+			err := session.Query("SELECT host_id FROM system.local").
+				SetHostID(expectedHost.HostID()).
+				Scan(&actualHostID)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			if expectedHost.HostID() != actualHostID {
+				t.Fatalf("Expected query to be executed on host %s, but it was executed on %s",
+					expectedHost.HostID(),
+					actualHostID,
+				)
+			}
+		}
+	}
+
+	// ensuring properly handled invalid host id
+	err := session.Query("SELECT host_id FROM system.local").
+		SetHostID("[invalid]").
+		Exec()
+	if !errors.Is(err, ErrNoPool) {
+		t.Fatalf("Expected error to be: %v, but got %v", ErrNoPool, err)
+	}
+
+	// ensuring that the driver properly handles the case
+	// when specified host for the query is down
+	host := hosts[0]
+	pool, _ := session.pool.getPoolByHostID(host.HostID())
+	// simulating specified host is down
+	pool.host.setState(NodeDown)
+	err = session.Query("SELECT host_id FROM system.local").
+		SetHostID(host.HostID()).
+		Exec()
+	if !errors.Is(err, ErrHostDown) {
+		t.Fatalf("Expected error to be: %v, but got %v", ErrHostDown, err)
+	}
+}

cluster.go

@@ -173,7 +173,7 @@ type ClusterConfig struct {
 
 	// Consistency for the serial part of queries, values can be either SERIAL or LOCAL_SERIAL.
 	// Default: unset
-	SerialConsistency SerialConsistency
+	SerialConsistency Consistency
 
 	// SslOpts configures TLS use when HostDialer is not set.
 	// SslOpts is ignored if HostDialer is set.
@@ -503,6 +503,10 @@ func (cfg *ClusterConfig) Validate() error {
 		cfg.Logger.Println("warning: enabling skipping metadata can lead to unpredictible results when executing query and altering columns involved in the query.")
 	}
 
+	if cfg.SerialConsistency > 0 && !cfg.SerialConsistency.IsSerial() {
+		return fmt.Errorf("the default SerialConsistency level is not allowed to be anything else but SERIAL or LOCAL_SERIAL. Recived value: %v", cfg.SerialConsistency)
+	}
+
 	return cfg.ValidateAndInitSSL()
 }
 

connectionpool.go

@@ -196,6 +196,13 @@ func (p *policyConnPool) getPool(host *HostInfo) (pool *hostConnPool, ok bool) {
 	return
 }
 
+func (p *policyConnPool) getPoolByHostID(hostID string) (pool *hostConnPool, ok bool) {
+	p.mu.RLock()
+	pool, ok = p.hostConnPools[hostID]
+	p.mu.RUnlock()
+	return
+}
+
 func (p *policyConnPool) Close() {
 	p.mu.Lock()
 	defer p.mu.Unlock()

frame.go

@@ -273,7 +273,6 @@ func (c *Consistency) UnmarshalText(text []byte) error {
 	default:
 		return fmt.Errorf("invalid consistency %q", string(text))
 	}
-
 	return nil
 }
 

policies.go

@@ -390,6 +390,34 @@ func (host selectedHost) Token() Token {
 
 func (host selectedHost) Mark(err error) {}
 
+func newSingleHost(info *HostInfo, maxRetries byte, retryDelay time.Duration) *singleHost {
+	return &singleHost{info: info, maxRetries: maxRetries, delay: retryDelay}
+}
+
+type singleHost struct {
+	retry      byte
+	maxRetries byte
+	delay      time.Duration
+	info       *HostInfo
+}
+
+func (s *singleHost) selectHost() SelectedHost {
+	if s.retry >= s.maxRetries {
+		return nil
+	}
+	if s.retry > 0 && s.delay > 0 {
+		time.Sleep(s.delay)
+	}
+	s.retry++
+	return s
+}
+
+func (s singleHost) Info() *HostInfo { return s.info }
+
+func (s singleHost) Token() Token { return nil }
+
+func (s singleHost) Mark(error) {}
+
 // NextHost is an iteration function over picked hosts
 type NextHost func() SelectedHost
 

policies_integration_test.go

@@ -4,7 +4,9 @@
 package gocql
 
 import (
+	"context"
 	"testing"
+	"time"
 )
 
 // Check if session fail to start if DC name provided in the policy is wrong
@@ -39,3 +41,64 @@ func TestDCValidationRackAware(t *testing.T) {
 		t.Fatal("createSession was expected to fail with wrong DC name provided.")
 	}
 }
+
+// This test ensures  that when all hosts are down, the query execution does not hang.
+func TestNoHangAllHostsDown(t *testing.T) {
+	cluster := createCluster()
+	session := createSessionFromCluster(cluster, t)
+
+	hosts := session.GetHosts()
+	dc := hosts[0].DataCenter()
+	rack := hosts[0].Rack()
+	session.Close()
+
+	policies := []HostSelectionPolicy{
+		DCAwareRoundRobinPolicy(dc),
+		DCAwareRoundRobinPolicy(dc, HostPolicyOptionDisableDCFailover),
+		TokenAwareHostPolicy(DCAwareRoundRobinPolicy(dc)),
+		TokenAwareHostPolicy(DCAwareRoundRobinPolicy(dc, HostPolicyOptionDisableDCFailover)),
+		RackAwareRoundRobinPolicy(dc, rack),
+		RackAwareRoundRobinPolicy(dc, rack, HostPolicyOptionDisableDCFailover),
+		TokenAwareHostPolicy(RackAwareRoundRobinPolicy(dc, rack)),
+		TokenAwareHostPolicy(RackAwareRoundRobinPolicy(dc, rack, HostPolicyOptionDisableDCFailover)),
+		nil,
+	}
+
+	for _, policy := range policies {
+		cluster = createCluster()
+		cluster.PoolConfig.HostSelectionPolicy = policy
+		session = createSessionFromCluster(cluster, t)
+		hosts = session.GetHosts()
+
+		// simulating hosts are down
+		for _, host := range hosts {
+			pool, _ := session.pool.getPoolByHostID(host.HostID())
+			pool.host.setState(NodeDown)
+			if policy != nil {
+				policy.AddHost(host)
+			}
+		}
+
+		ctx, _ := context.WithTimeout(context.Background(), 12*time.Second)
+		_ = session.Query("SELECT host_id FROM system.local").WithContext(ctx).Exec()
+		if ctx.Err() != nil {
+			t.Errorf("policy %T should be no hangups when all hosts are down", policy)
+		}
+
+		// remove all host except one
+		if policy != nil {
+			for i, host := range hosts {
+				if i != 0 {
+					policy.RemoveHost(host)
+				}
+			}
+		}
+
+		ctx, _ = context.WithTimeout(context.Background(), 12*time.Second)
+		_ = session.Query("SELECT host_id FROM system.local").WithContext(ctx).Exec()
+		if ctx.Err() != nil {
+			t.Errorf("policy %T should be no hangups when all hosts are down", policy)
+		}
+		session.Close()
+	}
+}

query_executor.go

@@ -27,6 +27,7 @@ package gocql
 import (
 	"context"
 	"errors"
+	"fmt"
 	"sync"
 	"time"
 )
@@ -44,6 +45,7 @@ type ExecutableQuery interface {
 	IsIdempotent() bool
 	IsLWT() bool
 	GetCustomPartitioner() Partitioner
+	GetHostID() string
 
 	withContext(context.Context) ExecutableQuery
 
@@ -88,12 +90,29 @@ func (q *queryExecutor) speculate(ctx context.Context, qry ExecutableQuery, sp S
 }
 
 func (q *queryExecutor) executeQuery(qry ExecutableQuery) (*Iter, error) {
-	hostIter := q.policy.Pick(qry)
+	var hostIter NextHost
+
+	// check if the hostID is specified for the query,
+	// if true  - the query execute at the specified host.
+	// if false - the query execute at the host picked by HostSelectionPolicy
+	if hostID := qry.GetHostID(); hostID != "" {
+		pool, ok := q.pool.getPoolByHostID(hostID)
+		if !ok {
+			// if the specified host ID have no connection pool we return error
+			return nil, fmt.Errorf("query is targeting unknown host id %s: %w", hostID, ErrNoPool)
+		} else if pool.Size() == 0 {
+			// if the pool have no connection we return error
+			return nil, fmt.Errorf("query is targeting host id %s that driver is not connected to: %w", hostID, ErrNoConnectionsInPool)
+		}
+		hostIter = newSingleHost(pool.host, 5, 200*time.Millisecond).selectHost
+	} else {
+		hostIter = q.policy.Pick(qry)
+	}
 
 	// check if the query is not marked as idempotent, if
 	// it is, we force the policy to NonSpeculative
 	sp := qry.speculativeExecutionPolicy()
-	if !qry.IsIdempotent() || sp.Attempts() == 0 {
+	if qry.GetHostID() != "" || !qry.IsIdempotent() || sp.Attempts() == 0 {
 		return q.do(qry.Context(), qry, hostIter), nil
 	}