From 95f1b02b730a17a5fedffa578197e5b308a9783c Mon Sep 17 00:00:00 2001
From: Igal Klebanov <igalklebanov@gmail.com>
Date: Sat, 11 Jan 2025 19:43:59 +0200
Subject: [PATCH] fix: `WithSchemaPlugin` is adding schema to table function
 arguments and causes db errors in `json_agg` and `to_json`. (#827)

---
 .../with-schema/with-schema-transformer.ts    | 41 ++++++++++++++++
 test/node/src/with-schema.test.ts             | 47 +++++++++++++++++++
 2 files changed, 88 insertions(+)

diff --git a/src/plugin/with-schema/with-schema-transformer.ts b/src/plugin/with-schema/with-schema-transformer.ts
index b326628b4..04c241d28 100644
--- a/src/plugin/with-schema/with-schema-transformer.ts
+++ b/src/plugin/with-schema/with-schema-transformer.ts
@@ -1,4 +1,6 @@
+import { AggregateFunctionNode } from '../../operation-node/aggregate-function-node.js'
 import { AliasNode } from '../../operation-node/alias-node.js'
+import { FunctionNode } from '../../operation-node/function-node.js'
 import { IdentifierNode } from '../../operation-node/identifier-node.js'
 import { OperationNodeTransformer } from '../../operation-node/operation-node-transformer.js'
 import { OperationNode } from '../../operation-node/operation-node.js'
@@ -34,6 +36,11 @@ const ROOT_OPERATION_NODES: Record<RootOperationNode['kind'], true> = freeze({
   MergeQueryNode: true,
 })
 
+const SCHEMALESS_FUNCTIONS: Record<string, true> = {
+  json_agg: true,
+  to_json: true,
+}
+
 export class WithSchemaTransformer extends OperationNodeTransformer {
   readonly #schema: string
   readonly #schemableIds = new Set<string>()
@@ -105,6 +112,40 @@ export class WithSchemaTransformer extends OperationNodeTransformer {
     }
   }
 
+  protected override transformAggregateFunction(
+    node: AggregateFunctionNode,
+  ): AggregateFunctionNode {
+    return {
+      ...super.transformAggregateFunction({ ...node, aggregated: [] }),
+      aggregated: this.#transformTableArgsWithoutSchemas(node, 'aggregated'),
+    }
+  }
+
+  protected override transformFunction(node: FunctionNode): FunctionNode {
+    return {
+      ...super.transformFunction({ ...node, arguments: [] }),
+      arguments: this.#transformTableArgsWithoutSchemas(node, 'arguments'),
+    }
+  }
+
+  #transformTableArgsWithoutSchemas<
+    A extends string,
+    N extends { func: string } & {
+      [K in A]: readonly OperationNode[]
+    },
+  >(node: N, argsKey: A): readonly OperationNode[] {
+    return SCHEMALESS_FUNCTIONS[node.func]
+      ? node[argsKey].map((arg) =>
+          !TableNode.is(arg) || arg.table.schema
+            ? this.transformNode(arg)
+            : {
+                ...arg,
+                table: this.transformIdentifier(arg.table.identifier),
+              },
+        )
+      : this.transformNodeList(node[argsKey])
+  }
+
   #isRootOperationNode(node: OperationNode): node is RootOperationNode {
     return node.kind in ROOT_OPERATION_NODES
   }
diff --git a/test/node/src/with-schema.test.ts b/test/node/src/with-schema.test.ts
index ad9556d5c..55e30889e 100644
--- a/test/node/src/with-schema.test.ts
+++ b/test/node/src/with-schema.test.ts
@@ -218,6 +218,53 @@ for (const dialect of DIALECTS.filter(
 
         await query.execute()
       })
+
+      if (dialect === 'postgres') {
+        it('should not add schema for json_agg parameters', async () => {
+          const query = ctx.db
+            .withSchema('mammals')
+            .selectFrom('pet')
+            .select((eb) => [
+              eb.fn.jsonAgg('pet').as('one'),
+              eb.fn.jsonAgg(eb.table('pet')).as('two'),
+              eb.fn.jsonAgg('pet').orderBy('pet.name', 'desc').as('three'),
+            ])
+
+          testSql(query, dialect, {
+            postgres: {
+              sql: 'select json_agg("pet") as "one", json_agg("pet") as "two", json_agg("pet" order by "mammals"."pet"."name" desc) as "three" from "mammals"."pet"',
+              parameters: [],
+            },
+            mysql: NOT_SUPPORTED,
+            mssql: NOT_SUPPORTED,
+            sqlite: NOT_SUPPORTED,
+          })
+
+          await query.execute()
+        })
+
+        it('should not add schema for to_json parameters', async () => {
+          const query = ctx.db
+            .withSchema('mammals')
+            .selectFrom('pet')
+            .select((eb) => [
+              eb.fn.toJson('pet').as('one'),
+              eb.fn.toJson(eb.table('pet')).as('two'),
+            ])
+
+          testSql(query, dialect, {
+            postgres: {
+              sql: 'select to_json("pet") as "one", to_json("pet") as "two" from "mammals"."pet"',
+              parameters: [],
+            },
+            mysql: NOT_SUPPORTED,
+            mssql: NOT_SUPPORTED,
+            sqlite: NOT_SUPPORTED,
+          })
+
+          await query.execute()
+        })
+      }
     })
 
     describe('insert into', () => {