server/storage/database.ts

Rama main · réplica de lectura del repositorio original.

import cluster from "node:cluster";
import path from "node:path";
import { DatabaseError, UniqueConstraintError } from "sequelize";
import type {
  InferAttributes,
  InferCreationAttributes,
  Transaction,
} from "sequelize";
import sequelizeStrictAttributes from "sequelize-strict-attributes";
import type { SequelizeOptions } from "sequelize-typescript";
import { Sequelize } from "sequelize-typescript";
import type { MigrationError } from "umzug";
import { Umzug, SequelizeStorage } from "umzug";
import { toError } from "@shared/utils/error";
import env from "@server/env";
import { ClientClosedRequestError } from "@server/errors";
import type Model from "@server/models/base/Model";
import Logger from "../logging/Logger";
import * as models from "../models";
import { requestContext } from "./requestContext";
import { getConnectionName } from "./utils";

/**
 * Returns database configuration for Sequelize constructor.
 * Either uses DATABASE_URL or constructs options from individual components.
 */
function getDatabaseConfig() {
  if (env.DATABASE_URL) {
    return env.DATABASE_URL;
  }

  // If using individual components, return Sequelize options object
  if (env.DATABASE_HOST && env.DATABASE_NAME && env.DATABASE_USER) {
    return {
      database: env.DATABASE_NAME,
      username: env.DATABASE_USER,
      password: env.DATABASE_PASSWORD || undefined,
      host: env.DATABASE_HOST,
      port: env.DATABASE_PORT || 5432,
      dialect: "postgres" as const,
    };
  }

  throw new Error(
    "DATABASE_URL is not set or individual database components (DATABASE_HOST, DATABASE_NAME, DATABASE_USER) are not properly configured."
  );
}

/** Postgres error code for `query_canceled`. */
const QueryCanceledErrorCode = "57014";

/** Headroom between the statement timeout and the HTTP request timeout. */
const StatementTimeoutMargin = 500;

/** Sequelize types the query generator as unknown; this narrows to the single
 * method used to begin a transaction. */
interface QueryGeneratorWithTransaction {
  startTransactionQuery(
    transaction: Transaction & { parent?: Transaction }
  ): string;
}

const isSSLDisabled = env.PGSSLMODE === "disable";
const poolMax = env.DATABASE_CONNECTION_POOL_MAX ?? 5;
const poolMin = env.DATABASE_CONNECTION_POOL_MIN ?? 0;
const databaseConfig = env.DATABASE_CONNECTION_POOL_URL || getDatabaseConfig();
const schema = env.DATABASE_SCHEMA;

const isApiProcess =
  (env.SERVICES.includes("web") ||
    env.SERVICES.includes("collaboration") ||
    env.SERVICES.includes("websockets") ||
    env.SERVICES.includes("admin")) &&
  !env.SERVICES.includes("worker") &&
  !env.SERVICES.includes("cron");

// Request-handling processes get a Postgres `statement_timeout` slightly under
// the HTTP request timeout, so a single slow query cannot hold a connection
// past the point at which its response could be delivered, and the cancelation
// still leaves enough headroom to write an error response before the socket is
// closed. Applied as `SET LOCAL` inside each transaction so the value is scoped
// to the transaction.
const statementTimeout =
  isApiProcess && cluster.isWorker
    ? Math.max(
        env.REQUEST_TIMEOUT - StatementTimeoutMargin,
        // Fall back to a fraction of the request timeout when it is shorter
        // than the margin, so the query is always canceled first.
        Math.round(env.REQUEST_TIMEOUT / 2)
      )
    : undefined;

/**
 * Whether an error was caused by Postgres canceling a query, most commonly
 * because it exceeded the configured `statement_timeout`.
 *
 * @param err the error to inspect.
 * @returns true if the query was canceled by the database.
 */
export function isQueryCanceledError(err: unknown): boolean {
  return (
    err instanceof DatabaseError &&
    !!err.parent &&
    "code" in err.parent &&
    err.parent.code === QueryCanceledErrorCode
  );
}

/**
 * Run the given function, retrying it when it fails due to a unique constraint
 * violation. Useful for operations that choose a unique value based on a prior
 * read, where a concurrent request may claim the same value first.
 *
 * @param fn the function to run, this should include any transaction as a
 * violation leaves the surrounding transaction unusable.
 * @param attempts the maximum number of times to run the function.
 * @returns the result of the function.
 * @throws the last error if it is not a unique constraint violation, or the
 * maximum number of attempts has been reached.
 */
