Improve reconnection policy and context handling in Queryx

scylladb · Jan 29, 2025 · b1bf7e7 · b1bf7e7
1 parent 2f79f86
commit b1bf7e7
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/gocqlxtest/gocqlxtest.go b/gocqlxtest/gocqlxtest.go
@@ -51,6 +51,10 @@ func CreateCluster() *gocql.ClusterConfig {
 	cluster.Timeout = *flagTimeout
 	cluster.Consistency = gocql.Quorum
 	cluster.MaxWaitSchemaAgreement = 2 * time.Minute // travis might be slow
+	cluster.ReconnectionPolicy = &gocql.ConstantReconnectionPolicy{
+		MaxRetries: 10,
+		Interval:   3 * time.Second,
+	}
 	if *flagRetry > 0 {
 		cluster.RetryPolicy = &gocql.SimpleRetryPolicy{NumRetries: *flagRetry}
 	}

diff --git a/queryx_wrap.go b/queryx_wrap.go
@@ -86,6 +86,9 @@ func (q *Queryx) RoutingKey(routingKey []byte) *Queryx {
 // query, queries will be canceled and return once the context is
 // canceled.
 func (q *Queryx) WithContext(ctx context.Context) *Queryx {
+	if ctx == nil {
+		ctx = context.Background()
+	}
 	q.Query = q.Query.WithContext(ctx)
 	return q
 }