Expose some internals of SelectStatement for inspection by QueryHandlers

patch by Sam Tunnicliffe; reviewed by Benjamin Lerer and Mick Semb Wever
for CASSANDRA-9532


Project: http://git-wip-us.apache.org/repos/asf/cassandra/repo
Commit: http://git-wip-us.apache.org/repos/asf/cassandra/commit/f32cff8e
Tree: http://git-wip-us.apache.org/repos/asf/cassandra/tree/f32cff8e
Diff: http://git-wip-us.apache.org/repos/asf/cassandra/diff/f32cff8e

Branch: refs/heads/cassandra-2.2
Commit: f32cff8e1fb69317219ffaee81b5861a54b83a1b
Parents: ad8047a
Author: Sam Tunnicliffe <s...@beobal.com>
Authored: Thu Jun 4 18:12:35 2015 +0100
Committer: Sam Tunnicliffe <s...@beobal.com>
Committed: Thu Jun 18 17:11:00 2015 +0100

----------------------------------------------------------------------
 CHANGES.txt                                     |   1 +
 .../cassandra/cql3/ColumnSpecification.java     |  22 ++
 .../cql3/statements/SelectStatement.java        |  37 ++-
 .../cassandra/cql3/statements/Selection.java    |  80 +++---
 .../cql3/statements/SelectionColumnMapping.java | 106 ++++++++
 .../cql3/statements/SelectionColumns.java       |  19 ++
 .../statements/SelectionColumnMappingTest.java  | 244 +++++++++++++++++++
 7 files changed, 476 insertions(+), 33 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/cassandra/blob/f32cff8e/CHANGES.txt
----------------------------------------------------------------------
diff --git a/CHANGES.txt b/CHANGES.txt
index 753fb1c..a235528 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -1,4 +1,5 @@
 2.0.16:
+ * Expose some internals of SelectStatement for inspection (CASSANDRA-9532)
  * ArrivalWindow should use primitives (CASSANDRA-9496)
  * Periodically submit background compaction tasks (CASSANDRA-9592)
  * Set HAS_MORE_PAGES flag to false when PagingState is null (CASSANDRA-9571)

http://git-wip-us.apache.org/repos/asf/cassandra/blob/f32cff8e/src/java/org/apache/cassandra/cql3/ColumnSpecification.java
----------------------------------------------------------------------
diff --git a/src/java/org/apache/cassandra/cql3/ColumnSpecification.java 
b/src/java/org/apache/cassandra/cql3/ColumnSpecification.java
index 4dae701..089a1c5 100644
--- a/src/java/org/apache/cassandra/cql3/ColumnSpecification.java
+++ b/src/java/org/apache/cassandra/cql3/ColumnSpecification.java
@@ -17,6 +17,8 @@
  */
 package org.apache.cassandra.cql3;
 
+import com.google.common.base.Objects;
+
 import org.apache.cassandra.db.marshal.AbstractType;
 
 public class ColumnSpecification
@@ -40,4 +42,24 @@ public class ColumnSpecification
         // Not fully conventional, but convenient (for error message to users 
in particular)
         return name.toString();
     }
+
+    public boolean equals(Object obj)
+    {
+        if (null == obj)
+            return false;
+
+        if(!(obj instanceof ColumnSpecification))
+            return false;
+
+        ColumnSpecification other = (ColumnSpecification)obj;
+        return Objects.equal(ksName, other.ksName)
+            && Objects.equal(cfName, other.cfName)
+            && Objects.equal(name, other.name)
+            && Objects.equal(type, other.type);
+    }
+
+    public int hashCode()
+    {
+        return Objects.hashCode(ksName, cfName, name, type);
+    }
 }

http://git-wip-us.apache.org/repos/asf/cassandra/blob/f32cff8e/src/java/org/apache/cassandra/cql3/statements/SelectStatement.java
----------------------------------------------------------------------
diff --git a/src/java/org/apache/cassandra/cql3/statements/SelectStatement.java 
b/src/java/org/apache/cassandra/cql3/statements/SelectStatement.java
index 95e0441..1c19760 100644
--- a/src/java/org/apache/cassandra/cql3/statements/SelectStatement.java
+++ b/src/java/org/apache/cassandra/cql3/statements/SelectStatement.java
@@ -60,7 +60,10 @@ import org.slf4j.LoggerFactory;
 /**
  * Encapsulates a completely parsed SELECT query, including the target
  * column family, expression, result count, and ordering clause.
- *
+ * A number of public methods here are only used internally. However,
+ * many of these are made accessible for the benefit of custom
+ * QueryHandler implementations, so before reducing their accessibility
+ * due consideration should be given.
  */
 public class SelectStatement implements CQLStatement, 
