implement stateful connection (#81)

Author: shiyuhang0

README.md

@@ -2,6 +2,12 @@
 
 This driver is for serverless and edge compute platforms that require HTTP external connections, such as Vercel Edge Functions or Cloudflare Workers.
 
+There are three ways to use the driver:
+
+1. Stateless connection (default): each query is independent, ideal for edge environments with short-lived, frequently created connections.
+2. Stateful connection (experimental): use it when you require session.
+3. Transaction (experimental): use it when you require interactive transaction.
+
 ## Usage
 
 **Install**
@@ -12,9 +18,9 @@ You can install the driver with npm:
 npm install @tidbcloud/serverless
 ```
 
-**Query**
+**Stateless Connection**
 
-To query from TiDB Serverless, you need to create a connection first. Then you can use the connection to execute raw SQL queries. For example:
+To query from TiDB Serverless, you need to create a connection first. Then you can use the connection to execute raw SQL queries.
 
 ```ts
 import { connect } from '@tidbcloud/serverless'
@@ -23,10 +29,38 @@ const conn = connect({url: 'mysql://username:password@host/database'})
 const results = await conn.execute('select * from test where id = ?',[1])
 ```
 
-**Transaction (Experimental)**
+**Stateful Connection (experimental)**
+
+If you want to keep session state across multiple queries, create a stateful connection. Remember to call `close()` to release the connection, or you may reach the connection limits.
+
+> **Note:**
+> 
+> Connections idle for 10 minutes will be closed automatically.
+> The Stateful connection is not concurrent-safe. You are not allowed to run SQLs parallel in the same stateful connection.
+
+```ts
+import { connect } from '@tidbcloud/serverless'
+
+const conn = connect({url: 'mysql://username:password@host/database'})
+const stateful = await conn.persist()
+
+try {
+  const r1 = await stateful.execute('use db2')
+  const r2 = await stateful.execute('select * from test where id = ?', [2])
+} finally {
+  await stateful.close()
+}
+```
+
+**Transaction (experimental)**
 
 You can also perform interactive transactions with the serverless driver. For example:
 
+> **Note:**
+> 
+> Transactions idle for 10 minutes will be rolled back automatically if it has not been committed or rolled back.
+> The transaction is not concurrent-safe. You are not allowed to run SQLs parallel in the same transaction.
+
 ```ts
 import { connect } from '@tidbcloud/serverless'
 
@@ -43,10 +77,6 @@ try {
 }
 ```
 
-> **Note:**
-> 
-> The transaction is not concurrent-safe. You are not allowed to run SQLs parallel in the same transaction.
-
 **Edge example**
 
 The serverless driver is suitable for the edge environments. See how to use it with Vercel Edge Functions:

integration-test/basic.test.ts

@@ -212,4 +212,34 @@ describe('basic', () => {
     await tx.commit()
     expect(result1.length + 1).toEqual(result2.rows?.length ?? result2.rowCount)
   })
+
+  test('stateful connection normal flow', async () => {
+    const conn = connect({ url: databaseURL, database: database, fetch, debug: true })
+    const stateful = await conn.persist()
+
+    await stateful.execute(`use mysql`)
+    await expect(stateful.execute(`select * from ${table} where emp_no = 0`)).rejects.toThrow()
+
+    await stateful.execute(`use ${database}`)
+    const r = (await stateful.execute(`select * from ${table} where emp_no = 0`)) as Row[]
+    expect(r.length).toEqual(1)
+
+    await stateful.close()
+  })
+
+  test('stateful connection use after close', async () => {
+    const conn = connect({ url: databaseURL, database: database, fetch, debug: true })
+    const stateful = await conn.persist()
+
+    const r1 = (await stateful.execute(`select * from ${table} where emp_no = 0`)) as Row[]
+    expect(r1.length).toEqual(1)
+
+    await stateful.close()
+
+    await expect(stateful.execute(`select * from ${table} where emp_no = 0`)).rejects.toThrow()
+
+    // original connection should still work
+    const r2 = (await conn.execute(`select * from ${table} where emp_no = 0`)) as Row[]
+    expect(r2.length).toEqual(1)
+  })
 })

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tidbcloud/serverless",
-  "version": "0.2.0",
+  "version": "0.3.0",
   "description": "TiDB Cloud Serverless Driver",
   "main": "./dist/index.cjs",
   "module": "./dist/index.js",

src/index.ts

@@ -45,10 +45,9 @@ export class Tx<T extends Config> {
   async execute<E extends ExecuteOptions>(
     query: string,
     args: ExecuteArgs = null,
-    options: E = defaultExecuteOptions as E,
-    txOptions: TxOptions = {}
+    options: E = defaultExecuteOptions as E
   ): Promise<ExecuteResult<E, T>> {
-    return this.conn.execute(query, args, options, txOptions)
+    return this.conn.execute(query, args, options)
   }
 
   async commit(): Promise<T['fullResult'] extends true ? FullResult : Row[]> {
@@ -104,15 +103,23 @@ export class Connection<T extends Config> {
   async begin(txOptions: TxOptions = {}) {
     const conn = new Connection<T>(this.config)
     const tx = new Tx<T>(conn)
-    await tx.execute<T>('BEGIN', undefined, undefined, txOptions)
+    await conn.execute('BEGIN', undefined, undefined, txOptions)
     return tx
   }
 
+  async persist() {
+    const conn = new Connection<T>(this.config)
+    await conn.execute('', null, defaultExecuteOptions as ExecuteOptions, {}, 'open')
+    const stateful = new StatefulConnection<T>(conn)
+    return stateful
+  }
+
   async execute<E extends ExecuteOptions>(
     query: string,
     args: ExecuteArgs = null,
     options: E = defaultExecuteOptions as E,
-    txOptions: TxOptions = {}
+    txOptions: TxOptions = {},
+    statefulAction?: 'open' | 'close'
   ): Promise<ExecuteResult<E, T>> {
     const sql = args ? format(query, args) : query
     const body = JSON.stringify({ query: sql })
@@ -125,7 +132,8 @@ export class Connection<T extends Config> {
       body,
       this.session ?? '',
       sql == 'BEGIN' ? txOptions.isolation : null,
-      debug
+      debug,
+      statefulAction
     )
 
     this.session = resp?.session ?? null
@@ -159,6 +167,26 @@ export class Connection<T extends Config> {
   }
 }
 
+export class StatefulConnection<T extends Config> {
+  private conn: Connection<T>
+
+  constructor(conn: Connection<T>) {
+    this.conn = conn
+  }
+
+  async execute<E extends ExecuteOptions>(
+    query: string,
+    args: ExecuteArgs = null,
+    options: E = defaultExecuteOptions as E
+  ): Promise<ExecuteResult<E, T>> {
+    return this.conn.execute(query, args, options)
+  }
+
+  async close(): Promise<void> {
+    await this.conn.execute('', null, defaultExecuteOptions as ExecuteOptions, {}, 'close')
+  }
+}
+
 export function connect<T extends Config>(config: T): Connection<T> {
   return new Connection<T>(config)
 }

src/serverless.ts

@@ -1,7 +1,7 @@
 import { Config } from './config.js'
 import { DatabaseError } from './error.js'
 import { Version } from './version.js'
-export async function postQuery<T>(config: Config, body, session = '', isolationLevel = null, debug): Promise<T> {
+export async function postQuery<T>(config: Config, body, session = '', isolationLevel = null, debug, statefulAction?: string): Promise<T> {
   let fetchCacheOption: Record<string, any> = { cache: 'no-store' }
   // Cloudflare Workers does not support cache now https://github.com/cloudflare/workerd/issues/69
   try {
@@ -32,6 +32,9 @@ export async function postQuery<T>(config: Config, body, session = '', isolation
   if (isolationLevel) {
     headers['TiDB-Isolation-Level'] = isolationLevel
   }
+  if (statefulAction) {
+    headers['TiDB-Stateful-Action'] = statefulAction
+  }
   const response = await fetch(url.toString(), {
     method: 'POST',
     body: body,

src/version.ts

@@ -1 +1 @@
-export const Version = '0.2.0'
+export const Version = '0.3.0'