Add authentication and port to cassandra tool (#278)

uber-archive · Aug 17, 2017 · 0e6a2b9 · 0e6a2b9
1 parent 39ddbc0
commit 0e6a2b9
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
 	}