jack 1 месяц назад
Родитель
Сommit
de34460249

+ 2 - 2
py/core/providers/database/chunks.py

@@ -671,7 +671,7 @@ class PostgresChunksHandler(Handler):
         query = f"""
         UPDATE {self._get_table_name(PostgresChunksHandler.TABLE_NAME)}
         SET collection_ids = array_append(collection_ids, $1)
-        WHERE document_id = $2 AND NOT ($1 = ANY(collection_ids));
+        WHERE document_id = $2 AND NOT (collection_ids @> ARRAY[$1]::uuid[]);
         """
         logger.info(query)
         logger.info(f"collection_id: {collection_id}")
@@ -702,7 +702,7 @@ class PostgresChunksHandler(Handler):
     async def delete_collection_vector(self, collection_id: UUID) -> None:
         query = f"""
          DELETE FROM {self._get_table_name(PostgresChunksHandler.TABLE_NAME)}
-         WHERE $1 = ANY(collection_ids)
+         WHERE collection_ids @> ARRAY[$1]::uuid[]
          RETURNING collection_ids
          """
         await self.connection_manager.fetchrow_query(query, (collection_id,))

+ 4 - 4
py/core/providers/database/collections.py

@@ -258,7 +258,7 @@ class PostgresCollectionsHandler(Handler):
         user_update_query = f"""
             UPDATE {self._get_table_name("users")}
             SET collection_ids = array_remove(collection_ids, $1)
-            WHERE $1 = ANY(collection_ids)
+            WHERE collection_ids @> ARRAY[$1]::uuid[]
         """
         logger.info(f"Deleting collection {collection_id} from users")
         await self.connection_manager.execute_query(
@@ -271,7 +271,7 @@ class PostgresCollectionsHandler(Handler):
             WITH updated AS (
                 UPDATE {self._get_table_name("documents")}
                 SET collection_ids = array_remove(collection_ids, $1)
-                WHERE $1 = ANY(collection_ids)
+                WHERE collection_ids @> ARRAY[$1]::uuid[]
                 RETURNING 1
             )
             SELECT COUNT(*) AS affected_rows FROM updated
@@ -468,7 +468,7 @@ class PostgresCollectionsHandler(Handler):
             assign_query = f"""
                 UPDATE {self._get_table_name("documents")}
                 SET collection_ids = array_append(collection_ids, $1)
-                WHERE id = $2 AND NOT ($1 = ANY(collection_ids))
+                WHERE id = $2 AND NOT (collection_ids @> ARRAY[$1]::uuid[])
                 RETURNING id
             """
             result = await self.connection_manager.fetchrow_query(
@@ -520,7 +520,7 @@ class PostgresCollectionsHandler(Handler):
         query = f"""
             UPDATE {self._get_table_name("documents")}
             SET collection_ids = array_remove(collection_ids, $1)
-            WHERE id = $2 AND $1 = ANY(collection_ids)
+            WHERE id = $2 AND collection_ids @> ARRAY[$1]::uuid[]
             RETURNING id
         """
         result = await self.connection_manager.fetchrow_query(

+ 3 - 3
py/core/providers/database/users.py

@@ -709,7 +709,7 @@ class PostgresUserHandler(Handler):
         query = f"""
             UPDATE {self._get_table_name(PostgresUserHandler.TABLE_NAME)}
             SET collection_ids = array_append(collection_ids, $1)
-            WHERE id = $2 AND NOT ($1 = ANY(collection_ids))
+            WHERE id = $2 AND NOT (collection_ids @> ARRAY[$1]::uuid[])
             RETURNING id
         """
         result = await self.connection_manager.fetchrow_query(
@@ -741,7 +741,7 @@ class PostgresUserHandler(Handler):
         query = f"""
             UPDATE {self._get_table_name(PostgresUserHandler.TABLE_NAME)}
             SET collection_ids = array_remove(collection_ids, $1)
-            WHERE id = $2 AND $1 = ANY(collection_ids)
+            WHERE id = $2 AND collection_ids @> ARRAY[$1]::uuid[]
             RETURNING id
         """
         result = await self.connection_manager.fetchrow_query(
@@ -785,7 +785,7 @@ class PostgresUserHandler(Handler):
                     "COUNT(*) OVER() AS total_entries",
                 ]
             )
-            .where("$1 = ANY(collection_ids)")
+            .where("collection_ids @> ARRAY[$1]::uuid[]")
             .order_by("name")
             .offset("$2")
             .limit("$3" if limit != -1 else None)