diff --git a/src/main/java/com/rockset/jdbc/RocksetResultSet.java b/src/main/java/com/rockset/jdbc/RocksetResultSet.java index a19f5367..98afe604 100644 --- a/src/main/java/com/rockset/jdbc/RocksetResultSet.java +++ b/src/main/java/com/rockset/jdbc/RocksetResultSet.java @@ -44,15 +44,21 @@ import java.time.temporal.ChronoField; import java.util.ArrayList; import java.util.Calendar; +import java.util.Collection; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; import java.util.TimeZone; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; + +import org.jetbrains.annotations.Nullable; import org.joda.time.DateTimeZone; import org.joda.time.LocalTime; import org.joda.time.format.DateTimeFormat; @@ -1567,55 +1573,49 @@ private static void log(String msg) { private static List getColumns(QueryResponse response) throws SQLException { List out = new ArrayList(); - ObjectMapper mapper = new ObjectMapper(); + try { - if (response.getResults().size() > 0) { - Set fieldNames = new HashSet<>(); - // Loop through all the rows to get the fields and (their first - // non-null) types. + if (!response.getResults().isEmpty()) { + List orderedFields = Collections.emptyList(); + // column fields will be null if a wildcard is used in the query + // otherwise, we want to return results in the order set by the projection + if(response.getColumnFields() != null) { + orderedFields = response.getColumnFields() + .stream() + .map(QueryFieldType::getName) + .collect(Collectors.toList()); + } + // Loop through all the rows to get the fields and types + Map fieldTypes = new HashMap<>(); for (int i = 0; i < response.getResults().size(); ++i) { log("Extracting column information from record " + i + " in resultset"); Object onedoc = response.getResults().get(i); - JsonNode docRootNode = mapper.readTree(mapper.writeValueAsString(onedoc)); - - Iterator> fields = docRootNode.fields(); - while (fields.hasNext()) { - Map.Entry field = fields.next(); - String fieldName = field.getKey(); - if (fieldNames.contains(fieldName)) { - // This fieldname was already found to have a non-null type - // in a previous record. - continue; - } - JsonNode value = field.getValue(); - Column.ColumnTypes type = Column.ColumnTypes.fromValue(value.getNodeType().toString()); - // Skip over the fields with null type unless all values for that - // field are null - if (type.equals(Column.ColumnTypes.NULL) && i != response.getResults().size() - 1) { - continue; - } - if (type.equals(Column.ColumnTypes.STRING)) { - java.time.format.DateTimeFormatter format = TIMESTAMP_PARSE_FORMAT; - try { - LocalDateTime.parse(value.asText(), format); - type = Column.ColumnTypes.TIMESTAMP; - } catch (DateTimeParseException e) { - // ignore - } + JsonNode docRootNode = OBJECT_MAPPER.readTree(OBJECT_MAPPER.writeValueAsString(onedoc)); + + for (Iterator> it = docRootNode.fields(); it.hasNext(); ) { + Map.Entry field = it.next(); + String fieldName = field.getKey(); + if (fieldTypes.containsKey(fieldName) && !fieldTypes.get(fieldName).equals(Column.ColumnTypes.NULL)) { + // This fieldname was already found to have a non-null type + // in a previous record. + continue; + } + JsonNode value = field.getValue(); + Column.ColumnTypes type = valueToColumnType(value); + log("getColumns::column name " + fieldName + " type: " + type.toString()); + fieldTypes.put(fieldName, type); } - if (type.equals(Column.ColumnTypes.OBJECT)) { - if (value.get("__rockset_type") != null) { - type = Column.ColumnTypes.fromValue(value.get("__rockset_type").asText()); - } - } - log("getColumns::column name " + fieldName + " type: " + type.toString()); - Column c = new Column(fieldName, type); - out.add(c); - fieldNames.add(fieldName); - } } - } else if (response.getColumnFields() != null && response.getColumnFields().size() > 0) { + + // If we know the desired field ordering through explicit projection, use that ordering. + // Otherwise, just iterate over in arbitrary ordering + Collection fields = orderedFields.isEmpty() ? fieldTypes.keySet() : orderedFields; + for(String field : fields){ + out.add(new Column(field, fieldTypes.get(field))); + } + + } else if (response.getColumnFields() != null && !response.getColumnFields().isEmpty()) { // If this is not a select star query, and has returned 0 rows. // Extrapolate the fields from query response's getColumnFields log("Extracting column information from explicit fields"); @@ -1637,6 +1637,26 @@ private static List getColumns(QueryResponse response) throws SQLExcepti } } + @Nullable + private static Column.ColumnTypes valueToColumnType(JsonNode value) { + Column.ColumnTypes type = Column.ColumnTypes.fromValue(value.getNodeType().toString()); + if (type.equals(Column.ColumnTypes.STRING)) { + java.time.format.DateTimeFormatter format = TIMESTAMP_PARSE_FORMAT; + try { + LocalDateTime.parse(value.asText(), format); + type = Column.ColumnTypes.TIMESTAMP; + } catch (DateTimeParseException e) { + // ignore + } + } + if (type.equals(Column.ColumnTypes.OBJECT)) { + if (value.get("__rockset_type") != null) { + type = Column.ColumnTypes.fromValue(value.get("__rockset_type").asText()); + } + } + return type; + } + private static Map getFieldMap(List columns) { Map map = new HashMap<>(); for (int i = 0; i < columns.size(); i++) { diff --git a/src/test/java/com/rockset/client/TestWorkspace.java b/src/test/java/com/rockset/client/TestWorkspace.java index 469ee75d..6da8587b 100644 --- a/src/test/java/com/rockset/client/TestWorkspace.java +++ b/src/test/java/com/rockset/client/TestWorkspace.java @@ -60,7 +60,8 @@ public void testDeleteWorkspace() throws Exception { // wait for collection to go away Awaitility.await("Waiting for collection to be cleaned up ") - .atMost(60, TimeUnit.SECONDS) + .atMost(3, TimeUnit.MINUTES) + .pollInterval(1, TimeUnit.SECONDS) .until( (Callable) () -> { @@ -70,7 +71,6 @@ public void testDeleteWorkspace() throws Exception { } catch (Exception e) { return true; // collection deleted } - Thread.sleep(1000); return false; }); diff --git a/src/test/java/com/rockset/jdbc/TestTable.java b/src/test/java/com/rockset/jdbc/TestTable.java index 8d5a367f..fc9f3058 100644 --- a/src/test/java/com/rockset/jdbc/TestTable.java +++ b/src/test/java/com/rockset/jdbc/TestTable.java @@ -6,6 +6,8 @@ import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.Collections2; +import com.google.common.collect.Lists; import com.rockset.client.RocksetClient; import com.rockset.client.model.Collection; import com.rockset.client.model.CreateCollectionRequest; @@ -20,6 +22,7 @@ import java.sql.DatabaseMetaData; import java.sql.Date; import java.sql.DriverManager; +import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; @@ -27,8 +30,15 @@ import java.sql.Time; import java.sql.Types; import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; import java.util.List; +import java.util.Map; import java.util.Properties; +import java.util.Set; +import java.util.stream.Collectors; + import okhttp3.MediaType; import okhttp3.MultipartBody; import okhttp3.OkHttpClient; @@ -38,6 +48,7 @@ import org.apache.commons.lang3.RandomStringUtils; import org.testng.Assert; import org.testng.annotations.BeforeSuite; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; // @@ -360,6 +371,68 @@ public void testGetColumnsAllTypes() throws Exception { } } + @Test(dataProvider = "columnOrderings") + public void testColumnOrdering(Map projectionTypes, String queryFormatString, String filePath) throws Exception { + List collections = generateCollectionNames(/*numCollections*/ 1); + Connection conn = null; + try { + createCollections(collections); + waitCollections(collections); + + String collection = collections.get(0); + uploadFile(collection, filePath, null); + waitNumberDocs(collection, 2); + + conn = DriverManager.getConnection(DB_URL, property); + + + + // When there is a wildcard in the projection, no ordering is guaranteed + String query = String.format("select * EXCEPT(_id, _meta, _event_time) from %s", collection); + try (PreparedStatement stmt = conn.prepareStatement(query); + ResultSet rs = stmt.executeQuery()) { + + ResultSetMetaData rsmd = rs.getMetaData(); + int cc = rsmd.getColumnCount(); + while (rs.next()) { + Set colNames = new HashSet<>(); + for (int i = 1; i <= cc; i++) { + colNames.add(rsmd.getColumnName(i)); + } + Assert.assertEquals(colNames, projectionTypes.keySet()); + } + } + + java.util.Collection> projectionOrderings = Collections2.permutations(new ArrayList<>(projectionTypes.keySet())); + + for(List projections : projectionOrderings) { + List expectedTypes = projections.stream().map(projectionTypes::get).collect(Collectors.toList()); + List queryStringParams = new ArrayList<>(projections); + queryStringParams.add(collection); + query = String.format(queryFormatString, queryStringParams.toArray()); + try (PreparedStatement stmt = conn.prepareStatement(query); + ResultSet rs = stmt.executeQuery()) { + + ResultSetMetaData rsmd = rs.getMetaData(); + int cc = rsmd.getColumnCount(); + List colNames = new ArrayList<>(); + List colTypes = new ArrayList<>(); + for (int i = 1; i <= cc; i++) { + colNames.add(rsmd.getColumnName(i)); + colTypes.add(rsmd.getColumnTypeName(i)); + } + Assert.assertEquals(colNames, Lists.newArrayList(projections)); + Assert.assertEquals(colTypes, expectedTypes); + } + } + + + } finally { + cleanup(collections, null, conn); + } + } + + private void assertNextEquals(ResultSet rs, String expectedColumnName, int expectedType) throws SQLException { int columnNameIndex = rs.findColumn("COLUMN_NAME"); @@ -527,4 +600,39 @@ private void waitNumberDocs(String collectionName, int expectedDocs) throws Exce Thread.sleep(1000); } } + + @DataProvider(name="columnOrderings") + Object[][] columnOrderings(){ + + Map nullProjectionTypes = new HashMap<>(); + nullProjectionTypes.put("a", "null"); + nullProjectionTypes.put("b", "null"); + + String nullProjectionQueryFormat = "select %s, %s from %s "; + String nullProjectionFile = "src/test/resources/all_nulls.json"; + + // id is non null in all records + // name is null in one of the records + // age is null in all the records + // mixed is a string or array, however we query such that string is the first non-null type in the result set + // the type in the column will be string + // When there is a single non-null field, the type should be of the first non-null field + Map mixedProjectionTypes = new HashMap<>(); + mixedProjectionTypes.put("age", "null"); + mixedProjectionTypes.put("id", "bigint"); + mixedProjectionTypes.put("name", "varchar"); + mixedProjectionTypes.put("mixed", "varchar"); + + String mixedProjectionQueryFormat = "select %s, %s, %s, %s from %s ORDER BY id ASC"; + String mixedProjectionFile = "src/test/resources/mixed_nulls.json"; + + + + return new Object[][]{{ + nullProjectionTypes, nullProjectionQueryFormat, nullProjectionFile + }, { + mixedProjectionTypes, mixedProjectionQueryFormat, mixedProjectionFile + }}; + + } } diff --git a/src/test/resources/all_nulls.json b/src/test/resources/all_nulls.json new file mode 100644 index 00000000..6f021157 --- /dev/null +++ b/src/test/resources/all_nulls.json @@ -0,0 +1,2 @@ +{"a": null, "b": null} +{"a": null, "b": null} diff --git a/src/test/resources/mixed_nulls.json b/src/test/resources/mixed_nulls.json new file mode 100644 index 00000000..d14bb6fb --- /dev/null +++ b/src/test/resources/mixed_nulls.json @@ -0,0 +1,2 @@ +{"id": 1, "name": null, "age": null, "mixed": "hello" } +{"id": 2, "name": "bob", "age": null, "mixed": ["a", "b"] }