Allow setting a custom Authenticator (apache#1242)

* Allow setting a custom Authenticator Signed-off-by: Alex Lourie <djay.il@gmail.com> * Review commit Signed-off-by: Alex Lourie <djay.il@gmail.com>
mpenick · Nov 24, 2018 · 70385f8 · 70385f8
1 parent 33c0e89
commit 70385f8
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 17 deletions.
diff --git a/cluster.go b/cluster.go
@@ -45,22 +45,23 @@ type ClusterConfig struct {
 	// highest supported protocol for the cluster. In clusters with nodes of different
 	// versions the protocol selected is not defined (ie, it can be any of the supported in the cluster)
 	ProtoVersion       int
-	Timeout            time.Duration      // connection timeout (default: 600ms)
-	ConnectTimeout     time.Duration      // initial connection timeout, used during initial dial to server (default: 600ms)
-	Port               int                // port (default: 9042)
-	Keyspace           string             // initial keyspace (optional)
-	NumConns           int                // number of connections per host (default: 2)
-	Consistency        Consistency        // default consistency level (default: Quorum)
-	Compressor         Compressor         // compression algorithm (default: nil)
-	Authenticator      Authenticator      // authenticator (default: nil)
-	RetryPolicy        RetryPolicy        // Default retry policy to use for queries (default: 0)
-	ConvictionPolicy   ConvictionPolicy   // Decide whether to mark host as down based on the error and host info (default: SimpleConvictionPolicy)
-	ReconnectionPolicy ReconnectionPolicy // Default reconnection policy to use for reconnecting before trying to mark host as down (default: see below)
-	SocketKeepalive    time.Duration      // The keepalive period to use, enabled if > 0 (default: 0)
-	MaxPreparedStmts   int                // Sets the maximum cache size for prepared statements globally for gocql (default: 1000)
-	MaxRoutingKeyInfo  int                // Sets the maximum cache size for query info about statements for each session (default: 1000)
-	PageSize           int                // Default page size to use for created sessions (default: 5000)
-	SerialConsistency  SerialConsistency  // Sets the consistency for the serial part of queries, values can be either SERIAL or LOCAL_SERIAL (default: unset)
+	Timeout            time.Duration                            // connection timeout (default: 600ms)
+	ConnectTimeout     time.Duration                            // initial connection timeout, used during initial dial to server (default: 600ms)
+	Port               int                                      // port (default: 9042)
+	Keyspace           string                                   // initial keyspace (optional)
+	NumConns           int                                      // number of connections per host (default: 2)
+	Consistency        Consistency                              // default consistency level (default: Quorum)
+	Compressor         Compressor                               // compression algorithm (default: nil)
+	Authenticator      Authenticator                            // authenticator (default: nil)
+	AuthProvider       func(h *HostInfo) (Authenticator, error) // an authenticator factory. Can be used to create alternative authenticators (default: nil)
+	RetryPolicy        RetryPolicy                              // Default retry policy to use for queries (default: 0)
+	ConvictionPolicy   ConvictionPolicy                         // Decide whether to mark host as down based on the error and host info (default: SimpleConvictionPolicy)
+	ReconnectionPolicy ReconnectionPolicy                       // Default reconnection policy to use for reconnecting before trying to mark host as down (default: see below)
+	SocketKeepalive    time.Duration                            // The keepalive period to use, enabled if > 0 (default: 0)
+	MaxPreparedStmts   int                                      // Sets the maximum cache size for prepared statements globally for gocql (default: 1000)
+	MaxRoutingKeyInfo  int                                      // Sets the maximum cache size for query info about statements for each session (default: 1000)
+	PageSize           int                                      // Default page size to use for created sessions (default: 5000)
+	SerialConsistency  SerialConsistency                        // Sets the consistency for the serial part of queries, values can be either SERIAL or LOCAL_SERIAL (default: unset)
 	SslOpts            *SslOptions
 	DefaultTimestamp   bool // Sends a client side timestamp for all requests which overrides the timestamp at which it arrives at the server. (default: true, only enabled for protocol 3 and above)
 	// PoolConfig configures the underlying connection pool, allowing the

diff --git a/conn.go b/conn.go
@@ -98,6 +98,7 @@ type ConnConfig struct {
 	ConnectTimeout time.Duration
 	Compressor     Compressor
 	Authenticator  Authenticator
+	AuthProvider   func(h *HostInfo) (Authenticator, error)
 	Keepalive      time.Duration
 
 	tlsConfig       *tls.Config
@@ -205,7 +206,6 @@ func (s *Session) dial(host *HostInfo, cfg *ConnConfig, errorHandler ConnErrorHa
 		addr:          conn.RemoteAddr().String(),
 		errorHandler:  errorHandler,
 		compressor:    cfg.Compressor,
-		auth:          cfg.Authenticator,
 		quit:          make(chan struct{}),
 		session:       s,
 		streams:       streams.New(cfg.ProtoVersion),
@@ -217,6 +217,15 @@ func (s *Session) dial(host *HostInfo, cfg *ConnConfig, errorHandler ConnErrorHa
 		},
 	}
 
+	if cfg.AuthProvider != nil {
+		c.auth, err = cfg.AuthProvider(host)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		c.auth = cfg.Authenticator
+	}
+
 	var (
 		ctx    context.Context
 		cancel func()

diff --git a/connectionpool.go b/connectionpool.go
@@ -96,6 +96,7 @@ func connConfig(cfg *ClusterConfig) (*ConnConfig, error) {
 		ConnectTimeout: cfg.ConnectTimeout,
 		Compressor:     cfg.Compressor,
 		Authenticator:  cfg.Authenticator,
+		AuthProvider:   cfg.AuthProvider,
 		Keepalive:      cfg.SocketKeepalive,
 		tlsConfig:      tlsConfig,
 	}, nil

diff --git a/session.go b/session.go
@@ -107,6 +107,11 @@ func NewSession(cfg ClusterConfig) (*Session, error) {
 		return nil, ErrNoHosts
 	}
 
+	// Check that either Authenticator is set or AuthProvider, not both
+	if cfg.Authenticator != nil && cfg.AuthProvider != nil {
+		return nil, errors.New("Can't use both Authenticator and AuthProvider in cluster config.")
+	}
+
 	s := &Session{
 		cons:            cfg.Consistency,
 		prefetch:        0.25,