From 0d521ee28620daeae6e36efb39890fe7dc5240f1 Mon Sep 17 00:00:00 2001
From: Marius Volkhart <marius@volkhart.com>
Date: Fri, 1 Sep 2023 00:00:39 -0400
Subject: [PATCH] Add Postgres JSON function definitions

It's not unreasonable for someone to use these functions in a SELECT, so we should have return types for them.

This does not cover any of the functions that return "setof", as those would essentially be a virtual table.
---
 .../postgresql/PostgreSqlTypeResolver.kt      | 22 ++++++++++++++++---
 .../sqldelight/postgresql/integration/Json.sq | 13 +++++++++++
 2 files changed, 32 insertions(+), 3 deletions(-)
 create mode 100644 sqldelight-gradle-plugin/src/test/integration-postgresql/src/main/sqldelight/app/cash/sqldelight/postgresql/integration/Json.sq

diff --git a/dialects/postgresql/src/main/kotlin/app/cash/sqldelight/dialects/postgresql/PostgreSqlTypeResolver.kt b/dialects/postgresql/src/main/kotlin/app/cash/sqldelight/dialects/postgresql/PostgreSqlTypeResolver.kt
index ae3595974f1..6bfcda1fb9e 100644
--- a/dialects/postgresql/src/main/kotlin/app/cash/sqldelight/dialects/postgresql/PostgreSqlTypeResolver.kt
+++ b/dialects/postgresql/src/main/kotlin/app/cash/sqldelight/dialects/postgresql/PostgreSqlTypeResolver.kt
@@ -2,8 +2,8 @@ package app.cash.sqldelight.dialects.postgresql
 
 import app.cash.sqldelight.dialect.api.DialectType
 import app.cash.sqldelight.dialect.api.IntermediateType
-import app.cash.sqldelight.dialect.api.PrimitiveType
 import app.cash.sqldelight.dialect.api.PrimitiveType.BLOB
+import app.cash.sqldelight.dialect.api.PrimitiveType.BOOLEAN
 import app.cash.sqldelight.dialect.api.PrimitiveType.INTEGER
 import app.cash.sqldelight.dialect.api.PrimitiveType.REAL
 import app.cash.sqldelight.dialect.api.PrimitiveType.TEXT
@@ -64,7 +64,7 @@ class PostgreSqlTypeResolver(private val parentResolver: TypeResolver) : TypeRes
           }
         }
         jsonDataType != null -> TEXT
-        booleanDataType != null -> PrimitiveType.BOOLEAN
+        booleanDataType != null -> BOOLEAN
         blobDataType != null -> BLOB
         else -> throw IllegalArgumentException("Unknown kotlin type for sql type ${this.text}")
       },
@@ -123,6 +123,22 @@ class PostgreSqlTypeResolver(private val parentResolver: TypeResolver) : TypeRes
     "regr_count" -> IntermediateType(BIG_INT).asNullable()
     "gen_random_uuid" -> IntermediateType(PostgreSqlType.UUID)
     "length", "character_length", "char_length" -> IntermediateType(PostgreSqlType.INTEGER).nullableIf(resolvedType(exprList[0]).javaType.isNullable)
+    "to_json", "to_jsonb",
+    "array_to_json", "row_to_json",
+    "json_build_array", "jsonb_build_array",
+    "json_build_object", "jsonb_build_object",
+    "json_object", "jsonb_object",
+    "json_extract_path", "jsonb_extract_path",
+    "json_extract_path_text", "jsonb_extract_path_text",
+    "jsonb_set", "jsonb_set_lax", "jsonb_insert",
+    "json_strip_nulls", "jsonb_strip_nulls",
+    "jsonb_path_query_array", "jsonb_path_query_first", "jsonb_path_query_array_tz", "jsonb_path_query_first_tz",
+    "jsonb_pretty",
+    "json_typeof", "jsonb_typeof",
+    "json_agg", "jsonb_agg", "json_object_agg", "jsonb_object_agg",
+    -> IntermediateType(TEXT)
+    "json_array_length", "jsonb_array_length" -> IntermediateType(INTEGER)
+    "jsonb_path_exists", "jsonb_path_match", "jsonb_path_exists_tz", "jsonb_path_match_tz" -> IntermediateType(BOOLEAN)
     else -> null
   }
 
@@ -166,7 +182,7 @@ class PostgreSqlTypeResolver(private val parentResolver: TypeResolver) : TypeRes
   private fun SqlExpr.postgreSqlType(): IntermediateType = when (this) {
     is SqlBinaryExpr -> {
       if (node.findChildByType(binaryExprChildTypesResolvingToBool) != null) {
-        IntermediateType(PrimitiveType.BOOLEAN)
+        IntermediateType(BOOLEAN)
       } else {
         encapsulatingType(
           exprList = getExprList(),
diff --git a/sqldelight-gradle-plugin/src/test/integration-postgresql/src/main/sqldelight/app/cash/sqldelight/postgresql/integration/Json.sq b/sqldelight-gradle-plugin/src/test/integration-postgresql/src/main/sqldelight/app/cash/sqldelight/postgresql/integration/Json.sq
new file mode 100644
index 00000000000..801bced9291
--- /dev/null
+++ b/sqldelight-gradle-plugin/src/test/integration-postgresql/src/main/sqldelight/app/cash/sqldelight/postgresql/integration/Json.sq
@@ -0,0 +1,13 @@
+CREATE TABLE myJson(
+  data JSON NOT NULL,
+  datab JSONB NOT NULL
+);
+
+insert:
+INSERT INTO myJson(data, datab) VALUES(
+  json_build_object(:key, :value),
+  jsonb_build_object('key', 'value')
+);
+
+buildJson:
+SELECT json_build_object('key', 'value'), jsonb_build_object('key', 'value');
\ No newline at end of file