openAgent

1/* oxlint-disable */
2import * as Context from "effect/Context"
3import * as Effect from "effect/Effect"
4import * as Exit from "effect/Exit"
5import * as Scope from "effect/Scope"
6import type { SqlClient } from "effect/unstable/sql/SqlClient"
7import type { SqlError } from "effect/unstable/sql/SqlError"
8import type { EffectCacheShape } from "drizzle-orm/cache/core/cache-effect"
9import type { WithCacheConfig } from "drizzle-orm/cache/core/types"
10import type { EffectDrizzleQueryError } from "drizzle-orm/effect-core/errors"
11import type { EffectLoggerShape } from "drizzle-orm/effect-core/logger"
12import type { QueryEffectHKTBase } from "drizzle-orm/effect-core/query-effect"
13import { entityKind } from "drizzle-orm/entity"
14import type { AnyRelations } from "drizzle-orm/relations"
15import type { RelationalQueryMapperConfig } from "drizzle-orm/relations"
16import type { Query } from "drizzle-orm/sql/sql"
17import type { SQLiteAsyncDialect } from "drizzle-orm/sqlite-core/dialect"
18import { SQLiteEffectPreparedQuery, SQLiteEffectSession, SQLiteEffectTransaction } from "../sqlite-core/effect/session"
19import type { SelectedFieldsOrdered } from "drizzle-orm/sqlite-core/query-builders/select.types"
20import type { PreparedQueryConfig, SQLiteExecuteMethod, SQLiteTransactionConfig } from "drizzle-orm/sqlite-core/session"
21 
22export interface EffectSQLiteQueryEffectHKT extends QueryEffectHKTBase {
23  readonly error: EffectDrizzleQueryError
24  readonly context: never
25}
26 
27export type EffectSQLiteRunResult = readonly never[]
28 
29export interface EffectSQLiteSessionOptions {
30  logger: EffectLoggerShape
31  cache: EffectCacheShape
32  useJitMappers?: boolean
33}
34 
35export class EffectSQLiteSession<TRelations extends AnyRelations> extends SQLiteEffectSession<
36  EffectSQLiteQueryEffectHKT,
37  EffectSQLiteRunResult,
38  TRelations
39> {
40  static override readonly [entityKind]: string = "EffectSQLiteSession"
41 
42  constructor(
43    private client: SqlClient,
44    dialect: SQLiteAsyncDialect,
45    protected relations: TRelations,
46    private options: EffectSQLiteSessionOptions,
47  ) {
48    super(dialect)
49  }
50 
51  override prepareQuery<T extends PreparedQueryConfig = PreparedQueryConfig>(
52    query: Query,
53    fields: SelectedFieldsOrdered | undefined,
54    executeMethod: SQLiteExecuteMethod,
55    customResultMapper?: (rows: unknown[][], mapColumnValue?: (value: unknown) => unknown) => unknown,
56    queryMetadata?: {
57      type: "select" | "update" | "delete" | "insert"
58      tables: string[]
59    },
60    cacheConfig?: WithCacheConfig,
61  ): SQLiteEffectPreparedQuery<T, EffectSQLiteQueryEffectHKT> {
62    return new SQLiteEffectPreparedQuery<T, EffectSQLiteQueryEffectHKT>(
63      (params, method) => this.execute(query, params, method),
64      query,
65      this.options.logger,
66      this.options.cache,
67      queryMetadata,
68      cacheConfig,
69      fields,
70      executeMethod,
71      this.options.useJitMappers,
72      customResultMapper,
73      undefined,
74      undefined,
75      this.isInTransaction(),
76    )
77  }
78 
79  override prepareRelationalQuery<T extends PreparedQueryConfig = PreparedQueryConfig>(
80    query: Query,
81    fields: SelectedFieldsOrdered | undefined,
82    executeMethod: SQLiteExecuteMethod,
83    customResultMapper: (rows: Record<string, unknown>[], mapColumnValue?: (value: unknown) => unknown) => unknown,
84    config: RelationalQueryMapperConfig,
85  ): SQLiteEffectPreparedQuery<T, EffectSQLiteQueryEffectHKT, true> {
86    return new SQLiteEffectPreparedQuery<T, EffectSQLiteQueryEffectHKT, true>(
87      (params, method) => this.execute(query, params, method),
88      query,
89      this.options.logger,
90      this.options.cache,
91      undefined,
92      undefined,
93      fields,
94      executeMethod,
95      this.options.useJitMappers,
96      customResultMapper,
97      true,
98      config,
99      this.isInTransaction(),
100    )
101  }
102 
103  private execute(query: Query, params: unknown[], method: SQLiteExecuteMethod | "values") {
104    const statement = this.client.unsafe(query.sql, params)
105    if (method === "values") return statement.values
106    if (method === "get") return statement.withoutTransform.pipe(Effect.map((rows) => rows[0]))
107    return statement.withoutTransform
108  }
109 
110  private isInTransaction() {
111    return Effect.serviceOption(this.client.transactionService).pipe(Effect.map((option) => option._tag === "Some"))
112  }
113 
114  private executeTransactionStatement(connection: Effect.Success<SqlClient["reserve"]>, query: string) {
115    return connection.executeUnprepared(query, [], undefined).pipe(Effect.asVoid)
116  }
117 
118  private withTransaction<A, E, R>(effect: Effect.Effect<A, E, R>, config: SQLiteTransactionConfig | undefined) {
119    return Effect.uninterruptibleMask((restore) =>
120      Effect.withFiber<A, E | SqlError, R>((fiber) => {
121        const services = fiber.context
122        const connectionOption = Context.getOption(services, this.client.transactionService)
123        const connection: Effect.Effect<
124          readonly [Scope.Closeable | undefined, Effect.Success<SqlClient["reserve"]>],
125          SqlError
126        > =
127          connectionOption._tag === "Some"
128            ? Effect.succeed([undefined, connectionOption.value[0]] as const)
129            : Scope.make().pipe(
130                Effect.flatMap((scope) =>
131                  Scope.provide(this.client.reserve, scope).pipe(
132                    Effect.map((connection) => [scope, connection] as const),
133                    Effect.catch((error) =>
134                      Scope.close(scope, Exit.fail(error)).pipe(Effect.andThen(Effect.fail(error))),
135                    ),
136                  ),
137                ),
138              )
139        const id = connectionOption._tag === "Some" ? connectionOption.value[1] + 1 : 0
140 
141        return connection.pipe(
142          Effect.flatMap(([scope, connection]) => {
143            const transaction = this.executeTransactionStatement(
144              connection,
145              id === 0 ? `begin ${config?.behavior ?? "deferred"}` : `savepoint effect_sql_${id}`,
146            ).pipe(
147              Effect.flatMap(() =>
148                Effect.provideContext(
149                  restore(effect),
150                  Context.add(services, this.client.transactionService, [connection, id]),
151                ).pipe(
152                  Effect.exit,
153                  Effect.flatMap((exit) => {
154                    const finalize = Exit.isSuccess(exit)
155                      ? id === 0
156                        ? this.executeTransactionStatement(connection, "commit").pipe(
157                            // SQLite keeps the transaction open after deferred constraint commit failures.
158                            Effect.catch((error) =>
159                              this.executeTransactionStatement(connection, "rollback").pipe(
160                                Effect.catch(() => Effect.void),
161                                Effect.andThen(Effect.fail(error)),
162                              ),
163                            ),
164                          )
165                        : this.executeTransactionStatement(connection, `release savepoint effect_sql_${id}`)
166                      : id === 0
167                        ? this.executeTransactionStatement(connection, "rollback")
168                        : this.executeTransactionStatement(connection, `rollback to savepoint effect_sql_${id}`).pipe(
169                            Effect.andThen(
170                              this.executeTransactionStatement(connection, `release savepoint effect_sql_${id}`),
171                            ),
172                          )
173 
174                    return finalize.pipe(Effect.flatMap(() => exit))
175                  }),
176                ),
177              ),
178            )
179 
180            return scope === undefined
181              ? transaction
182              : transaction.pipe(Effect.onExit((exit) => Scope.close(scope, exit)))
183          }),
184        )
185      }),
186    )
187  }
188 
189  override transaction<A, E, R>(
190    transaction: (tx: EffectSQLiteTransaction<TRelations>) => Effect.Effect<A, E, R>,
191    config?: SQLiteTransactionConfig,
192  ): Effect.Effect<A, E | SqlError, R> {
193    const { dialect, relations } = this
194 
195    return this.withTransaction(
196      Effect.gen({ self: this }, function* () {
197        const tx = new EffectSQLiteTransaction<TRelations>(dialect, this, relations)
198 
199        return yield* transaction(tx)
200      }),
201      config,
202    )
203  }
204}
205 
206export class EffectSQLiteTransaction<TRelations extends AnyRelations> extends SQLiteEffectTransaction<
207  EffectSQLiteQueryEffectHKT,
208  EffectSQLiteRunResult,
209  TRelations
210> {
211  static override readonly [entityKind]: string = "EffectSQLiteTransaction"
212 
213  override transaction: <A, E, R>(
214    transaction: (
215      tx: SQLiteEffectTransaction<EffectSQLiteQueryEffectHKT, EffectSQLiteRunResult, TRelations>,
216    ) => Effect.Effect<A, E, R>,
217  ) => Effect.Effect<A, SqlError | E, R> = (tx) => this.session.transaction(tx)
218}
219