1
0
Fork 0
mirror of https://github.com/immich-app/immich.git synced 2024-12-28 22:51:59 +00:00

refactor(server): startup checks for vector extension (#11559)

* update update logic

refactor

* update tests

* get version range through repo method, make tests more static

* move "should work" test
This commit is contained in:
Mert 2024-08-05 21:00:25 -04:00 committed by GitHub
parent 9765ccb5a7
commit d5b23373c7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 390 additions and 366 deletions

View file

@ -28,6 +28,11 @@ export const EXTENSION_NAMES: Record<DatabaseExtension, string> = {
vectors: 'pgvecto.rs',
} as const;
export interface ExtensionVersion {
availableVersion: string | null;
installedVersion: string | null;
}
export interface VectorUpdateResult {
restartRequired: boolean;
}
@ -35,9 +40,10 @@ export interface VectorUpdateResult {
export const IDatabaseRepository = 'IDatabaseRepository';
export interface IDatabaseRepository {
getExtensionVersion(extensionName: string): Promise<string | undefined>;
getAvailableExtensionVersion(extension: DatabaseExtension): Promise<string | undefined>;
getExtensionVersion(extension: DatabaseExtension): Promise<ExtensionVersion>;
getExtensionVersionRange(extension: VectorExtension): string;
getPostgresVersion(): Promise<string>;
getPostgresVersionRange(): string;
createExtension(extension: DatabaseExtension): Promise<void>;
updateExtension(extension: DatabaseExtension, version?: string): Promise<void>;
updateVectorExtension(extension: VectorExtension, version?: string): Promise<VectorUpdateResult>;

View file

@ -2,11 +2,13 @@ import { Inject, Injectable } from '@nestjs/common';
import { InjectDataSource } from '@nestjs/typeorm';
import AsyncLock from 'async-lock';
import semver from 'semver';
import { POSTGRES_VERSION_RANGE, VECTOR_VERSION_RANGE, VECTORS_VERSION_RANGE } from 'src/constants';
import { getVectorExtension } from 'src/database.config';
import {
DatabaseExtension,
DatabaseLock,
EXTENSION_NAMES,
ExtensionVersion,
IDatabaseRepository,
VectorExtension,
VectorIndex,
@ -29,20 +31,18 @@ export class DatabaseRepository implements IDatabaseRepository {
this.logger.setContext(DatabaseRepository.name);
}
async getExtensionVersion(extension: DatabaseExtension): Promise<string | undefined> {
const res = await this.dataSource.query(`SELECT extversion FROM pg_extension WHERE extname = $1`, [extension]);
return res[0]?.['extversion'];
}
async getAvailableExtensionVersion(extension: DatabaseExtension): Promise<string | undefined> {
const res = await this.dataSource.query(
`
SELECT version FROM pg_available_extension_versions
WHERE name = $1 AND installed = false
ORDER BY version DESC`,
async getExtensionVersion(extension: DatabaseExtension): Promise<ExtensionVersion> {
const [res]: ExtensionVersion[] = await this.dataSource.query(
`SELECT default_version as "availableVersion", installed_version as "installedVersion"
FROM pg_available_extensions
WHERE name = $1`,
[extension],
);
return res[0]?.['version'];
return res ?? { availableVersion: null, installedVersion: null };
}
getExtensionVersionRange(extension: VectorExtension): string {
return extension === DatabaseExtension.VECTORS ? VECTORS_VERSION_RANGE : VECTOR_VERSION_RANGE;
}
async getPostgresVersion(): Promise<string> {
@ -50,6 +50,10 @@ export class DatabaseRepository implements IDatabaseRepository {
return version;
}
getPostgresVersionRange(): string {
return POSTGRES_VERSION_RANGE;
}
async createExtension(extension: DatabaseExtension): Promise<void> {
await this.dataSource.query(`CREATE EXTENSION IF NOT EXISTS ${extension}`);
}
@ -59,28 +63,34 @@ export class DatabaseRepository implements IDatabaseRepository {
}
async updateVectorExtension(extension: VectorExtension, targetVersion?: string): Promise<VectorUpdateResult> {
const currentVersion = await this.getExtensionVersion(extension);
if (!currentVersion) {
const { availableVersion, installedVersion } = await this.getExtensionVersion(extension);
if (!installedVersion) {
throw new Error(`${EXTENSION_NAMES[extension]} extension is not installed`);
}
if (!availableVersion) {
throw new Error(`No available version for ${EXTENSION_NAMES[extension]} extension`);
}
targetVersion ??= availableVersion;
const isVectors = extension === DatabaseExtension.VECTORS;
let restartRequired = false;
await this.dataSource.manager.transaction(async (manager) => {
await this.setSearchPath(manager);
const isSchemaUpgrade = targetVersion && semver.satisfies(targetVersion, '0.1.1 || 0.1.11');
if (isVectors && installedVersion === '0.1.1') {
await this.setExtVersion(manager, DatabaseExtension.VECTORS, '0.1.11');
}
const isSchemaUpgrade = semver.satisfies(installedVersion, '0.1.1 || 0.1.11');
if (isSchemaUpgrade && isVectors) {
await this.updateVectorsSchema(manager, currentVersion);
await this.updateVectorsSchema(manager);
}
await manager.query(`ALTER EXTENSION ${extension} UPDATE${targetVersion ? ` TO '${targetVersion}'` : ''}`);
await manager.query(`ALTER EXTENSION ${extension} UPDATE TO '${targetVersion}'`);
if (!isSchemaUpgrade) {
return;
}
if (isVectors) {
const diff = semver.diff(installedVersion, targetVersion);
if (isVectors && diff && ['minor', 'major'].includes(diff)) {
await manager.query('SELECT pgvectors_upgrade()');
restartRequired = true;
} else {
@ -96,24 +106,24 @@ export class DatabaseRepository implements IDatabaseRepository {
try {
await this.dataSource.query(`REINDEX INDEX ${index}`);
} catch (error) {
if (getVectorExtension() === DatabaseExtension.VECTORS) {
this.logger.warn(`Could not reindex index ${index}. Attempting to auto-fix.`);
const table = index === VectorIndex.CLIP ? 'smart_search' : 'face_search';
const dimSize = await this.getDimSize(table);
await this.dataSource.manager.transaction(async (manager) => {
await this.setSearchPath(manager);
await manager.query(`DROP INDEX IF EXISTS ${index}`);
await manager.query(`ALTER TABLE ${table} ALTER COLUMN embedding SET DATA TYPE real[]`);
await manager.query(`ALTER TABLE ${table} ALTER COLUMN embedding SET DATA TYPE vector(${dimSize})`);
await manager.query(`SET vectors.pgvector_compatibility=on`);
await manager.query(`
CREATE INDEX IF NOT EXISTS ${index} ON ${table}
USING hnsw (embedding vector_cosine_ops)
WITH (ef_construction = 300, m = 16)`);
});
} else {
if (getVectorExtension() !== DatabaseExtension.VECTORS) {
throw error;
}
this.logger.warn(`Could not reindex index ${index}. Attempting to auto-fix.`);
const table = await this.getIndexTable(index);
const dimSize = await this.getDimSize(table);
await this.dataSource.manager.transaction(async (manager) => {
await this.setSearchPath(manager);
await manager.query(`DROP INDEX IF EXISTS ${index}`);
await manager.query(`ALTER TABLE ${table} ALTER COLUMN embedding SET DATA TYPE real[]`);
await manager.query(`ALTER TABLE ${table} ALTER COLUMN embedding SET DATA TYPE vector(${dimSize})`);
await manager.query(`SET vectors.pgvector_compatibility=on`);
await manager.query(`
CREATE INDEX IF NOT EXISTS ${index} ON ${table}
USING hnsw (embedding vector_cosine_ops)
WITH (ef_construction = 300, m = 16)`);
});
}
}
@ -123,13 +133,8 @@ export class DatabaseRepository implements IDatabaseRepository {
}
try {
const res = await this.dataSource.query(
`
SELECT idx_status
FROM pg_vector_index_stat
WHERE indexname = $1`,
[name],
);
const query = `SELECT idx_status FROM pg_vector_index_stat WHERE indexname = $1`;
const res = await this.dataSource.query(query, [name]);
return res[0]?.['idx_status'] === 'UPGRADE';
} catch (error) {
const message: string = (error as any).message;
@ -146,19 +151,27 @@ export class DatabaseRepository implements IDatabaseRepository {
await manager.query(`SET search_path TO "$user", public, vectors`);
}
private async updateVectorsSchema(manager: EntityManager, currentVersion: string): Promise<void> {
await manager.query('CREATE SCHEMA IF NOT EXISTS vectors');
await manager.query(`UPDATE pg_catalog.pg_extension SET extversion = $1 WHERE extname = $2`, [
currentVersion,
DatabaseExtension.VECTORS,
]);
await manager.query('UPDATE pg_catalog.pg_extension SET extrelocatable = true WHERE extname = $1', [
DatabaseExtension.VECTORS,
]);
private async setExtVersion(manager: EntityManager, extName: DatabaseExtension, version: string): Promise<void> {
const query = `UPDATE pg_catalog.pg_extension SET extversion = $1 WHERE extname = $2`;
await manager.query(query, [version, extName]);
}
private async getIndexTable(index: VectorIndex): Promise<string> {
const tableQuery = `SELECT relname FROM pg_stat_all_indexes WHERE indexrelname = $1`;
const [res]: { relname: string | null }[] = await this.dataSource.manager.query(tableQuery, [index]);
const table = res?.relname;
if (!table) {
throw new Error(`Could not find table for index ${index}`);
}
return table;
}
private async updateVectorsSchema(manager: EntityManager): Promise<void> {
const extension = DatabaseExtension.VECTORS;
await manager.query(`CREATE SCHEMA IF NOT EXISTS ${extension}`);
await manager.query('UPDATE pg_catalog.pg_extension SET extrelocatable = true WHERE extname = $1', [extension]);
await manager.query('ALTER EXTENSION vectors SET SCHEMA vectors');
await manager.query('UPDATE pg_catalog.pg_extension SET extrelocatable = false WHERE extname = $1', [
DatabaseExtension.VECTORS,
]);
await manager.query('UPDATE pg_catalog.pg_extension SET extrelocatable = false WHERE extname = $1', [extension]);
}
private async getDimSize(table: string, column = 'embedding'): Promise<number> {

View file

@ -1,4 +1,4 @@
import { DatabaseExtension, IDatabaseRepository } from 'src/interfaces/database.interface';
import { DatabaseExtension, EXTENSION_NAMES, IDatabaseRepository } from 'src/interfaces/database.interface';
import { ILoggerRepository } from 'src/interfaces/logger.interface';
import { DatabaseService } from 'src/services/database.service';
import { newDatabaseRepositoryMock } from 'test/repositories/database.repository.mock';
@ -9,15 +9,33 @@ describe(DatabaseService.name, () => {
let sut: DatabaseService;
let databaseMock: Mocked<IDatabaseRepository>;
let loggerMock: Mocked<ILoggerRepository>;
let extensionRange: string;
let versionBelowRange: string;
let minVersionInRange: string;
let updateInRange: string;
let versionAboveRange: string;
beforeEach(() => {
delete process.env.DB_SKIP_MIGRATIONS;
delete process.env.DB_VECTOR_EXTENSION;
databaseMock = newDatabaseRepositoryMock();
loggerMock = newLoggerRepositoryMock();
sut = new DatabaseService(databaseMock, loggerMock);
databaseMock.getExtensionVersion.mockResolvedValue('0.2.0');
extensionRange = '0.2.x';
databaseMock.getExtensionVersionRange.mockReturnValue(extensionRange);
versionBelowRange = '0.1.0';
minVersionInRange = '0.2.0';
updateInRange = '0.2.1';
versionAboveRange = '0.3.0';
databaseMock.getExtensionVersion.mockResolvedValue({
installedVersion: minVersionInRange,
availableVersion: minVersionInRange,
});
});
afterEach(() => {
delete process.env.DB_SKIP_MIGRATIONS;
delete process.env.DB_VECTOR_EXTENSION;
});
it('should work', () => {
@ -32,264 +50,238 @@ describe(DatabaseService.name, () => {
expect(databaseMock.getPostgresVersion).toHaveBeenCalledTimes(1);
});
it(`should start up successfully with pgvectors`, async () => {
databaseMock.getPostgresVersion.mockResolvedValue('14.0.0');
describe.each([
{ extension: DatabaseExtension.VECTOR, extensionName: EXTENSION_NAMES[DatabaseExtension.VECTOR] },
{ extension: DatabaseExtension.VECTORS, extensionName: EXTENSION_NAMES[DatabaseExtension.VECTORS] },
])('should work with $extensionName', ({ extension, extensionName }) => {
beforeEach(() => {
process.env.DB_VECTOR_EXTENSION = extensionName;
});
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
it(`should start up successfully with ${extension}`, async () => {
databaseMock.getPostgresVersion.mockResolvedValue('14.0.0');
databaseMock.getExtensionVersion.mockResolvedValue({
installedVersion: null,
availableVersion: minVersionInRange,
});
expect(databaseMock.getPostgresVersion).toHaveBeenCalled();
expect(databaseMock.createExtension).toHaveBeenCalledWith(DatabaseExtension.VECTORS);
expect(databaseMock.createExtension).toHaveBeenCalledTimes(1);
expect(databaseMock.getExtensionVersion).toHaveBeenCalled();
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
it(`should start up successfully with pgvector`, async () => {
process.env.DB_VECTOR_EXTENSION = 'pgvector';
databaseMock.getPostgresVersion.mockResolvedValue('14.0.0');
databaseMock.getExtensionVersion.mockResolvedValue('0.5.0');
expect(databaseMock.getPostgresVersion).toHaveBeenCalled();
expect(databaseMock.createExtension).toHaveBeenCalledWith(extension);
expect(databaseMock.createExtension).toHaveBeenCalledTimes(1);
expect(databaseMock.getExtensionVersion).toHaveBeenCalled();
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
it(`should throw an error if the ${extension} extension is not installed`, async () => {
databaseMock.getExtensionVersion.mockResolvedValue({ installedVersion: null, availableVersion: null });
const message = `The ${extensionName} extension is not available in this Postgres instance.
If using a container image, ensure the image has the extension installed.`;
await expect(sut.onBootstrapEvent()).rejects.toThrow(message);
expect(databaseMock.createExtension).toHaveBeenCalledWith(DatabaseExtension.VECTOR);
expect(databaseMock.createExtension).toHaveBeenCalledTimes(1);
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
expect(databaseMock.createExtension).not.toHaveBeenCalled();
expect(databaseMock.runMigrations).not.toHaveBeenCalled();
});
it(`should throw an error if the pgvecto.rs extension is not installed`, async () => {
databaseMock.getExtensionVersion.mockResolvedValue('');
await expect(sut.onBootstrapEvent()).rejects.toThrow(`Unexpected: The pgvecto.rs extension is not installed.`);
it(`should throw an error if the ${extension} extension version is below minimum supported version`, async () => {
databaseMock.getExtensionVersion.mockResolvedValue({
installedVersion: versionBelowRange,
availableVersion: versionBelowRange,
});
expect(databaseMock.createExtension).toHaveBeenCalledTimes(1);
expect(databaseMock.runMigrations).not.toHaveBeenCalled();
});
await expect(sut.onBootstrapEvent()).rejects.toThrow(
`The ${extensionName} extension version is ${versionBelowRange}, but Immich only supports ${extensionRange}`,
);
it(`should throw an error if the pgvector extension is not installed`, async () => {
process.env.DB_VECTOR_EXTENSION = 'pgvector';
databaseMock.getExtensionVersion.mockResolvedValue('');
await expect(sut.onBootstrapEvent()).rejects.toThrow(`Unexpected: The pgvector extension is not installed.`);
expect(databaseMock.runMigrations).not.toHaveBeenCalled();
});
expect(databaseMock.createExtension).toHaveBeenCalledTimes(1);
expect(databaseMock.runMigrations).not.toHaveBeenCalled();
});
it(`should throw an error if ${extension} extension version is a nightly`, async () => {
databaseMock.getExtensionVersion.mockResolvedValue({ installedVersion: '0.0.0', availableVersion: '0.0.0' });
it(`should throw an error if the pgvecto.rs extension version is below minimum supported version`, async () => {
databaseMock.getExtensionVersion.mockResolvedValue('0.1.0');
await expect(sut.onBootstrapEvent()).rejects.toThrow(
`The ${extensionName} extension version is 0.0.0, which means it is a nightly release.`,
);
await expect(sut.onBootstrapEvent()).rejects.toThrow(
'The pgvecto.rs extension version is 0.1.0, but Immich only supports 0.2.x.',
);
expect(databaseMock.createExtension).not.toHaveBeenCalled();
expect(databaseMock.updateVectorExtension).not.toHaveBeenCalled();
expect(databaseMock.runMigrations).not.toHaveBeenCalled();
});
expect(databaseMock.runMigrations).not.toHaveBeenCalled();
});
it(`should do in-range update for ${extension} extension`, async () => {
databaseMock.getExtensionVersion.mockResolvedValue({
availableVersion: updateInRange,
installedVersion: minVersionInRange,
});
databaseMock.updateVectorExtension.mockResolvedValue({ restartRequired: false });
it(`should throw an error if the pgvector extension version is below minimum supported version`, async () => {
process.env.DB_VECTOR_EXTENSION = 'pgvector';
databaseMock.getExtensionVersion.mockResolvedValue('0.1.0');
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
await expect(sut.onBootstrapEvent()).rejects.toThrow(
'The pgvector extension version is 0.1.0, but Immich only supports >=0.5 <1',
);
expect(databaseMock.updateVectorExtension).toHaveBeenCalledWith(extension, updateInRange);
expect(databaseMock.updateVectorExtension).toHaveBeenCalledTimes(1);
expect(databaseMock.getExtensionVersion).toHaveBeenCalled();
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
expect(databaseMock.runMigrations).not.toHaveBeenCalled();
});
it(`should not upgrade ${extension} if same version`, async () => {
databaseMock.getExtensionVersion.mockResolvedValue({
availableVersion: minVersionInRange,
installedVersion: minVersionInRange,
});
it(`should throw an error if pgvecto.rs extension version is a nightly`, async () => {
databaseMock.getExtensionVersion.mockResolvedValue('0.0.0');
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
await expect(sut.onBootstrapEvent()).rejects.toThrow(
'The pgvecto.rs extension version is 0.0.0, which means it is a nightly release.',
);
expect(databaseMock.updateVectorExtension).not.toHaveBeenCalled();
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
expect(databaseMock.createExtension).toHaveBeenCalledTimes(1);
expect(databaseMock.runMigrations).not.toHaveBeenCalled();
});
it(`should throw error if ${extension} available version is below range`, async () => {
databaseMock.getExtensionVersion.mockResolvedValue({
availableVersion: versionBelowRange,
installedVersion: null,
});
it(`should throw an error if pgvector extension version is a nightly`, async () => {
process.env.DB_VECTOR_EXTENSION = 'pgvector';
databaseMock.getExtensionVersion.mockResolvedValue('0.0.0');
await expect(sut.onBootstrapEvent()).rejects.toThrow();
await expect(sut.onBootstrapEvent()).rejects.toThrow(
'The pgvector extension version is 0.0.0, which means it is a nightly release.',
);
expect(databaseMock.createExtension).not.toHaveBeenCalled();
expect(databaseMock.updateVectorExtension).not.toHaveBeenCalled();
expect(databaseMock.runMigrations).not.toHaveBeenCalled();
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
expect(databaseMock.createExtension).toHaveBeenCalledTimes(1);
expect(databaseMock.runMigrations).not.toHaveBeenCalled();
});
it(`should throw error if ${extension} available version is above range`, async () => {
databaseMock.getExtensionVersion.mockResolvedValue({
availableVersion: versionAboveRange,
installedVersion: minVersionInRange,
});
it(`should throw error if pgvecto.rs extension could not be created`, async () => {
databaseMock.createExtension.mockRejectedValue(new Error('Failed to create extension'));
await expect(sut.onBootstrapEvent()).rejects.toThrow();
await expect(sut.onBootstrapEvent()).rejects.toThrow('Failed to create extension');
expect(databaseMock.createExtension).not.toHaveBeenCalled();
expect(databaseMock.updateVectorExtension).not.toHaveBeenCalled();
expect(databaseMock.runMigrations).not.toHaveBeenCalled();
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
expect(loggerMock.fatal).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal.mock.calls[0][0]).toContain(
'Alternatively, if your Postgres instance has pgvector, you may use this instead',
);
expect(databaseMock.createExtension).toHaveBeenCalledTimes(1);
expect(databaseMock.runMigrations).not.toHaveBeenCalled();
it('should throw error if available version is below installed version', async () => {
databaseMock.getExtensionVersion.mockResolvedValue({
availableVersion: minVersionInRange,
installedVersion: updateInRange,
});
await expect(sut.onBootstrapEvent()).rejects.toThrow(
`The database currently has ${extensionName} ${updateInRange} activated, but the Postgres instance only has ${minVersionInRange} available.`,
);
expect(databaseMock.updateVectorExtension).not.toHaveBeenCalled();
expect(databaseMock.runMigrations).not.toHaveBeenCalled();
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
it(`should raise error if ${extension} extension upgrade failed`, async () => {
databaseMock.getExtensionVersion.mockResolvedValue({
availableVersion: updateInRange,
installedVersion: minVersionInRange,
});
databaseMock.updateVectorExtension.mockRejectedValue(new Error('Failed to update extension'));
await expect(sut.onBootstrapEvent()).rejects.toThrow('Failed to update extension');
expect(loggerMock.warn.mock.calls[0][0]).toContain(
`The ${extensionName} extension can be updated to ${updateInRange}.`,
);
expect(loggerMock.fatal).not.toHaveBeenCalled();
expect(databaseMock.updateVectorExtension).toHaveBeenCalledWith(extension, updateInRange);
expect(databaseMock.runMigrations).not.toHaveBeenCalled();
});
it(`should warn if ${extension} extension update requires restart`, async () => {
databaseMock.getExtensionVersion.mockResolvedValue({
availableVersion: updateInRange,
installedVersion: minVersionInRange,
});
databaseMock.updateVectorExtension.mockResolvedValue({ restartRequired: true });
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
expect(loggerMock.warn).toHaveBeenCalledTimes(1);
expect(loggerMock.warn.mock.calls[0][0]).toContain(extensionName);
expect(databaseMock.updateVectorExtension).toHaveBeenCalledWith(extension, updateInRange);
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
it(`should reindex ${extension} indices if needed`, async () => {
databaseMock.shouldReindex.mockResolvedValue(true);
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
expect(databaseMock.shouldReindex).toHaveBeenCalledTimes(2);
expect(databaseMock.reindex).toHaveBeenCalledTimes(2);
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
it(`should not reindex ${extension} indices if not needed`, async () => {
databaseMock.shouldReindex.mockResolvedValue(false);
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
expect(databaseMock.shouldReindex).toHaveBeenCalledTimes(2);
expect(databaseMock.reindex).toHaveBeenCalledTimes(0);
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
it('should skip migrations if DB_SKIP_MIGRATIONS=true', async () => {
process.env.DB_SKIP_MIGRATIONS = 'true';
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
expect(databaseMock.runMigrations).not.toHaveBeenCalled();
});
});
it(`should throw error if pgvector extension could not be created`, async () => {
process.env.DB_VECTOR_EXTENSION = 'pgvector';
databaseMock.getExtensionVersion.mockResolvedValue('0.0.0');
databaseMock.getExtensionVersion.mockResolvedValue({
installedVersion: null,
availableVersion: minVersionInRange,
});
databaseMock.updateVectorExtension.mockResolvedValue({ restartRequired: false });
databaseMock.createExtension.mockRejectedValue(new Error('Failed to create extension'));
await expect(sut.onBootstrapEvent()).rejects.toThrow('Failed to create extension');
expect(loggerMock.fatal).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal.mock.calls[0][0]).toContain(
'Alternatively, if your Postgres instance has pgvecto.rs, you may use this instead',
`Alternatively, if your Postgres instance has pgvecto.rs, you may use this instead`,
);
expect(databaseMock.createExtension).toHaveBeenCalledTimes(1);
expect(databaseMock.updateVectorExtension).not.toHaveBeenCalled();
expect(databaseMock.runMigrations).not.toHaveBeenCalled();
});
for (const version of ['0.2.1', '0.2.0', '0.2.9']) {
it(`should update the pgvecto.rs extension to ${version}`, async () => {
databaseMock.getAvailableExtensionVersion.mockResolvedValue(version);
databaseMock.getExtensionVersion.mockResolvedValueOnce(void 0);
databaseMock.getExtensionVersion.mockResolvedValue(version);
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
expect(databaseMock.updateVectorExtension).toHaveBeenCalledWith('vectors', version);
expect(databaseMock.updateVectorExtension).toHaveBeenCalledTimes(1);
expect(databaseMock.getExtensionVersion).toHaveBeenCalled();
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
it(`should throw error if pgvecto.rs extension could not be created`, async () => {
databaseMock.getExtensionVersion.mockResolvedValue({
installedVersion: null,
availableVersion: minVersionInRange,
});
}
databaseMock.updateVectorExtension.mockResolvedValue({ restartRequired: false });
databaseMock.createExtension.mockRejectedValue(new Error('Failed to create extension'));
for (const version of ['0.5.1', '0.6.0', '0.7.10']) {
it(`should update the pgvectors extension to ${version}`, async () => {
process.env.DB_VECTOR_EXTENSION = 'pgvector';
databaseMock.getAvailableExtensionVersion.mockResolvedValue(version);
databaseMock.getExtensionVersion.mockResolvedValueOnce(void 0);
databaseMock.getExtensionVersion.mockResolvedValue(version);
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
expect(databaseMock.updateVectorExtension).toHaveBeenCalledWith('vector', version);
expect(databaseMock.updateVectorExtension).toHaveBeenCalledTimes(1);
expect(databaseMock.getExtensionVersion).toHaveBeenCalled();
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
}
for (const version of ['0.1.0', '0.3.0', '1.0.0']) {
it(`should not upgrade pgvecto.rs to ${version}`, async () => {
databaseMock.getAvailableExtensionVersion.mockResolvedValue(version);
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
expect(databaseMock.updateVectorExtension).not.toHaveBeenCalled();
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
}
for (const version of ['0.4.0', '0.7.1', '0.7.2', '1.0.0']) {
it(`should not upgrade pgvector to ${version}`, async () => {
process.env.DB_VECTOR_EXTENSION = 'pgvector';
databaseMock.getExtensionVersion.mockResolvedValue('0.7.2');
databaseMock.getAvailableExtensionVersion.mockResolvedValue(version);
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
expect(databaseMock.updateVectorExtension).not.toHaveBeenCalled();
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
}
it(`should warn if the pgvecto.rs extension upgrade failed`, async () => {
process.env.DB_VECTOR_EXTENSION = 'pgvector';
databaseMock.getExtensionVersion.mockResolvedValue('0.5.0');
databaseMock.getAvailableExtensionVersion.mockResolvedValue('0.5.2');
databaseMock.updateVectorExtension.mockRejectedValue(new Error('Failed to update extension'));
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
expect(loggerMock.warn.mock.calls[0][0]).toContain('The pgvector extension can be updated to 0.5.2.');
expect(loggerMock.error).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
expect(databaseMock.updateVectorExtension).toHaveBeenCalledWith('vector', '0.5.2');
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
});
it(`should warn if the pgvector extension upgrade failed`, async () => {
databaseMock.getAvailableExtensionVersion.mockResolvedValue('0.2.1');
databaseMock.updateVectorExtension.mockRejectedValue(new Error('Failed to update extension'));
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
expect(loggerMock.warn.mock.calls[0][0]).toContain('The pgvecto.rs extension can be updated to 0.2.1.');
expect(loggerMock.error).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
expect(databaseMock.updateVectorExtension).toHaveBeenCalledWith('vectors', '0.2.1');
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
});
it(`should warn if the pgvecto.rs extension update requires restart`, async () => {
databaseMock.getAvailableExtensionVersion.mockResolvedValue('0.2.1');
databaseMock.updateVectorExtension.mockResolvedValue({ restartRequired: true });
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
expect(loggerMock.warn).toHaveBeenCalledTimes(1);
expect(loggerMock.warn.mock.calls[0][0]).toContain('pgvecto.rs');
expect(databaseMock.updateVectorExtension).toHaveBeenCalledWith('vectors', '0.2.1');
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
it(`should warn if the pgvector extension update requires restart`, async () => {
process.env.DB_VECTOR_EXTENSION = 'pgvector';
databaseMock.getExtensionVersion.mockResolvedValue('0.5.0');
databaseMock.getAvailableExtensionVersion.mockResolvedValue('0.5.1');
databaseMock.updateVectorExtension.mockResolvedValue({ restartRequired: true });
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
expect(loggerMock.warn).toHaveBeenCalledTimes(1);
expect(loggerMock.warn.mock.calls[0][0]).toContain('pgvector');
expect(databaseMock.updateVectorExtension).toHaveBeenCalledWith('vector', '0.5.1');
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
it('should reindex if needed', async () => {
databaseMock.shouldReindex.mockResolvedValue(true);
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
expect(databaseMock.shouldReindex).toHaveBeenCalledTimes(2);
expect(databaseMock.reindex).toHaveBeenCalledTimes(2);
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
it('should not reindex if not needed', async () => {
databaseMock.shouldReindex.mockResolvedValue(false);
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
expect(databaseMock.shouldReindex).toHaveBeenCalledTimes(2);
expect(databaseMock.reindex).toHaveBeenCalledTimes(0);
expect(databaseMock.runMigrations).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal).not.toHaveBeenCalled();
});
it('should skip migrations if DB_SKIP_MIGRATIONS=true', async () => {
process.env.DB_SKIP_MIGRATIONS = 'true';
databaseMock.getExtensionVersion.mockResolvedValue('0.2.0');
await expect(sut.onBootstrapEvent()).resolves.toBeUndefined();
await expect(sut.onBootstrapEvent()).rejects.toThrow('Failed to create extension');
expect(loggerMock.fatal).toHaveBeenCalledTimes(1);
expect(loggerMock.fatal.mock.calls[0][0]).toContain(
`Alternatively, if your Postgres instance has pgvector, you may use this instead`,
);
expect(databaseMock.createExtension).toHaveBeenCalledTimes(1);
expect(databaseMock.updateVectorExtension).not.toHaveBeenCalled();
expect(databaseMock.runMigrations).not.toHaveBeenCalled();
});
});

View file

@ -1,6 +1,5 @@
import { Inject, Injectable } from '@nestjs/common';
import semver from 'semver';
import { POSTGRES_VERSION_RANGE, VECTORS_VERSION_RANGE, VECTOR_VERSION_RANGE } from 'src/constants';
import { getVectorExtension } from 'src/database.config';
import { EventHandlerOptions } from 'src/decorators';
import {
@ -8,6 +7,7 @@ import {
DatabaseLock,
EXTENSION_NAMES,
IDatabaseRepository,
VectorExtension,
VectorIndex,
} from 'src/interfaces/database.interface';
import { OnEvents } from 'src/interfaces/event.interface';
@ -18,50 +18,46 @@ type UpdateFailedArgs = { name: string; extension: string; availableVersion: str
type RestartRequiredArgs = { name: string; availableVersion: string };
type NightlyVersionArgs = { name: string; extension: string; version: string };
type OutOfRangeArgs = { name: string; extension: string; version: string; range: string };
const EXTENSION_RANGES = {
[DatabaseExtension.VECTOR]: VECTOR_VERSION_RANGE,
[DatabaseExtension.VECTORS]: VECTORS_VERSION_RANGE,
};
type InvalidDowngradeArgs = { name: string; extension: string; installedVersion: string; availableVersion: string };
const messages = {
notInstalled: (name: string) => `Unexpected: The ${name} extension is not installed.`,
notInstalled: (name: string) =>
`The ${name} extension is not available in this Postgres instance.
If using a container image, ensure the image has the extension installed.`,
nightlyVersion: ({ name, extension, version }: NightlyVersionArgs) => `
The ${name} extension version is ${version}, which means it is a nightly release.
The ${name} extension version is ${version}, which means it is a nightly release.
Please run 'DROP EXTENSION IF EXISTS ${extension}' and switch to a release version.
See https://immich.app/docs/guides/database-queries for how to query the database.`,
outOfRange: ({ name, extension, version, range }: OutOfRangeArgs) => `
The ${name} extension version is ${version}, but Immich only supports ${range}.
Please run 'DROP EXTENSION IF EXISTS ${extension}' and switch to a release version.
See https://immich.app/docs/guides/database-queries for how to query the database.`,
outOfRange: ({ name, version, range }: OutOfRangeArgs) =>
`The ${name} extension version is ${version}, but Immich only supports ${range}.
Please change ${name} to a compatible version in the Postgres instance.`,
createFailed: ({ name, extension, otherName }: CreateFailedArgs) =>
`Failed to activate ${name} extension.
Please ensure the Postgres instance has ${name} installed.
If the Postgres instance already has a compatible version installed, Immich may not have the necessary permissions to activate it.
In this case, please run 'ALTER EXTENSION UPDATE ${extension}' manually as a superuser.
See https://immich.app/docs/guides/database-queries for how to query the database.
If the Postgres instance already has ${name} installed, Immich may not have the necessary permissions to activate it.
In this case, please run 'CREATE EXTENSION IF NOT EXISTS ${extension}' manually as a superuser.
See https://immich.app/docs/guides/database-queries for how to query the database.
Otherwise, please update the version of ${name} in the Postgres instance to a compatible version.`,
createFailed: ({ name, extension, otherName }: CreateFailedArgs) => `
Failed to activate ${name} extension.
Please ensure the Postgres instance has ${name} installed.
Alternatively, if your Postgres instance has ${otherName}, you may use this instead by setting the environment variable 'DB_VECTOR_EXTENSION=${otherName}'.
Note that switching between the two extensions after a successful startup is not supported.
The exception is if your version of Immich prior to upgrading was 1.90.2 or earlier.
In this case, you may set either extension now, but you will not be able to switch to the other extension following a successful startup.`,
updateFailed: ({ name, extension, availableVersion }: UpdateFailedArgs) =>
`The ${name} extension can be updated to ${availableVersion}.
Immich attempted to update the extension, but failed to do so.
This may be because Immich does not have the necessary permissions to update the extension.
If the Postgres instance already has ${name} installed, Immich may not have the necessary permissions to activate it.
In this case, please run 'CREATE EXTENSION IF NOT EXISTS ${extension}' manually as a superuser.
See https://immich.app/docs/guides/database-queries for how to query the database.
Alternatively, if your Postgres instance has ${otherName}, you may use this instead by setting the environment variable 'DB_VECTOR_EXTENSION=${otherName}'.
Note that switching between the two extensions after a successful startup is not supported.
The exception is if your version of Immich prior to upgrading was 1.90.2 or earlier.
In this case, you may set either extension now, but you will not be able to switch to the other extension following a successful startup.
`,
updateFailed: ({ name, extension, availableVersion }: UpdateFailedArgs) => `
The ${name} extension can be updated to ${availableVersion}.
Immich attempted to update the extension, but failed to do so.
This may be because Immich does not have the necessary permissions to update the extension.
Please run 'ALTER EXTENSION ${extension} UPDATE' manually as a superuser.
See https://immich.app/docs/guides/database-queries for how to query the database.`,
restartRequired: ({ name, availableVersion }: RestartRequiredArgs) => `
The ${name} extension has been updated to ${availableVersion}.
Please restart the Postgres instance to complete the update.`,
Please run 'ALTER EXTENSION ${extension} UPDATE' manually as a superuser.
See https://immich.app/docs/guides/database-queries for how to query the database.`,
restartRequired: ({ name, availableVersion }: RestartRequiredArgs) =>
`The ${name} extension has been updated to ${availableVersion}.
Please restart the Postgres instance to complete the update.`,
invalidDowngrade: ({ name, installedVersion, availableVersion }: InvalidDowngradeArgs) =>
`The database currently has ${name} ${installedVersion} activated, but the Postgres instance only has ${availableVersion} available.
This most likely means the extension was downgraded.
If ${name} ${installedVersion} is compatible with Immich, please ensure the Postgres instance has this available.`,
};
@Injectable()
@ -77,74 +73,90 @@ export class DatabaseService implements OnEvents {
async onBootstrapEvent() {
const version = await this.databaseRepository.getPostgresVersion();
const current = semver.coerce(version);
if (!current || !semver.satisfies(current, POSTGRES_VERSION_RANGE)) {
const postgresRange = this.databaseRepository.getPostgresVersionRange();
if (!current || !semver.satisfies(current, postgresRange)) {
throw new Error(
`Invalid PostgreSQL version. Found ${version}, but needed ${POSTGRES_VERSION_RANGE}. Please use a supported version.`,
`Invalid PostgreSQL version. Found ${version}, but needed ${postgresRange}. Please use a supported version.`,
);
}
await this.databaseRepository.withLock(DatabaseLock.Migrations, async () => {
const extension = getVectorExtension();
const otherExtension =
extension === DatabaseExtension.VECTORS ? DatabaseExtension.VECTOR : DatabaseExtension.VECTORS;
const otherName = EXTENSION_NAMES[otherExtension];
const name = EXTENSION_NAMES[extension];
const extensionRange = EXTENSION_RANGES[extension];
const extensionRange = this.databaseRepository.getExtensionVersionRange(extension);
try {
await this.databaseRepository.createExtension(extension);
} catch (error) {
this.logger.fatal(messages.createFailed({ name, extension, otherName }));
throw error;
}
const initialVersion = await this.databaseRepository.getExtensionVersion(extension);
const availableVersion = await this.databaseRepository.getAvailableExtensionVersion(extension);
const isAvailable = availableVersion && semver.satisfies(availableVersion, extensionRange);
if (isAvailable && (!initialVersion || semver.gt(availableVersion, initialVersion))) {
try {
this.logger.log(`Updating ${name} extension to ${availableVersion}`);
const { restartRequired } = await this.databaseRepository.updateVectorExtension(extension, availableVersion);
if (restartRequired) {
this.logger.warn(messages.restartRequired({ name, availableVersion }));
}
} catch (error) {
this.logger.warn(messages.updateFailed({ name, extension, availableVersion }));
this.logger.error(error);
}
}
const version = await this.databaseRepository.getExtensionVersion(extension);
if (!version) {
const { availableVersion, installedVersion } = await this.databaseRepository.getExtensionVersion(extension);
if (!availableVersion) {
throw new Error(messages.notInstalled(name));
}
if (semver.eq(version, '0.0.0')) {
throw new Error(messages.nightlyVersion({ name, extension, version }));
if ([availableVersion, installedVersion].some((version) => version && semver.eq(version, '0.0.0'))) {
throw new Error(messages.nightlyVersion({ name, extension, version: '0.0.0' }));
}
if (!semver.satisfies(version, extensionRange)) {
throw new Error(messages.outOfRange({ name, extension, version, range: extensionRange }));
if (!semver.satisfies(availableVersion, extensionRange)) {
throw new Error(messages.outOfRange({ name, extension, version: availableVersion, range: extensionRange }));
}
try {
if (await this.databaseRepository.shouldReindex(VectorIndex.CLIP)) {
await this.databaseRepository.reindex(VectorIndex.CLIP);
}
if (await this.databaseRepository.shouldReindex(VectorIndex.FACE)) {
await this.databaseRepository.reindex(VectorIndex.FACE);
}
} catch (error) {
this.logger.warn(
'Could not run vector reindexing checks. If the extension was updated, please restart the Postgres instance.',
);
throw error;
if (!installedVersion) {
await this.createExtension(extension);
}
if (installedVersion && semver.gt(availableVersion, installedVersion)) {
await this.updateExtension(extension, availableVersion);
} else if (installedVersion && !semver.satisfies(installedVersion, extensionRange)) {
throw new Error(messages.outOfRange({ name, extension, version: installedVersion, range: extensionRange }));
} else if (installedVersion && semver.lt(availableVersion, installedVersion)) {
throw new Error(messages.invalidDowngrade({ name, extension, availableVersion, installedVersion }));
}
await this.checkReindexing();
if (process.env.DB_SKIP_MIGRATIONS !== 'true') {
await this.databaseRepository.runMigrations();
}
});
}
private async createExtension(extension: DatabaseExtension) {
try {
await this.databaseRepository.createExtension(extension);
} catch (error) {
const otherExtension =
extension === DatabaseExtension.VECTORS ? DatabaseExtension.VECTOR : DatabaseExtension.VECTORS;
const name = EXTENSION_NAMES[extension];
this.logger.fatal(messages.createFailed({ name, extension, otherName: EXTENSION_NAMES[otherExtension] }));
throw error;
}
}
private async updateExtension(extension: VectorExtension, availableVersion: string) {
this.logger.log(`Updating ${EXTENSION_NAMES[extension]} extension to ${availableVersion}`);
try {
const { restartRequired } = await this.databaseRepository.updateVectorExtension(extension, availableVersion);
if (restartRequired) {
this.logger.warn(messages.restartRequired({ name: EXTENSION_NAMES[extension], availableVersion }));
}
} catch (error) {
this.logger.warn(messages.updateFailed({ name: EXTENSION_NAMES[extension], extension, availableVersion }));
throw error;
}
}
private async checkReindexing() {
try {
if (await this.databaseRepository.shouldReindex(VectorIndex.CLIP)) {
await this.databaseRepository.reindex(VectorIndex.CLIP);
}
if (await this.databaseRepository.shouldReindex(VectorIndex.FACE)) {
await this.databaseRepository.reindex(VectorIndex.FACE);
}
} catch (error) {
this.logger.warn(
'Could not run vector reindexing checks. If the extension was updated, please restart the Postgres instance.',
);
throw error;
}
}
}

View file

@ -4,8 +4,9 @@ import { Mocked, vitest } from 'vitest';
export const newDatabaseRepositoryMock = (): Mocked<IDatabaseRepository> => {
return {
getExtensionVersion: vitest.fn(),
getAvailableExtensionVersion: vitest.fn(),
getExtensionVersionRange: vitest.fn(),
getPostgresVersion: vitest.fn().mockResolvedValue('14.10 (Debian 14.10-1.pgdg120+1)'),
getPostgresVersionRange: vitest.fn().mockReturnValue('>=14.0.0'),
createExtension: vitest.fn().mockResolvedValue(void 0),
updateExtension: vitest.fn(),
updateVectorExtension: vitest.fn(),