django-lambda-tasks 0.4.1__tar.gz → 0.4.3__tar.gz

{django_lambda_tasks-0.4.1 → django_lambda_tasks-0.4.3}/.kiro/steering/product.md

@@ -136,10 +136,11 @@ class SQSLambdaTaskMessage(BaseModel):
 - Returns `{"batchItemFailures": [...]}` for partial-batch failure reporting
 - Only pre-execution failures (malformed message, import error, misconfiguration) are reported as `batchItemFailures` — task logic failures are caught and recorded as `FAILED` TaskRecords without raising
 - Recommended SQS queue settings: `maxReceiveCount=1` with a DLQ configured; automatic retries are not useful since task failures are not re-driven by design
-- Cold-start sequence runs inside the handler on the first invocation (not at module import time) to avoid Lambda init-duration timeouts: a temporary `StreamHandler` is attached to the `lambda_tasks` logger, then `resolve_environment()` → `resolve_secrets_into_env()` (handler removed) → conditional `django.setup()`
+- Cold-start sequence runs inside the handler on the first invocation (not at module import time) to avoid Lambda init-duration timeouts: memory limit is set, a temporary `StreamHandler` is attached to the `lambda_tasks` logger, then `resolve_environment()` → `resolve_secrets_into_env()` (handler removed) → conditional `django.setup()`
 - A module-level `_cold_start_done` sentinel ensures the sequence runs only once; subsequent warm invocations skip it
 - Both loaders run unconditionally (outside the `DJANGO_SETTINGS_MODULE` check) — the environment secret may provide that var, and individual secrets may depend on environment-loaded vars
 - A temporary `StreamHandler` is attached to the `lambda_tasks` logger for the duration of the loaders so their log output is visible before Django's `LOGGING` dictConfig has run; it is removed immediately after so that Django's configuration is the sole authority on logging from that point on
+- `resource.setrlimit(RLIMIT_DATA)` is set from `AWS_LAMBDA_FUNCTION_MEMORY_SIZE` so that excessive allocation raises `MemoryError` instead of triggering the OOM killer
 
 ## Environment Loader
 

{django_lambda_tasks-0.4.1 → django_lambda_tasks-0.4.3}/.kiro/steering/structure.md

@@ -42,7 +42,7 @@ django-lambda-tasks/
 - `decorators.py` — defines `@lambda_task`; enforces kwargs-only at decoration time
 - `models.py` — `TaskRecord` (Django ORM), `SQSLambdaTaskMessage` (Pydantic, SQS schema + execution logic), `SQSLambdaTask` (Pydantic, holds message + routing; `_execute()` publishes to SQS, executes eagerly, or submits to the local process pool; `execute_on_commit()` registers `_execute` with `transaction.on_commit`)
 - `local_executor.py` — `ProcessPoolExecutor`-based async local execution; `get_pool()` lazily creates a module-level pool; `submit_task()` fire-and-forget submission; `_execute_in_worker()` deserializes and runs the task in a child process; `_pool_initializer()` calls `django.setup()` per worker
-- `handler.py` — Lambda entry point; cold-start init (temporary log handler → `resolve_environment()` → `resolve_secrets_into_env()` → handler removed → conditional `django.setup()`) runs inside the handler on first invocation, guarded by `_cold_start_done` sentinel; processes SQS records independently; returns `batchItemFailures`
+- `handler.py` — Lambda entry point; cold-start init (memory limit → temporary log handler → `resolve_environment()` → `resolve_secrets_into_env()` → handler removed → conditional `django.setup()`) runs inside the handler on first invocation, guarded by `_cold_start_done` sentinel; processes SQS records independently; returns `batchItemFailures`
 - `environment_loader.py` — loads env vars from a Secrets Manager secret at cold start; validates flat JSON format; idempotent via `_loaded` sentinel
 - `secret_loader.py` — resolves `LAMBDA_TASKS_SECRET_*` env vars from Secrets Manager before Django starts; validates format, detects conflicts, batches API calls; idempotent via `_loaded` sentinel
 - `logging.py` — `task_logger` singleton; `message_id` set/cleared around each task execution

