RD-14922: Added HSTORE support to SqlCompilerService (#518)

raw-labs · Sep 20, 2024 · 439c913 · 439c913
1 parent 4251a72
commit 439c913
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 1 deletion.
diff --git a/sql-compiler/src/main/scala/com/rawlabs/sql/compiler/SqlTypesUtils.scala b/sql-compiler/src/main/scala/com/rawlabs/sql/compiler/SqlTypesUtils.scala
@@ -138,7 +138,8 @@ object SqlTypesUtils {
         val otherTypeMap: Map[String, RawType] = Map(
           "interval" -> RawIntervalType(false, false),
           "json" -> RawAnyType(),
-          "jsonb" -> RawAnyType()
+          "jsonb" -> RawAnyType(),
+          "hstore" -> RawAnyType()
         )
         otherTypeMap.get(pgTypeName) match {
           case Some(t) => Right(t)

diff --git a/sql-compiler/src/main/scala/com/rawlabs/sql/compiler/writers/TypedResultSetJsonWriter.scala b/sql-compiler/src/main/scala/com/rawlabs/sql/compiler/writers/TypedResultSetJsonWriter.scala
@@ -13,6 +13,7 @@
 package com.rawlabs.sql.compiler.writers
 
 import com.fasterxml.jackson.core.{JsonEncoding, JsonFactory, JsonParser}
+import com.fasterxml.jackson.databind.node.ObjectNode
 import com.fasterxml.jackson.databind.{JsonNode, ObjectMapper}
 import com.rawlabs.compiler.{
   RawAnyType,
@@ -129,6 +130,14 @@ class TypedResultSetJsonWriter(os: OutputStream, maxRows: Option[Long]) {
               val json = mapper.readTree(data)
               writeRawJson(json)
             }
+          case "hstore" =>
+            val hstoreMap = v.getObject(i).asInstanceOf[java.util.Map[String, String]]
+            if (v.wasNull()) gen.writeNull()
+            else {
+              // Convert hstore to JSON-like structure
+              val json = mapper.valueToTree[ObjectNode](hstoreMap)
+              writeRawJson(json)
+            }
           case _ => throw new IOException("unsupported type")
         }
       case _: RawDateType =>

diff --git a/sql-compiler/src/test/resources/example.sql b/sql-compiler/src/test/resources/example.sql
@@ -1,5 +1,7 @@
 CREATE SCHEMA example;
 
+CREATE EXTENSION hstore;
+
 CREATE TABLE example.airports (
     airport_id integer,
     name character varying(256),

diff --git a/sql-compiler/src/test/scala/com/rawlabs/sql/compiler/TestSqlCompilerServiceAirports.scala b/sql-compiler/src/test/scala/com/rawlabs/sql/compiler/TestSqlCompilerServiceAirports.scala
@@ -1153,4 +1153,61 @@ class TestSqlCompilerServiceAirports
     }
   }
 
+  test("SELECT 'a=>1,b=>tralala'::hstore AS r -- JSON") { t =>
+    val baos = new ByteArrayOutputStream()
+    assert(
+      compilerService.execute(
+        t.q,
+        asJson(),
+        None,
+        baos
+      ) == ExecutionSuccess(true)
+    )
+    assert(baos.toString() === """[{"r":{"a":"1","b":"tralala"}}]""")
+  }
+
+  test("SELECT NULL::hstore AS r -- JSON") { t =>
+    val baos = new ByteArrayOutputStream()
+    assert(
+      compilerService.execute(
+        t.q,
+        asJson(),
+        None,
+        baos
+      ) == ExecutionSuccess(true)
+    )
+    assert(baos.toString() === """[{"r":null}]""")
+  }
+
+  // TODO What do we do in CSV?
+  ignore("SELECT 'a=>1,b=>tralala'::hstore AS r -- CSV") { t =>
+    val baos = new ByteArrayOutputStream()
+    assert(
+      compilerService.execute(
+        t.q,
+        asCsv(),
+        None,
+        baos
+      ) == ExecutionSuccess(true)
+    )
+    assert(
+      baos.toString() ===
+        """r
+          |{"a":"1","b":"tralala"}
+          |""".stripMargin
+    )
+  }
+
+  ignore("SELECT NULL::hstore AS r -- CSV") { t =>
+    val baos = new ByteArrayOutputStream()
+    assert(
+      compilerService.execute(
+        t.q,
+        asCsv(),
+        None,
+        baos
+      ) == ExecutionSuccess(true)
+    )
+    assert(baos.toString() === """[{"r":{"a":"1","b":"tralala"}}]""")
+  }
 }