Fix result when complex type contains user defined type in Cassandra

trinodb · Jan 20, 2023 · 6acfc5f · 6acfc5f
1 parent d8b4817
commit 6acfc5f
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 1 deletion.
diff --git a/plugin/trino-cassandra/src/main/java/io/trino/plugin/cassandra/CassandraTypeManager.java b/plugin/trino-cassandra/src/main/java/io/trino/plugin/cassandra/CassandraTypeManager.java
@@ -472,8 +472,9 @@ private String objectToJson(Object cassandraValue, DataType dataType)
             case INET:
             case VARINT:
             case TUPLE:
-            case UDT:
                 return quoteStringLiteralForJson(cassandraValue.toString());
+            case UDT:
+                return quoteStringLiteralForJson(((UdtValue) cassandraValue).getFormattedContents());
 
             case BLOB:
             case CUSTOM:

diff --git a/...n/trino-cassandra/src/test/java/io/trino/plugin/cassandra/TestCassandraConnectorTest.java b/...n/trino-cassandra/src/test/java/io/trino/plugin/cassandra/TestCassandraConnectorTest.java
@@ -63,6 +63,7 @@
 import static io.trino.testing.MaterializedResult.resultBuilder;
 import static io.trino.testing.QueryAssertions.assertContains;
 import static io.trino.testing.QueryAssertions.assertContainsEventually;
+import static io.trino.testing.TestingNames.randomNameSuffix;
 import static io.trino.type.IpAddressType.IPADDRESS;
 import static java.lang.String.format;
 import static java.util.Comparator.comparing;
@@ -993,6 +994,54 @@ public void testColumnNameAmbiguity()
         session.execute("DROP KEYSPACE keyspace_5");
     }
 
+    @Test
+    public void testUserDefinedTypeInArray()
+    {
+        String tableName = "test_udt_in_array" + randomNameSuffix();
+        String userDefinedTypeName = "test_udt" + randomNameSuffix();
+
+        session.execute("CREATE TYPE tpch." + userDefinedTypeName + "(udt_field bigint)");
+        session.execute("CREATE TABLE tpch." + tableName + "(id bigint, col list<frozen<tpch." + userDefinedTypeName + ">>, primary key (id))");
+        session.execute("INSERT INTO tpch." + tableName + "(id, col) values (1, [{udt_field: 10}])");
+
+        assertQuery("SELECT * FROM " + tableName, "VALUES (1, '[\"{udt_field:10}\"]')");
+
+        session.execute("DROP TABLE tpch." + tableName);
+        session.execute("DROP TYPE tpch." + userDefinedTypeName);
+    }
+
+    @Test
+    public void testUserDefinedTypeInMap()
+    {
+        String tableName = "test_udt_in_map" + randomNameSuffix();
+        String userDefinedTypeName = "test_udt" + randomNameSuffix();
+
+        session.execute("CREATE TYPE tpch." + userDefinedTypeName + "(udt_field bigint)");
+        session.execute("CREATE TABLE tpch." + tableName + "(id bigint, col map<frozen<tpch." + userDefinedTypeName + ">, frozen<tpch." + userDefinedTypeName + ">>, primary key (id))");
+        session.execute("INSERT INTO tpch." + tableName + "(id, col) values (1, {{udt_field: 10}: {udt_field: -10}})");
+
+        assertQuery("SELECT * FROM " + tableName, "VALUES (1, '{\"{udt_field:10}\":\"{udt_field:-10}\"}')");
+
+        session.execute("DROP TABLE tpch." + tableName);
+        session.execute("DROP TYPE tpch." + userDefinedTypeName);
+    }
+
+    @Test
+    public void testUserDefinedTypeInSet()
+    {
+        String tableName = "test_udt_in_set" + randomNameSuffix();
+        String userDefinedTypeName = "test_udt" + randomNameSuffix();
+
+        session.execute("CREATE TYPE tpch." + userDefinedTypeName + "(udt_field bigint)");
+        session.execute("CREATE TABLE tpch." + tableName + "(id bigint, col set<frozen<tpch." + userDefinedTypeName + ">>, primary key (id))");
+        session.execute("INSERT INTO tpch." + tableName + "(id, col) values (1, {{udt_field: 10}})");
+
+        assertQuery("SELECT * FROM " + tableName, "VALUES (1, '[\"{udt_field:10}\"]')");
+
+        session.execute("DROP TABLE tpch." + tableName);
+        session.execute("DROP TYPE tpch." + userDefinedTypeName);
+    }
+
     @Test
     public void testUnsupportedColumnType()
     {