export async function retryOnUniqueConstraintError<T>(
  fn: () => Promise<T>,
  attempts = 3
): Promise<T> {
  for (let attempt = 1; ; attempt++) {
    try {
      return await fn();
    } catch (err) {
      if (!(err instanceof UniqueConstraintError) || attempt >= attempts) {
        throw err;
      }
      Logger.info("database", "Retrying after unique constraint violation", {
        attempt,
        fields: err.fields,
      });
    }
  }
}

export function createDatabaseInstance(
  databaseConfig: string | object,
  input: {
    [key: string]: typeof Model<
      InferAttributes<Model>,
      InferCreationAttributes<Model>
    >;
  },
  options?: { readOnly?: boolean }
): Sequelize {
  try {
    let instance;
    const isReadOnly = options?.readOnly ?? false;

    // Common options for both URL and object configurations
    const commonOptions: SequelizeOptions = {
      logging: (msg) =>
        process.env.DEBUG?.includes("database") &&
        Logger.debug("database", msg),
      typeValidation: true,
      logQueryParameters: env.isDevelopment,
      dialectOptions: {
        application_name: getConnectionName(),
        ssl:
          env.isProduction && !isSSLDisabled
            ? {
                // Ref.: https://github.com/brianc/node-postgres/issues/2009
                rejectUnauthorized: false,
              }
            : false,
      },
      models: Object.values(input),
      pool: {
        // Read-only connections can have larger pools since there's no write contention
        max: isReadOnly ? poolMax * 2 : poolMax,
        min: poolMin,
        acquire: 30000,
        idle: 10000,
      },
      // Only retry on deadlocks for write connections
      retry: isReadOnly
        ? undefined
        : {
            match: [/deadlock/i],
            max: 3,
            backoffBase: 200,
            backoffExponent: 1.5,
          },
      schema,
    };

    // If databaseConfig is a string, it's a URL; if it's an object, merge with common options
    if (typeof databaseConfig === "string") {
      instance = new Sequelize(databaseConfig, commonOptions);
    } else {
      instance = new Sequelize({ ...databaseConfig, ...commonOptions });
    }

    sequelizeStrictAttributes(instance);

    if (statementTimeout) {
      instance = applyStatementTimeoutToTransactions(
        instance,
        Number(statementTimeout)
      );
    }

    if (env.isTest) {
      instance = monkeyPatchSequelizeErrorsForTests(instance);
    }

    // Skip queries when the originating HTTP request socket has been destroyed
    // (e.g. client disconnected or server timeout). This avoids wasting database
    // resources on work whose response can never be delivered.
    const assertConnectionOpen = () => {
      const store = requestContext.getStore();
      if (store?.req.socket.destroyed) {
        throw ClientClosedRequestError();
      }
    };
    instance.addHook("beforeFind", assertConnectionOpen);
    instance.addHook("beforeCount", assertConnectionOpen);

    // Add hooks to warn about write operations on read-only connections
    if (isReadOnly) {
      const warnWriteOperation = (operation: string) => {
        Logger.warn(
          `Attempted ${operation} operation on read-only database connection`
        );
      };

      instance.addHook("beforeCreate", () => warnWriteOperation("CREATE"));
      instance.addHook("beforeUpdate", () => warnWriteOperation("UPDATE"));
      instance.addHook("beforeDestroy", () => warnWriteOperation("DELETE"));
      instance.addHook("beforeBulkCreate", () =>
        warnWriteOperation("BULK CREATE")
      );
      instance.addHook("beforeBulkUpdate", () =>
        warnWriteOperation("BULK UPDATE")
      );
      instance.addHook("beforeBulkDestroy", () =>
        warnWriteOperation("BULK DELETE")
      );
    }

    return instance;
  } catch (_err) {
    Logger.fatal(
      "Could not connect to database",
      typeof databaseConfig === "string"
        ? new Error(
            `Failed to parse: "${databaseConfig}". Ensure special characters in database URL are encoded`
          )
        : new Error(
            `Failed to connect using database credentials. Please check DATABASE_HOST, DATABASE_NAME, DATABASE_USER configuration`
          )
    );
    // To satisfy TypeScript that a Sequelize instance is always returned
    throw _err;
  }
}

/**
 * This function is used to test the database connection on startup. It will
 * throw a descriptive error if the connection fails.
 */
