nucliadb 6.9.1.post5180__py3-none-any.whl → 6.9.2.post5282__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/manager.py

@@ -43,8 +43,6 @@ from nucliadb_protos import knowledgebox_pb2, writer_pb2
 from nucliadb_telemetry import errors
 from nucliadb_utils.utilities import get_storage
 
-from .settings import settings
-
 logger = logging.getLogger(__name__)
 
 
@@ -113,6 +111,8 @@ class KBShardManager:
         self,
         txn: Transaction,
         kbid: str,
+        *,
+        prewarm_enabled: bool,
     ) -> writer_pb2.ShardObject:
         kb_shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid, for_update=True)
         if kb_shards is None:
@@ -133,6 +133,7 @@ class KBShardManager:
             req = NewShardRequest(
                 kbid=kbid,
                 vectorsets_configs=vectorsets,
+                prewarm_enabled=prewarm_enabled,
             )
 
             resp = await nidx_api.NewShard(req)  # type: ignore
@@ -232,23 +233,6 @@ class KBShardManager:
         indexpb.shard = shard.nidx_shard_id
         await nidx.index(indexpb)
 
-    def should_create_new_shard(self, num_paragraphs: int) -> bool:
-        return num_paragraphs > settings.max_shard_paragraphs
-
-    async def maybe_create_new_shard(
-        self,
-        kbid: str,
-        num_paragraphs: int,
-    ):
-        if not self.should_create_new_shard(num_paragraphs):
-            return
-
-        logger.info({"message": "Adding shard", "kbid": kbid})
-
-        async with datamanagers.with_transaction() as txn:
-            await self.create_shard_by_kbid(txn, kbid)
-            await txn.commit()
-
     async def create_vectorset(self, kbid: str, config: knowledgebox_pb2.VectorSetConfig):
         """Create a new vectorset in all KB shards."""