MeasurableForPreparedCache
 {
@@ -184,6 +187,14 @@ public class SelectStatement implements CQLStatement, 
MeasurableForPreparedCache
         return boundTerms;
     }
 
+    /**
+     * May be used by custom QueryHandler implementations
+     */
+    public Selection getSelection()
+    {
+        return selection;
+    }
+
     public void checkAccess(ClientState state) throws InvalidRequestException, 
UnauthorizedException
     {
         state.hasColumnFamilyAccess(keyspace(), columnFamily(), 
Permission.SELECT);
@@ -580,7 +591,10 @@ public class SelectStatement implements CQLStatement, 
MeasurableForPreparedCache
         return new SliceQueryFilter(slices, isReversed, limit, toGroup);
     }
 
-    private int getLimit(List<ByteBuffer> variables) throws 
InvalidRequestException
+    /**
+     * May be used by custom QueryHandler implementations
+     */
+    public int getLimit(List<ByteBuffer> variables) throws 
InvalidRequestException
     {
         int l = Integer.MAX_VALUE;
         if (limit != null)
@@ -1067,6 +1081,9 @@ public class SelectStatement implements CQLStatement, 
MeasurableForPreparedCache
                           variables);
     }
 
+    /**
+     * May be used by custom QueryHandler implementations
+     */
     public List<IndexExpression> getIndexExpressions(List<ByteBuffer> 
variables) throws InvalidRequestException
     {
         if (!usesSecondaryIndexing || restrictedNames.isEmpty())
@@ -1446,7 +1463,21 @@ public class SelectStatement implements CQLStatement, 
MeasurableForPreparedCache
         return true;
     }
 
