nucliadb 6.9.2.post5276__py3-none-any.whl → 6.9.3.post5290__py3-none-any.whl

migrations/pg/0010_shards_index.py

@@ -0,0 +1,34 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+#
+
+from nucliadb.common.maindb.pg import PGTransaction
+
+
+async def migrate(txn: PGTransaction) -> None:
+    # Concurrent index must be created outside of a transaction but psycopg automatically
+    # creates transactions. We temporarily disable this for building indexes.
+    await txn.connection.commit()
+    try:
+        await txn.connection.set_autocommit(True)
+        await txn.connection.execute(
+            "CREATE INDEX CONCURRENTLY ON resources (key, value) WHERE key ~ '/kbs/[^/]*/r/[^/]*/shard$';"
+        )
+    finally:
+        await txn.connection.set_autocommit(False)

nucliadb/common/cluster/rebalance.py

@@ -22,9 +22,8 @@ import dataclasses
 import logging
 import math
 import random
-from typing import Optional
+from typing import Optional, cast
 
-import aioitertools
 from grpc import StatusCode
 from grpc.aio import AioRpcError
 from nidx_protos import nodereader_pb2, noderesources_pb2
@@ -32,7 +31,8 @@ from nidx_protos import nodereader_pb2, noderesources_pb2
 from nucliadb.common import datamanagers, locking
 from nucliadb.common.cluster.utils import get_shard_manager
 from nucliadb.common.context import ApplicationContext
-from nucliadb.common.datamanagers.resources import KB_RESOURCE_SHARD
+from nucliadb.common.maindb.driver import Driver
+from nucliadb.common.maindb.pg import PGDriver
 from nucliadb.common.nidx import get_nidx_api_client, get_nidx_searcher_client
 from nucliadb_protos import writer_pb2
 from nucliadb_telemetry import errors
@@ -68,7 +68,7 @@ class Rebalancer:
         self.context = context
         self.kbid = kbid
         self.kb_shards: Optional[writer_pb2.Shards] = None
-        self.index: dict[str, set[str]] = {}
+        self.index: dict[str, int] = {}
 
     async def get_rebalance_shards(self) -> list[RebalanceShard]:
         """
@@ -93,16 +93,7 @@ class Rebalancer:
         )
 
     async def build_shard_resources_index(self):
-        async with datamanagers.with_ro_transaction() as txn:
-            iterable = datamanagers.resources.iterate_resource_ids(kbid=self.kbid)
-            async for resources_batch in aioitertools.batched(iterable, n=200):
-                shards = await txn.batch_get(
-                    keys=[KB_RESOURCE_SHARD.format(kbid=self.kbid, uuid=rid) for rid in resources_batch],
-                    for_update=False,
-                )
-                for rid, shard_bytes in zip(resources_batch, shards):
-                    if shard_bytes is not None:
-                        self.index.setdefault(shard_bytes.decode(), set()).add(rid)
+        self.index = await build_shard_resources_index(self.context.kv_driver, self.kbid)
 
     async def move_paragraphs(
         self, from_shard: RebalanceShard, to_shard: RebalanceShard, max_paragraphs: int
@@ -113,13 +104,19 @@ class Rebalancer:
         """
         moved_paragraphs = 0
 
+        resources_batch: list[str] = []
+
         while moved_paragraphs < max_paragraphs:
+            if len(resources_batch) == 0:
+                resources_batch = await get_resources_from_shard(
+                    self.context.kv_driver, self.kbid, from_shard.id, n=50
+                )
+                if len(resources_batch) == 0:
+                    # No more resources to move or shard not found
+                    break
+
             # Take a random resource to move
-            try:
-                resource_id = random.choice(tuple(self.index[from_shard.id]))
-            except (KeyError, IndexError):
-                # No more resources in shard or shard not found
-                break
+            resource_id = random.choice(resources_batch)
 
             assert self.kb_shards is not None
             from_shard_obj = next(s for s in self.kb_shards.shards if s.shard == from_shard.id)
@@ -129,8 +126,9 @@ class Rebalancer:
                 self.context, self.kbid, resource_id, from_shard_obj, to_shard_obj
             )
             if moved:
