nucliadb 6.3.1.post3492__py3-none-any.whl → 6.3.1.post3496__py3-none-any.whl

nucliadb/export_import/datamanager.py

@@ -19,7 +19,7 @@
 #
 import json
 from datetime import datetime
-from typing import AsyncGenerator, Union
+from typing import AsyncGenerator, Type, Union, cast
 
 from nucliadb.common.maindb.driver import Driver
 from nucliadb.export_import import logger
@@ -59,11 +59,12 @@ class ExportImportDataManager:
         if data is None or data == b"":
             raise MetadataNotFound()
         decoded = data.decode("utf-8")
+        model_type: Union[Type[ExportMetadata], Type[ImportMetadata]]
         if type == "export":
             model_type = ExportMetadata
         elif type == "import":
-            model_type = ImportMetadata  # type: ignore
-        else:
+            model_type = ImportMetadata
+        else:  # pragma: no cover
             raise ValueError(f"Invalid type: {type}")
         json_decoded = json.loads(decoded)
 
@@ -73,9 +74,11 @@ class ExportImportDataManager:
             json_decoded["total"] = 0
         if json_decoded.get("processed") is None:
             json_decoded["processed"] = 0
-
         return model_type.model_validate(json_decoded)
 
+    async def get_export_metadata(self, kbid: str, id: str) -> ExportMetadata:
+        return cast(ExportMetadata, await self.get_metadata("export", kbid, id))
+
     async def set_metadata(
         self,
         type: str,

nucliadb/export_import/exporter.py

@@ -76,8 +76,8 @@ async def export_kb_to_blob_storage(context: ApplicationContext, msg: NatsTaskMe
     """
     kbid, export_id = msg.kbid, msg.id
     dm = ExportImportDataManager(context.kv_driver, context.blob_storage)
-    metadata = await dm.get_metadata(type="export", kbid=kbid, id=export_id)
-    iterator = export_kb(context, kbid, metadata)  # type: ignore
+    metadata = await dm.get_export_metadata(kbid=kbid, id=export_id)
+    iterator = export_kb(context, kbid, metadata)
 
     retry_handler = TaskRetryHandler("export", dm, metadata)
 

nucliadb/export_import/tasks.py

@@ -27,40 +27,48 @@ from nucliadb.tasks.producer import NatsTaskProducer
 from nucliadb_utils import const
 
 
-def get_exports_consumer() -> NatsTaskConsumer:
+def get_exports_consumer() -> NatsTaskConsumer[NatsTaskMessage]:
     return create_consumer(
         name="exports_consumer",
-        stream=const.Streams.KB_EXPORTS,  # type: ignore
-        callback=export_kb_to_blob_storage,  # type: ignore
+        stream=const.Streams.KB_EXPORTS.name,
+        stream_subjects=[const.Streams.KB_EXPORTS.subject],
+        consumer_subject=const.Streams.KB_EXPORTS.subject,
+        callback=export_kb_to_blob_storage,
         msg_type=NatsTaskMessage,
         max_concurrent_messages=10,
     )
 
 
-async def get_exports_producer(context: ApplicationContext) -> NatsTaskProducer:
+async def get_exports_producer(context: ApplicationContext) -> NatsTaskProducer[NatsTaskMessage]:
     producer = create_producer(
         name="exports_producer",
-        stream=const.Streams.KB_EXPORTS,  # type: ignore
+        stream=const.Streams.KB_EXPORTS.name,
+        stream_subjects=[const.Streams.KB_EXPORTS.subject],
+        producer_subject=const.Streams.KB_EXPORTS.subject,
         msg_type=NatsTaskMessage,
     )
     await producer.initialize(context)
     return producer
 
 
-def get_imports_consumer() -> NatsTaskConsumer:
+def get_imports_consumer() -> NatsTaskConsumer[NatsTaskMessage]:
     return create_consumer(
         name="imports_consumer",
-        stream=const.Streams.KB_IMPORTS,  # type: ignore
-        callback=import_kb_from_blob_storage,  # type: ignore
+        stream=const.Streams.KB_IMPORTS.name,
+        stream_subjects=[const.Streams.KB_IMPORTS.subject],
+        consumer_subject=const.Streams.KB_IMPORTS.subject,
+        callback=import_kb_from_blob_storage,
         msg_type=NatsTaskMessage,
         max_concurrent_messages=10,
     )
 
 
-async def get_imports_producer(context: ApplicationContext) -> NatsTaskProducer:
+async def get_imports_producer(context: ApplicationContext) -> NatsTaskProducer[NatsTaskMessage]:
     producer = create_producer(
         name="imports_producer",
-        stream=const.Streams.KB_IMPORTS,  # type: ignore
+        stream=const.Streams.KB_IMPORTS.name,
+        stream_subjects=[const.Streams.KB_IMPORTS.subject],
+        producer_subject=const.Streams.KB_IMPORTS.subject,
         msg_type=NatsTaskMessage,
     )
     await producer.initialize(context)

nucliadb/ingest/app.py

@@ -214,6 +214,8 @@ def run_subscriber_workers() -> None:  # pragma: no cover
         - shard creator subscriber
         - audit counter subscriber
         - audit fields subscriber
+        - export/import subscriber
+        - materializer subscriber
     """
     setup_configuration()
     asyncio.run(main_subscriber_workers())

nucliadb/tasks/__init__.py

@@ -17,6 +17,5 @@
 # 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 .consumer import NatsTaskConsumer, create_consumer, start_consumer  # noqa
-from .producer import NatsTaskProducer, create_producer, get_producer  # noqa
-from .registry import register_task  # noqa
+from .consumer import NatsTaskConsumer, create_consumer  # noqa
+from .producer import NatsTaskProducer, create_producer  # noqa

nucliadb/tasks/consumer.py

@@ -19,7 +19,7 @@
 #
 
 import asyncio
-from typing import Optional
+from typing import Generic, Optional, Type
 
 import nats
 import pydantic
@@ -27,42 +27,41 @@ from nats.aio.client import Msg
 
 from nucliadb.common.context import ApplicationContext
 from nucliadb.tasks.logger import logger
-from nucliadb.tasks.models import MsgType
-from nucliadb.tasks.registry import get_registered_task
-from nucliadb.tasks.utils import TaskCallback, create_nats_stream_if_not_exists
+from nucliadb.tasks.models import Callback, MsgType
+from nucliadb.tasks.utils import create_nats_stream_if_not_exists
 from nucliadb_telemetry import errors
-from nucliadb_utils import const
 from nucliadb_utils.nats import MessageProgressUpdater
 from nucliadb_utils.settings import nats_consumer_settings
 
 BEFORE_NAK_SLEEP_SECONDS = 2
 
 
-class NatsTaskConsumer:
+class NatsTaskConsumer(Generic[MsgType]):
     def __init__(
         self,
         name: str,
-        stream: const.Streams,
-        callback: TaskCallback,
-        msg_type: MsgType,
+        stream: str,
+        stream_subjects: list[str],
+        consumer_subject: str,
+        callback: Callback,
+        msg_type: Type[MsgType],
         max_concurrent_messages: Optional[int] = None,
     ):
         self.name = name
         self.stream = stream
+        self.stream_subjects = stream_subjects
+        self.consumer_subject = consumer_subject
         self.callback = callback
         self.msg_type = msg_type
         self.max_concurrent_messages = max_concurrent_messages
         self.initialized = False
-        self.context: Optional[ApplicationContext] = None
         self.running_tasks: list[asyncio.Task] = []
         self.subscription = None
 
     async def initialize(self, context: ApplicationContext):
         self.context = context
         await create_nats_stream_if_not_exists(
-            self.context,
-            self.stream.name,  # type: ignore
-            subjects=[self.stream.subject],  # type: ignore
+            context, stream_name=self.stream, subjects=self.stream_subjects
         )
         await self._setup_nats_subscription()
         self.initialized = True
@@ -81,9 +80,8 @@ class NatsTaskConsumer:
 
     async def _setup_nats_subscription(self):
         # Nats push consumer
-        subject = self.stream.subject
-        group = self.stream.group
-        stream = self.stream.name
+        stream = self.stream
+        subject = group = self.consumer_subject
         max_ack_pending = (
             self.max_concurrent_messages
             if self.max_concurrent_messages
@@ -146,7 +144,7 @@ class NatsTaskConsumer:
 
             logger.info(f"Starting task consumption", extra={"consumer_name": self.name})
             try:
-                await self.callback(self.context, task_msg)  # type: ignore
+                await self.callback(self.context, task_msg)
             except asyncio.CancelledError:
                 logger.debug(
                     f"Task cancelled. Naking and exiting...",
@@ -180,38 +178,23 @@ class NatsTaskConsumer:
 
 def create_consumer(
     name: str,
-    stream: const.Streams,
-    callback: TaskCallback,
-    msg_type: MsgType,
+    stream: str,
+    stream_subjects: list[str],
+    consumer_subject: str,
+    callback: Callback,
+    msg_type: Type[MsgType],
     max_concurrent_messages: Optional[int] = None,
-) -> NatsTaskConsumer:
+) -> NatsTaskConsumer[MsgType]:
     """
     Returns a non-initialized consumer
     """
     consumer = NatsTaskConsumer(
         name=name,
         stream=stream,
+        stream_subjects=stream_subjects,
+        consumer_subject=consumer_subject,
         callback=callback,
         msg_type=msg_type,
         max_concurrent_messages=max_concurrent_messages,
     )
     return consumer
-
-
-async def start_consumer(task_name: str, context: ApplicationContext) -> NatsTaskConsumer:
-    """
-    Returns an initialized consumer for the given task name, ready to consume messages from the task stream.
-    """
-    try:
-        task = get_registered_task(task_name)
-    except KeyError:
-        raise ValueError(f"Task {task_name} not registered")
-    consumer = create_consumer(
-        name=f"{task_name}_consumer",
-        stream=task.stream,
-        callback=task.callback,  # type: ignore
-        msg_type=task.msg_type,
-        max_concurrent_messages=task.max_concurrent_messages,
-    )
-    await consumer.initialize(context)
-    return consumer

nucliadb/tasks/models.py

@@ -17,23 +17,13 @@
 # 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 dataclasses import dataclass
-from typing import Any, Callable, Coroutine, Optional, Type
+from typing import Any, Callable, Coroutine, TypeVar
 
 import pydantic
 
 from nucliadb.common.context import ApplicationContext
-from nucliadb_utils import const
 
-MsgType = Type[pydantic.BaseModel]
+MsgType = TypeVar("MsgType", bound=pydantic.BaseModel)
 
 # async def callback(context: ApplicationContext, msg: MyPydanticModel):
 Callback = Callable[[ApplicationContext, MsgType], Coroutine[Any, Any, Any]]
-
-
-@dataclass
-class RegisteredTask:
-    stream: const.Streams
-    callback: Callback
-    msg_type: MsgType
-    max_concurrent_messages: Optional[int] = None

nucliadb/tasks/producer.py

@@ -17,26 +17,28 @@
 # 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 typing import Optional
+from typing import Generic, Optional, Type
 
 from nucliadb.common.context import ApplicationContext
 from nucliadb.tasks.logger import logger
 from nucliadb.tasks.models import MsgType
-from nucliadb.tasks.registry import get_registered_task
 from nucliadb.tasks.utils import create_nats_stream_if_not_exists
 from nucliadb_telemetry import errors
-from nucliadb_utils import const
 
 
-class NatsTaskProducer:
+class NatsTaskProducer(Generic[MsgType]):
     def __init__(
         self,
         name: str,
-        stream: const.Streams,
-        msg_type: MsgType,
+        stream: str,
+        stream_subjects: list[str],
+        producer_subject: str,
+        msg_type: Type[MsgType],
     ):
         self.name = name
         self.stream = stream
+        self.stream_subjects = stream_subjects
+        self.producer_subject = producer_subject
         self.msg_type = msg_type
         self.context: Optional[ApplicationContext] = None
         self.initialized = False
@@ -45,12 +47,12 @@ class NatsTaskProducer:
         self.context = context
         await create_nats_stream_if_not_exists(
             self.context,
-            self.stream.name,  # type: ignore
-            subjects=[self.stream.subject],  # type: ignore
+            self.stream,
+            subjects=self.stream_subjects,
         )
         self.initialized = True
 
-    async def __call__(self, msg: MsgType) -> int:
+    async def send(self, msg: MsgType) -> int:
         """
         Publish message to the producer's nats stream.
         Returns the sequence number of the published message.
@@ -59,7 +61,7 @@ class NatsTaskProducer:
             raise RuntimeError("NatsTaskProducer not initialized")
         try:
             pub_ack = await self.context.nats_manager.js.publish(  # type: ignore
-                self.stream.subject,  # type: ignore
+                self.producer_subject,
                 msg.model_dump_json().encode("utf-8"),  # type: ignore
             )
             logger.info(
@@ -79,23 +81,19 @@ class NatsTaskProducer:
 
 def create_producer(
     name: str,
-    stream: const.Streams,
-    msg_type: MsgType,
-) -> NatsTaskProducer:
+    stream: str,
+    stream_subjects: list[str],
+    producer_subject: str,
+    msg_type: Type[MsgType],
+) -> NatsTaskProducer[MsgType]:
     """
     Returns a non-initialized producer.
     """
-    return NatsTaskProducer(name=name, stream=stream, msg_type=msg_type)
-
-
-async def get_producer(task_name: str, context: ApplicationContext) -> NatsTaskProducer:
-    """
-    Returns a producer for the given task name, ready to be used to send messages to the task stream.
-    """
-    try:
-        task = get_registered_task(task_name)
-    except KeyError:
-        raise ValueError(f"Task {task_name} not registered")
-    producer = create_producer(name=f"{task_name}_producer", stream=task.stream, msg_type=task.msg_type)
-    await producer.initialize(context)
+    producer = NatsTaskProducer[MsgType](
+        name=name,
+        stream=stream,
+        stream_subjects=stream_subjects,
+        producer_subject=producer_subject,
+        msg_type=msg_type,
+    )
     return producer

nucliadb/tasks/utils.py

@@ -18,15 +18,10 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import Awaitable, Callable
-
 import nats
-from pydantic import BaseModel
 
 from nucliadb.common.context import ApplicationContext
 
-TaskCallback = Callable[[ApplicationContext, BaseModel], Awaitable[None]]
-
 
 async def create_nats_stream_if_not_exists(
     context: ApplicationContext, stream_name: str, subjects: list[str]

nucliadb/writer/api/v1/export_import.py

@@ -197,7 +197,7 @@ async def start_export_task(context: ApplicationContext, kbid: str, export_id: s
     try:
         producer = await get_exports_producer(context)
         msg = NatsTaskMessage(kbid=kbid, id=export_id)
-        seqid = await producer(msg)  # type: ignore
+        seqid = await producer.send(msg)
         logger.info(f"Export task produced. seqid={seqid} kbid={kbid} export_id={export_id}")
     except Exception as e:
         errors.capture_exception(e)
@@ -214,7 +214,7 @@ async def start_import_task(context: ApplicationContext, kbid: str, import_id: s
     try:
         producer = await get_imports_producer(context)
         msg = NatsTaskMessage(kbid=kbid, id=import_id)
-        seqid = await producer(msg)  # type: ignore
+        seqid = await producer.send(msg)
         logger.info(f"Import task produced. seqid={seqid} kbid={kbid} import_id={import_id}")
     except Exception as e:
         errors.capture_exception(e)

{nucliadb-6.3.1.post3492.dist-info → nucliadb-6.3.1.post3496.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: nucliadb
-Version: 6.3.1.post3492
+Version: 6.3.1.post3496
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License: AGPL
@@ -20,11 +20,11 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Requires-Python: <4,>=3.9
 Description-Content-Type: text/markdown
-Requires-Dist: nucliadb-telemetry[all]>=6.3.1.post3492
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.3.1.post3492
-Requires-Dist: nucliadb-protos>=6.3.1.post3492
-Requires-Dist: nucliadb-models>=6.3.1.post3492
-Requires-Dist: nidx-protos>=6.3.1.post3492
+Requires-Dist: nucliadb-telemetry[all]>=6.3.1.post3496
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.3.1.post3496
+Requires-Dist: nucliadb-protos>=6.3.1.post3496
+Requires-Dist: nucliadb-models>=6.3.1.post3496
+Requires-Dist: nidx-protos>=6.3.1.post3496
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn

{nucliadb-6.3.1.post3492.dist-info → nucliadb-6.3.1.post3496.dist-info}/RECORD

@@ -95,15 +95,15 @@ nucliadb/common/models_utils/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJ
 nucliadb/common/models_utils/from_proto.py,sha256=d-q6-JWvjVVxtI57itQVKRIY8Hzv-tBhDt-tr4aIKNM,15979
 nucliadb/common/models_utils/to_proto.py,sha256=97JvOR_3odu50YvzLa2CERfEN3w_QPmAVcCJwJB5m5A,2438
 nucliadb/export_import/__init__.py,sha256=y-Is0Bxa8TMV6UiOW0deC_D3U465P65CQ5RjBjIWnow,932
-nucliadb/export_import/datamanager.py,sha256=ZL96M6WuyrFLcGdoM6Cb6hQwrGFRtSeRlsKuvDjV3sY,6722
+nucliadb/export_import/datamanager.py,sha256=xL8b0xvk45q6wx1l7J32JgPmpyjnF6fKiJi2F2B_UQY,6968
 nucliadb/export_import/exceptions.py,sha256=Dw8WqfG4r6MPJc5TPfbjMvCgXXWTcTOecGHRVU1h3kM,1949
-nucliadb/export_import/exporter.py,sha256=kgbW-B7FNW7mlc9rBVEfwkkFTqD58TWSTDe9zkmEnBc,7098
+nucliadb/export_import/exporter.py,sha256=k2QVx1EjqFlDYiggriWiEJzwtMXzHbldsqWdpGQM3_U,7074
 nucliadb/export_import/importer.py,sha256=v5cq9Nn8c2zrY_K_00mydR52f8mdFxR7tLdtNLQ0qvk,4229
 nucliadb/export_import/models.py,sha256=dbjScNkiMRv4X3Ktudy1JRliD25bfoDTy3JmEZgQSCc,2121
-nucliadb/export_import/tasks.py,sha256=fpCBeFYPReyLIdk38LDM9Tpnw_VczeMrobT4n1RAIp4,2507
+nucliadb/export_import/tasks.py,sha256=yPNdBdvTD7eGc7zvV9Rp7UZ0-mDhA34OOsLqHvns_v0,2975
 nucliadb/export_import/utils.py,sha256=t7xLA3f5W3zGq3HNXe3noOQnY7gRO8TAoe8S8BG3_so,19733
 nucliadb/ingest/__init__.py,sha256=fsw3C38VP50km3R-nHL775LNGPpJ4JxqXJ2Ib1f5SqE,1011
-nucliadb/ingest/app.py,sha256=n5PlZyrf-bJsrc6jMRSO4DTjGyZLb4q6EA6Y4x6TYfQ,7028
+nucliadb/ingest/app.py,sha256=ErNd3q44xbBcUOl-Ae_lvcKPAsfFMSzb8zqWAjekrM4,7097
 nucliadb/ingest/cache.py,sha256=w7jMMzamOmQ7gwXna6Dqm6isRNBVv6l5BTBlTxaYWjE,1005
 nucliadb/ingest/partitions.py,sha256=2NIhMYbNT0TNBL6bX1UMSi7vxFGICstCKEqsB0TXHOE,2410
 nucliadb/ingest/processing.py,sha256=gg1DqbMFwqdOsmCSGsZc2abRdYz86xOZJun9vrHOCzs,20618
@@ -260,13 +260,12 @@ nucliadb/standalone/versions.py,sha256=8CxNMNt2NgWM8ct50UsR4d44-ae7wtQI-sV-yGiFq
 nucliadb/standalone/static/favicon.ico,sha256=96pKGp6Sx457JkTfjy1dtApMhkitixfU6invCUGAYOU,2285
 nucliadb/standalone/static/index.html,sha256=PEZfuEQFYnYACAL1ceN8xC0im8lBrUx838RkE8tbvgA,3833
 nucliadb/standalone/static/logo.svg,sha256=-wQqSvPGTdlKjUP6pHE6kiq005pgYjDzp9nPl0X71Mk,2639
-nucliadb/tasks/__init__.py,sha256=ie5S4LUZ3Pu16FKOva9CM3JLwHqT7g41y4n1-sJ_oB4,1037
-nucliadb/tasks/consumer.py,sha256=u4c9gonpu1OuwX1G1gz3lalK5JjTA7A0p_wywpOHS2I,7632
+nucliadb/tasks/__init__.py,sha256=vruCOMmCu0BcAZQrEinlgiQtiR1WYxSxvI5UsydAopc,963
+nucliadb/tasks/consumer.py,sha256=x-999Nsw6lBcKvyGyCGPiGP_naANVYMfl9M-u0U3mhY,7052
 nucliadb/tasks/logger.py,sha256=C7keOEO_mjLVp5VbqAZ2QXfqVB2Hot7NgBlUP_SDSMw,924
-nucliadb/tasks/models.py,sha256=hzbPqaq-SZUz4cTRSgIy0G8hBOYBaYBCNgH2KtLibVg,1378
-nucliadb/tasks/producer.py,sha256=_aTEulet1ebGhQMhoD3SlfHtfhZHNCbBfOf8uKyNpWk,3512
-nucliadb/tasks/registry.py,sha256=tKvv_Py_O3peuLshq4bB3w197E33O7P7B63_6loLMrA,1753
-nucliadb/tasks/utils.py,sha256=abXcE_bjAkO4UwrKXhZykFPsRwbtwnuc0reNWhYImjw,1360
+nucliadb/tasks/models.py,sha256=qrZKi5DNDQ07waMsp5L4_Fi7WRs57YiO-kmXlrBzEAA,1168
+nucliadb/tasks/producer.py,sha256=w4R1YXgXtmCPGcoNNOr3qkqJYcHJtSmix-xjt7vsPqk,3261
+nucliadb/tasks/utils.py,sha256=6tQVckqyzxv8PhVAd3ZqcMYpGcn73ZY6p1cpm1FxagA,1214
 nucliadb/tests/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/tests/config.py,sha256=JN_Jhgj-fwM9_8IeO9pwxr6C1PiwRDrXxm67Y38rU30,2080
 nucliadb/tests/vectors.py,sha256=CcNKx-E8LPpyvRyljbmb-Tn_wST9Juw2CBoogWrKiTk,62843
@@ -315,7 +314,7 @@ nucliadb/writer/api/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20
 nucliadb/writer/api/constants.py,sha256=qWEDjFUycrEZnSJyLnNK4PQNodU2oVmkO4NycaEZtio,1738
 nucliadb/writer/api/utils.py,sha256=wIQHlU8RQiIGVLI72suvyVIKlCU44Unh0Ae0IiN6Qwo,1313
 nucliadb/writer/api/v1/__init__.py,sha256=akI9A_jloNLb0dU4T5zjfdyvmSAiDeIdjAlzNx74FlU,1128
-nucliadb/writer/api/v1/export_import.py,sha256=Cv4DNtqoR_x2LMNx80C9ehyaCaFIxhq2eutEk2uzcZg,8249
+nucliadb/writer/api/v1/export_import.py,sha256=FxGY_rofj3AH-HQReoMDOYfEpACeWB5bsmxrT03QUtQ,8227
 nucliadb/writer/api/v1/field.py,sha256=OsWOYA0WQ6onE5Rkl20QIEdtrSi7Jgnu62fUt90Ziy8,17503
 nucliadb/writer/api/v1/knowledgebox.py,sha256=MLeIuym4jPrJgfy1NTcN9CpUGwuBiqDHMcx0hY9DR7g,9530
 nucliadb/writer/api/v1/learning_config.py,sha256=CKBjqcbewkfPwGUPLDWzZSpro6XkmCaVppe5Qtpu5Go,3117
@@ -340,8 +339,8 @@ nucliadb/writer/tus/local.py,sha256=7jYa_w9b-N90jWgN2sQKkNcomqn6JMVBOVeDOVYJHto,
 nucliadb/writer/tus/s3.py,sha256=vF0NkFTXiXhXq3bCVXXVV-ED38ECVoUeeYViP8uMqcU,8357
 nucliadb/writer/tus/storage.py,sha256=ToqwjoYnjI4oIcwzkhha_MPxi-k4Jk3Lt55zRwaC1SM,2903
 nucliadb/writer/tus/utils.py,sha256=MSdVbRsRSZVdkaum69_0wku7X3p5wlZf4nr6E0GMKbw,2556
-nucliadb-6.3.1.post3492.dist-info/METADATA,sha256=xDXk1LjHJfFz-tVPOziUfXKnHVpl3EU2gao7cD8xwYI,4291
-nucliadb-6.3.1.post3492.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
-nucliadb-6.3.1.post3492.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.3.1.post3492.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.3.1.post3492.dist-info/RECORD,,
+nucliadb-6.3.1.post3496.dist-info/METADATA,sha256=5MWHuNYen2cYlCYbpGOrnkEXINxqHOxutzk5cR0iimk,4291
+nucliadb-6.3.1.post3496.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
+nucliadb-6.3.1.post3496.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.3.1.post3496.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.3.1.post3496.dist-info/RECORD,,

nucliadb/tasks/registry.py

@@ -1,59 +0,0 @@
-# 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/>.
-#
-import functools
-from typing import Optional
-
-import pydantic
-
-from nucliadb.tasks.models import RegisteredTask
-from nucliadb_utils import const
-
-REGISTRY: dict[str, RegisteredTask] = dict()
-
-
-def register_task(
-    name: str,
-    stream: const.Streams,
-    msg_type: pydantic.BaseModel,
-    max_concurrent_messages: Optional[int] = None,
-):
-    """
-    Decorator to register a task in the registry.
-    """
-
-    def decorator_register(func):
-        REGISTRY[name] = RegisteredTask(
-            callback=func,
-            stream=stream,
-            msg_type=msg_type,
-            max_concurrent_messages=max_concurrent_messages,
-        )
-
-        @functools.wraps(func)
-        async def wrapper_register(*args, **kwargs):
-            return await func(*args, **kwargs)
-
-        return wrapper_register
-
-    return decorator_register
-
-
-def get_registered_task(name: str) -> RegisteredTask:
-    return REGISTRY[name]