chore(server): startup check for pgvecto.rs (#5815)

* startup check for pgvecto.rs

* prefilter after assertion

* formatting

* add assert to migration

* more specific import

* use runner
This commit is contained in:
Mert
2023-12-18 11:38:25 -05:00
committed by GitHub
parent fade8b627f
commit de1514a441
5 changed files with 53 additions and 15 deletions

View File

@@ -1,4 +1,4 @@
import { DataSource } from 'typeorm';
import { DataSource, QueryRunner } from 'typeorm';
import { PostgresConnectionOptions } from 'typeorm/driver/postgres/PostgresConnectionOptions';
const url = process.env.DB_URL;
@@ -26,9 +26,50 @@ export const databaseConfig: PostgresConnectionOptions = {
// this export is used by TypeORM commands in package.json#scripts
export const dataSource = new DataSource(databaseConfig);
export async function enablePrefilter() {
export async function databaseChecks() {
if (!dataSource.isInitialized) {
await dataSource.initialize();
}
await dataSource.query(`SET vectors.enable_prefilter = on`);
await assertVectors(dataSource);
await enablePrefilter(dataSource);
await dataSource.runMigrations();
}
export async function enablePrefilter(runner: DataSource | QueryRunner) {
await runner.query(`SET vectors.enable_prefilter = on`);
}
export async function getExtensionVersion(extName: string, runner: DataSource | QueryRunner): Promise<string | null> {
const res = await runner.query(`SELECT extversion FROM pg_extension WHERE extname = $1`, [extName]);
return res[0]?.['extversion'] ?? null;
}
export async function getPostgresVersion(runner: DataSource | QueryRunner): Promise<string> {
const res = await runner.query(`SHOW server_version`);
return res[0]['server_version'].split('.')[0];
}
export async function assertVectors(runner: DataSource | QueryRunner) {
const postgresVersion = await getPostgresVersion(runner);
const expected = ['0.1.1', '0.1.11'];
const image = `tensorchord/pgvecto-rs:pg${postgresVersion}-v${expected[expected.length - 1]}`;
await runner.query('CREATE EXTENSION IF NOT EXISTS vectors').catch((err) => {
console.error(
'Failed to create pgvecto.rs extension. ' +
`If you have not updated your Postgres instance to an image that supports pgvecto.rs (such as '${image}'), please do so. ` +
'See the v1.91.0 release notes for more info: https://github.com/immich-app/immich/releases/tag/v1.91.0',
);
throw err;
});
const version = await getExtensionVersion('vectors', runner);
if (version != null && !expected.includes(version)) {
throw new Error(
`The pgvecto.rs extension version is ${version} instead of the expected version ${
expected[expected.length - 1]
}.` + `If you're using the 'latest' tag, please switch to '${image}'.`,
);
}
}