Merge pull request #222 from sylwiaszunejko/decouple_schema_fetch

Schema queries with configurable server-side timeouts

cluster.go

@@ -255,6 +255,9 @@ type ClusterConfig struct {
 	// If not specified, defaults to the global gocql.Logger.
 	Logger StdLogger
 
+	// The timeout for the requests to the schema tables. (default: 60s)
+	MetadataSchemaRequestTimeout time.Duration
+
 	// internal config for testing
 	disableControlConn bool
 	disableInit        bool
@@ -275,23 +278,24 @@ type Dialer interface {
 // the same host, and will not mark the node being down or up from events.
 func NewCluster(hosts ...string) *ClusterConfig {
 	cfg := &ClusterConfig{
-		Hosts:                  hosts,
-		CQLVersion:             "3.0.0",
-		Timeout:                11 * time.Second,
-		ConnectTimeout:         11 * time.Second,
-		Port:                   9042,
-		NumConns:               2,
-		Consistency:            Quorum,
-		MaxPreparedStmts:       defaultMaxPreparedStmts,
-		MaxRoutingKeyInfo:      1000,
-		PageSize:               5000,
-		DefaultTimestamp:       true,
-		MaxWaitSchemaAgreement: 60 * time.Second,
-		ReconnectInterval:      60 * time.Second,
-		ConvictionPolicy:       &SimpleConvictionPolicy{},
-		ReconnectionPolicy:     &ConstantReconnectionPolicy{MaxRetries: 3, Interval: 1 * time.Second},
-		SocketKeepalive:        15 * time.Second,
-		WriteCoalesceWaitTime:  200 * time.Microsecond,
+		Hosts:                        hosts,
+		CQLVersion:                   "3.0.0",
+		Timeout:                      11 * time.Second,
+		ConnectTimeout:               11 * time.Second,
+		Port:                         9042,
+		NumConns:                     2,
+		Consistency:                  Quorum,
+		MaxPreparedStmts:             defaultMaxPreparedStmts,
+		MaxRoutingKeyInfo:            1000,
+		PageSize:                     5000,
+		DefaultTimestamp:             true,
+		MaxWaitSchemaAgreement:       60 * time.Second,
+		ReconnectInterval:            60 * time.Second,
+		ConvictionPolicy:             &SimpleConvictionPolicy{},
+		ReconnectionPolicy:           &ConstantReconnectionPolicy{MaxRetries: 3, Interval: 1 * time.Second},
+		SocketKeepalive:              15 * time.Second,
+		WriteCoalesceWaitTime:        200 * time.Microsecond,
+		MetadataSchemaRequestTimeout: 60 * time.Second,
 	}
 
 	return cfg

conn.go

@@ -1769,9 +1769,13 @@ func (c *Conn) query(ctx context.Context, statement string, values ...interface{
 }
 
 func (c *Conn) querySystemPeers(ctx context.Context, version cassVersion) *Iter {
-	const (
-		peerSchema    = "SELECT * FROM system.peers"
-		peerV2Schemas = "SELECT * FROM system.peers_v2"
+	usingClause := ""
+	if c.session.control != nil {
+		usingClause = c.session.usingTimeoutClause
+	}
+	var (
+		peerSchema    = "SELECT * FROM system.peers" + usingClause
+		peerV2Schemas = "SELECT * FROM system.peers_v2" + usingClause
 	)
 
 	c.mu.Lock()
@@ -1804,11 +1808,19 @@ func (c *Conn) querySystemPeers(ctx context.Context, version cassVersion) *Iter
 }
 
 func (c *Conn) querySystemLocal(ctx context.Context) *Iter {
-	return c.query(ctx, "SELECT * FROM system.local WHERE key='local'")
+	usingClause := ""
+	if c.session.control != nil {
+		usingClause = c.session.usingTimeoutClause
+	}
+	return c.query(ctx, "SELECT * FROM system.local WHERE key='local'"+usingClause)
 }
 
 func (c *Conn) awaitSchemaAgreement(ctx context.Context) (err error) {
-	const localSchemas = "SELECT schema_version FROM system.local WHERE key='local'"
+	usingClause := ""
+	if c.session.control != nil {
+		usingClause = c.session.usingTimeoutClause
+	}
+	var localSchemas = "SELECT schema_version FROM system.local WHERE key='local'" + usingClause
 
 	var versions map[string]struct{}
 	var schemaVersion string

control.go

@@ -50,10 +50,11 @@ type controlConn struct {
 }
 
 func createControlConn(session *Session) *controlConn {
+
 	control := &controlConn{
-		session: session,
-		quit:    make(chan struct{}),
-		retry:   &SimpleRetryPolicy{NumRetries: 3},
+		session:            session,
+		quit:               make(chan struct{}),
+		retry:              &SimpleRetryPolicy{NumRetries: 3},
 	}
 
 	control.conn.Store((*connHost)(nil))

host_source.go

@@ -642,7 +642,7 @@ type ringDescriber struct {
 
 // Returns true if we are using system_schema.keyspaces instead of system.schema_keyspaces
 func checkSystemSchema(control *controlConn) (bool, error) {
-	iter := control.query("SELECT * FROM system_schema.keyspaces")
+	iter := control.query("SELECT * FROM system_schema.keyspaces" + control.session.usingTimeoutClause)
 	if err := iter.err; err != nil {
 		if errf, ok := err.(*errorFrame); ok {
 			if errf.code == ErrCodeSyntax {

metadata_cassandra.go

@@ -557,7 +557,7 @@ func getKeyspaceMetadata(session *Session, keyspaceName string) (*KeyspaceMetada
 
 		var replication map[string]string
 
-		iter := session.control.query(stmt, keyspaceName)
+		iter := session.control.query(stmt+session.usingTimeoutClause, keyspaceName)
 		if iter.NumRows() == 0 {
 			return nil, ErrKeyspaceDoesNotExist
 		}
@@ -583,7 +583,7 @@ func getKeyspaceMetadata(session *Session, keyspaceName string) (*KeyspaceMetada
 
 		var strategyOptionsJSON []byte
 
-		iter := session.control.query(stmt, keyspaceName)
+		iter := session.control.query(stmt+session.usingTimeoutClause, keyspaceName)
 		if iter.NumRows() == 0 {
 			return nil, ErrKeyspaceDoesNotExist
 		}
@@ -631,7 +631,7 @@ func getTableMetadata(session *Session, keyspaceName string) ([]TableMetadata, e
 					view_name
 				FROM system_schema.views
 				WHERE keyspace_name = ?`
-			iter = session.control.query(stmt, keyspaceName)
+			iter = session.control.query(stmt+session.usingTimeoutClause, keyspaceName)
 			return iter
 		}
 
@@ -693,7 +693,7 @@ func getTableMetadata(session *Session, keyspaceName string) ([]TableMetadata, e
 		}
 	}
 
-	iter = session.control.query(stmt, keyspaceName)
+	iter = session.control.query(stmt+session.usingTimeoutClause, keyspaceName)
 
 	tables := []TableMetadata{}
 	table := TableMetadata{Keyspace: keyspaceName}
@@ -756,7 +756,7 @@ func (s *Session) scanColumnMetadataV1(keyspace string) ([]ColumnMetadata, error
 
 	var columns []ColumnMetadata
 
-	rows := s.control.query(stmt, keyspace).Scanner()
+	rows := s.control.query(stmt+s.usingTimeoutClause, keyspace).Scanner()
 	for rows.Next() {
 		var (
 			column           = ColumnMetadata{Keyspace: keyspace}
@@ -817,7 +817,7 @@ func (s *Session) scanColumnMetadataV2(keyspace string) ([]ColumnMetadata, error
 
 	var columns []ColumnMetadata
 
-	rows := s.control.query(stmt, keyspace).Scanner()
+	rows := s.control.query(stmt+s.usingTimeoutClause, keyspace).Scanner()
 	for rows.Next() {
 		var (
 			column           = ColumnMetadata{Keyspace: keyspace}
@@ -875,7 +875,7 @@ func (s *Session) scanColumnMetadataSystem(keyspace string) ([]ColumnMetadata, e
 
 	var columns []ColumnMetadata
 
-	rows := s.control.query(stmt, keyspace).Scanner()
+	rows := s.control.query(stmt+s.usingTimeoutClause, keyspace).Scanner()
 	for rows.Next() {
 		column := ColumnMetadata{Keyspace: keyspace}
 

metadata_scylla.go

@@ -522,7 +522,7 @@ func getKeyspaceMetadata(session *Session, keyspaceName string) (*KeyspaceMetada
 
 	var replication map[string]string
 
-	iter := session.control.query(stmt, keyspaceName)
+	iter := session.control.query(stmt+session.usingTimeoutClause, keyspaceName)
 	if iter.NumRows() == 0 {
 		return nil, ErrKeyspaceDoesNotExist
 	}
@@ -550,7 +550,7 @@ func getTableMetadata(session *Session, keyspaceName string) ([]TableMetadata, e
 	}
 
 	stmt := `SELECT * FROM system_schema.tables WHERE keyspace_name = ?`
-	iter := session.control.query(stmt, keyspaceName)
+	iter := session.control.query(stmt+session.usingTimeoutClause, keyspaceName)
 
 	var tables []TableMetadata
 	table := TableMetadata{Keyspace: keyspaceName}
@@ -582,7 +582,7 @@ func getTableMetadata(session *Session, keyspaceName string) ([]TableMetadata, e
 
 	stmt = `SELECT * FROM system_schema.scylla_tables WHERE keyspace_name = ? AND table_name = ?`
 	for i, t := range tables {
-		iter := session.control.query(stmt, keyspaceName, t.Name)
+		iter := session.control.query(stmt+session.usingTimeoutClause, keyspaceName, t.Name)
 
 		table := TableMetadata{}
 		if iter.MapScan(map[string]interface{}{
@@ -610,7 +610,7 @@ func getColumnMetadata(session *Session, keyspaceName string) ([]ColumnMetadata,
 
 	var columns []ColumnMetadata
 
-	iter := session.control.query(stmt, keyspaceName)
+	iter := session.control.query(stmt+session.usingTimeoutClause, keyspaceName)
 	column := ColumnMetadata{Keyspace: keyspaceName}
 
 	for iter.MapScan(map[string]interface{}{
@@ -639,7 +639,7 @@ func getTypeMetadata(session *Session, keyspaceName string) ([]TypeMetadata, err
 	}
 
 	stmt := `SELECT * FROM system_schema.types WHERE keyspace_name = ?`
-	iter := session.control.query(stmt, keyspaceName)
+	iter := session.control.query(stmt+session.usingTimeoutClause, keyspaceName)
 
 	var types []TypeMetadata
 	tm := TypeMetadata{Keyspace: keyspaceName}
@@ -670,7 +670,7 @@ func getFunctionsMetadata(session *Session, keyspaceName string) ([]FunctionMeta
 	var functions []FunctionMetadata
 	function := FunctionMetadata{Keyspace: keyspaceName}
 
-	iter := session.control.query(stmt, keyspaceName)
+	iter := session.control.query(stmt+session.usingTimeoutClause, keyspaceName)
 	for iter.MapScan(map[string]interface{}{
 		"function_name":        &function.Name,
 		"argument_types":       &function.ArgumentTypes,
@@ -702,7 +702,7 @@ func getAggregatesMetadata(session *Session, keyspaceName string) ([]AggregateMe
 	var aggregates []AggregateMetadata
 	aggregate := AggregateMetadata{Keyspace: keyspaceName}
 
-	iter := session.control.query(stmt, keyspaceName)
+	iter := session.control.query(stmt+session.usingTimeoutClause, keyspaceName)
 	for iter.MapScan(map[string]interface{}{
 		"aggregate_name": &aggregate.Name,
 		"argument_types": &aggregate.ArgumentTypes,
@@ -734,7 +734,7 @@ func getIndexMetadata(session *Session, keyspaceName string) ([]IndexMetadata, e
 	var indexes []IndexMetadata
 	index := IndexMetadata{}
 
-	iter := session.control.query(stmt, keyspaceName)
+	iter := session.control.query(stmt+session.usingTimeoutClause, keyspaceName)
 	for iter.MapScan(map[string]interface{}{
 		"index_name":    &index.Name,
 		"keyspace_name": &index.KeyspaceName,
@@ -790,7 +790,7 @@ func getViewMetadata(session *Session, keyspaceName string) ([]ViewMetadata, err
 
 	stmt := `SELECT * FROM system_schema.views WHERE keyspace_name = ?`
 
-	iter := session.control.query(stmt, keyspaceName)
+	iter := session.control.query(stmt+session.usingTimeoutClause, keyspaceName)
 
 	var views []ViewMetadata
 	view := ViewMetadata{KeyspaceName: keyspaceName}

schema_queries_test.go

@@ -0,0 +1,20 @@
+//go:build integration && scylla
+// +build integration,scylla
+
+package gocql
+
+import (
+	"testing"
+)
+
+func TestSchemaQueries(t *testing.T) {
+	cluster := createCluster()
+
+	fallback := RoundRobinHostPolicy()
+	cluster.PoolConfig.HostSelectionPolicy = TokenAwareHostPolicy(fallback)
+
+	session := createSessionFromCluster(cluster, t)
+	defer session.Close()
+
+	assertTrue(t, "keyspace present in schemaDescriber", session.schemaDescriber.cache["gocql_test"].Name == "gocql_test")
+}

session.go

@@ -12,6 +12,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"strconv"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -85,6 +86,8 @@ type Session struct {
 	logger StdLogger
 
 	tabletsRoutingV1 bool
+
+	usingTimeoutClause string
 }
 
 var queryPool = &sync.Pool{
@@ -233,9 +236,13 @@ func (s *Session) init() error {
 		if err := s.control.connect(hosts); err != nil {
 			return err
 		}
-		s.control.getConn().conn.mu.Lock()
-		s.tabletsRoutingV1 = s.control.getConn().conn.isTabletSupported()
-		s.control.getConn().conn.mu.Unlock()
+		conn := s.control.getConn().conn
+		conn.mu.Lock()
+		s.tabletsRoutingV1 = conn.isTabletSupported()
+		if s.cfg.MetadataSchemaRequestTimeout > time.Duration(0) && isScyllaConn(conn) {
+			s.usingTimeoutClause = " USING TIMEOUT " + strconv.FormatInt(int64(s.cfg.MetadataSchemaRequestTimeout.Milliseconds()), 10) + "ms"
+		}
+		conn.mu.Unlock()
 
 		if !s.cfg.DisableInitialHostLookup {
 			var partitioner string