add FETCH clause support. (#822)

src/dialect/mssql/mssql-query-compiler.ts

@@ -1,6 +1,7 @@
 import { AddColumnNode } from '../../operation-node/add-column-node.js'
 import { AlterTableColumnAlterationNode } from '../../operation-node/alter-table-node.js'
 import { DropColumnNode } from '../../operation-node/drop-column-node.js'
+import { OffsetNode } from '../../operation-node/offset-node.js'
 import { MergeQueryNode } from '../../operation-node/merge-query-node.js'
 import { DefaultQueryCompiler } from '../../query-compiler/default-query-compiler.js'
 
@@ -9,6 +10,11 @@ export class MssqlQueryCompiler extends DefaultQueryCompiler {
     return `@${this.numParameters}`
   }
 
+  protected override visitOffset(node: OffsetNode): void {
+    super.visitOffset(node)
+    this.append(' rows')
+  }
+
   // mssql allows multi-column alterations in a single statement,
   // but you can only use the command keyword/s once.
   // it also doesn't support multiple kinds of commands in the same

src/index.ts

@@ -201,6 +201,7 @@ export * from './operation-node/json-operator-chain-node.js'
 export * from './operation-node/tuple-node.js'
 export * from './operation-node/merge-query-node.js'
 export * from './operation-node/matched-node.js'
+export * from './operation-node/fetch-node.js'
 
 export * from './util/column-type.js'
 export * from './util/compilable.js'

src/operation-node/fetch-node.ts

@@ -0,0 +1,27 @@
+import { OperationNode } from './operation-node.js'
+import { ValueNode } from './value-node.js'
+
+export type FetchModifier = 'only' | 'with ties'
+
+export interface FetchNode extends OperationNode {
+  readonly kind: 'FetchNode'
+  readonly rowCount: ValueNode
+  readonly modifier: FetchModifier
+}
+
+/**
+ * @internal
+ */
+export const FetchNode = {
+  is(node: OperationNode): node is FetchNode {
+    return node.kind === 'FetchNode'
+  },
+
+  create(rowCount: number | bigint, modifier: FetchModifier): FetchNode {
+    return {
+      kind: 'FetchNode',
+      rowCount: ValueNode.create(rowCount),
+      modifier,
+    }
+  },
+}

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

@@ -91,6 +91,7 @@ import { MergeQueryNode } from './merge-query-node.js'
 import { MatchedNode } from './matched-node.js'
 import { AddIndexNode } from './add-index-node.js'
 import { CastNode } from './cast-node.js'
+import { FetchNode } from './fetch-node.js'
 
 /**
  * Transforms an operation node tree into another one.
@@ -216,6 +217,7 @@ export class OperationNodeTransformer {
     MatchedNode: this.transformMatched.bind(this),
     AddIndexNode: this.transformAddIndex.bind(this),
     CastNode: this.transformCast.bind(this),
+    FetchNode: this.transformFetch.bind(this),
   })
 
   transformNode<T extends OperationNode | undefined>(node: T): T {
@@ -262,6 +264,7 @@ export class OperationNodeTransformer {
       having: this.transformNode(node.having),
       explain: this.transformNode(node.explain),
       setOperations: this.transformNodeList(node.setOperations),
+      fetch: this.transformNode(node.fetch),
     })
   }
 
@@ -1027,6 +1030,14 @@ export class OperationNodeTransformer {
     })
   }
 
+  protected transformFetch(node: FetchNode): FetchNode {
+    return requireAllProps<FetchNode>({
+      kind: 'FetchNode',
+      rowCount: this.transformNode(node.rowCount),
+      modifier: node.modifier,
+    })
+  }
+
   protected transformDataType(node: DataTypeNode): DataTypeNode {
     // An Object.freezed leaf node. No need to clone.
     return node

src/operation-node/operation-node-visitor.ts

@@ -93,6 +93,7 @@ import { MergeQueryNode } from './merge-query-node.js'
 import { MatchedNode } from './matched-node.js'
 import { AddIndexNode } from './add-index-node.js'
 import { CastNode } from './cast-node.js'
+import { FetchNode } from './fetch-node.js'
 
 export abstract class OperationNodeVisitor {
   protected readonly nodeStack: OperationNode[] = []
@@ -193,6 +194,7 @@ export abstract class OperationNodeVisitor {
     MatchedNode: this.visitMatched.bind(this),
     AddIndexNode: this.visitAddIndex.bind(this),
     CastNode: this.visitCast.bind(this),
+    FetchNode: this.visitFetch.bind(this),
   })
 
   protected readonly visitNode = (node: OperationNode): void => {
@@ -301,4 +303,5 @@ export abstract class OperationNodeVisitor {
   protected abstract visitMatched(node: MatchedNode): void
   protected abstract visitAddIndex(node: AddIndexNode): void
   protected abstract visitCast(node: CastNode): void
+  protected abstract visitFetch(node: FetchNode): void
 }

src/operation-node/operation-node.ts

@@ -89,6 +89,7 @@ export type OperationNodeKind =
   | 'MatchedNode'
   | 'AddIndexNode'
   | 'CastNode'
+  | 'FetchNode'
 
 export interface OperationNode {
   readonly kind: OperationNodeKind

src/operation-node/select-query-node.ts

@@ -15,6 +15,7 @@ import { WithNode } from './with-node.js'
 import { SelectModifierNode } from './select-modifier-node.js'
 import { ExplainNode } from './explain-node.js'
 import { SetOperationNode } from './set-operation-node.js'
+import { FetchNode } from './fetch-node.js'
 
 export interface SelectQueryNode extends OperationNode {
   readonly kind: 'SelectQueryNode'
@@ -33,6 +34,7 @@ export interface SelectQueryNode extends OperationNode {
   readonly having?: HavingNode
   readonly explain?: ExplainNode
   readonly setOperations?: ReadonlyArray<SetOperationNode>
+  readonly fetch?: FetchNode
 }
 
 /**
@@ -153,6 +155,16 @@ export const SelectQueryNode = freeze({
     })
   },
 
+  cloneWithFetch(
+    selectNode: SelectQueryNode,
+    fetch: FetchNode
+  ): SelectQueryNode {
+    return freeze({
+      ...selectNode,
+      fetch,
+    })
+  },
+
   cloneWithHaving(
     selectNode: SelectQueryNode,
     operation: OperationNode,

src/parser/fetch-parser.ts

@@ -0,0 +1,21 @@
+import { FetchModifier, FetchNode } from '../operation-node/fetch-node.js'
+import { isBigInt, isNumber } from '../util/object-utils.js'
+
+export function parseFetch(
+  rowCount: number | bigint,
+  modifier: FetchModifier,
+): FetchNode {
+  if (!isNumber(rowCount) && !isBigInt(rowCount)) {
+    throw new Error(`Invalid fetch row count: ${rowCount}`)
+  }
+
+  if (!isFetchModifier(modifier)) {
+    throw new Error(`Invalid fetch modifier: ${modifier}`)
+  }
+
+  return FetchNode.create(rowCount, modifier)
+}
+
+function isFetchModifier(value: any): value is FetchModifier {
+  return value === 'only' || value === 'with ties'
+}

src/query-builder/select-query-builder.ts

@@ -79,6 +79,8 @@ import {
   ValueExpression,
   parseValueExpression,
 } from '../parser/value-parser.js'
+import { FetchModifier } from '../operation-node/fetch-node.js'
+import { parseFetch } from '../parser/fetch-parser.js'
 
 export interface SelectQueryBuilder<DB, TB extends keyof DB, O>
   extends WhereInterface<DB, TB>,
@@ -1045,10 +1047,12 @@ export interface SelectQueryBuilder<DB, TB extends keyof DB, O>
    *   .limit(10)
    * ```
    */
-  limit(limit: ValueExpression<DB, TB, number>): SelectQueryBuilder<DB, TB, O>
+  limit(
+    limit: ValueExpression<DB, TB, number | bigint>,
+  ): SelectQueryBuilder<DB, TB, O>
 
   /**
-   * Adds an offset clause to the query.
+   * Adds an `offset` clause to the query.
    *
    * ### Examples
    *
@@ -1058,11 +1062,45 @@ export interface SelectQueryBuilder<DB, TB extends keyof DB, O>
    * return await db
    *   .selectFrom('person')
    *   .select('first_name')
-   *   .offset(10)
    *   .limit(10)
+   *   .offset(10)
    * ```
    */
-  offset(offset: ValueExpression<DB, TB, number>): SelectQueryBuilder<DB, TB, O>
+  offset(
+    offset: ValueExpression<DB, TB, number | bigint>,
+  ): SelectQueryBuilder<DB, TB, O>
+
+  /**
+   * Adds a `fetch` clause to the query.
+   *
+   * This clause is only supported by some dialects like PostgreSQL or MS SQL Server.
+   *
+   * ### Examples
+   *
+   * ```ts
+   * return await db
+   *   .selectFrom('person')
+   *   .select('first_name')
+   *   .orderBy('first_name')
+   *   .offset(0)
+   *   .fetch(10)
+   *   .execute()
+   * ```
+   *
+   * The generated SQL (MS SQL Server):
+   *
+   * ```sql
+   * select "first_name"
+   * from "person"
+   * order by "first_name"
+   * offset 0 rows
+   * fetch next 10 rows only
+   * ```
+   */
+  fetch(
+    rowCount: number | bigint,
+    modifier?: FetchModifier,
+  ): SelectQueryBuilder<DB, TB, O>
 
   /**
    * Combines another select query or raw expression to this query using `union`.
@@ -1966,7 +2004,9 @@ class SelectQueryBuilderImpl<DB, TB extends keyof DB, O>
     })
   }
 
-  limit(limit: ValueExpression<DB, TB, number>): SelectQueryBuilder<DB, TB, O> {
+  limit(
+    limit: ValueExpression<DB, TB, number | bigint>,
+  ): SelectQueryBuilder<DB, TB, O> {
     return new SelectQueryBuilderImpl({
       ...this.#props,
       queryNode: SelectQueryNode.cloneWithLimit(
@@ -1977,7 +2017,7 @@ class SelectQueryBuilderImpl<DB, TB extends keyof DB, O>
   }
 
   offset(
-    offset: ValueExpression<DB, TB, number>,
+    offset: ValueExpression<DB, TB, number | bigint>,
   ): SelectQueryBuilder<DB, TB, O> {
     return new SelectQueryBuilderImpl({
       ...this.#props,
@@ -1988,6 +2028,19 @@ class SelectQueryBuilderImpl<DB, TB extends keyof DB, O>
     })
   }
 
+  fetch(
+    rowCount: number | bigint,
+    modifier: FetchModifier = 'only',
+  ): SelectQueryBuilder<DB, TB, O> {
+    return new SelectQueryBuilderImpl({
+      ...this.#props,
+      queryNode: SelectQueryNode.cloneWithFetch(
+        this.#props.queryNode,
+        parseFetch(rowCount, modifier),
+      ),
+    })
+  }
+
   union(
     expression: SetOperandExpression<DB, O>,
   ): SelectQueryBuilder<DB, TB, O> {

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

@@ -108,6 +108,7 @@ import { MergeQueryNode } from '../operation-node/merge-query-node.js'
 import { MatchedNode } from '../operation-node/matched-node.js'
 import { AddIndexNode } from '../operation-node/add-index-node.js'
 import { CastNode } from '../operation-node/cast-node.js'
+import { FetchNode } from '../operation-node/fetch-node.js'
 
 export class DefaultQueryCompiler
   extends OperationNodeVisitor
@@ -222,6 +223,11 @@ export class DefaultQueryCompiler
       this.visitNode(node.offset)
     }
 
+    if (node.fetch) {
+      this.append(' ')
+      this.visitNode(node.fetch)
+    }
+
     if (node.endModifiers?.length) {
       this.append(' ')
       this.compileList(this.sortSelectModifiers([...node.endModifiers]), ' ')
@@ -1526,6 +1532,12 @@ export class DefaultQueryCompiler
     this.visitNode(node.dataType)
     this.append(')')
   }
+
+  protected override visitFetch(node: FetchNode): void {
+    this.append('fetch next ')
+    this.visitNode(node.rowCount)
+    this.append(` rows ${node.modifier}`)
+  }
 
   protected append(str: string): void {
     this.#sql += str