Revisit getToken/setToken (JAVA-312).

olim7t · olim7t · commit 3c319c30a891 · 2015-02-20T18:43:55.000+01:00
diff --git a/driver-core/src/main/java/com/datastax/driver/core/ArrayBackedRow.java b/driver-core/src/main/java/com/datastax/driver/core/ArrayBackedRow.java
@@ -256,19 +256,19 @@ public Token getToken(int i) {
     }
 
     public Token getToken(String name) {
-        // first, try with `token(name)`, preserving case sensitivity
-        String tokenName = (name.length() >= 2 && name.charAt(0) == '"' && name.charAt(name.length() - 1) == '"')
-            ? "\"token(" + name.substring(1, name.length() - 1) + ")\""
-            : "token(" + name + ")";
-        int[] indexes = metadata.findAllIdx(tokenName);
-        if (indexes != null) {
-            return getToken(indexes[0]);
-        }
-
-        // Otherwise, use normal name
         return getToken(metadata.getFirstIdx(name));
     }
 
+    public Token getPartitionKeyToken() {
+        int i = 0;
+        for (ColumnDefinitions.Definition column : metadata) {
+            if (column.getName().matches("token(.*)"))
+                return getToken(i);
+            i++;
+        }
+        throw new IllegalStateException("Found no column named 'token(...)'. If the column is aliased, use getToken(String).");
+    }
+
     @SuppressWarnings("unchecked")
     public <T> List<T> getList(int i, Class<T> elementsClass) {
         DataType type = metadata.getType(i);
diff --git a/driver-core/src/main/java/com/datastax/driver/core/BoundStatement.java b/driver-core/src/main/java/com/datastax/driver/core/BoundStatement.java
@@ -527,8 +527,8 @@ public BoundStatement setDouble(String name, double v) {
      */
     public BoundStatement setString(int i, String v) {
         DataType.Name type = metadata().checkType(i, DataType.Name.VARCHAR,
-                                                     DataType.Name.TEXT,
-                                                     DataType.Name.ASCII);
+            DataType.Name.TEXT,
+            DataType.Name.ASCII);
         switch (type) {
             case ASCII:
                 return setValue(i, v == null ? null : TypeCodec.StringCodec.asciiInstance.serialize(v));
@@ -742,7 +742,7 @@ public BoundStatement setDecimal(String name, BigDecimal v) {
      */
     public BoundStatement setUUID(int i, UUID v) {
         DataType.Name type = metadata().checkType(i, DataType.Name.UUID,
-                                                       DataType.Name.TIMEUUID);
+            DataType.Name.TIMEUUID);
 
         if (v == null)
             return setValue(i, null);
@@ -823,6 +823,9 @@ public BoundStatement setInet(String name, InetAddress v) {
 
     /**
      * Sets the {@code i}th value to the provided {@link Token}.
+     * <p>
+     * {@link #setPartitionKeyToken(Token)} should generally be preferred if you
+     * have a single token variable.
      *
      * @param i the index of the variable to set.
      * @param v the value to set.
@@ -839,6 +842,18 @@ public BoundStatement setToken(int i, Token v) {
     /**
      * Sets the value for (all occurrences of) variable {@code name} to the
      * provided token.
+     * <p>
+     * {@link #setPartitionKeyToken(Token)} should generally be preferred if you
+     * have a single token variable.
+     * <p>
+     * If you have multiple token variables, use positional binding ({@link #setToken(int, Token)},
+     * or named bind markers:
+     * <pre>
+     * {@code
+     * PreparedStatement pst = session.prepare("SELECT * FROM my_table WHERE token(k) > :min AND token(k) <= :max");
+     * BoundStatement b = pst.bind().setToken("min", minToken).setToken("max", maxToken);
+     * }
+     * </pre>
      *
      * @param name the name of the variable to set; if multiple variables
      * {@code name} are prepared, all of them are set.
@@ -857,6 +872,40 @@ public BoundStatement setToken(String name, Token v) {
         return this;
     }
 
+    /**
+     * Sets the value for (all occurrences of) variable "{@code partition key token}"
+     * to the provided token (this is the name generated by Cassandra for markers
+     * corresponding to a {@code token(...)} call).
+     * <p>
+     * This method is a shorthand for statements with a single token variable:
+     * <pre>
+     * {@code
+     * Token token = ...
+     * PreparedStatement pst = session.prepare("SELECT * FROM my_table WHERE token(k) = ?");
+     * BoundStatement b = pst.bind().setPartitionKeyToken(token);
+     * }
+     * </pre>
+     * If you have multiple token variables, use positional binding ({@link #setToken(int, Token)},
+     * or named bind markers:
+     * <pre>
+     * {@code
+     * PreparedStatement pst = session.prepare("SELECT * FROM my_table WHERE token(k) > :min AND token(k) <= :max");
+     * BoundStatement b = pst.bind().setToken("min", minToken).setToken("max", maxToken);
+     * }
+     * </pre>
+     *
+     * @param v the value to set.
+     * @return this BoundStatement.
+     *
+     * @throws IllegalArgumentException if {@code name} is not a prepared
+     * variable, that is, if {@code !this.preparedStatement().variables().names().contains(name)}.
+     * @throws InvalidTypeException if (any occurrence of) {@code name} is
+     * not of the type of the token's value.
+     */
+    public BoundStatement setPartitionKeyToken(Token v) {
+        return setToken("partition key token", v);
+    }
+
     /**
      * Sets the {@code i}th value to the provided list.
      * <p>
diff --git a/driver-core/src/main/java/com/datastax/driver/core/Row.java b/driver-core/src/main/java/com/datastax/driver/core/Row.java
@@ -406,10 +406,12 @@ public interface Row {
 
     /**
      * Returns the {@code i}th value of this row as a {@link Token}.
+     * <p>
+     * {@link #getPartitionKeyToken()} should generally be preferred to this method (unless the
+     * token column is aliased).
      *
      * @param i the index ({@code 0 <= i < size()}) of the column to retrieve.
      * @return the value of the {@code i}th column in this row as an Token.
-     * If the value is NULL, {@code null} is returned.
      *
      * @throws IndexOutOfBoundsException if {@code i < 0 || i >= this.columns().size()}.
      * @throws InvalidTypeException if column {@code i} is not of the type of token values
@@ -420,26 +422,11 @@ public interface Row {
     /**
      * Returns the value of column {@code name} as a {@link Token}.
      * <p>
-     * This method will first try with {@code token(name)}, because in most cases the column
-     * is named that way since it results from the application of a CQL function:
-     * <pre>
-     * {@code
-     * ResultSet rs = session.execute("SELECT token(k) FROM my_table WHERE k = 1");
-     * Token token = rs.one().getToken("k"); // retrieves token(k)
-     * }
-     * </pre>
-     * If that fails, it will try will {@code name}. This handles the case where the column
-     * was aliased:
-     * <pre>
-     * {@code
-     * ResultSet rs = session.execute("SELECT token(k) AS t FROM my_table WHERE k = 1");
-     * Token token = rs.one().getToken("t");
-     * }
-     * </pre>
+     * {@link #getPartitionKeyToken()} should generally be preferred to this method (unless the
+     * token column is aliased).
      *
      * @param name the name of the column to retrieve.
      * @return the value of column {@code name} as a Token.
-     * If the value is NULL, {@code null} is returned.
      *
      * @throws IllegalArgumentException if {@code name} is not part of the
      * ResultSet this row is part of, i.e. if {@code !this.columns().names().contains(name)}.
@@ -448,6 +435,29 @@ public interface Row {
      */
     public Token getToken(String name);
 
+    /**
+     * Returns the value of the first column containing a {@link Token}.
+     * <p>
+     * This method is a shorthand for queries returning a single token in an unaliased
+     * column. It will look for the first name matching {@code token(...)}:
+     * <pre>
+     * {@code
+     * ResultSet rs = session.execute("SELECT token(k) FROM my_table WHERE k = 1");
+     * Token token = rs.one().getPartitionKeyToken(); // retrieves token(k)
+     * }
+     * </pre>
+     * If that doesn't work for you (for example, if you're using an alias), use
+     * {@link #getToken(int)} or {@link #getToken(String)}.
+     *
+     * @return the value of column {@code name} as a Token.
+     *
+     * @throws IllegalStateException if no column named {@code token(...)} exists in this
+     * ResultSet.
+     * @throws InvalidTypeException if the first column named {@code token(...)} is not of
+     * the type of token values for this cluster (this depends on the configured partitioner).
+     */
+    public Token getPartitionKeyToken();
+
     /**
      * Returns the {@code i}th value of this row as a list.
      *
diff --git a/driver-core/src/test/java/com/datastax/driver/core/TokenIntegrationTest.java b/driver-core/src/test/java/com/datastax/driver/core/TokenIntegrationTest.java
@@ -166,21 +166,16 @@ public void should_get_token_from_row_and_set_token_in_query() {
         Token token = row.getToken(0);
         assertThat(token.getType()).isEqualTo(expectedTokenType);
 
-        // get by "base" column name (will add token() around it):
         assertThat(
-            row.getToken("i")
+            row.getPartitionKeyToken()
         ).isEqualTo(token);
 
-        // get by name:
+        // get by name when column is aliased:
         assertThat(
             session.execute("SELECT token(i) AS t FROM test.foo WHERE i = 1").one()
                 .getToken("t")
         ).isEqualTo(token);
 
-        // get by "base" column name when it's case-sensitive
-        session.execute("SELECT token(\"caseSensitiveKey\") FROM test.foo2 WHERE \"caseSensitiveKey\" = 1").one()
-            .getToken("\"caseSensitiveKey\"");
-
 
         PreparedStatement pst = session.prepare("SELECT * FROM test.foo WHERE token(i) = ?");
         row = session.execute(pst.bind(token)).one();
@@ -189,6 +184,9 @@ public void should_get_token_from_row_and_set_token_in_query() {
         row = session.execute(pst.bind().setToken(0, token)).one();
         assertThat(row.getInt(0)).isEqualTo(1);
 
+        row = session.execute(pst.bind().setPartitionKeyToken(token)).one();
+        assertThat(row.getInt(0)).isEqualTo(1);
+
         PreparedStatement pst2 = session.prepare("SELECT * FROM test.foo WHERE token(i) = :myToken");
         row = session.execute(pst2.bind().setToken("myToken", token)).one();
         assertThat(row.getInt(0)).isEqualTo(1);
