From 0e6a2b933be217c941d9acd05a80f0ee8c4c9b60 Mon Sep 17 00:00:00 2001
From: Guang Yang <kobe@uber.com>
Date: Thu, 17 Aug 2017 16:18:24 -0700
Subject: [PATCH] Add authentication and port to cassandra tool (#278)

---
 tools/cassandra/config.go       |  3 +++
 tools/cassandra/cqlclient.go    | 17 +++++++++++++----
 tools/cassandra/updateschema.go |  9 +++++++++
 tools/cassandra/updatetask.go   |  2 +-
 4 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/tools/cassandra/config.go b/tools/cassandra/config.go
index 6649d51f..66744daa 100644
--- a/tools/cassandra/config.go
+++ b/tools/cassandra/config.go
@@ -25,7 +25,10 @@ type (
 	// for the schema update task
 	SchemaUpdaterConfig struct {
 		Keyspace         string
+		Username         string
+		Password         string
 		HostsCsv         string
+		Port             int
 		IsDryRun         bool
 		SkipVersionCheck bool
 		ManifestFilePath string
diff --git a/tools/cassandra/cqlclient.go b/tools/cassandra/cqlclient.go
index eab3df1a..e0bae06b 100644
--- a/tools/cassandra/cqlclient.go
+++ b/tools/cassandra/cqlclient.go
@@ -66,16 +66,25 @@ const (
 )
 
 // newCQLClient returns a new instance of CQLClient
-func newCQLClient(hostsCsv string, keyspace string, protoVersion int) (CQLClient, error) {
-	hosts := parseHosts(hostsCsv)
+func newCQLClient(config *SchemaUpdaterConfig) (CQLClient, error) {
+	hosts := parseHosts(config.HostsCsv)
 	if len(hosts) == 0 {
 		return nil, errNoHosts
 	}
 	clusterCfg := gocql.NewCluster(hosts...)
-	clusterCfg.Keyspace = keyspace
+	clusterCfg.Port = config.Port
+	clusterCfg.Keyspace = config.Keyspace
 	clusterCfg.Timeout = defaultTimeout
-	clusterCfg.ProtoVersion = protoVersion
+	clusterCfg.ProtoVersion = config.ProtoVersion
 	clusterCfg.Consistency = gocql.ParseConsistency(defaultConsistency)
+
+	if config.Username != "" && config.Password != "" {
+		clusterCfg.Authenticator = gocql.PasswordAuthenticator{
+			Username: config.Username,
+			Password: config.Password,
+		}
+	}
+
 	cqlClient := new(cqlClient)
 	cqlClient.clusterConfig = clusterCfg
 	var err error
diff --git a/tools/cassandra/updateschema.go b/tools/cassandra/updateschema.go
index 4b546d0b..d65976cd 100644
--- a/tools/cassandra/updateschema.go
+++ b/tools/cassandra/updateschema.go
@@ -48,6 +48,9 @@ func parseSchemaUpadaterConfig(args []string) *SchemaUpdaterConfig {
 	var config SchemaUpdaterConfig
 	cmd := flag.NewFlagSet("update-schema", flag.ExitOnError)
 	cmd.StringVar(&config.HostsCsv, "h", "", "Path to json file containing cassandra hosts")
+	cmd.IntVar(&config.Port, "port", 9042, "Port to cassandra hosts")
+	cmd.StringVar(&config.Username, "u", "", "Cassandra username")
+	cmd.StringVar(&config.Password, "pw", "", "Cassandra password")
 	cmd.StringVar(&config.Keyspace, "k", "", "Cassandra keyspace")
 	cmd.BoolVar(&config.IsDryRun, "d", true, "Dry run")
 	cmd.BoolVar(&config.SkipVersionCheck, "v", false, "Skip previous version check")
@@ -77,6 +80,12 @@ func printHelp() {
 	helpMessage := `updateSchema -h [cassandra_host1,host2..] -k [keyspace] -m [/tmp/manifest.json] [-d false] [-v false]
     -h
         List of cassandra hosts to connect to
+    -port
+        Port to cassandra cluster, 9042 by default
+    -u
+        Cassandra username
+    -pw
+        Cassandra password
     -k
         Cassandra keyspace
     -m
diff --git a/tools/cassandra/updatetask.go b/tools/cassandra/updatetask.go
index 2260a076..c2c6ee36 100644
--- a/tools/cassandra/updatetask.go
+++ b/tools/cassandra/updatetask.go
@@ -57,7 +57,7 @@ var (
 
 // NewSchemaUpdateTask creates and returns a new SchemaUpdateTask
 func NewSchemaUpdateTask(config *SchemaUpdaterConfig) (*SchemaUpdateTask, error) {
-	cqlClient, err := newCQLClient(config.HostsCsv, config.Keyspace, config.ProtoVersion)
+	cqlClient, err := newCQLClient(config)
 	if err != nil {
 		return nil, err
 	}