-                self.index[from_shard.id].remove(resource_id)
-                self.index.setdefault(to_shard.id, set()).add(resource_id)
+                resources_batch.remove(resource_id)
+                self.index[from_shard.id] = self.index.get(from_shard.id, 1) - 1
+                self.index[to_shard.id] = self.index.get(to_shard.id, 0) + 1
                 moved_paragraphs += paragraphs_count
 
         return moved_paragraphs
@@ -261,7 +259,7 @@ class Rebalancer:
         empty_shard = False
 
         for _ in range(MAX_MOVES_PER_SHARD):
-            resources_count = len(self.index.get(shard_to_merge.id, []))
+            resources_count = self.index.get(shard_to_merge.id, 0)
             if resources_count == 0:
                 logger.info(
                     "Shard is now empty",
@@ -306,14 +304,14 @@ class Rebalancer:
         if empty_shard:
             # Build the index again, and make sure there is no resource assigned to this shard
             await self.build_shard_resources_index()
-            shard_resources = self.index.get(shard_to_merge.id, set())
-            if len(shard_resources) > 0:
+            shard_resources = self.index.get(shard_to_merge.id, 0)
+            if shard_resources > 0:
                 logger.error(
                     f"Shard expected to be empty, but it isn't. Won't be deleted.",
                     extra={
                         "kbid": self.kbid,
                         "shard": shard_to_merge.id,
-                        "resources": list(shard_resources)[:30],
+                        "resources": shard_resources,
                     },
                 )
                 return
@@ -356,6 +354,40 @@ class Rebalancer:
                 )
 
 
+async def build_shard_resources_index(driver: Driver, kbid: str) -> dict[str, int]:
+    index: dict[str, int] = {}
+    driver = cast(PGDriver, driver)
+    async with driver._get_connection() as conn:
+        cur = conn.cursor("")
+        await cur.execute(
+            """
+            SELECT encode(value, 'escape'), COUNT(*) FROM resources WHERE key ~ '/kbs/[^/]*/r/[^/]*/shard$' AND key ~ %s GROUP BY value;
+            """,
+            (f"/kbs/{kbid}/r/[^/]*/shard$",),
+        )
+        records = await cur.fetchall()
+        shard: str
+        resources_count: int
+        for shard, resources_count in records:
+            index[shard] = resources_count
+        return index
+
+
+async def get_resources_from_shard(driver: Driver, kbid: str, shard_id: str, n: int) -> list[str]:
+    driver = cast(PGDriver, driver)
+    async with driver._get_connection() as conn:
+        cur = conn.cursor("")
+        await cur.execute(
+            """
+            SELECT split_part(key, '/', 5) FROM resources WHERE key ~ '/kbs/[^/]*/r/[^/]*/shard$' AND key ~ %s AND encode(value, 'escape') LIKE %s limit %s;
+            """,
+            (f"/kbs/{kbid}/r/[^/]*/shard$", shard_id, n),
+        )
+        records = await cur.fetchall()
+        rids: list[str] = [r[0] for r in records]
+        return rids
+
+
 async def get_resource_paragraphs_count(resource_id: str, nidx_shard_id: str) -> int:
     # Do a search on the fields (paragraph) index and return the number of paragraphs this resource has
     try:

nucliadb/ingest/fields/conversation.py

@@ -25,7 +25,7 @@ from nucliadb_protos.resources_pb2 import CloudFile, FieldConversation, SplitMet
 from nucliadb_protos.resources_pb2 import Conversation as PBConversation
 from nucliadb_utils.storages.storage import StorageField
 
-MAX_CONVERSATION_MESSAGES = 50 * 1024
+MAX_CONVERSATION_MESSAGES = None  # No limit
 
 PAGE_SIZE = 200
 

nucliadb/writer/resource/field.py