{django_lambda_tasks-0.4.1 → django_lambda_tasks-0.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: django-lambda-tasks
-Version: 0.4.1
+Version: 0.4.3
 Summary: Run async tasks in a lambda function
 Requires-Python: >=3.10
 Requires-Dist: awslambdaric

{django_lambda_tasks-0.4.1 → django_lambda_tasks-0.4.3}/lambda_tasks/admin.py

@@ -3,8 +3,13 @@
 from django.contrib import admin
 from django.db.models import DurationField, ExpressionWrapper, F, QuerySet
 from django.http import HttpRequest
+from django.utils.module_loading import import_string
 
-from lambda_tasks.models import TaskRecord
+from lambda_tasks.models import (
+    SQSLambdaTask,
+    SQSLambdaTaskMessage,
+    TaskRecord,
+)
 
 
 @admin.register(TaskRecord)
@@ -33,6 +38,7 @@ class TaskRecordAdmin(admin.ModelAdmin):
         "result",
         "traceback",
     )
+    actions = ["replay_tasks"]
 
     def get_queryset(self, request: HttpRequest) -> QuerySet:
         return (
@@ -52,3 +58,18 @@ class TaskRecordAdmin(admin.ModelAdmin):
             return None
         total_seconds = d.total_seconds()
         return f"{total_seconds:.3f}s"
+
+    @admin.action(description="Replay selected tasks", permissions=("change",))
+    def replay_tasks(self, request: HttpRequest, queryset: QuerySet) -> None:
+        for record in queryset:
+            queue = import_string(record.task_name).queue
+            task = SQSLambdaTask(
+                message=SQSLambdaTaskMessage(
+                    task_name=record.task_name,
+                    kwargs=record.kwargs,
+                    n_retries=0,
+                ),
+                delay=0,
+                queue=queue,
+            )
+            task.execute_on_commit()

{django_lambda_tasks-0.4.1 → django_lambda_tasks-0.4.3}/lambda_tasks/handler.py

@@ -14,6 +14,7 @@ warm invocations skip it.
 
 import logging
 import os
+import resource
 
 import django
 from django.apps import apps
@@ -26,6 +27,36 @@ logger = logging.getLogger(__name__)
 _cold_start_done: bool = False
 
 
+_MEMORY_RESERVED_MB = 128
+_MEMORY_MINIMUM_LIMIT_MB = 64
+
+
+def _set_memory_limit() -> None:
+    """Set RLIMIT_DATA from AWS_LAMBDA_FUNCTION_MEMORY_SIZE if available.
+
+    When running in Lambda, this causes Python to raise MemoryError on
+    excessive allocation instead of being killed by the OOM killer.
+
+    128 MB is reserved for the Python runtime, shared libraries, and OS
+    overhead. The limit is floored at 64 MB so that even small Lambdas
+    get some protection.
+    """
+    memory_mb = os.environ.get("AWS_LAMBDA_FUNCTION_MEMORY_SIZE")
+
+    if memory_mb is None:
+        return
+
+    limit_bytes = (
+        max(int(memory_mb) - _MEMORY_RESERVED_MB, _MEMORY_MINIMUM_LIMIT_MB)
+        * 1024
+        * 1024
+    )
+    resource.setrlimit(resource.RLIMIT_AS, (limit_bytes, limit_bytes))
+    logger.info(
+        f"Set RLIMIT_AS to {limit_bytes} bytes ({memory_mb} MB - {_MEMORY_RESERVED_MB} MB reserved)"
+    )
+
+
 def _perform_cold_start() -> None:
     """Run one-time initialisation: env loading, secrets, Django setup.
 
@@ -49,6 +80,8 @@ def _perform_cold_start() -> None:
     lambda_tasks_logger.addHandler(boot_handler)
     lambda_tasks_logger.setLevel(logging.DEBUG)
 
+    _set_memory_limit()
+
     try:
         resolve_environment()
         resolve_secrets_into_env()

{django_lambda_tasks-0.4.1 → django_lambda_tasks-0.4.3}/pyproject.toml

@@ -7,7 +7,7 @@ packages = ["lambda_tasks"]
 
 [project]
 name = "django-lambda-tasks"
-version = "0.4.1"
+version = "0.4.3"
 description = "Run async tasks in a lambda function"
 readme = "README.md"
 requires-python = ">=3.10"

django_lambda_tasks-0.4.3/tests/test_admin.py

@@ -0,0 +1,282 @@
+from unittest.mock import patch
+
+import pytest
+from django.contrib import admin
+from django.contrib.auth.models import Permission, User
+from django.contrib.contenttypes.models import ContentType
+from django.test import RequestFactory
+from django.utils.timezone import now
+
+import lambda_tasks.admin  # noqa: F401 — triggers @admin.register side-effect
+from lambda_tasks.models import SQSLambdaTask, TaskRecord, TaskStatus
+
+
+def test_task_record_registered_in_admin():
+    assert TaskRecord in admin.site._registry
+
+
+def test_task_record_admin_list_display():
+    assert admin.site._registry[TaskRecord].list_display == (
+        "pk",
+        "task_name",
+        "status",
+        "start_time",
+        "end_time",
+        "n_retries",
+        "duration",
+        "result",
+    )
+
+
+@pytest.fixture()
+def admin_user(db: None) -> User:
+    return User.objects.create_superuser(
+        username="admin", password="password", email="admin@example.com"
+    )
+
+
+@pytest.fixture()
+def task_record(db: None) -> TaskRecord:
+    return TaskRecord.objects.create(
+        id="00000000-0000-0000-0000-000000000001",
+        task_name="myapp.tasks.my_task",
+        kwargs={"user_id": 42},
+        n_retries=3,
+        status=TaskStatus.FAILED,
+        start_time=now(),
+        end_time=now(),
+    )
+
+
+@pytest.fixture()
+def _mock_import_string():
+    class _FakeWrapper:
+        queue = "default"
+
+    with patch("lambda_tasks.admin.import_string", return_value=_FakeWrapper()):
+        yield
+
+
+@pytest.mark.django_db()
+class TestReplayAction:
+    def test_replay_action_is_registered(self) -> None:
+        model_admin = admin.site._registry[TaskRecord]
+        request = RequestFactory().post("/")
+        request.user = User(is_staff=True, is_superuser=True)
+        action_names = list(model_admin.get_actions(request).keys())
+        assert "replay_tasks" in action_names
+
+    def test_replay_enqueues_task_with_original_kwargs(
+        self, admin_user: User, task_record: TaskRecord, _mock_import_string: None
+    ) -> None:
+        model_admin = admin.site._registry[TaskRecord]
+        request = RequestFactory().post("/")
+        request.user = admin_user
+
+        queryset = TaskRecord.objects.filter(pk=task_record.pk)
+
+        with patch(
+            "lambda_tasks.admin.SQSLambdaTask.execute_on_commit"
+        ) as mock_execute:
+            model_admin.replay_tasks(request, queryset)
+
+        mock_execute.assert_called_once()
+
+    def test_replay_resets_n_retries_to_zero(
+        self, admin_user: User, task_record: TaskRecord, _mock_import_string: None
+    ) -> None:
+        model_admin = admin.site._registry[TaskRecord]
+        request = RequestFactory().post("/")
+        request.user = admin_user
+
+        queryset = TaskRecord.objects.filter(pk=task_record.pk)
+
+        built_tasks: list[SQSLambdaTask] = []
+
+        def capture_execute(self: SQSLambdaTask) -> None:
+            built_tasks.append(self)
+
+        with patch.object(SQSLambdaTask, "execute_on_commit", capture_execute):
+            model_admin.replay_tasks(request, queryset)
+
+        assert len(built_tasks) == 1
+        assert built_tasks[0].message.n_retries == 0
+
+    def test_replay_preserves_task_name(
+        self, admin_user: User, task_record: TaskRecord, _mock_import_string: None
+    ) -> None:
+        model_admin = admin.site._registry[TaskRecord]
+        request = RequestFactory().post("/")
+        request.user = admin_user
+
+        queryset = TaskRecord.objects.filter(pk=task_record.pk)
+
+        built_tasks: list[SQSLambdaTask] = []
+
+        def capture_execute(self: SQSLambdaTask) -> None:
+            built_tasks.append(self)
+
+        with patch.object(SQSLambdaTask, "execute_on_commit", capture_execute):
+            model_admin.replay_tasks(request, queryset)
+
+        assert built_tasks[0].message.task_name == "myapp.tasks.my_task"
+        assert built_tasks[0].message.kwargs == {"user_id": 42}
+
+    def test_replay_uses_delay_zero(
+        self, admin_user: User, task_record: TaskRecord, _mock_import_string: None
+    ) -> None:
+        model_admin = admin.site._registry[TaskRecord]
+        request = RequestFactory().post("/")
+        request.user = admin_user
+
+        queryset = TaskRecord.objects.filter(pk=task_record.pk)
+
+        built_tasks: list[SQSLambdaTask] = []
+
+        def capture_execute(self: SQSLambdaTask) -> None:
+            built_tasks.append(self)
+
+        with patch.object(SQSLambdaTask, "execute_on_commit", capture_execute):
+            model_admin.replay_tasks(request, queryset)
+
+        assert built_tasks[0].delay == 0
+
+    def test_replay_resolves_queue_from_wrapper(
+        self, admin_user: User, db: None
+    ) -> None:
+        record = TaskRecord.objects.create(
+            id="00000000-0000-0000-0000-000000000002",
+            task_name="myapp.tasks.custom_queue_task",
+            kwargs={"x": 1},
+            n_retries=0,
+            status=TaskStatus.FAILED,
+            start_time=now(),
+            end_time=now(),
+        )
+
+        model_admin = admin.site._registry[TaskRecord]
+        request = RequestFactory().post("/")
+        request.user = admin_user
+
+        queryset = TaskRecord.objects.filter(pk=record.pk)
+
+        built_tasks: list[SQSLambdaTask] = []
+
+        def capture_execute(self: SQSLambdaTask) -> None:
+            built_tasks.append(self)
+
+        class _FakeWrapper:
+            queue = "high-priority"
+
+        with patch.object(SQSLambdaTask, "execute_on_commit", capture_execute):
+            with patch(
+                "lambda_tasks.admin.import_string",
+                return_value=_FakeWrapper(),
+            ):
+                model_admin.replay_tasks(request, queryset)
+
+        assert built_tasks[0].queue == "high-priority"
+
+    def test_replay_raises_on_import_error(
+        self, admin_user: User, task_record: TaskRecord
+    ) -> None:
+        model_admin = admin.site._registry[TaskRecord]
+        request = RequestFactory().post("/")
+        request.user = admin_user
+
+        queryset = TaskRecord.objects.filter(pk=task_record.pk)
+
+        with patch(
+            "lambda_tasks.admin.import_string",
+            side_effect=ImportError("no such module"),
+        ):
+            with pytest.raises(ImportError, match="no such module"):
+                model_admin.replay_tasks(request, queryset)
+
+    def test_replay_raises_when_not_wrapper(
+        self, admin_user: User, task_record: TaskRecord
+    ) -> None:
+        model_admin = admin.site._registry[TaskRecord]
+        request = RequestFactory().post("/")
+        request.user = admin_user
+
+        queryset = TaskRecord.objects.filter(pk=task_record.pk)
+
+        with patch(
+            "lambda_tasks.admin.import_string",
+            return_value="not a wrapper",
+        ):
+            with pytest.raises(AttributeError):
+                model_admin.replay_tasks(request, queryset)
+
+    def test_replay_multiple_records(
+        self, admin_user: User, db: None, _mock_import_string: None
+    ) -> None:
+        records = [
+            TaskRecord.objects.create(
+                id=f"00000000-0000-0000-0000-00000000000{i}",
+                task_name=f"myapp.tasks.task_{i}",
+                kwargs={"key": i},
+                n_retries=0,
+                status=TaskStatus.FAILED,
+                start_time=now(),
+                end_time=now(),
+            )
+            for i in range(3)
+        ]
+
+        model_admin = admin.site._registry[TaskRecord]
+        request = RequestFactory().post("/")
+        request.user = admin_user
+
+        queryset = TaskRecord.objects.filter(pk__in=[r.pk for r in records])
+
+        built_tasks: list[SQSLambdaTask] = []
+
+        def capture_execute(self: SQSLambdaTask) -> None:
+            built_tasks.append(self)
+
+        with patch.object(SQSLambdaTask, "execute_on_commit", capture_execute):
+            model_admin.replay_tasks(request, queryset)
+
+        assert len(built_tasks) == 3
+        task_names = {t.message.task_name for t in built_tasks}
+        assert task_names == {
+            "myapp.tasks.task_0",
+            "myapp.tasks.task_1",
+            "myapp.tasks.task_2",
+        }
+
+    def test_replay_action_hidden_without_change_permission(self, db: None) -> None:
+        user = User.objects.create_user(
+            username="viewer", password="password", is_staff=True
+        )
+        content_type = ContentType.objects.get_for_model(TaskRecord)
+        view_perm = Permission.objects.get(
+            codename="view_taskrecord", content_type=content_type
+        )
+        user.user_permissions.add(view_perm)
+
+        model_admin = admin.site._registry[TaskRecord]
+        request = RequestFactory().get("/")
+        request.user = user
+
+        action_names = list(model_admin.get_actions(request).keys())
+        assert "replay_tasks" not in action_names
+
+    def test_replay_action_visible_with_change_permission(self, db: None) -> None:
+        user = User.objects.create_user(
+            username="editor", password="password", is_staff=True
+        )
+        content_type = ContentType.objects.get_for_model(TaskRecord)
+        change_perm = Permission.objects.get(
+            codename="change_taskrecord", content_type=content_type
+        )
+        user.user_permissions.add(change_perm)
+
+        model_admin = admin.site._registry[TaskRecord]
+        request = RequestFactory().get("/")
+        request.user = user
+
+        action_names = list(model_admin.get_actions(request).keys())
+        assert "replay_tasks" in action_names

django_lambda_tasks-0.4.3/tests/test_memory_limit.py

@@ -0,0 +1,131 @@
+"""
+Unit tests for memory limit enforcement in the Lambda handler.
+
+The handler sets resource.RLIMIT_AS during cold start so that runaway
+memory allocation raises MemoryError instead of triggering the Lambda OOM
+killer.
+"""
+
+import resource
+from unittest.mock import patch
+
+import lambda_tasks.handler as handler_module
+
+
+class TestMemoryLimitSetDuringColdStart:
+    def test_rlimit_as_set_from_env_var(self, monkeypatch):
+        """When AWS_LAMBDA_FUNCTION_MEMORY_SIZE is set, RLIMIT_AS is configured with 128 MB reserved."""
+        monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
+        monkeypatch.setenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", "512")
+        monkeypatch.setattr("lambda_tasks.handler.resolve_environment", lambda: None)
+        monkeypatch.setattr(
+            "lambda_tasks.handler.resolve_secrets_into_env", lambda: None
+        )
+
+        calls: list[tuple] = []
+
+        def fake_setrlimit(which: int, limits: tuple[int, int]) -> None:
+            calls.append((which, limits))
+
+        monkeypatch.setattr(resource, "setrlimit", fake_setrlimit)
+
+        handler_module._cold_start_done = False
+        handler_module.handler(event={"Records": []}, context=None)
+
+        expected = (512 - 128) * 1024 * 1024
+        assert calls == [(resource.RLIMIT_AS, (expected, expected))]
+
+    def test_rlimit_as_not_set_when_env_var_missing(self, monkeypatch):
+        """When AWS_LAMBDA_FUNCTION_MEMORY_SIZE is not set, RLIMIT_AS is unchanged."""
+        monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
+        monkeypatch.delenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", raising=False)
+        monkeypatch.setattr("lambda_tasks.handler.resolve_environment", lambda: None)
+        monkeypatch.setattr(
+            "lambda_tasks.handler.resolve_secrets_into_env", lambda: None
+        )
+
+        calls: list[tuple] = []
+
+        def fake_setrlimit(which: int, limits: tuple[int, int]) -> None:
+            calls.append((which, limits))
+
+        monkeypatch.setattr(resource, "setrlimit", fake_setrlimit)
+
+        handler_module._cold_start_done = False
+        handler_module.handler(event={"Records": []}, context=None)
+
+        assert calls == []
+
+    def test_memory_limit_set_before_loaders(self, monkeypatch):
+        """The memory limit is set before resolve_environment runs."""
+        monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
+        monkeypatch.setenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", "256")
+
+        call_order: list[str] = []
+
+        def fake_setrlimit(which: int, limits: tuple[int, int]) -> None:
+            call_order.append("setrlimit")
+
+        def fake_resolve_environment() -> None:
+            call_order.append("resolve_environment")
+
+        def fake_resolve_secrets() -> None:
+            call_order.append("resolve_secrets_into_env")
+
+        monkeypatch.setattr(resource, "setrlimit", fake_setrlimit)
+        monkeypatch.setattr(
+            "lambda_tasks.handler.resolve_environment", fake_resolve_environment
+        )
+        monkeypatch.setattr(
+            "lambda_tasks.handler.resolve_secrets_into_env", fake_resolve_secrets
+        )
+
+        handler_module._cold_start_done = False
+        handler_module.handler(event={"Records": []}, context=None)
+
+        assert call_order == [
+            "setrlimit",
+            "resolve_environment",
+            "resolve_secrets_into_env",
+        ]
+
+    def test_memory_limit_is_logged(self, monkeypatch, caplog):
+        """Setting the memory limit emits an info log."""
+        import logging
+
+        monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
+        monkeypatch.setenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", "1024")
+        monkeypatch.setattr("lambda_tasks.handler.resolve_environment", lambda: None)
+        monkeypatch.setattr(
+            "lambda_tasks.handler.resolve_secrets_into_env", lambda: None
+        )
+        monkeypatch.setattr(resource, "setrlimit", lambda *args: None)
+
+        handler_module._cold_start_done = False
+
+        with caplog.at_level(logging.INFO, logger="lambda_tasks.handler"):
+            handler_module.handler(event={"Records": []}, context=None)
+
+        assert any("RLIMIT_AS" in msg for msg in caplog.messages)
+
+    def test_rlimit_floored_at_minimum_when_memory_too_small(self, monkeypatch):
+        """When AWS_LAMBDA_FUNCTION_MEMORY_SIZE <= reserved, RLIMIT_AS is set to the 64 MB minimum."""
+        monkeypatch.delenv("DJANGO_SETTINGS_MODULE", raising=False)
+        monkeypatch.setenv("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", "128")
+        monkeypatch.setattr("lambda_tasks.handler.resolve_environment", lambda: None)
+        monkeypatch.setattr(
+            "lambda_tasks.handler.resolve_secrets_into_env", lambda: None
+        )
+
+        calls: list[tuple] = []
+
+        def fake_setrlimit(which: int, limits: tuple[int, int]) -> None:
+            calls.append((which, limits))
+
+        monkeypatch.setattr(resource, "setrlimit", fake_setrlimit)
+
+        handler_module._cold_start_done = False
+        handler_module.handler(event={"Records": []}, context=None)
+
+        expected = 64 * 1024 * 1024
+        assert calls == [(resource.RLIMIT_AS, (expected, expected))]

django_lambda_tasks-0.4.1/tests/test_admin.py

@@ -1,21 +0,0 @@
-from django.contrib import admin
-
-import lambda_tasks.admin  # noqa: F401 — triggers @admin.register side-effect
-from lambda_tasks.models import TaskRecord
-
-
-def test_task_record_registered_in_admin():
-    assert TaskRecord in admin.site._registry
-
-
-def test_task_record_admin_list_display():
-    assert admin.site._registry[TaskRecord].list_display == (
-        "pk",
-        "task_name",
-        "status",
-        "start_time",
-        "end_time",
-        "n_retries",
-        "duration",
-        "result",
-    )