savepoints.

kysely-org · Apr 23, 2024 · 119e846 · 119e846
1 parent a5aa83c
commit 119e846
Show file tree

Hide file tree

Showing 8 changed files with 282 additions and 11 deletions.
diff --git a/src/dialect/mssql/mssql-driver.ts b/src/dialect/mssql/mssql-driver.ts
@@ -30,6 +30,8 @@ import { CompiledQuery } from '../../query-compiler/compiled-query.js'
 import { extendStackTrace } from '../../util/stack-trace-utils.js'
 import { randomString } from '../../util/random-string.js'
 import { Deferred } from '../../util/deferred.js'
+import { parseSavepointCommand } from '../../parser/savepoint-parser.js'
+import { QueryCompiler } from '../../query-compiler/query-compiler.js'
 
 const PRIVATE_RELEASE_METHOD = Symbol()
 const PRIVATE_DESTROY_METHOD = Symbol()
@@ -87,6 +89,34 @@ export class MssqlDriver implements Driver {
     await connection.rollbackTransaction()
   }
 
+  async savepoint(
+    connection: DatabaseConnection,
+    savepointName: string,
+    compileQuery: QueryCompiler['compileQuery'],
+  ): Promise<void> {
+    await connection.executeQuery(
+      compileQuery(parseSavepointCommand('save transaction', savepointName)),
+    )
+  }
+
+  async rollbackToSavepoint(
+    connection: DatabaseConnection,
+    savepointName: string,
+    compileQuery: QueryCompiler['compileQuery'],
+  ): Promise<void> {
+    await connection.executeQuery(
+      compileQuery(
+        parseSavepointCommand('rollback transaction', savepointName),
+      ),
+    )
+  }
+
+  async releaseSavepoint(): Promise<void> {
+    throw new Error(
+      'MS SQL Server (mssql) does not support releasing savepoints',
+    )
+  }
+
   async releaseConnection(connection: MssqlConnection): Promise<void> {
     await connection[PRIVATE_RELEASE_METHOD]()
     this.#pool.release(connection)

diff --git a/src/dialect/mysql/mysql-driver.ts b/src/dialect/mysql/mysql-driver.ts
@@ -3,7 +3,9 @@ import {
   QueryResult,
 } from '../../driver/database-connection.js'
 import { Driver, TransactionSettings } from '../../driver/driver.js'
+import { parseSavepointCommand } from '../../parser/savepoint-parser.js'
 import { CompiledQuery } from '../../query-compiler/compiled-query.js'
+import { QueryCompiler } from '../../query-compiler/query-compiler.js'
 import { isFunction, isObject, freeze } from '../../util/object-utils.js'
 import { extendStackTrace } from '../../util/stack-trace-utils.js'
 import {
@@ -86,6 +88,36 @@ export class MysqlDriver implements Driver {
     await connection.executeQuery(CompiledQuery.raw('rollback'))
   }
 
+  async savepoint(
+    connection: DatabaseConnection,
+    savepointName: string,
+    compileQuery: QueryCompiler['compileQuery'],
+  ): Promise<void> {
+    await connection.executeQuery(
+      compileQuery(parseSavepointCommand('savepoint', savepointName)),
+    )
+  }
+
+  async rollbackToSavepoint(
+    connection: DatabaseConnection,
+    savepointName: string,
+    compileQuery: QueryCompiler['compileQuery'],
+  ): Promise<void> {
+    await connection.executeQuery(
+      compileQuery(parseSavepointCommand('rollback to', savepointName)),
+    )
+  }
+
+  async releaseSavepoint(
+    connection: DatabaseConnection,
+    savepointName: string,
+    compileQuery: QueryCompiler['compileQuery'],
+  ): Promise<void> {
+    await connection.executeQuery(
+      compileQuery(parseSavepointCommand('release savepoint', savepointName)),
+    )
+  }
+
   async releaseConnection(connection: MysqlConnection): Promise<void> {
     connection[PRIVATE_RELEASE_METHOD]()
   }

diff --git a/src/dialect/postgres/postgres-driver.ts b/src/dialect/postgres/postgres-driver.ts
@@ -3,7 +3,14 @@ import {
   QueryResult,
 } from '../../driver/database-connection.js'
 import { Driver, TransactionSettings } from '../../driver/driver.js'
+import { IdentifierNode } from '../../operation-node/identifier-node.js'
+import { RawNode } from '../../operation-node/raw-node.js'
+import { parseSavepointCommand } from '../../parser/savepoint-parser.js'
 import { CompiledQuery } from '../../query-compiler/compiled-query.js'
+import {
+  QueryCompiler,
+  RootOperationNode,
+} from '../../query-compiler/query-compiler.js'
 import { isFunction, freeze } from '../../util/object-utils.js'
 import { extendStackTrace } from '../../util/stack-trace-utils.js'
 import {
@@ -74,6 +81,36 @@ export class PostgresDriver implements Driver {
     await connection.executeQuery(CompiledQuery.raw('rollback'))
   }
 
+  async savepoint(
+    connection: DatabaseConnection,
+    savepointName: string,
+    compileQuery: QueryCompiler['compileQuery'],
+  ): Promise<void> {
+    await connection.executeQuery(
+      compileQuery(parseSavepointCommand('savepoint', savepointName)),
+    )
+  }
+
+  async rollbackToSavepoint(
+    connection: DatabaseConnection,
+    savepointName: string,
+    compileQuery: QueryCompiler['compileQuery'],
+  ): Promise<void> {
+    await connection.executeQuery(
+      compileQuery(parseSavepointCommand('rollback to', savepointName)),
+    )
+  }
+
+  async releaseSavepoint(
+    connection: DatabaseConnection,
+    savepointName: string,
+    compileQuery: QueryCompiler['compileQuery'],
+  ): Promise<void> {
+    await connection.executeQuery(
+      compileQuery(parseSavepointCommand('release', savepointName)),
+    )
+  }
+
   async releaseConnection(connection: PostgresConnection): Promise<void> {
     connection[PRIVATE_RELEASE_METHOD]()
   }

diff --git a/src/dialect/sqlite/sqlite-driver.ts b/src/dialect/sqlite/sqlite-driver.ts
@@ -4,7 +4,9 @@ import {
 } from '../../driver/database-connection.js'
 import { Driver } from '../../driver/driver.js'
 import { SelectQueryNode } from '../../operation-node/select-query-node.js'
+import { parseSavepointCommand } from '../../parser/savepoint-parser.js'
 import { CompiledQuery } from '../../query-compiler/compiled-query.js'
+import { QueryCompiler } from '../../query-compiler/query-compiler.js'
 import { freeze, isFunction } from '../../util/object-utils.js'
 import { SqliteDatabase, SqliteDialectConfig } from './sqlite-dialect-config.js'
 
@@ -50,6 +52,36 @@ export class SqliteDriver implements Driver {
     await connection.executeQuery(CompiledQuery.raw('rollback'))
   }
 
+  async savepoint(
+    connection: DatabaseConnection,
+    savepointName: string,
+    compileQuery: QueryCompiler['compileQuery'],
+  ): Promise<void> {
+    await connection.executeQuery(
+      compileQuery(parseSavepointCommand('savepoint', savepointName)),
+    )
+  }
+
+  async rollbackToSavepoint(
+    connection: DatabaseConnection,
+    savepointName: string,
+    compileQuery: QueryCompiler['compileQuery'],
+  ): Promise<void> {
+    await connection.executeQuery(
+      compileQuery(parseSavepointCommand('rollback to', savepointName)),
+    )
+  }
+
+  async releaseSavepoint(
+    connection: DatabaseConnection,
+    savepointName: string,
+    compileQuery: QueryCompiler['compileQuery'],
+  ): Promise<void> {
+    await connection.executeQuery(
+      compileQuery(parseSavepointCommand('release', savepointName)),
+    )
+  }
+
   async releaseConnection(): Promise<void> {
     this.#connectionMutex.unlock()
   }

diff --git a/src/driver/driver.ts b/src/driver/driver.ts
@@ -1,3 +1,4 @@
+import { QueryCompiler } from '../query-compiler/query-compiler.js'
 import { ArrayItemType } from '../util/type-utils.js'
 import { DatabaseConnection } from './database-connection.js'
 
@@ -37,6 +38,33 @@ export interface Driver {
    */
   rollbackTransaction(connection: DatabaseConnection): Promise<void>
 
+  /**
+   * Establishses a new savepoint within a transaction.
+   */
+  savepoint?(
+    connection: DatabaseConnection,
+    savepointName: string,
+    compileQuery: QueryCompiler['compileQuery'],
+  ): Promise<void>
+
+  /**
+   * Rolls back to a savepoint within a transaction.
+   */
+  rollbackToSavepoint?(
+    connection: DatabaseConnection,
+    savepointName: string,
+    compileQuery: QueryCompiler['compileQuery'],
+  ): Promise<void>
+
+  /**
+   * Releases a savepoint within a transaction.
+   */
+  releaseSavepoint?(
+    connection: DatabaseConnection,
+    savepointName: string,
+    compileQuery: QueryCompiler['compileQuery'],
+  ): Promise<void>
+
   /**
    * Releases a connection back to the pool.
    */

diff --git a/src/driver/runtime-driver.ts b/src/driver/runtime-driver.ts
@@ -1,4 +1,5 @@
 import { CompiledQuery } from '../query-compiler/compiled-query.js'
+import { QueryCompiler } from '../query-compiler/query-compiler.js'
 import { Log } from '../util/log.js'
 import { performanceNow } from '../util/performance-now.js'
 import { DatabaseConnection, QueryResult } from './database-connection.js'
@@ -85,6 +86,50 @@ export class RuntimeDriver implements Driver {
     return this.#driver.rollbackTransaction(connection)
   }
 
+  savepoint(
+    connection: DatabaseConnection,
+    savepointName: string,
+    compileQuery: QueryCompiler['compileQuery'],
+  ): Promise<void> {
+    if (this.#driver.savepoint) {
+      return this.#driver.savepoint(connection, savepointName, compileQuery)
+    }
+
+    throw new Error('savepoints are not supported by this driver')
+  }
+
+  rollbackToSavepoint(
+    connection: DatabaseConnection,
+    savepointName: string,
+    compileQuery: QueryCompiler['compileQuery'],
+  ): Promise<void> {
+    if (this.#driver.rollbackToSavepoint) {
+      return this.#driver.rollbackToSavepoint(
+        connection,
+        savepointName,
+        compileQuery,
+      )
+    }
+
+    throw new Error('savepoints are not supported by this driver')
+  }
+
+  releaseSavepoint(
+    connection: DatabaseConnection,
+    savepointName: string,
+    compileQuery: QueryCompiler['compileQuery'],
+  ): Promise<void> {
+    if (this.#driver.releaseSavepoint) {
+      return this.#driver.releaseSavepoint(
+        connection,
+        savepointName,
+        compileQuery,
+      )
+    }
+
+    throw new Error('savepoints are not supported by this driver')
+  }
+
   async destroy(): Promise<void> {
     if (!this.#initPromise) {
       return