From 439c913bec3e033931bfac593ccfb94afbd3045c Mon Sep 17 00:00:00 2001
From: Benjamin Gaidioz <ben@raw-labs.com>
Date: Fri, 20 Sep 2024 11:16:01 +0200
Subject: [PATCH] RD-14922: Added HSTORE support to SqlCompilerService (#518)

---
 .../rawlabs/sql/compiler/SqlTypesUtils.scala  |  3 +-
 .../writers/TypedResultSetJsonWriter.scala    |  9 +++
 sql-compiler/src/test/resources/example.sql   |  2 +
 .../TestSqlCompilerServiceAirports.scala      | 57 +++++++++++++++++++
 4 files changed, 70 insertions(+), 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
index 62ea20486..f47773a87 100644
--- 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
index f49d2f67f..930f59f0d 100644
--- 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
index 6495927dd..73d7f6e8c 100644
--- 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
index 5939d12c2..e181395ae 100644
--- 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"}}]""")
+  }
 }