Skip to content

Commit

Permalink
fix(server): face search migration sometimes failing (#10827)
Browse files Browse the repository at this point in the history
* turn it off and back on

* handle missing smart search embedding column

* handle missing face embedding column

* simplify

* Revert "simplify"

This reverts commit 8322af0.

* fix migration
  • Loading branch information
mertalev authored Jul 4, 2024
1 parent 2b89420 commit 0b88bef
Showing 1 changed file with 43 additions and 18 deletions.
61 changes: 43 additions & 18 deletions server/src/migrations/1718486162779-AddFaceSearchRelation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,25 +9,50 @@ export class AddFaceSearchRelation1718486162779 implements MigrationInterface {
await queryRunner.query(`SET vectors.pgvector_compatibility=on`);
}

const hasEmbeddings = async (tableName: string): Promise<boolean> => {
const columns = await queryRunner.query(
`SELECT column_name as name
FROM information_schema.columns
WHERE table_name = '${tableName}'`);
return columns.some((column: { name: string }) => column.name === 'embedding');
}

const hasAssetEmbeddings = await hasEmbeddings('smart_search');
if (!hasAssetEmbeddings) {
await queryRunner.query(`TRUNCATE smart_search`);
await queryRunner.query(`ALTER TABLE smart_search ADD COLUMN IF NOT EXISTS embedding vector(512) NOT NULL`);
}

await queryRunner.query(`
CREATE TABLE face_search (
"faceId" uuid PRIMARY KEY REFERENCES asset_faces(id) ON DELETE CASCADE,
embedding vector(512) NOT NULL )`);
CREATE TABLE face_search (
"faceId" uuid PRIMARY KEY REFERENCES asset_faces(id) ON DELETE CASCADE,
embedding vector(512) NOT NULL )`);

await queryRunner.query(`ALTER TABLE face_search ALTER COLUMN embedding SET STORAGE EXTERNAL`);
await queryRunner.query(`ALTER TABLE smart_search ALTER COLUMN embedding SET STORAGE EXTERNAL`);

await queryRunner.query(`
INSERT INTO face_search("faceId", embedding)
SELECT id, embedding
FROM asset_faces faces`);
const hasFaceEmbeddings = await hasEmbeddings('asset_faces')
if (hasFaceEmbeddings) {
await queryRunner.query(`
INSERT INTO face_search("faceId", embedding)
SELECT id, embedding
FROM asset_faces faces`);
}

await queryRunner.query(`ALTER TABLE asset_faces DROP COLUMN IF EXISTS embedding`);

await queryRunner.query(`ALTER TABLE asset_faces DROP COLUMN "embedding"`);
await queryRunner.query(`ALTER TABLE face_search ALTER COLUMN embedding SET DATA TYPE real[]`);
await queryRunner.query(`ALTER TABLE face_search ALTER COLUMN embedding SET DATA TYPE vector(512)`);

await queryRunner.query(`
CREATE INDEX IF NOT EXISTS clip_index ON smart_search
USING hnsw (embedding vector_cosine_ops)
WITH (ef_construction = 300, m = 16)`);

await queryRunner.query(`
CREATE INDEX face_index ON face_search
USING hnsw (embedding vector_cosine_ops)
WITH (ef_construction = 300, m = 16)`);
CREATE INDEX face_index ON face_search
USING hnsw (embedding vector_cosine_ops)
WITH (ef_construction = 300, m = 16)`);
}

public async down(queryRunner: QueryRunner): Promise<void> {
Expand All @@ -40,15 +65,15 @@ export class AddFaceSearchRelation1718486162779 implements MigrationInterface {
await queryRunner.query(`ALTER TABLE face_search ALTER COLUMN embedding SET STORAGE DEFAULT`);
await queryRunner.query(`ALTER TABLE smart_search ALTER COLUMN embedding SET STORAGE DEFAULT`);
await queryRunner.query(`
UPDATE asset_faces
SET embedding = fs.embedding
FROM face_search fs
WHERE id = fs."faceId"`);
UPDATE asset_faces
SET embedding = fs.embedding
FROM face_search fs
WHERE id = fs."faceId"`);
await queryRunner.query(`DROP TABLE face_search`);

await queryRunner.query(`
CREATE INDEX face_index ON asset_faces
USING hnsw (embedding vector_cosine_ops)
WITH (ef_construction = 300, m = 16)`);
CREATE INDEX face_index ON asset_faces
USING hnsw (embedding vector_cosine_ops)
WITH (ef_construction = 300, m = 16)`);
}
}

0 comments on commit 0b88bef

Please sign in to comment.