Merge pull request #477 from Zariel/session-bind-token-policy

fix tokenaware policy when using session.Bind

cassandra_test.go

@@ -2157,3 +2157,35 @@ func TestLexicalUUIDType(t *testing.T) {
 		t.Errorf("got %s, expected %s", gotUUID, column1)
 	}
 }
+
+// Issue 475
+func TestSessionBindRoutingKey(t *testing.T) {
+	cluster := createCluster()
+	cluster.ConnPoolType = NewTokenAwareConnPool
+
+	session := createSessionFromCluster(cluster, t)
+	defer session.Close()
+
+	if err := createTable(session, `CREATE TABLE test_bind_routing_key (
+			key     varchar,
+			value   int,
+			PRIMARY KEY (key)
+		)`); err != nil {
+
+		t.Fatal(err)
+	}
+
+	const (
+		key   = "routing-key"
+		value = 5
+	)
+
+	fn := func(info *QueryInfo) ([]interface{}, error) {
+		return []interface{}{key, value}, nil
+	}
+
+	q := session.Bind("INSERT INTO test_bind_routing_key(key, value) VALUES(?, ?)", fn)
+	if err := q.Exec(); err != nil {
+		t.Fatal(err)
+	}
+}

policies.go

@@ -161,6 +161,11 @@ func (t *tokenAwareHostPolicy) resetTokenRing() {
 func (t *tokenAwareHostPolicy) Pick(qry *Query) NextHost {
 	if qry == nil {
 		return t.fallback.Pick(qry)
+	} else if qry.binding != nil && len(qry.values) == 0 {
+		// If this query was created using session.Bind we wont have the query
+		// values yet, so we have to pass down to the next policy.
+		// TODO: Remove this and handle this case
+		return t.fallback.Pick(qry)
 	}
 
 	routingKey, err := qry.GetRoutingKey()

session.go

@@ -518,6 +518,7 @@ func (q *Query) GetRoutingKey() ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
+
 	if routingKeyInfo == nil {
 		return nil, nil
 	}