From 0b88bef1570f023a7ffd5bbffcb5d06a3697f51d Mon Sep 17 00:00:00 2001
From: Mert <101130780+mertalev@users.noreply.github.com>
Date: Wed, 3 Jul 2024 22:02:43 -0400
Subject: [PATCH] fix(server): face search migration sometimes failing (#10827)

* turn it off and back on

* handle missing smart search embedding column

* handle missing face embedding column

* simplify

* Revert "simplify"

This reverts commit 8322af0baf2be8d59f868db9868b9694f999e615.

* fix migration
---
 .../1718486162779-AddFaceSearchRelation.ts    | 61 +++++++++++++------
 1 file changed, 43 insertions(+), 18 deletions(-)

diff --git a/server/src/migrations/1718486162779-AddFaceSearchRelation.ts b/server/src/migrations/1718486162779-AddFaceSearchRelation.ts
index 5bf3fcd97b..c8e02ec0c5 100644
--- a/server/src/migrations/1718486162779-AddFaceSearchRelation.ts
+++ b/server/src/migrations/1718486162779-AddFaceSearchRelation.ts
@@ -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 "embedding"`);
+    await queryRunner.query(`ALTER TABLE asset_faces DROP COLUMN IF EXISTS 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 face_index ON face_search
-            USING hnsw (embedding vector_cosine_ops)
-            WITH (ef_construction = 300, m = 16)`);
+      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)`);
   }
 
   public async down(queryRunner: QueryRunner): Promise<void> {
@@ -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)`);
   }
 }