fixes #537

src/plugin/with-schema/with-schema-transformer.ts

@@ -36,6 +36,7 @@ const ROOT_OPERATION_NODES: Record<RootOperationNode['kind'], true> = freeze({
 export class WithSchemaTransformer extends OperationNodeTransformer {
   readonly #schema: string
   readonly #schemableIds = new Set<string>()
+  readonly #ctes = new Set<string>()
 
   constructor(schema: string) {
     super()
@@ -47,6 +48,12 @@ export class WithSchemaTransformer extends OperationNodeTransformer {
       return super.transformNodeImpl(node)
     }
 
+    const ctes = this.#collectCTEs(node)
+
+    for (const cte of ctes) {
+      this.#ctes.add(cte)
+    }
+
     const tables = this.#collectSchemableIds(node)
 
     for (const table of tables) {
@@ -59,6 +66,10 @@ export class WithSchemaTransformer extends OperationNodeTransformer {
       this.#schemableIds.delete(table)
     }
 
+    for (const cte of ctes) {
+      this.#ctes.delete(cte)
+    }
+
     return transformed
   }
 
@@ -124,11 +135,17 @@ export class WithSchemaTransformer extends OperationNodeTransformer {
       }
     }
 
+    return schemableIds
+  }
+
+  #collectCTEs(node: RootOperationNode): Set<string> {
+    const ctes = new Set<string>()
+
     if ('with' in node && node.with) {
-      this.#removeCommonTableExpressionTables(node.with, schemableIds)
+      this.#collectCTEIds(node.with, ctes)
     }
 
-    return schemableIds
+    return ctes
   }
 
   #collectSchemableIdsFromTableExpr(
@@ -150,17 +167,20 @@ export class WithSchemaTransformer extends OperationNodeTransformer {
     node: SchemableIdentifierNode,
     schemableIds: Set<string>
   ): void {
-    if (!this.#schemableIds.has(node.identifier.name)) {
-      schemableIds.add(node.identifier.name)
+    const id = node.identifier.name
+
+    if (!this.#schemableIds.has(id) && !this.#ctes.has(id)) {
+      schemableIds.add(id)
     }
   }
 
-  #removeCommonTableExpressionTables(
-    node: WithNode,
-    schemableIds: Set<string>
-  ) {
+  #collectCTEIds(node: WithNode, ctes: Set<string>): void {
     for (const expr of node.expressions) {
-      schemableIds.delete(expr.name.table.table.identifier.name)
+      const cteId = expr.name.table.table.identifier.name
+
+      if (!this.#ctes.has(cteId)) {
+        ctes.add(cteId)
+      }
     }
   }
 }

test/node/src/with-schema.test.ts

@@ -282,6 +282,31 @@ if (DIALECTS.includes('postgres')) {
 
         await query.execute()
       })
+
+      it('should not add schema for common table expression names in subqueries', async () => {
+        const query = ctx.db
+          .withSchema('mammals')
+          .with('doggo', (qb) =>
+            qb.selectFrom('pet').where('name', '=', 'Doggo').select('pet.id')
+          )
+          .selectFrom('pet')
+          .select((eb) => [
+            'pet.id',
+            eb.selectFrom('doggo').select('id').as('doggo_id'),
+          ])
+          .selectAll()
+
+        testSql(query, dialect, {
+          postgres: {
+            sql: 'with "doggo" as (select "mammals"."pet"."id" from "mammals"."pet" where "name" = $1) select "mammals"."pet"."id", (select "id" from "doggo") as "doggo_id", * from "mammals"."pet"',
+            parameters: ['Doggo'],
+          },
+          mysql: NOT_SUPPORTED,
+          sqlite: NOT_SUPPORTED,
+        })
+
+        await query.execute()
+      })
     })
 
     describe('create table', () => {