@@ -564,7 +564,10 @@ async def _conversation_append_checks(
 
         # Make sure that the max number of messages is not exceeded
         current_message_count = (await conv.get_metadata()).total
-        if len(input.messages) + current_message_count > MAX_CONVERSATION_MESSAGES:
+        if (
+            MAX_CONVERSATION_MESSAGES is not None
+            and (len(input.messages) + current_message_count) > MAX_CONVERSATION_MESSAGES
+        ):
             raise HTTPException(
                 status_code=422,
                 detail=f"Conversation fields cannot have more than {MAX_CONVERSATION_MESSAGES} messages.",

{nucliadb-6.9.2.post5276.dist-info → nucliadb-6.9.3.post5290.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.9.2.post5276
+Version: 6.9.3.post5290
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License-Expression: AGPL-3.0-or-later
@@ -18,11 +18,11 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Requires-Python: <4,>=3.10
 Description-Content-Type: text/markdown
-Requires-Dist: nucliadb-telemetry[all]>=6.9.2.post5276
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.9.2.post5276
-Requires-Dist: nucliadb-protos>=6.9.2.post5276
-Requires-Dist: nucliadb-models>=6.9.2.post5276
-Requires-Dist: nidx-protos>=6.9.2.post5276
+Requires-Dist: nucliadb-telemetry[all]>=6.9.3.post5290
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.9.3.post5290
+Requires-Dist: nucliadb-protos>=6.9.3.post5290
+Requires-Dist: nucliadb-models>=6.9.3.post5290
+Requires-Dist: nidx-protos>=6.9.3.post5290
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.50.0
 Requires-Dist: uvicorn[standard]

{nucliadb-6.9.2.post5276.dist-info → nucliadb-6.9.3.post5290.dist-info}/RECORD

@@ -46,6 +46,7 @@ migrations/pg/0006_catalog_title_indexes.py,sha256=n2OGxwE4oeCwHAYaxBkja4t10BmwT
 migrations/pg/0007_catalog_slug.py,sha256=mArzZCBO-RD5DkWxRIyDKgEzrnAcis1TOGvSNUe7Kgg,1150
 migrations/pg/0008_catalog_facets.py,sha256=dxIUdHJHtI_Gyk2dpP7tjHEnL2iPzAufi6ajYm2FVMI,1595
 migrations/pg/0009_extract_facets_safety.py,sha256=k9Appx7ipp3wDyLy70qgw9oLjN7N6BEadE-N5Fhan-4,1066
+migrations/pg/0010_shards_index.py,sha256=7s7c2s5768BkOqexJ2CedzJq8SigGKC9rfOigtQOraA,1417
 migrations/pg/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/__init__.py,sha256=_abCmDJ_0ku483Os4UAjPX7Nywm39cQgAV_DiyjsKeQ,891
 nucliadb/health.py,sha256=UIxxA4oms4HIsCRZM_SZsdkIZIlgzmOxw-qSHLlWuak,3465
@@ -86,7 +87,7 @@ nucliadb/common/cluster/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIX
 nucliadb/common/cluster/exceptions.py,sha256=t7v_l93t44l2tQpdQXgO_w-c4YZRcaayOz1A2i0w4RQ,1258
 nucliadb/common/cluster/grpc_node_dummy.py,sha256=JkufazWzMA4KFEU8EBkMbiiDW4C8lLcRhiiCxP7aCQY,2949
 nucliadb/common/cluster/manager.py,sha256=p-haaGEnCa-20t-I2XEo4QJ5ZC1QgJ6p2jzXFYVB6nQ,12346
-nucliadb/common/cluster/rebalance.py,sha256=0RIyrv8kkTHJFpqC5H3fPukypF0HpRNiN5f4WRiFvh8,22541
+nucliadb/common/cluster/rebalance.py,sha256=XT8APmvvjwr5pbMPqnVcFj1ZehSQwOP5US-cKkEyEag,23559
 nucliadb/common/cluster/rollover.py,sha256=kmVCdyjJ1dilnSodHMqf0EUkTjPC5H0aA4JqW9KsEa4,27168
 nucliadb/common/cluster/settings.py,sha256=f6Y5K0PGahkedwe5wtkWMnbqwDFJgOOwX_MOIGwH9Dg,2271
 nucliadb/common/cluster/utils.py,sha256=E4GqidwTKczJX_lTnncBCof2fL4CFVVF1eLiz9NWjlc,5494
@@ -155,7 +156,7 @@ nucliadb/ingest/consumer/shard_creator.py,sha256=qUEpxZLE1etw1nL8L3O9HvZBx5NNql7
 nucliadb/ingest/consumer/utils.py,sha256=jpX8D4lKzuPCpArQLZeX_Zczq3pfen_zAf8sPJfOEZU,2642
 nucliadb/ingest/fields/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/ingest/fields/base.py,sha256=Yk6b7OQ96YAAaDEyckry6-dBhzXASzuJucjTpEvFUZ4,23715
-nucliadb/ingest/fields/conversation.py,sha256=KkOvNM1rZFQRg2RsfGd3Jrz3lpx0HpGpN1cmlpz_mZw,9563
+nucliadb/ingest/fields/conversation.py,sha256=QbHBtjOV_KWeJu_nrP8DTQ6uTc-yDAXliiJ1f2LienY,9570
 nucliadb/ingest/fields/exceptions.py,sha256=sZBk21BSrXFdOdo1qUdCAyD-9YMYakSLdn4_WdIPCIQ,1217
 nucliadb/ingest/fields/file.py,sha256=1v4jLg3balUua2VmSV8hHkAwPFShTUCOzufZvIUQcQw,4740
 nucliadb/ingest/fields/generic.py,sha256=elgtqv15aJUq3zY7X_g0bli_2BpcwPArVvzhe54Y4Ig,1547
@@ -374,7 +375,7 @@ nucliadb/writer/api/v1/vectorsets.py,sha256=F3iMViL5G95_Tns4aO2SOA0DwAzxK2_P8MXx
 nucliadb/writer/resource/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/writer/resource/audit.py,sha256=FvxMZPzrNHtd31HgpZEvxzwAkbxJTZRhPLqRYYJi3tA,1426
 nucliadb/writer/resource/basic.py,sha256=44GK8M9EEVoAUfGiabdLrrpENqeFwNn7qwxF2AHhQGg,10504
-nucliadb/writer/resource/field.py,sha256=kJFxOgmizGbEuTRPb5o0cNqonZ8sa9ehVlSfRk-ektY,22866
+nucliadb/writer/resource/field.py,sha256=FF1cDuvDPZN6TrwiMQiyGqxNU4bGxXYpVcHsM-BNpKs,22946
 nucliadb/writer/resource/origin.py,sha256=pvhUDdU0mlWPUcpoQi4LDUJaRtfjzVVrA8XcGVI_N8k,2021
 nucliadb/writer/tus/__init__.py,sha256=Kera0BtxoDX0ngPftXiMjNgjrhtQ3l2XFc5nJqSBOJY,5498
 nucliadb/writer/tus/azure.py,sha256=yxoRi4PhGDikTqVK3PiuVyguy8H9DOS66JpZCY4hpUY,4177
@@ -385,8 +386,8 @@ nucliadb/writer/tus/local.py,sha256=7jYa_w9b-N90jWgN2sQKkNcomqn6JMVBOVeDOVYJHto,
 nucliadb/writer/tus/s3.py,sha256=vu1BGg4VqJ_x2P1u2BxqPKlSfw5orT_a3R-Ln5oPUpU,8483
 nucliadb/writer/tus/storage.py,sha256=ToqwjoYnjI4oIcwzkhha_MPxi-k4Jk3Lt55zRwaC1SM,2903
 nucliadb/writer/tus/utils.py,sha256=MSdVbRsRSZVdkaum69_0wku7X3p5wlZf4nr6E0GMKbw,2556
-nucliadb-6.9.2.post5276.dist-info/METADATA,sha256=92ehiGl-d963I9_9rxJo1lyNxd7om9lcGXMtunuwhc8,4118
-nucliadb-6.9.2.post5276.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nucliadb-6.9.2.post5276.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.9.2.post5276.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.9.2.post5276.dist-info/RECORD,,
+nucliadb-6.9.3.post5290.dist-info/METADATA,sha256=4vQ4BF0XZCMBSLDcbFox0HLi4sYZxqM3sOzsdM1xH5w,4118
+nucliadb-6.9.3.post5290.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nucliadb-6.9.3.post5290.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.9.3.post5290.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.9.3.post5290.dist-info/RECORD,,