From 3c012305d145c3c6da430748f2c3f05e6f980c62 Mon Sep 17 00:00:00 2001
From: Gian Merlino <gianmerlino@gmail.com>
Date: Sat, 7 Jan 2017 12:33:53 -0800
Subject: [PATCH] SqlResource: Fix incorrect labeling of aliased columns.
 (#3829)

---
 .../java/io/druid/sql/http/SqlResource.java   |  2 +-
 .../sql/avatica/DruidAvaticaHandlerTest.java  | 19 +++++++++--
 .../druid/sql/calcite/CalciteQueryTest.java   | 25 +++++++++++++++
 .../sql/calcite/http/SqlResourceTest.java     | 32 +++++++++++++++++++
 4 files changed, 75 insertions(+), 3 deletions(-)

diff --git a/sql/src/main/java/io/druid/sql/http/SqlResource.java b/sql/src/main/java/io/druid/sql/http/SqlResource.java
index 527d2e09686..cb6105c54f5 100644
--- a/sql/src/main/java/io/druid/sql/http/SqlResource.java
+++ b/sql/src/main/java/io/druid/sql/http/SqlResource.java
@@ -106,7 +106,7 @@ public class SqlResource
                       value = resultSet.getObject(i + 1);
                     }
 
-                    jsonGenerator.writeObjectField(metaData.getColumnName(i + 1), value);
+                    jsonGenerator.writeObjectField(metaData.getColumnLabel(i + 1), value);
                   }
                   jsonGenerator.writeEndObject();
                 }
diff --git a/sql/src/test/java/io/druid/sql/avatica/DruidAvaticaHandlerTest.java b/sql/src/test/java/io/druid/sql/avatica/DruidAvaticaHandlerTest.java
index 2784450fbed..246026bd311 100644
--- a/sql/src/test/java/io/druid/sql/avatica/DruidAvaticaHandlerTest.java
+++ b/sql/src/test/java/io/druid/sql/avatica/DruidAvaticaHandlerTest.java
@@ -124,6 +124,21 @@ public class DruidAvaticaHandlerTest
     );
   }
 
+  @Test
+  public void testFieldAliasingSelect() throws Exception
+  {
+    final ResultSet resultSet = client.createStatement().executeQuery(
+        "SELECT dim2 AS \"x\", dim2 AS \"y\" FROM druid.foo LIMIT 1"
+    );
+    final List<Map<String, Object>> rows = getRows(resultSet);
+    Assert.assertEquals(
+        ImmutableList.of(
+            ImmutableMap.of("x", "a", "y", "a")
+        ),
+        rows
+    );
+  }
+
   @Test
   public void testExplainSelectCount() throws Exception
   {
@@ -244,8 +259,8 @@ public class DruidAvaticaHandlerTest
       while (resultSet.next()) {
         final Map<String, Object> row = Maps.newHashMap();
         for (int i = 0; i < metaData.getColumnCount(); i++) {
-          if (returnKeys == null || returnKeys.contains(metaData.getColumnName(i + 1))) {
-            row.put(metaData.getColumnName(i + 1), resultSet.getObject(i + 1));
+          if (returnKeys == null || returnKeys.contains(metaData.getColumnLabel(i + 1))) {
+            row.put(metaData.getColumnLabel(i + 1), resultSet.getObject(i + 1));
           }
         }
         rows.add(row);
diff --git a/sql/src/test/java/io/druid/sql/calcite/CalciteQueryTest.java b/sql/src/test/java/io/druid/sql/calcite/CalciteQueryTest.java
index 2f65f1535b2..dc8d9696f38 100644
--- a/sql/src/test/java/io/druid/sql/calcite/CalciteQueryTest.java
+++ b/sql/src/test/java/io/druid/sql/calcite/CalciteQueryTest.java
@@ -328,6 +328,31 @@ public class CalciteQueryTest
     );
   }
 
+  @Test
+  public void testSelectSingleColumnTwice() throws Exception
+  {
+    testQuery(
+        "SELECT dim2 x, dim2 y FROM druid.foo LIMIT 2",
+        ImmutableList.<Query>of(
+            Druids.newSelectQueryBuilder()
+                  .dataSource(CalciteTests.DATASOURCE)
+                  .intervals(QSS(Filtration.eternity()))
+                  .dimensionSpecs(DIMS(
+                      new DefaultDimensionSpec("dim2", "d1"),
+                      new DefaultDimensionSpec("dim2", "d2")
+                  ))
+                  .granularity(QueryGranularities.ALL)
+                  .descending(false)
+                  .pagingSpec(FIRST_PAGING_SPEC)
+                  .build()
+        ),
+        ImmutableList.of(
+            new Object[]{"a", "a"},
+            new Object[]{"", ""}
+        )
+    );
+  }
+
   @Test
   public void testSelectSingleColumnWithLimitDescending() throws Exception
   {
diff --git a/sql/src/test/java/io/druid/sql/calcite/http/SqlResourceTest.java b/sql/src/test/java/io/druid/sql/calcite/http/SqlResourceTest.java
index 625631b99c2..12a96394596 100644
--- a/sql/src/test/java/io/druid/sql/calcite/http/SqlResourceTest.java
+++ b/sql/src/test/java/io/druid/sql/calcite/http/SqlResourceTest.java
@@ -109,6 +109,38 @@ public class SqlResourceTest
     );
   }
 
+  @Test
+  public void testFieldAliasingSelect() throws Exception
+  {
+    final List<Map<String, Object>> rows = doPost(
+        new SqlQuery("SELECT dim2 \"x\", dim2 \"y\" FROM druid.foo LIMIT 1")
+    );
+
+    Assert.assertEquals(
+        ImmutableList.of(
+            ImmutableMap.of("x", "a", "y", "a")
+        ),
+        rows
+    );
+  }
+
+  @Test
+  public void testFieldAliasingGroupBy() throws Exception
+  {
+    final List<Map<String, Object>> rows = doPost(
+        new SqlQuery("SELECT dim2 \"x\", dim2 \"y\" FROM druid.foo GROUP BY dim2")
+    );
+
+    Assert.assertEquals(
+        ImmutableList.of(
+            ImmutableMap.of("x", "", "y", ""),
+            ImmutableMap.of("x", "a", "y", "a"),
+            ImmutableMap.of("x", "abc", "y", "abc")
+        ),
+        rows
+    );
+  }
+
   @Test
   public void testExplainCountStar() throws Exception
   {