export const checkConnection = async (db: Sequelize) => {
  try {
    await db.authenticate();
  } catch (error) {
    if (
      error instanceof Error &&
      error.message.includes("does not support SSL")
    ) {
      Logger.fatal(
        "The database does not support SSL connections. Set the `PGSSLMODE` environment variable to `disable` or enable SSL on your database server.",
        error
      );
    } else {
      Logger.fatal("Failed to connect to database", toError(error));
    }
  }
};

export function createMigrationRunner(
  db: Sequelize,
  glob:
    | string
    | [
        string,
        {
          cwd?: string | undefined;
          ignore?: string | string[] | undefined;
        },
      ]
) {
  return new Umzug({
    migrations: {
      glob,
      resolve: ({ name, path, context }) => {
        // oxlint-disable-next-line @typescript-eslint/no-require-imports
        const migration = require(path as string);
        return {
          name,
          up: async () => migration.up(context, Sequelize),
          down: async () => migration.down(context, Sequelize),
        };
      },
    },
    context: db.getQueryInterface(),
    storage: new SequelizeStorage({ sequelize: db }),
    logger: {
      warn: (params) => Logger.warn("database", params),
      error: (params: Record<string, unknown> & MigrationError) =>
        Logger.error(params.message, params),
      info: (params) =>
        Logger.info(
          "database",
          params.event === "migrating"
            ? `Migrating ${String(params.name)}…`
            : `Migrated ${String(params.name)} in ${String(params.durationSeconds)}s`
        ),
      debug: (params) =>
        Logger.debug(
          "database",
          params.event === "migrating"
            ? `Migrating ${String(params.name)}…`
            : `Migrated ${String(params.name)} in ${String(params.durationSeconds)}s`
        ),
    },
  });
}

/**
 * Appends `SET LOCAL statement_timeout` to the statement that begins each
 * transaction, so that both are sent in a single round trip. Using `SET LOCAL`
 * scopes the value to the transaction, preventing it from leaking to other
 * consumers (e.g. background workers) sharing the same underlying connection
 * via pgbouncer's transaction pooling.
 *
 * @param instance the Sequelize instance to modify.
 * @param timeoutMs the statement timeout in milliseconds.
 * @returns the modified Sequelize instance.
 */
export function applyStatementTimeoutToTransactions(
  instance: Sequelize,
  timeoutMs: number
) {
  const queryGenerator = instance.getQueryInterface()
    .queryGenerator as QueryGeneratorWithTransaction;
  const origStartTransactionQuery =
    queryGenerator.startTransactionQuery.bind(queryGenerator);

  queryGenerator.startTransactionQuery = (transaction) => {
    const sql = origStartTransactionQuery(transaction);

    // Savepoints inherit the value set when the outer transaction began.
    if (transaction.parent) {
      return sql;
    }
    return `${sql} SET LOCAL statement_timeout = ${timeoutMs};`;
  };

  return instance;
}

/**
 * Fixed in Sequelize v7, but hasn't been back-ported to Sequelize v6.
 * See https://github.com/sequelize/sequelize/issues/14807#issuecomment-1854398131
 */
export function monkeyPatchSequelizeErrorsForTests(instance: Sequelize) {
  const sequelizeVersion = (Sequelize as unknown as { version: string })
    .version;
  const major = sequelizeVersion.split(".").map(Number)[0];

  if (major >= 7) {
    Logger.fatal(
      "Redundant patch",
      new Error(
        "This patch was made redundant in Sequelize v7, you should check!"
      )
    );
  }

  const origQueryFunc = instance.query.bind(instance);
  instance.query = (async (...args: Parameters<typeof origQueryFunc>) => {
    try {
      return await origQueryFunc(...args);
    } catch (err) {
      // Ensure error appears in test output, not swallowed by Sequelize internals
      const error = err as Error & { parent?: Error };
      Logger.error(error.message, error.parent ?? error);
      throw err;
    }
  }) as typeof instance.query;

  return instance;
}

export const sequelize = createDatabaseInstance(databaseConfig, models);

/**
 * Read-only database connection for read replicas.
 * Falls back to the main connection if DATABASE_READ_ONLY_URL is not set, and
 * in the test environment, where DATABASE_READ_ONLY_URL would not point at
 * the isolated per-worker test database.
 */
export const sequelizeReadOnly =
  env.DATABASE_READ_ONLY_URL && !env.isTest
    ? createDatabaseInstance(
        env.DATABASE_READ_ONLY_URL,
        {},
        {
          readOnly: true,
        }
      )
    : sequelize;

export const migrations = createMigrationRunner(sequelize, [
  "migrations/*.js",
  { cwd: path.resolve("server") },
]);