nucliadb 6.3.3.post3608__py3-none-any.whl → 6.3.3.post3611__py3-none-any.whl

nucliadb/purge/__init__.py

@@ -37,6 +37,7 @@ from nucliadb.ingest.orm.knowledgebox import (
     RESOURCE_TO_DELETE_STORAGE_BASE,
     KnowledgeBox,
 )
+from nucliadb.tasks.retries import purge_metadata as purge_task_metadata
 from nucliadb_protos.knowledgebox_pb2 import VectorSetConfig, VectorSetPurge
 from nucliadb_telemetry import errors
 from nucliadb_telemetry.logs import setup_logging
@@ -270,6 +271,7 @@ async def main():
         service_name=SERVICE_NAME,
     )
     try:
+        purge_task_metadata_task = asyncio.create_task(purge_task_metadata(driver))
         purge_resources_storage_task = asyncio.create_task(
             purge_deleted_resource_storage(driver, storage)
         )
@@ -277,6 +279,7 @@ async def main():
         await purge_kb_storage(driver, storage)
         await purge_kb_vectorsets(driver, storage)
         await purge_resources_storage_task
+        await purge_task_metadata_task
     except Exception as ex:  # pragma: no cover
         logger.exception("Unhandled exception on purge command")
         errors.capture_exception(ex)

nucliadb/tasks/retries.py

@@ -19,12 +19,15 @@
 #
 import functools
 import logging
+from datetime import datetime, timezone
 from enum import Enum
-from typing import Callable, Optional
+from typing import Callable, Optional, cast
 
 from pydantic import BaseModel
 
 from nucliadb.common.context import ApplicationContext
+from nucliadb.common.maindb.driver import Driver
+from nucliadb.common.maindb.pg import PGDriver, PGTransaction
 
 logger = logging.getLogger(__name__)
 
@@ -41,6 +44,7 @@ class TaskMetadata(BaseModel):
     status: Status
     retries: int = 0
     error_messages: list[str] = []
+    last_modified: Optional[datetime] = None
 
 
 class TaskRetryHandler:
@@ -84,16 +88,10 @@ class TaskRetryHandler:
         )
 
     async def get_metadata(self) -> Optional[TaskMetadata]:
-        async with self.context.kv_driver.transaction(read_only=True) as txn:
-            metadata = await txn.get(self.metadata_key)
-            if metadata is None:
-                return None
-            return TaskMetadata.model_validate_json(metadata)
+        return await _get_metadata(self.context.kv_driver, self.metadata_key)
 
     async def set_metadata(self, metadata: TaskMetadata) -> None:
-        async with self.context.kv_driver.transaction() as txn:
-            await txn.set(self.metadata_key, metadata.model_dump_json().encode())
-            await txn.commit()
+        await _set_metadata(self.context.kv_driver, self.metadata_key, metadata)
 
     def wrap(self, func: Callable) -> Callable:
         @functools.wraps(func)
@@ -106,6 +104,7 @@ class TaskRetryHandler:
                     task_id=self.task_id,
                     status=TaskMetadata.Status.RUNNING,
                     retries=0,
+                    last_modified=datetime.now(timezone.utc),
                 )
                 await self.set_metadata(metadata)
 
@@ -123,6 +122,7 @@ class TaskRetryHandler:
                     f"Task reached max retries. Setting to FAILED state",
                     extra={"kbid": self.kbid, "task_type": self.task_type, "task_id": self.task_id},
                 )
+                metadata.last_modified = datetime.now(timezone.utc)
                 await self.set_metadata(metadata)
                 return
             try:
@@ -144,6 +144,91 @@ class TaskRetryHandler:
                 metadata.status = TaskMetadata.Status.COMPLETED
                 return func_result
             finally:
+                metadata.last_modified = datetime.now(timezone.utc)
                 await self.set_metadata(metadata)
 
         return wrapper
+
+
+async def _get_metadata(kv_driver: Driver, metadata_key: str) -> Optional[TaskMetadata]:
+    async with kv_driver.transaction(read_only=True) as txn:
+        metadata = await txn.get(metadata_key)
+        if metadata is None:
+            return None
+        return TaskMetadata.model_validate_json(metadata)
+
+
+async def _set_metadata(kv_driver: Driver, metadata_key: str, metadata: TaskMetadata) -> None:
+    async with kv_driver.transaction() as txn:
+        await txn.set(metadata_key, metadata.model_dump_json().encode())
+        await txn.commit()
+
+
+async def purge_metadata(kv_driver: Driver) -> int:
+    """
+    Purges old task metadata records that are in a final state and older than 15 days.
+    Returns the total number of records purged.
+    """
+    if not isinstance(kv_driver, PGDriver):
+        return 0
+
+    total_purged = 0
+    start: Optional[str] = ""
+    while True:
+        start, purged = await purge_batch(kv_driver, start)
+        total_purged += purged
+        if start is None:
+            break
+    return total_purged
+
+
+async def purge_batch(
+    kv_driver: PGDriver, start: Optional[str] = None, batch_size: int = 200
+) -> tuple[Optional[str], int]:
+    """
+    Returns the next start key and the number of purged records. If start is None, it means there are no more records to purge.
+    """
+    async with kv_driver.transaction() as txn:
+        txn = cast(PGTransaction, txn)
+        async with txn.connection.cursor() as cur:
+            await cur.execute(
+                """
+                SELECT key from resources
+                WHERE key  ~ '^/kbs/[^/]*/tasks/[^/]*/[^/]*$'
+                AND key > %s
+                ORDER BY key
+                LIMIT %s
+                """,
+                (start, batch_size),
+            )
+            records = await cur.fetchall()
+            keys = [r[0] for r in records]
+
+    if not keys:
+        # No more records to purge
+        return None, 0
+
+    to_delete = []
+    for key in keys:
+        metadata = await _get_metadata(kv_driver, key)
+        if metadata is None:  # pragma: no cover
+            continue
+        task_finished = metadata.status in (TaskMetadata.Status.COMPLETED, TaskMetadata.Status.FAILED)
+        old_task = (
+            metadata.last_modified is None
+            or (datetime.now(timezone.utc) - metadata.last_modified).days >= 15
+        )
+        if task_finished and old_task:
+            to_delete.append(key)
+
+    n_to_delete = len(to_delete)
+    delete_batch_size = 50
+    while len(to_delete) > 0:
+        batch = to_delete[:delete_batch_size]
+        to_delete = to_delete[delete_batch_size:]
+        async with kv_driver.transaction() as txn:
+            for key in batch:
+                logger.info("Purging task metadata", extra={"key": key})
+                await txn.delete(key)
+            await txn.commit()
+    return keys[-1], n_to_delete