-    private boolean hasClusteringColumnsRestriction()
+    /**
+     * May be used by custom QueryHandler implementations
+     */
+    public boolean hasPartitionKeyRestriction()
+    {
+        for (int i = 0; i < keyRestrictions.length; i++)
+            if (keyRestrictions[i] != null)
+                return true;
+        return false;
+    }
+
+    /**
+     * May be used by custom QueryHandler implementations
+     */
+    public boolean hasClusteringColumnsRestriction()
     {
         for (int i = 0; i < columnRestrictions.length; i++)
             if (columnRestrictions[i] != null)

http://git-wip-us.apache.org/repos/asf/cassandra/blob/f32cff8e/src/java/org/apache/cassandra/cql3/statements/Selection.java
----------------------------------------------------------------------
diff --git a/src/java/org/apache/cassandra/cql3/statements/Selection.java 
b/src/java/org/apache/cassandra/cql3/statements/Selection.java
index 223f698..50a34bf 100644
--- a/src/java/org/apache/cassandra/cql3/statements/Selection.java
+++ b/src/java/org/apache/cassandra/cql3/statements/Selection.java
@@ -18,8 +18,9 @@
 package org.apache.cassandra.cql3.statements;
 
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.List;
+import java.util.*;
+
+import com.google.common.collect.*;
 
 import org.apache.cassandra.cql3.*;
 import org.apache.cassandra.cql3.functions.Function;
@@ -37,14 +38,14 @@ import org.apache.cassandra.utils.ByteBufferUtil;
 public abstract class Selection
 {
     private final List<CFDefinition.Name> columns;
-    private final List<ColumnSpecification> metadata;
+    private final SelectionColumns columnMapping;
     private final boolean collectTimestamps;
     private final boolean collectTTLs;
 
-    protected Selection(List<CFDefinition.Name> columns, 
List<ColumnSpecification> metadata, boolean collectTimestamps, boolean 
collectTTLs)
+    protected Selection(List<CFDefinition.Name> columns, SelectionColumns 
columnMapping, boolean collectTimestamps, boolean collectTTLs)
     {
         this.columns = columns;
-        this.metadata = metadata;
+        this.columnMapping = columnMapping;
         this.collectTimestamps = collectTimestamps;
         this.collectTTLs = collectTTLs;
     }
@@ -57,7 +58,7 @@ public abstract class Selection
 
     public ResultSet.Metadata getResultMetadata()
     {
-        return new ResultSet.Metadata(metadata);
+        return new ResultSet.Metadata(columnMapping.getColumnSpecifications());
     }
 
     public static Selection wildcard(CFDefinition cfDef)
@@ -94,21 +95,28 @@ public abstract class Selection
         return idx;
     }
 
-    private static Selector makeSelector(CFDefinition cfDef, RawSelector raw, 
List<CFDefinition.Name> names, List<ColumnSpecification> metadata) throws 
InvalidRequestException
+    private static Selector makeSelector(CFDefinition cfDef,
+                                         RawSelector raw,
+                                         List<CFDefinition.Name> names,
+                                         SelectionColumnMapping columnMapping) 
throws InvalidRequestException
     {
         Selectable selectable = raw.selectable.prepare(cfDef.cfm);
-        return makeSelector(cfDef, selectable, raw.alias, names, metadata);
+        return makeSelector(cfDef, selectable, raw.alias, names, 
columnMapping);
     }
 
-    private static Selector makeSelector(CFDefinition cfDef, Selectable 
selectable, ColumnIdentifier alias, List<CFDefinition.Name> names, 
List<ColumnSpecification> metadata) throws InvalidRequestException
+    private static Selector makeSelector(CFDefinition cfDef,
+                                         Selectable selectable,
+                                         ColumnIdentifier alias,
+                                         List<CFDefinition.Name> names,
+                                         SelectionColumnMapping columnMapping) 
throws InvalidRequestException
     {
         if (selectable instanceof ColumnIdentifier)
         {
-            CFDefinition.Name name = cfDef.get((ColumnIdentifier)selectable);
+            CFDefinition.Name name = cfDef.get((ColumnIdentifier) selectable);
             if (name == null)
                 throw new InvalidRequestException(String.format("Undefined 
name %s in selection clause", selectable));
-            if (metadata != null)
-                metadata.add(alias == null ? name : makeAliasSpec(cfDef, 
name.type, alias));
+            if (columnMapping != null)
+                columnMapping.addMapping(alias == null ? name : 
makeAliasSpec(cfDef, name.type, alias), name);
             return new SimpleSelector(name.toString(), addAndGetIndex(name, 
names), name.type);
         }
         else if (selectable instanceof Selectable.WritetimeOrTTL)
@@ -121,25 +129,26 @@ public abstract class Selection
                 throw new InvalidRequestException(String.format("Cannot use 
selection function %s on PRIMARY KEY part %s", tot.isWritetime ? "writeTime" : 
"ttl", name));
             if (name.type.isCollection())
                 throw new InvalidRequestException(String.format("Cannot use 
selection function %s on collections", tot.isWritetime ? "writeTime" : "ttl"));
-
-            if (metadata != null)
-                metadata.add(makeWritetimeOrTTLSpec(cfDef, tot, alias));
+            if (columnMapping != null)
+                columnMapping.addMapping(makeWritetimeOrTTLSpec(cfDef, tot, 
alias), name);
             return new WritetimeOrTTLSelector(name.toString(), 
addAndGetIndex(name, names), tot.isWritetime);
         }
         else
         {
             Selectable.WithFunction withFun = 
(Selectable.WithFunction)selectable;
             List<Selector> args = new ArrayList<Selector>(withFun.args.size());
+            // use a temporary column mapping to collate the columns used by 
all the function args
+            SelectionColumnMapping tmpMapping = 
SelectionColumnMapping.newMapping();
             for (Selectable rawArg : withFun.args)
-                args.add(makeSelector(cfDef, rawArg, null, names, null));
+                args.add(makeSelector(cfDef, rawArg, null, names, tmpMapping));
 
             AbstractType<?> returnType = 
Functions.getReturnType(withFun.functionName, cfDef.cfm.ksName, 
cfDef.cfm.cfName);
             if (returnType == null)
                 throw new InvalidRequestException(String.format("Unknown 
function '%s'", withFun.functionName));
             ColumnSpecification spec = makeFunctionSpec(cfDef, withFun, 
returnType, alias);
             Function fun = Functions.get(withFun.functionName, args, spec);
-            if (metadata != null)
-                metadata.add(spec);
+            if (columnMapping != null)
+                columnMapping.addMapping(spec, 
tmpMapping.getMappings().values());
             return new FunctionSelector(fun, args);
         }
     }
@@ -178,23 +187,23 @@ public abstract class Selection
         if (needsProcessing)
         {
             List<CFDefinition.Name> names = new ArrayList<CFDefinition.Name>();
-            List<ColumnSpecification> metadata = new 
ArrayList<ColumnSpecification>(rawSelectors.size());
+            SelectionColumnMapping columnMapping = 
SelectionColumnMapping.newMapping();
             List<Selector> selectors = new 
ArrayList<Selector>(rawSelectors.size());
             boolean collectTimestamps = false;
             boolean collectTTLs = false;
             for (RawSelector rawSelector : rawSelectors)
             {
-                Selector selector = makeSelector(cfDef, rawSelector, names, 
metadata);
+                Selector selector = makeSelector(cfDef, rawSelector, names, 
columnMapping);
                 selectors.add(selector);
                 collectTimestamps |= selector.usesTimestamps();
                 collectTTLs |= selector.usesTTLs();
             }
-            return new SelectionWithProcessing(names, metadata, selectors, 
collectTimestamps, collectTTLs);
+            return new SelectionWithProcessing(names, columnMapping, 
selectors, collectTimestamps, collectTTLs);
         }
         else
         {
             List<CFDefinition.Name> names = new 
ArrayList<CFDefinition.Name>(rawSelectors.size());
-            List<ColumnSpecification> metadata = new 
ArrayList<ColumnSpecification>(rawSelectors.size());
+            SelectionColumnMapping columnMapping = 
SelectionColumnMapping.newMapping();
             for (RawSelector rawSelector : rawSelectors)
             {
                 assert rawSelector.selectable instanceof ColumnIdentifier.Raw;
@@ -203,9 +212,12 @@ public abstract class Selection
                 if (name == null)
                     throw new InvalidRequestException(String.format("Undefined 
name %s in selection clause", id));
                 names.add(name);
-                metadata.add(rawSelector.alias == null ? name : 
makeAliasSpec(cfDef, name.type, rawSelector.alias));
+                columnMapping.addMapping(rawSelector.alias == null ? name : 
makeAliasSpec(cfDef,
+                                                                               
           name.type,
+                                                                               
           rawSelector.alias),
+                                         name);
             }
-            return new SimpleSelection(names, metadata, false);
+            return new SimpleSelection(names, columnMapping, false);
         }
     }
 
@@ -233,6 +245,14 @@ public abstract class Selection
         return columns;
     }
 
+    /**
+     * @return the mappings between resultset columns and the underlying 
columns
+     */
+    public SelectionColumns getColumnMapping()
+    {
+        return columnMapping;
+    }
+
     public ResultSetBuilder resultSetBuilder(long now)
     {
         return new ResultSetBuilder(now);
@@ -264,7 +284,7 @@ public abstract class Selection
 
         private ResultSetBuilder(long now)
         {
-            this.resultSet = new ResultSet(metadata);
+            this.resultSet = new 
ResultSet(columnMapping.getColumnSpecifications());
             this.timestamps = collectTimestamps ? new long[columns.size()] : 
null;
             this.ttls = collectTTLs ? new int[columns.size()] : null;
             this.now = now;
@@ -321,17 +341,17 @@ public abstract class Selection
 
         public SimpleSelection(List<CFDefinition.Name> columns, boolean 
isWildcard)
         {
-            this(columns, new ArrayList<ColumnSpecification>(columns), 
isWildcard);
+            this(columns, SelectionColumnMapping.simpleMapping(columns), 
isWildcard);
         }
 
-        public SimpleSelection(List<CFDefinition.Name> columns, 
List<ColumnSpecification> metadata, boolean isWildcard)
+        public SimpleSelection(List<CFDefinition.Name> columns, 
SelectionColumnMapping columnMapping, boolean isWildcard)
         {
             /*
              * In theory, even a simple selection could have multiple time the 
same column, so we
              * could filter those duplicate out of columns. But since we're 
very unlikely to
              * get much duplicate in practice, it's more efficient not to 
bother.
              */
-            super(columns, metadata, false, false);
+            super(columns, columnMapping, false, false);
             this.isWildcard = isWildcard;
         }
 
@@ -351,9 +371,9 @@ public abstract class Selection
     {
         private final List<Selector> selectors;
 
-        public SelectionWithProcessing(List<CFDefinition.Name> columns, 
List<ColumnSpecification> metadata, List<Selector> selectors, boolean 
collectTimestamps, boolean collectTTLs)
+        public SelectionWithProcessing(List<CFDefinition.Name> columns, 
SelectionColumns columnMapping, List<Selector> selectors, boolean 
collectTimestamps, boolean collectTTLs)
         {
-            super(columns, metadata, collectTimestamps, collectTTLs);
+            super(columns, columnMapping, collectTimestamps, collectTTLs);
             this.selectors = selectors;
         }
 

http://git-wip-us.apache.org/repos/asf/cassandra/blob/f32cff8e/src/java/org/apache/cassandra/cql3/statements/SelectionColumnMapping.java
----------------------------------------------------------------------
diff --git 
a/src/java/org/apache/cassandra/cql3/statements/SelectionColumnMapping.java 
b/src/java/org/apache/cassandra/cql3/statements/SelectionColumnMapping.java
new file mode 100644
index 0000000..d09612f
--- /dev/null
+++ b/src/java/org/apache/cassandra/cql3/statements/SelectionColumnMapping.java
@@ -0,0 +1,106 @@
+package org.apache.cassandra.cql3.statements;
+
+import java.util.*;
+
+import com.google.common.base.Function;
+import com.google.common.base.Joiner;
+import com.google.common.collect.*;
+
+import org.apache.cassandra.cql3.CFDefinition;
+import org.apache.cassandra.cql3.ColumnSpecification;
+
+public class SelectionColumnMapping implements SelectionColumns
+{
+
+    // Uses LinkedHashMultimap because ordering of keys must be maintained
+    private final LinkedHashMultimap<ColumnSpecification, CFDefinition.Name> 
columnMappings;
+
+    private SelectionColumnMapping()
+    {
+        this.columnMappings = LinkedHashMultimap.create();
+    }
+
+    protected static SelectionColumnMapping newMapping()
+    {
+        return new SelectionColumnMapping();
+    }
+
+    protected static SelectionColumnMapping 
simpleMapping(List<CFDefinition.Name> columnDefinitions)
+    {
+        SelectionColumnMapping mapping = new SelectionColumnMapping();
+        for (CFDefinition.Name def: columnDefinitions)
+            mapping.addMapping(def, def);
+        return mapping;
+    }
+
+    protected SelectionColumnMapping addMapping(ColumnSpecification colSpec, 
CFDefinition.Name column)
+    {
+        columnMappings.put(colSpec, column);
+        return this;
+    }
+
+    protected SelectionColumnMapping addMapping(ColumnSpecification colSpec, 
Iterable<CFDefinition.Name> columns)
+    {
+        columnMappings.putAll(colSpec, columns);
+        return this;
+    }
+
+    public List<ColumnSpecification> getColumnSpecifications()
+    {
+        // return a mutable copy as we may add extra columns
+        // for ordering (CASSANDRA-4911 & CASSANDRA-8286)
+        return new ArrayList(columnMappings.keySet());
+    }
+
+    public Multimap<ColumnSpecification, CFDefinition.Name> getMappings()
+    {
+        return Multimaps.unmodifiableMultimap(columnMappings);
+    }
+
+    public boolean equals(Object obj)
+    {
+        if (obj == null)
+            return false;
+
+        if (!(obj instanceof SelectionColumns))
+            return false;
+
+        return Objects.equals(columnMappings, ((SelectionColumns) 
obj).getMappings());
+    }
+
+    public int hashCode()
+    {
+        return Objects.hashCode(columnMappings);
+    }
+
+    public String toString()
+    {
+        final Function<CFDefinition.Name, String> getDefName = new 
Function<CFDefinition.Name, String>()
+        {
+            public String apply(CFDefinition.Name name)
+            {
+                return name.toString();
+            }
+        };
+        Function<Map.Entry<ColumnSpecification, 
Collection<CFDefinition.Name>>, String> mappingEntryToString =
+        new Function<Map.Entry<ColumnSpecification, 
Collection<CFDefinition.Name>>, String>(){
+            public String apply(Map.Entry<ColumnSpecification, 
Collection<CFDefinition.Name>> entry)
+            {
+                StringBuilder builder = new StringBuilder();
+                builder.append(entry.getKey().name.toString());
+                builder.append(":[");
+                
builder.append(Joiner.on(',').join(Iterables.transform(entry.getValue(), 
getDefName)));
+                builder.append("]");
+                return builder.toString();
+            }
+        };
+
+        StringBuilder builder = new StringBuilder();
+        builder.append("{ ");
+        builder.append(Joiner.on(", ")
+                             
.join(Iterables.transform(columnMappings.asMap().entrySet(),
+                                                       mappingEntryToString)));
+        builder.append(" }");
+        return builder.toString();
+    }
+}

http://git-wip-us.apache.org/repos/asf/cassandra/blob/f32cff8e/src/java/org/apache/cassandra/cql3/statements/SelectionColumns.java
----------------------------------------------------------------------
diff --git 
a/src/java/org/apache/cassandra/cql3/statements/SelectionColumns.java 
b/src/java/org/apache/cassandra/cql3/statements/SelectionColumns.java
new file mode 100644
index 0000000..3053f99
--- /dev/null
+++ b/src/java/org/apache/cassandra/cql3/statements/SelectionColumns.java
@@ -0,0 +1,19 @@
+package org.apache.cassandra.cql3.statements;
+
+import java.util.List;
+
+import com.google.common.collect.Multimap;
+
+import org.apache.cassandra.cql3.CFDefinition;
+import org.apache.cassandra.cql3.ColumnSpecification;
+
+/**
+ * Represents a mapping between the actual columns used to satisfy a Selection
+ * and the column definitions included in the resultset metadata for the query.
+ */
+public interface SelectionColumns
+{
+    List<ColumnSpecification> getColumnSpecifications();
+    Multimap<ColumnSpecification, CFDefinition.Name> getMappings();
+}
+

http://git-wip-us.apache.org/repos/asf/cassandra/blob/f32cff8e/test/unit/org/apache/cassandra/cql3/statements/SelectionColumnMappingTest.java
----------------------------------------------------------------------
diff --git 
a/test/unit/org/apache/cassandra/cql3/statements/SelectionColumnMappingTest.java
 
b/test/unit/org/apache/cassandra/cql3/statements/SelectionColumnMappingTest.java
new file mode 100644
index 0000000..9c31653
--- /dev/null
+++ 
b/test/unit/org/apache/cassandra/cql3/statements/SelectionColumnMappingTest.java
@@ -0,0 +1,244 @@
+package org.apache.cassandra.cql3.statements;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+import org.apache.cassandra.SchemaLoader;
+import org.apache.cassandra.config.Schema;
+import org.apache.cassandra.cql3.*;
+import org.apache.cassandra.db.ConsistencyLevel;
+import org.apache.cassandra.db.marshal.*;
+import org.apache.cassandra.exceptions.RequestValidationException;
+import org.apache.cassandra.service.ClientState;
+
+import static org.apache.cassandra.cql3.QueryProcessor.process;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class SelectionColumnMappingTest
+{
+    static String KEYSPACE = "selection_column_mapping_test_ks";
+    String tableName = "test_table";
+
+    @BeforeClass
+    public static void setupSchema() throws Throwable
+    {
+        SchemaLoader.loadSchema();
+        executeSchemaChange(String.format("CREATE KEYSPACE IF NOT EXISTS %s " +
+                                          "WITH replication = {'class': 
'SimpleStrategy', " +
+                                          "                    
'replication_factor': '1'}",
+                                          KEYSPACE));
+    }
+
+    @Test
+    public void testSelectionColumnMapping() throws Throwable
+    {
+        // Organised as a single test to avoid the overhead of
+        // table creation for each variant
+        tableName = "table1";
+        createTable("CREATE TABLE %s (" +
+                    " k int PRIMARY KEY," +
+                    " v1 int," +
+                    " v2 ascii)");
+        testSimpleTypes();
+        testWildcard();
+        testSimpleTypesWithAliases();
+        testWritetimeAndTTL();
+        testWritetimeAndTTLWithAliases();
+        testFunction();
+        testFunctionWithAlias();
+        testMultipleAliasesOnSameColumn();
+        testMixedColumnTypes();
+    }
+
+    @Test
+    public void testMultipleArgumentFunction() throws Throwable
+    {
+        // token() is currently the only function which accepts multiple 
arguments
+        tableName = "table2";
+        createTable("CREATE TABLE %s (a int, b text, PRIMARY KEY ((a, b)))");
+        ColumnSpecification tokenSpec = columnSpecification("token(a, b)", 
BytesType.instance);
+        SelectionColumns expected = SelectionColumnMapping.newMapping()
+                                                          
.addMapping(tokenSpec, columnDefinitions("a", "b"));
+
+        assertEquals(expected, extractColumnMappingFromSelect("SELECT 
token(a,b) FROM %s"));
+    }
+
+    private void testSimpleTypes() throws Throwable
+    {
+        // simple column identifiers without aliases are represented in
+        // ResultSet.Metadata by the underlying ColumnDefinition
+        CFDefinition.Name kDef = columnDefinition("k");
+        CFDefinition.Name v1Def = columnDefinition("v1");
+        CFDefinition.Name v2Def = columnDefinition("v2");
+        SelectionColumns expected = SelectionColumnMapping.newMapping()
+                                                          .addMapping(kDef, 
columnDefinition("k"))
+                                                          .addMapping(v1Def, 
columnDefinition("v1"))
+                                                          .addMapping(v2Def, 
columnDefinition("v2"));
+
+        assertEquals(expected, extractColumnMappingFromSelect("SELECT k, v1, 
v2 FROM %s"));
+    }
+
+    private void testWildcard() throws Throwable
+    {
+        // Wildcard select should behave just as though we had
+        // explicitly selected each column
+        CFDefinition.Name kDef = columnDefinition("k");
+        CFDefinition.Name v1Def = columnDefinition("v1");
+        CFDefinition.Name v2Def = columnDefinition("v2");
+        SelectionColumns expected = SelectionColumnMapping.newMapping()
+                                                          .addMapping(kDef, 
columnDefinition("k"))
+                                                          .addMapping(v1Def, 
columnDefinition("v1"))
+                                                          .addMapping(v2Def, 
columnDefinition("v2"));
+
+        assertEquals(expected, extractColumnMappingFromSelect("SELECT * FROM 
%s"));
+    }
+
+    private void testSimpleTypesWithAliases() throws Throwable
+    {
+        // simple column identifiers with aliases are represented in 
ResultSet.Metadata
+        // by a ColumnSpecification based on the underlying ColumnDefinition
+        ColumnSpecification kSpec = columnSpecification("k_alias", 
Int32Type.instance);
+        ColumnSpecification v1Spec = columnSpecification("v1_alias", 
Int32Type.instance);
+        ColumnSpecification v2Spec = columnSpecification("v2_alias", 
AsciiType.instance);
+        SelectionColumns expected = SelectionColumnMapping.newMapping()
+                                                          .addMapping(kSpec, 
columnDefinition("k"))
+                                                          .addMapping(v1Spec, 
columnDefinition("v1"))
+                                                          .addMapping(v2Spec, 
columnDefinition("v2"));
+
+        assertEquals(expected, extractColumnMappingFromSelect(
+                                                             "SELECT k AS 
k_alias, v1 AS v1_alias, v2 AS v2_alias FROM %s"));
+    }
+
+    private void testWritetimeAndTTL() throws Throwable
+    {
+        // writetime and ttl are represented in ResultSet.Metadata by a 
ColumnSpecification
+        // with the function name plus argument and a long or int type 
respectively
+        ColumnSpecification wtSpec = columnSpecification("writetime(v1)", 
LongType.instance);
+        ColumnSpecification ttlSpec = columnSpecification("ttl(v2)", 
Int32Type.instance);
+        SelectionColumns expected = SelectionColumnMapping.newMapping()
+                                                          .addMapping(wtSpec, 
columnDefinition("v1"))
+                                                          .addMapping(ttlSpec, 
columnDefinition("v2"));
+
+        assertEquals(expected, extractColumnMappingFromSelect("SELECT 
writetime(v1), ttl(v2) FROM %s"));
+    }
+
+    private void testWritetimeAndTTLWithAliases() throws Throwable
+    {
+        // writetime and ttl with aliases are represented in ResultSet.Metadata
+        // by a ColumnSpecification with the alias name and the appropriate 
numeric type
+        ColumnSpecification wtSpec = columnSpecification("wt_alias", 
LongType.instance);
+        ColumnSpecification ttlSpec = columnSpecification("ttl_alias", 
Int32Type.instance);
+        SelectionColumns expected = SelectionColumnMapping.newMapping()
+                                                          .addMapping(wtSpec, 
columnDefinition("v1"))
+                                                          .addMapping(ttlSpec, 
columnDefinition("v2"));
+
+        assertEquals(expected, extractColumnMappingFromSelect("SELECT 
writetime(v1) AS wt_alias, ttl(v2) AS ttl_alias FROM %s"));
+    }
+
+    private void testFunction() throws Throwable
+    {
+        // a function such as intasblob(<col>) is represented in 
ResultSet.Metadata
+        // by a ColumnSpecification with the function name plus args and the 
type set
+        // to the function's return type
+        ColumnSpecification fnSpec = columnSpecification("intasblob(v1)", 
BytesType.instance);
+        SelectionColumns expected = SelectionColumnMapping.newMapping()
+                                                          .addMapping(fnSpec, 
columnDefinition("v1"));
+
+        assertEquals(expected, extractColumnMappingFromSelect("SELECT 
intasblob(v1) FROM %s"));
+    }
+
+    private void testFunctionWithAlias() throws Throwable
+    {
+        // a function with an alias is represented in ResultSet.Metadata by a
+        // ColumnSpecification with the alias and the type set to the 
function's
+        // return type
+        ColumnSpecification fnSpec = columnSpecification("fn_alias", 
BytesType.instance);
+        SelectionColumns expected = SelectionColumnMapping.newMapping()
+                                                          .addMapping(fnSpec, 
columnDefinition("v1"));
+
+        assertEquals(expected, extractColumnMappingFromSelect("SELECT 
intasblob(v1) AS fn_alias FROM %s"));
+    }
+
+    private void testMultipleAliasesOnSameColumn() throws Throwable
+    {
+        // Multiple result columns derived from the same underlying column are
+        // represented by ColumnSpecifications
+        ColumnSpecification alias1 = columnSpecification("alias_1", 
Int32Type.instance);
+        ColumnSpecification alias2 = columnSpecification("alias_2", 
Int32Type.instance);
+        SelectionColumns expected = SelectionColumnMapping.newMapping()
+                                                          .addMapping(alias1, 
columnDefinition("v1"))
+                                                          .addMapping(alias2, 
columnDefinition("v1"));
+
+        assertEquals(expected, extractColumnMappingFromSelect("SELECT v1 AS 
alias_1, v1 AS alias_2 FROM %s"));
+    }
+
+    private void testMixedColumnTypes() throws Throwable
+    {
+        ColumnSpecification kSpec = columnSpecification("k_alias", 
Int32Type.instance);
+        ColumnSpecification v1Spec = columnSpecification("writetime(v1)", 
LongType.instance);
+        ColumnSpecification v2Spec = columnSpecification("ttl_alias", 
Int32Type.instance);
+
+        SelectionColumns expected = SelectionColumnMapping.newMapping()
+                                                          .addMapping(kSpec, 
columnDefinition("k"))
+                                                          .addMapping(v1Spec, 
columnDefinition("v1"))
+                                                          .addMapping(v2Spec, 
columnDefinition("v2"));
+
+        assertEquals(expected, extractColumnMappingFromSelect("SELECT k AS 
k_alias," +
+                                                              "       
writetime(v1)," +
+                                                              "       ttl(v2) 
as ttl_alias" +
+                                                              " FROM %s"));
+    }
+
+    private SelectionColumns extractColumnMappingFromSelect(String query) 
throws RequestValidationException
+    {
+        CQLStatement statement = 
QueryProcessor.getStatement(String.format(query, KEYSPACE + "." + tableName),
+                                                             
ClientState.forInternalCalls()).statement;
+        assertTrue(statement instanceof SelectStatement);
+        return ((SelectStatement)statement).getSelection().getColumnMapping();
+    }
+
+    private CFDefinition.Name columnDefinition(String name)
+    {
+        return Schema.instance.getCFMetaData(KEYSPACE, tableName)
+                              .getCfDef()
+                              .get(new ColumnIdentifier(name, true));
+
+    }
+
+    private Iterable<CFDefinition.Name> columnDefinitions(String...name)
+    {
+        List<CFDefinition.Name> list = new ArrayList<>();
+        for (String n : name)
+            list.add(columnDefinition(n));
+        return list;
+    }
+
+    private ColumnSpecification columnSpecification(String name, 
AbstractType<?> type)
+    {
+        return new ColumnSpecification(KEYSPACE,
+                                       tableName,
+                                       new ColumnIdentifier(name, true),
+                                       type);
+    }
+
+    private void createTable(String query) throws Throwable
+    {
+        executeSchemaChange(String.format(query, KEYSPACE + "." + tableName));
+    }
+
+    private static void executeSchemaChange(String query) throws Throwable
+    {
+        try
+        {
+            process(query, ConsistencyLevel.ONE);
+        }
+        catch (RuntimeException exc)
+        {
+            throw exc.getCause();
+        }
+    }
+}

Reply via email to