Support materialized CTEs using CTE builder

src/operation-node/common-table-expression-node.ts

@@ -2,9 +2,15 @@ import { freeze } from '../util/object-utils.js'
 import { CommonTableExpressionNameNode } from './common-table-expression-name-node.js'
 import { OperationNode } from './operation-node.js'
 
+type CommonTableExpressionNodeProps = Pick<
+  CommonTableExpressionNode,
+  'materialized'
+>
+
 export interface CommonTableExpressionNode extends OperationNode {
   readonly kind: 'CommonTableExpressionNode'
   readonly name: CommonTableExpressionNameNode
+  readonly materialized?: boolean
   readonly expression: OperationNode
 }
 
@@ -26,4 +32,14 @@ export const CommonTableExpressionNode = freeze({
       expression,
     })
   },
+
+  cloneWith(
+    node: CommonTableExpressionNode,
+    props: CommonTableExpressionNodeProps
+  ): CommonTableExpressionNode {
+    return freeze({
+      ...node,
+      ...props,
+    })
+  },
 })

src/operation-node/operation-node-transformer.ts

@@ -642,6 +642,7 @@ export class OperationNodeTransformer {
     return requireAllProps<CommonTableExpressionNode>({
       kind: 'CommonTableExpressionNode',
       name: this.transformNode(node.name),
+      materialized: node.materialized,
       expression: this.transformNode(node.expression),
     })
   }

src/parser/join-parser.ts