{nucliadb-6.3.3.post3608.dist-info → nucliadb-6.3.3.post3611.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: nucliadb
-Version: 6.3.3.post3608
+Version: 6.3.3.post3611
 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.3.post3608
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.3.3.post3608
-Requires-Dist: nucliadb-protos>=6.3.3.post3608
-Requires-Dist: nucliadb-models>=6.3.3.post3608
-Requires-Dist: nidx-protos>=6.3.3.post3608
+Requires-Dist: nucliadb-telemetry[all]>=6.3.3.post3611
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.3.3.post3611
+Requires-Dist: nucliadb-protos>=6.3.3.post3611
+Requires-Dist: nucliadb-models>=6.3.3.post3611
+Requires-Dist: nidx-protos>=6.3.3.post3611
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn

{nucliadb-6.3.3.post3608.dist-info → nucliadb-6.3.3.post3611.dist-info}/RECORD

@@ -165,7 +165,7 @@ nucliadb/migrator/settings.py,sha256=jOUX0ZMunCXN8HpF9xXN0aunJYRhu4Vdr_ffjRIqwtw
 nucliadb/migrator/utils.py,sha256=NgUreUvON8_nWEzTxELBMWlfV7E6-6qi-g0DMEbVEz4,2885
 nucliadb/models/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/models/responses.py,sha256=qnuOoc7TrVSUnpikfTwHLKez47_DE4mSFzpxrwtqijA,1599
-nucliadb/purge/__init__.py,sha256=BphuNvsJ1aSwuVXUcSOaK4nj9pDcpuKRBf_QAcoRj-A,11787
+nucliadb/purge/__init__.py,sha256=p-DxJ3vGtEcCuiSsWY22bQn99vTMHk32BQdQJVy5b0k,11983
 nucliadb/purge/orphan_shards.py,sha256=AU1Jfc4qtQFasB6dkuGz0a_Zjs-i7liUTf7Xpl8R8ng,7939
 nucliadb/reader/__init__.py,sha256=C5Efic7WlGm2U2C5WOyquMFbIj2Pojwe_8mwzVYnOzE,1304
 nucliadb/reader/app.py,sha256=Se-BFTE6d1v1msLzQn4q5XIhjnSxa2ckDSHdvm7NRf8,3096
@@ -276,7 +276,7 @@ nucliadb/tasks/consumer.py,sha256=4CWfBdXVr2a25n7seldbQ0PaK0FcxJZuWgosU6aODS8,69
 nucliadb/tasks/logger.py,sha256=C7keOEO_mjLVp5VbqAZ2QXfqVB2Hot7NgBlUP_SDSMw,924
 nucliadb/tasks/models.py,sha256=qrZKi5DNDQ07waMsp5L4_Fi7WRs57YiO-kmXlrBzEAA,1168
 nucliadb/tasks/producer.py,sha256=UnpJAzhj_GElsCoO5G6T4m6MshsgOaqR2tVzJmEta64,2625
-nucliadb/tasks/retries.py,sha256=tLNtwAutGEfFV3Adr4a9ew-Wg4vuxBd72GGDE_Mma8s,5160
+nucliadb/tasks/retries.py,sha256=L8jbnkyNe6LVoXAmpdtl4QLoMNI8NUXOna7u0sSM0_Y,8157
 nucliadb/tasks/utils.py,sha256=tV1AbWdFc3qfIULX44Veqj41FCD1B6XYjG6brULBeiw,1459
 nucliadb/tests/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/tests/config.py,sha256=JN_Jhgj-fwM9_8IeO9pwxr6C1PiwRDrXxm67Y38rU30,2080
@@ -351,8 +351,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.3.post3608.dist-info/METADATA,sha256=D39vPx8WWL01vqSIJ7XLuZ5j4QTwhJWsYECLccPs6lk,4291
-nucliadb-6.3.3.post3608.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
-nucliadb-6.3.3.post3608.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.3.3.post3608.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.3.3.post3608.dist-info/RECORD,,
+nucliadb-6.3.3.post3611.dist-info/METADATA,sha256=k4CK1ALRNK5NsSVw0Lo-VG8RNfvYELnUJbd2pxziQzE,4291
+nucliadb-6.3.3.post3611.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
+nucliadb-6.3.3.post3611.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.3.3.post3611.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.3.3.post3611.dist-info/RECORD,,