diff --git a/server/src/infra/database.config.ts b/server/src/infra/database.config.ts index 089fd6878..7e197b51e 100644 --- a/server/src/infra/database.config.ts +++ b/server/src/infra/database.config.ts @@ -23,5 +23,30 @@ export const databaseConfig: PostgresConnectionOptions = { ...urlOrParts, }; +export const databaseConfigVector: PostgresConnectionOptions = { + ...databaseConfig, + migrationsRun: false, + migrations: [__dirname + '/migrations/vector/*.{js,ts}'], +}; + +export async function initDataSource(): Promise { + const dataSource = await new DataSource(databaseConfig).initialize(); + + const hasVectorExtension = (await dataSource.query( + `SELECT * FROM pg_extension WHERE name = 'vectors'`, + )).length > 0; + + if (hasVectorExtension) { + const dataSourceVector = await new DataSource(databaseConfigVector).initialize(); + await dataSourceVector.runMigrations(); + await dataSourceVector.destroy(); + } + + return dataSource; +} + // this export is used by TypeORM commands in package.json#scripts -export const dataSource = new DataSource(databaseConfig); +export let dataSource: DataSource; +(async () => { + dataSource = await initDataSource(); +})(); diff --git a/server/src/infra/migrations/1699746198141-UsePgVector.ts b/server/src/infra/migrations/1699746198141-UsePgVector.ts deleted file mode 100644 index 362e547e2..000000000 --- a/server/src/infra/migrations/1699746198141-UsePgVector.ts +++ /dev/null @@ -1,47 +0,0 @@ -import { MigrationInterface, QueryRunner } from 'typeorm'; - -export class UsePgVector1699746198141 implements MigrationInterface { - name = 'UsePgVector1699746198141'; - - public async up(queryRunner: QueryRunner): Promise { - await queryRunner.query('DROP EXTENSION IF EXISTS vectors'); - await queryRunner.query('CREATE EXTENSION vectors'); - - const faceDimQuery = await queryRunner.query(` - SELECT CARDINALITY(embedding) as dimsize - FROM asset_faces - LIMIT 1`); - const clipDimQuery = await queryRunner.query(` - SELECT CARDINALITY("clipEmbedding") as dimsize - FROM smart_info - LIMIT 1`); - - const faceDimSize = faceDimQuery?.[0]?.['dimsize'] ?? 512; - const clipDimSize = clipDimQuery?.[0]?.['dimsize'] ?? 512; - - await queryRunner.query(`ALTER TABLE asset_faces ALTER COLUMN embedding TYPE vector(${faceDimSize})`); - await queryRunner.query(`CREATE TABLE smart_search ( - "assetId" uuid PRIMARY KEY NOT NULL REFERENCES assets(id) ON DELETE CASCADE, - embedding vector(${clipDimSize}) NOT NULL )`); - await queryRunner.query(` - INSERT INTO smart_search("assetId", embedding) - SELECT si."assetId", si."clipEmbedding" - FROM smart_info si - WHERE "clipEmbedding" IS NOT NULL - `); - await queryRunner.query(`ALTER TABLE smart_info DROP COLUMN "clipEmbedding"`); - } - - public async down(queryRunner: QueryRunner): Promise { - await queryRunner.query(`ALTER TABLE asset_faces ALTER COLUMN embedding TYPE real array`); - await queryRunner.query(`ALTER TABLE smart_info ADD COLUMN IF NOT EXISTS "clipEmbedding" TYPE real array`); - await queryRunner.query(` - INSERT INTO smart_info - ("assetId", "clipEmbedding") - SELECT s."assetId", s.embedding - FROM smart_search s - ON CONFLICT (s."assetId") DO UPDATE SET "clipEmbedding" = s.embedding - `); - await queryRunner.query(`DROP TABLE IF EXISTS smart_search`); - } -} diff --git a/server/src/infra/migrations/vector/1699746198141-UsePgVectors.ts b/server/src/infra/migrations/vector/1699746198141-UsePgVectors.ts new file mode 100644 index 000000000..16e85ee30 --- /dev/null +++ b/server/src/infra/migrations/vector/1699746198141-UsePgVectors.ts @@ -0,0 +1,61 @@ +import { MigrationInterface, QueryRunner } from 'typeorm'; + +export class UsePgVectors1699746198141 implements MigrationInterface { + name = 'UsePgVectors1699746198141'; + + public async up(queryRunner: QueryRunner): Promise { + const faceDimQuery = await queryRunner.query(` + SELECT CARDINALITY(embedding) as dimsize + FROM asset_faces + LIMIT 1`); + const clipDimQuery = await queryRunner.query(` + SELECT CARDINALITY("clipEmbedding") as dimsize + FROM smart_info + LIMIT 1`); + + const faceDimSize = faceDimQuery?.[0]?.['dimsize'] ?? 512; + const clipDimSize = clipDimQuery?.[0]?.['dimsize'] ?? 512; + + await queryRunner.query('DROP EXTENSION IF EXISTS vectors'); + await queryRunner.query('CREATE EXTENSION vectors'); + + await queryRunner.query(` + BEGIN; + + ALTER TABLE asset_faces ALTER COLUMN embedding TYPE vector(${faceDimSize}); + + CREATE TABLE smart_search ( + "assetId" uuid PRIMARY KEY NOT NULL REFERENCES assets(id) ON DELETE CASCADE, + embedding vector(${clipDimSize}) NOT NULL ); + + INSERT INTO smart_search("assetId", embedding) + SELECT si."assetId", si."clipEmbedding" + FROM smart_info si + WHERE "clipEmbedding" IS NOT NULL; + + ALTER TABLE smart_info DROP COLUMN "clipEmbedding"; + + COMMIT; + `); + } + + public async down(queryRunner: QueryRunner): Promise { + await queryRunner.query(` + BEGIN; + + ALTER TABLE asset_faces ALTER COLUMN embedding TYPE real array; + + ALTER TABLE smart_info ADD COLUMN IF NOT EXISTS "clipEmbedding" TYPE real array; + + INSERT INTO smart_info + ("assetId", "clipEmbedding") + SELECT s."assetId", s.embedding + FROM smart_search s + ON CONFLICT (s."assetId") DO UPDATE SET "clipEmbedding" = s.embedding; + + DROP TABLE IF EXISTS smart_search; + + COMMIT; + `); + } +} diff --git a/server/src/infra/migrations/1699746301742-AddCLIPEmbeddingIndex.ts b/server/src/infra/migrations/vector/1699746301742-AddCLIPEmbeddingIndex.ts similarity index 100% rename from server/src/infra/migrations/1699746301742-AddCLIPEmbeddingIndex.ts rename to server/src/infra/migrations/vector/1699746301742-AddCLIPEmbeddingIndex.ts diff --git a/server/src/infra/migrations/1699746444644-AddFaceEmbeddingIndex.ts b/server/src/infra/migrations/vector/1699746444644-AddFaceEmbeddingIndex.ts similarity index 100% rename from server/src/infra/migrations/1699746444644-AddFaceEmbeddingIndex.ts rename to server/src/infra/migrations/vector/1699746444644-AddFaceEmbeddingIndex.ts