@@ -43,8 +43,7 @@ function parseCallbackJoin(
   from: TableExpression<any, any>,
   callback: JoinCallbackExpression<any, any, any>
 ): JoinNode {
-  const joinBuilder = callback(createJoinBuilder(joinType, from))
-  return joinBuilder.toOperationNode()
+  return callback(createJoinBuilder(joinType, from)).toOperationNode()
 }
 
 function parseSingleOnJoin(

src/parser/with-parser.ts

@@ -1,12 +1,15 @@
 import { UpdateQueryBuilder } from '../query-builder/update-query-builder.js'
 import { DeleteQueryBuilder } from '../query-builder/delete-query-builder.js'
 import { InsertQueryBuilder } from '../query-builder/insert-query-builder.js'
-import { CommonTableExpressionNode } from '../operation-node/common-table-expression-node.js'
 import { CommonTableExpressionNameNode } from '../operation-node/common-table-expression-name-node.js'
 import { QueryCreator } from '../query-creator.js'
-import { createQueryCreator } from './parse-utils.js'
 import { Expression } from '../expression/expression.js'
 import { ShallowRecord } from '../util/type-utils.js'
+import { OperationNode } from '../operation-node/operation-node.js'
+import { createQueryCreator } from './parse-utils.js'
+import { isFunction } from '../util/object-utils.js'
+import { CTEBuilder, CTEBuilderCallback } from '../query-builder/cte-builder.js'
+import { CommonTableExpressionNode } from '../operation-node/common-table-expression-node.js'
 
 export type CommonTableExpression<DB, CN extends string> = (
   creator: QueryCreator<DB>
@@ -91,23 +94,39 @@ type ExtractColumnNamesFromColumnList<R extends string> =
     : R
 
 export function parseCommonTableExpression(
-  name: string,
-  expression: CommonTableExpression<any, any>
+  nameOrBuilderCallback: string | CTEBuilderCallback<string>,
+  expression: CommonTableExpression<any, string>
 ): CommonTableExpressionNode {
-  const builder = expression(createQueryCreator())
+  const expressionNode = expression(createQueryCreator()).toOperationNode()
+
+  if (isFunction(nameOrBuilderCallback)) {
+    return nameOrBuilderCallback(
+      cteBuilderFactory(expressionNode)
+    ).toOperationNode()
+  }
 
   return CommonTableExpressionNode.create(
-    parseCommonTableExpressionName(name),
-    builder.toOperationNode()
+    parseCommonTableExpressionName(nameOrBuilderCallback),
+    expressionNode
   )
 }
 
+function cteBuilderFactory(expressionNode: OperationNode) {
+  return (name: string) => {
+    return new CTEBuilder({
+      node: CommonTableExpressionNode.create(
+        parseCommonTableExpressionName(name),
+        expressionNode
+      ),
+    })
+  }
+}
+
 function parseCommonTableExpressionName(
   name: string
 ): CommonTableExpressionNameNode {
   if (name.includes('(')) {
     const parts = name.split(/[\(\)]/)
-
     const table = parts[0]
     const columns = parts[1].split(',').map((it) => it.trim())
 

src/query-builder/cte-builder.ts

@@ -0,0 +1,48 @@
+import { OperationNodeSource } from '../operation-node/operation-node-source.js'
+import { CommonTableExpressionNode } from '../operation-node/common-table-expression-node.js'
+import { preventAwait } from '../util/prevent-await.js'
+import { freeze } from '../util/object-utils.js'
+
+export class CTEBuilder<N extends string> implements OperationNodeSource {
+  readonly #props: CTEBuilderProps
+
+  constructor(props: CTEBuilderProps) {
+    this.#props = freeze(props)
+  }
+
+  asMaterialized(): CTEBuilder<N> {
+    return new CTEBuilder({
+      ...this.#props,
+      node: CommonTableExpressionNode.cloneWith(this.#props.node, {
+        materialized: true,
+      }),
+    })
+  }
+
+  asNotMaterialized(): CTEBuilder<N> {
+    return new CTEBuilder({
+      ...this.#props,
+      node: CommonTableExpressionNode.cloneWith(this.#props.node, {
+        materialized: false,
+      }),
+    })
+  }
+
+  toOperationNode(): CommonTableExpressionNode {
+    return this.#props.node
+  }
+}
+
+preventAwait(
+  CTEBuilder,
+  "don't await CTEBuilder instances. They are never executed directly and are always just a part of a query."
+)
+
+interface CTEBuilderProps {
+  readonly node: CommonTableExpressionNode
+}
+
+export type CTEBuilderCallback<N extends string> = (
+  // N2 is needed for proper inference. Don't remove it.
+  cte: <N2 extends string>(name: N2) => CTEBuilder<N2>
+) => CTEBuilder<N>

src/query-compiler/default-query-compiler.ts

@@ -938,6 +938,15 @@ export class DefaultQueryCompiler
   ): void {
     this.visitNode(node.name)
     this.append(' as ')
+
+    if (isBoolean(node.materialized)) {
+      if (node.materialized) {
+        this.append('materialized ')
+      } else {
+        this.append('not materialized ')
+      }
+    }
+
     this.visitNode(node.expression)
   }
 

src/query-creator.ts

@@ -23,9 +23,9 @@ import {
 import { QueryExecutor } from './query-executor/query-executor.js'
 import {
   CommonTableExpression,
-  parseCommonTableExpression,
   QueryCreatorWithCommonTableExpression,
   RecursiveCommonTableExpression,
+  parseCommonTableExpression,
 } from './parser/with-parser.js'
 import { WithNode } from './operation-node/with-node.js'
 import { createQueryId } from './util/query-id.js'
@@ -35,6 +35,7 @@ import { InsertResult } from './query-builder/insert-result.js'
 import { DeleteResult } from './query-builder/delete-result.js'
 import { UpdateResult } from './query-builder/update-result.js'
 import { KyselyPlugin } from './plugin/kysely-plugin.js'
+import { CTEBuilderCallback } from './query-builder/cte-builder.js'
 
 export class QueryCreator<DB> {
   readonly #props: QueryCreatorProps
@@ -455,10 +456,10 @@ export class QueryCreator<DB> {
    * ```
    */
   with<N extends string, E extends CommonTableExpression<DB, N>>(
-    name: N,
+    nameOrBuilder: N | CTEBuilderCallback<N>,
     expression: E
   ): QueryCreatorWithCommonTableExpression<DB, N, E> {
-    const cte = parseCommonTableExpression(name, expression)
+    const cte = parseCommonTableExpression(nameOrBuilder, expression)
 
     return new QueryCreator({
       ...this.#props,
@@ -476,8 +477,11 @@ export class QueryCreator<DB> {
   withRecursive<
     N extends string,
     E extends RecursiveCommonTableExpression<DB, N>
-  >(name: N, expression: E): QueryCreatorWithCommonTableExpression<DB, N, E> {
-    const cte = parseCommonTableExpression(name, expression)
+  >(
+    nameOrBuilder: N | CTEBuilderCallback<N>,
+    expression: E
+  ): QueryCreatorWithCommonTableExpression<DB, N, E> {
+    const cte = parseCommonTableExpression(nameOrBuilder, expression)
 
     return new QueryCreator({
       ...this.#props,

test/node/src/with.test.ts

@@ -178,6 +178,60 @@ for (const dialect of DIALECTS) {
           ])
         })
       })
+
+      it('should create a CTE with `as materialized`', async () => {
+        const query = ctx.db
+          .with(
+            (cte) => cte('person_name').asMaterialized(),
+            (qb) => qb.selectFrom('person').select('first_name')
+          )
+          .selectFrom('person_name')
+          .select('person_name.first_name')
+          .orderBy('first_name')
+
+        testSql(query, dialect, {
+          postgres: {
+            sql: 'with "person_name" as materialized (select "first_name" from "person") select "person_name"."first_name" from "person_name" order by "first_name"',
+            parameters: [],
+          },
+          mysql: NOT_SUPPORTED,
+          sqlite: NOT_SUPPORTED,
+        })
+
+        const result = await query.execute()
+        expect(result).to.eql([
+          { first_name: 'Arnold' },
+          { first_name: 'Jennifer' },
+          { first_name: 'Sylvester' },
+        ])
+      })
+
+      it('should create a CTE with `as not materialized`', async () => {
+        const query = ctx.db
+          .with(
+            (cte) => cte('person_name').asNotMaterialized(),
+            (qb) => qb.selectFrom('person').select('first_name')
+          )
+          .selectFrom('person_name')
+          .select('person_name.first_name')
+          .orderBy('first_name')
+
+        testSql(query, dialect, {
+          postgres: {
+            sql: 'with "person_name" as not materialized (select "first_name" from "person") select "person_name"."first_name" from "person_name" order by "first_name"',
+            parameters: [],
+          },
+          mysql: NOT_SUPPORTED,
+          sqlite: NOT_SUPPORTED,
+        })
+
+        const result = await query.execute()
+        expect(result).to.eql([
+          { first_name: 'Arnold' },
+          { first_name: 'Jennifer' },
+          { first_name: 'Sylvester' },
+        ])
+      })
     }
 
     if (dialect !== 'mysql') {

test/typings/test-d/with.test-d.ts

@@ -64,6 +64,27 @@ async function testWith(db: Kysely<Database>) {
     }[]
   >(r3)
 
+  // Using CTE builder.
+  const r4 = await db
+    .with(
+      (cte) => cte('jennifers').asMaterialized(),
+      (db) =>
+        db
+          .selectFrom('person')
+          .where('first_name', '=', 'Jennifer')
+          .select(['first_name', 'last_name as ln', 'gender'])
+    )
+    .selectFrom('jennifers')
+    .select(['first_name', 'ln'])
+    .execute()
+
+  expectType<
+    {
+      first_name: string
+      ln: string | null
+    }[]
+  >(r4)
+
   // Different columns in expression and CTE name.
   expectError(
     db
@@ -76,6 +97,22 @@ async function testWith(db: Kysely<Database>) {
       .selectFrom('jennifers')
       .select(['first_name', 'last_name'])
   )
+
+  // Unknown CTE name when using the CTE builder.
+  expectError(
+    db
+      .with(
+        (cte) => cte('jennifers').asMaterialized(),
+        (db) =>
+          db
+            .selectFrom('person')
+            .where('first_name', '=', 'Jennifer')
+            .select(['first_name', 'last_name as ln', 'gender'])
+      )
+      .selectFrom('lollifers')
+      .select(['first_name', 'ln'])
+      .execute()
+  )
 }
 
 async function testManyWith(db: Kysely<Database>) {