kysely-replication/dist/index.js

Replication-aware Kysely query execution

// src/connection.ts
import {
  SelectQueryNode
} from "kysely";
var PRIMARY_OPERATION_NODE_KINDS = {
  AlterTableNode: true,
  CreateIndexNode: true,
  CreateSchemaNode: true,
  CreateTableNode: true,
  CreateTypeNode: true,
  CreateViewNode: true,
  DeleteQueryNode: true,
  DropIndexNode: true,
  DropSchemaNode: true,
  DropTableNode: true,
  DropTypeNode: true,
  DropViewNode: true,
  InsertQueryNode: true,
  MergeQueryNode: true,
  RawNode: true,
  UpdateQueryNode: true
};
var KyselyReplicationConnection = class {
  #primaryDriver;
  #getReplicaDriver;
  #onReplicaTransaction;
  #connection;
  #driver;
  constructor(primary, getReplica, onReplicaTransaction) {
    this.#primaryDriver = primary;
    this.#getReplicaDriver = getReplica;
    this.#onReplicaTransaction = onReplicaTransaction;
    this.#connection = null;
    this.#driver = null;
  }
  async executeQuery(compiledQuery) {
    const { connection } = await this.#acquireDriverAndConnection(compiledQuery);
    return await connection.executeQuery(compiledQuery);
  }
  async *streamQuery(compiledQuery, chunkSize) {
    const { connection } = await this.#acquireDriverAndConnection(compiledQuery);
    for await (const result of connection.streamQuery(
      compiledQuery,
      chunkSize
    )) {
      yield result;
    }
  }
  async beginTransaction(settings) {
    const { connection, driver } = await this.#acquireDriverAndConnection("transaction");
    if (driver !== this.#primaryDriver) {
      const message = "KyselyReplication: transaction started with replica connection!";
      if (this.#onReplicaTransaction === "error") {
        throw new Error(message);
      }
      if (this.#onReplicaTransaction === "warn") {
        console.warn(message);
      }
    }
    await driver.beginTransaction(connection, settings);
  }
  async commitTransaction() {
    if (!this.#connection) {
      throw new Error("commitTransaction called without a transaction");
    }
    await this.#driver?.commitTransaction(this.#connection);
  }
  async rollbackTransaction() {
    if (!this.#connection) {
      throw new Error("rollbackTransaction called without a transaction");
    }
    await this.#driver?.rollbackTransaction(this.#connection);
  }
  async release() {
    if (!this.#connection) return;
    await this.#driver?.releaseConnection(this.#connection);
  }
  async #acquireDriverAndConnection(compiledQueryOrContext) {
    if (this.#connection && this.#driver) {
      return { connection: this.#connection, driver: this.#driver };
    }
    this.#driver = compiledQueryOrContext === "transaction" || this.#isQueryForPrimary(compiledQueryOrContext) ? this.#primaryDriver : await this.#getReplicaDriver(compiledQueryOrContext);
    this.#connection = await this.#driver.acquireConnection();
    return { connection: this.#connection, driver: this.#driver };
  }
  #isQueryForPrimary(compiledQuery) {
    const { query } = compiledQuery;
    if ("__dialect__" in query) {
      return query.__dialect__ === "primary";
    }
    return this.#isOperationNodeForPrimary(query) || SelectQueryNode.is(query) && Boolean(
      query.with?.expressions.some(
        (e) => this.#isOperationNodeForPrimary(e.expression)
      )
    );
  }
  #isOperationNodeForPrimary(node) {
    return PRIMARY_OPERATION_NODE_KINDS[node.kind];
  }
};

// src/driver.ts
var KyselyReplicationDriver = class {
  #primaryDriver;
  #replicaDrivers;
  #replicaStrategy;
  constructor(primaryDriver, replicaDrivers, replicaStrategy) {
    this.#primaryDriver = primaryDriver;
    this.#replicaDrivers = replicaDrivers;
    this.#replicaStrategy = replicaStrategy;
  }
  async acquireConnection() {
    return new KyselyReplicationConnection(
      this.#primaryDriver,
      async (compiledQuery) => {
        const replicaIndex = "__replicaIndex__" in compiledQuery.query ? compiledQuery.query.__replicaIndex__ : await this.#replicaStrategy.next(this.#replicaDrivers.length);
        const replicaDriver = this.#replicaDrivers[replicaIndex];
        if (!replicaDriver) {
          throw new Error(
            `KyselyReplication: no replicas found at index ${replicaIndex}!`
          );
        }
        return replicaDriver;
      },
      this.#replicaStrategy.onTransaction || "error"
    );
  }
  async beginTransaction(connection, settings) {
    await connection.beginTransaction(settings);
  }
  async commitTransaction(connection) {
    await connection.commitTransaction();
  }
  async destroy() {
    const results = await Promise.allSettled([
      this.#primaryDriver.destroy(),
      ...this.#replicaDrivers.map((replica) => replica.destroy())
    ]);
    const errors = this.#compileErrors(results);
    if (errors.length) {
      throw new AggregateError(
        errors,
        "KyselyReplicationDriver.destroy failed!"
      );
    }
  }
  async init() {
    const results = await Promise.allSettled([
      this.#primaryDriver.init(),
      ...this.#replicaDrivers.map((replica) => replica.init())
    ]);
    const errors = this.#compileErrors(results);
    if (errors.length) {
      throw new AggregateError(errors, "KyselyReplicationDriver.init failed!");
    }
  }
  async releaseConnection(connection) {
    await connection.release();
  }
  async rollbackTransaction(connection) {
    await connection.rollbackTransaction();
  }
  #compileErrors(results) {
    return results.map(
      (result, index) => result.status === "fulfilled" ? null : `${!index ? "primary" : `replica-${index - 1}`}: ${result.reason}`
    ).filter(Boolean);
  }
};

// src/dialect.ts
var KyselyReplicationDialect = class {
  #config;
  constructor(config) {
    this.#config = {
      ...config,
      replicaDialects: [...config.replicaDialects]
    };
  }
  createAdapter() {
    return this.#config.primaryDialect.createAdapter();
  }
  createDriver() {
    return new KyselyReplicationDriver(
      this.#config.primaryDialect.createDriver(),
      this.#config.replicaDialects.map((replica) => replica.createDriver()),
      this.#config.replicaStrategy
    );
  }
  createIntrospector(db) {
    return this.#config.primaryDialect.createIntrospector(db);
  }
  createQueryCompiler() {
    return this.#config.primaryDialect.createQueryCompiler();
  }
};
export {
  KyselyReplicationDialect,
  KyselyReplicationDriver
};
//# sourceMappingURL=index.js.map