From a75f0d6591c74c00dd8bf49bf1c17d5b314c6e84 Mon Sep 17 00:00:00 2001
From: Jeff Mitchell <jeffrey.mitchell@gmail.com>
Date: Fri, 21 Jun 2019 17:04:50 -0400
Subject: [PATCH] Pass context to Cassandra queries (#6954)

---
 plugins/database/cassandra/cassandra.go           | 8 ++++----
 plugins/database/cassandra/connection_producer.go | 8 ++++----
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/plugins/database/cassandra/cassandra.go b/plugins/database/cassandra/cassandra.go
index 001570ffb297..886e7c8b3196 100644
--- a/plugins/database/cassandra/cassandra.go
+++ b/plugins/database/cassandra/cassandra.go
@@ -129,7 +129,7 @@ func (c *Cassandra) CreateUser(ctx context.Context, statements dbplugin.Statemen
 			err = session.Query(dbutil.QueryHelper(query, map[string]string{
 				"username": username,
 				"password": password,
-			})).Exec()
+			})).WithContext(ctx).Exec()
 			if err != nil {
 				for _, stmt := range rollbackCQL {
 					for _, query := range strutil.ParseArbitraryStringSlice(stmt, ";") {
@@ -140,7 +140,7 @@ func (c *Cassandra) CreateUser(ctx context.Context, statements dbplugin.Statemen
 
 						session.Query(dbutil.QueryHelper(query, map[string]string{
 							"username": username,
-						})).Exec()
+						})).WithContext(ctx).Exec()
 					}
 				}
 				return "", "", err
@@ -185,7 +185,7 @@ func (c *Cassandra) RevokeUser(ctx context.Context, statements dbplugin.Statemen
 
 			err := session.Query(dbutil.QueryHelper(query, map[string]string{
 				"username": username,
-			})).Exec()
+			})).WithContext(ctx).Exec()
 
 			result = multierror.Append(result, err)
 		}
@@ -225,7 +225,7 @@ func (c *Cassandra) RotateRootCredentials(ctx context.Context, statements []stri
 			err := session.Query(dbutil.QueryHelper(query, map[string]string{
 				"username": c.Username,
 				"password": password,
-			})).Exec()
+			})).WithContext(ctx).Exec()
 
 			result = multierror.Append(result, err)
 		}
diff --git a/plugins/database/cassandra/connection_producer.go b/plugins/database/cassandra/connection_producer.go
index bebff4ef059f..a34dd4e45c00 100644
--- a/plugins/database/cassandra/connection_producer.go
+++ b/plugins/database/cassandra/connection_producer.go
@@ -137,7 +137,7 @@ func (c *cassandraConnectionProducer) Init(ctx context.Context, conf map[string]
 	return conf, nil
 }
 
-func (c *cassandraConnectionProducer) Connection(_ context.Context) (interface{}, error) {
+func (c *cassandraConnectionProducer) Connection(ctx context.Context) (interface{}, error) {
 	if !c.Initialized {
 		return nil, connutil.ErrNotInitialized
 	}
@@ -147,7 +147,7 @@ func (c *cassandraConnectionProducer) Connection(_ context.Context) (interface{}
 		return c.session, nil
 	}
 
-	session, err := c.createSession()
+	session, err := c.createSession(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -172,7 +172,7 @@ func (c *cassandraConnectionProducer) Close() error {
 	return nil
 }
 
-func (c *cassandraConnectionProducer) createSession() (*gocql.Session, error) {
+func (c *cassandraConnectionProducer) createSession(ctx context.Context) (*gocql.Session, error) {
 	hosts := strings.Split(c.Hosts, ",")
 	clusterConfig := gocql.NewCluster(hosts...)
 	clusterConfig.Authenticator = gocql.PasswordAuthenticator{
@@ -256,7 +256,7 @@ func (c *cassandraConnectionProducer) createSession() (*gocql.Session, error) {
 	}
 
 	// Verify the info
-	err = session.Query(`LIST ALL`).Exec()
+	err = session.Query(`LIST ALL`).WithContext(ctx).Exec()
 	if err != nil && len(c.Username) != 0 && strings.Contains(err.Error(), "not authorized") {
 		rowNum := session.Query(dbutil.QueryHelper(`LIST CREATE ON ALL ROLES OF '{{username}}';`, map[string]string{
 			"username": c.Username,