django-lambda-tasks 0.1.0__py3-none-any.whl

lambda_tasks/models.py

@@ -0,0 +1,240 @@
+"""Django model for persisting task execution records."""
+
+import random
+import traceback
+
+import boto3
+from django.core.exceptions import ImproperlyConfigured
+from django.db import models, transaction
+from django.db.models import Q
+from django.utils.module_loading import import_string
+from django.utils.timezone import now
+from pydantic import BaseModel, ConfigDict, Field
+
+from lambda_tasks.logging import task_logger
+from lambda_tasks.settings import LambdaTasksSettings
+from lambda_tasks.timeouts import TimeoutContext
+
+
+class MaxRetriesExceededError(Exception):
+    def __init__(self, *, task_name: str, n_retries: int) -> None:
+        super().__init__(
+            f"Task '{task_name}' exceeded the maximum retry limit ({n_retries} retries)."
+        )
+
+
+class TaskStatus(models.TextChoices):
+    RUNNING = "RUNNING"
+    SUCCESS = "SUCCESS"
+    FAILED = "FAILED"
+    RETRYING = "RETRYING"
+
+
+class TaskRecord(models.Model):
+    TaskStatus = TaskStatus
+
+    task_name = models.CharField(max_length=255, editable=False)
+    invocation_id = models.UUIDField(unique=True, editable=False)
+    kwargs = models.JSONField(editable=False)
+    n_retries = models.PositiveSmallIntegerField(editable=False)
+    status = models.CharField(
+        max_length=10,
+        choices=TaskStatus,
+        editable=False,
+    )
+    start_time = models.DateTimeField(null=True, editable=False)
+    end_time = models.DateTimeField(null=True, editable=False)
+    result = models.JSONField(null=True, editable=False)
+    traceback = models.TextField(null=True, editable=False)
+
+    @property
+    def duration(self) -> object:
+        return self.end_time - self.start_time  # type: ignore[operator]
+
+    class Meta:
+        ordering = ["-start_time"]
+        indexes = [
+            models.Index(fields=["task_name"]),
+            models.Index(fields=["invocation_id"]),
+            models.Index(fields=["status"]),
+            models.Index(fields=["-start_time"]),
+        ]
+        constraints = [
+            models.CheckConstraint(
+                condition=Q(status__in=TaskStatus.values),
+                name="taskrecord_status_valid",
+            ),
+        ]
+
+
+class SQSLambdaTaskMessage(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    task_name: str
+    invocation_id: str
+    kwargs: dict
+    n_retries: int = Field(default=0, ge=0)
+
+    def execute_immediately(self) -> None:
+        """Execute a background task described.
+
+        Creates a TaskRecord, resolves timeouts, validates configuration,
+        runs the task inside an atomic block with timeout enforcement, and
+        persists the outcome.
+        """
+        task_logger.invocation_id = self.invocation_id
+
+        try:
+            task_logger.info(f"Received {self.task_name}")
+
+            # Local import to avoid circular dependency
+            from lambda_tasks.decorators import LambdaTaskWrapper
+
+            wrapper = import_string(self.task_name)
+
+            if not isinstance(wrapper, LambdaTaskWrapper):
+                raise TypeError(
+                    f"import_string('{self.task_name}') returned {type(wrapper)!r}, "
+                    f"expected LambdaTaskWrapper."
+                )
+
+            # Resolve before creating the TaskRecord — ConfigurationError here means
+            # misconfigured settings, not a task failure, so no record should be written.
+            soft_timeout, hard_timeout = wrapper.resolved_timeouts
+
+            record, created = (
+                TaskRecord.objects.get_or_create(  # ty: ignore[unresolved-attribute]
+                    invocation_id=self.invocation_id,
+                    defaults={
+                        "task_name": self.task_name,
+                        "kwargs": self.kwargs,
+                        "n_retries": self.n_retries,
+                        "status": TaskRecord.TaskStatus.RUNNING,
+                        "start_time": now(),
+                    },
+                )
+            )
+
+            if not created:
+                task_logger.warning(
+                    f"Skipping duplicate delivery (existing record status={record.status})"
+                )
+                return
+
+            ignored_exception: BaseException | None = None
+            ignored_traceback: str | None = None
+            result = None
+
+            try:
+                with transaction.atomic():
+                    with TimeoutContext(
+                        soft_timeout=soft_timeout, hard_timeout=hard_timeout
+                    ):
+                        result = wrapper(**self.kwargs)
+            except Exception as error:
+                if wrapper.ignore_errors and isinstance(error, wrapper.ignore_errors):
+                    ignored_exception = error
+                    ignored_traceback = traceback.format_exc()
+                elif wrapper.retry_on and isinstance(error, wrapper.retry_on):
+                    conf = LambdaTasksSettings()
+                    if self.n_retries >= conf.MAX_RETRIES:
+                        record.status = TaskRecord.TaskStatus.FAILED
+                        record.traceback = traceback.format_exc()
+                        record.end_time = now()
+                        record.save(update_fields=["status", "traceback", "end_time"])
+
+                        task_logger.warning(
+                            f"Failed due to MaxRetriesExceededError in {record.duration}"
+                        )
+
+                        raise MaxRetriesExceededError(
+                            task_name=self.task_name, n_retries=self.n_retries
+                        )
+                    else:
+                        record.status = TaskRecord.TaskStatus.RETRYING
+                        record.traceback = traceback.format_exc()
+                        record.end_time = now()
+                        record.save(update_fields=["status", "traceback", "end_time"])
+
+                        task_logger.warning(
+                            f"Retrying (due to {type(ignored_exception).__name__}) after {record.duration}"
+                        )
+
+                        delay = (
+                            wrapper._delay
+                            if wrapper._delay != 0
+                            else round(random.uniform(1, 5))
+                        )
+                        wrapper.execute_on_commit(
+                            **self.kwargs,
+                            _delay=delay,
+                            _n_retries=self.n_retries + 1,
+                        )
+                        return
+                else:
+                    task_logger.error(error, exc_info=True)
+
+                    record.status = TaskRecord.TaskStatus.FAILED
+                    record.traceback = traceback.format_exc()
+                    record.end_time = now()
+                    record.save(update_fields=["status", "traceback", "end_time"])
+
+                    task_logger.warning(f"Failed in {record.duration}")
+                    return
+
+            if ignored_exception is None:
+                record.status = TaskRecord.TaskStatus.SUCCESS
+                record.result = result
+                record.end_time = now()
+                record.save(update_fields=["status", "result", "end_time"])
+
+                task_logger.info(f"Succeeded in {record.duration}")
+            else:
+                record.status = TaskRecord.TaskStatus.SUCCESS
+                record.traceback = ignored_traceback
+                record.end_time = now()
+                record.save(update_fields=["status", "traceback", "end_time"])
+
+                task_logger.info(
+                    f"Succeeded (ignored {type(ignored_exception).__name__}) in {record.duration}"
+                )
+
+        finally:
+            task_logger.invocation_id = None
+
+
+class SQSLambdaTask(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    message: SQSLambdaTaskMessage
+    delay: int
+    queue: str
+
+    def execute_on_commit(self) -> None:
+        """Enqueue this task after the current transaction commits."""
+        transaction.on_commit(self._execute)
+
+    def _execute(self) -> None:
+        """Send this task to SQS (or execute eagerly).
+
+        Raises:
+            ImproperlyConfigured: if the queue name is not found in settings.
+            Any boto3 exception: propagated directly to the caller.
+        """
+        conf = LambdaTasksSettings()
+
+        if conf.EAGER:
+            self.message.execute_immediately()
+        else:
+            try:
+                queue_url = conf.QUEUES[self.queue]
+            except KeyError:
+                raise ImproperlyConfigured(
+                    f"Queue '{self.queue}' is not defined in LAMBDA_TASKS_QUEUES."
+                )
+            client = boto3.client("sqs")
+            client.send_message(
+                QueueUrl=queue_url,
+                MessageBody=self.message.model_dump_json(),
+                DelaySeconds=self.delay,
+            )

lambda_tasks/secret_loader.py

@@ -0,0 +1,183 @@
+"""
+Resolves environment variables that reference AWS Secrets Manager ARNs.
+
+Any env var prefixed with ``AWS_SECRETS_MANAGER_`` is treated as a pointer
+to a secret value.  The unprefixed name is the target env var to populate.
+
+Required value format
+---------------------
+Every reference must follow the full dynamic reference syntax::
+
+    AWS_SECRETS_MANAGER_DJANGO_ADMIN_URL=arn:aws:secretsmanager:eu-west-1:123:secret:my-secret:DJANGO_ADMIN_URL:AWSCURRENT:v1
+
+That is: ``<arn>:<json-key>:<version-stage>:<version-id>``
+
+All four suffix segments must be present and non-empty.
+A malformed reference raises ``ValueError`` immediately so the Lambda
+container fails at cold start rather than silently misconfiguring Django.
+
+It is a configuration error to set both ``AWS_SECRETS_MANAGER_FOO`` and
+``FOO`` — use one or the other.  Having both raises ``ValueError`` at cold
+start so the misconfiguration is caught immediately.
+
+Calls are batched by unique (ARN, version-stage, version-id) combination —
+one ``GetSecretValue`` call per unique combination, regardless of how many
+env vars reference it.  Results are cached in-process so repeated calls to
+``resolve_secrets_into_env`` are free after the first.
+"""
+
+import json
+import logging
+import os
+from typing import NamedTuple
+
+import boto3
+
+logger = logging.getLogger(__name__)
+
+_PREFIX = "AWS_SECRETS_MANAGER_"
+
+# Module-level cache: (arn, version_stage, version_id) → raw secret string.
+# Populated on first call; reused for the lifetime of the Lambda container.
+_secret_cache: dict[tuple[str, str, str], dict[str, str]] = {}
+
+
+class _SecretReference(NamedTuple):
+    arn: str
+    json_key: str
+    version_stage: str
+    version_id: str
+
+
+def _parse_reference(*, env_var: str, value: str) -> _SecretReference:
+    """Parse and validate a Secrets Manager reference.
+
+    Expected format::
+
+        <arn>:<json-key>:<version-stage>:<version-id>
+
+    The ARN itself is 7 colon-separated segments, plus 3 suffix segments
+    (json-key, version-stage, version-id) = 10 total.
+    All four suffix fields must be non-empty.
+
+    Raises ``ValueError`` if the format is invalid — this is intentional so
+    the Lambda container fails at cold start rather than starting with a
+    misconfigured environment.
+    """
+    # ARN has exactly 7 colon-separated parts:
+    # arn : aws : secretsmanager : region : account : secret : name
+    # Plus 3 suffix parts: json-key : version-stage : version-id  → 10 total
+    parts = value.split(":")
+
+    if len(parts) != 10:
+        raise ValueError(
+            f"{env_var} has an invalid Secrets Manager reference format. "
+            "Expected <arn>:<json-key>:<version-stage>:<version-id> "
+            f"(10 colon-separated segments), got {len(parts)}: {value!r}"
+        )
+
+    arn = ":".join(parts[:7])
+    json_key, version_stage, version_id = parts[7], parts[8], parts[9]
+
+    for field, field_value in (
+        ("json-key", json_key),
+        ("version-stage", version_stage),
+        ("version-id", version_id),
+    ):
+        if not field_value:
+            raise ValueError(
+                f"{env_var} is missing the {field} segment in its Secrets Manager "
+                f"reference: {value!r}"
+            )
+
+    return _SecretReference(
+        arn=arn,
+        json_key=json_key,
+        version_stage=version_stage,
+        version_id=version_id,
+    )
+
+
+def _fetch_secret(*, client: object, ref: _SecretReference) -> dict[str, str]:
+    """Fetch a secret string from Secrets Manager, using the module cache.
+
+    The cache key is ``(arn, version_stage, version_id)`` so that references
+    to different versions of the same secret are fetched independently.
+    """
+    cache_key = (ref.arn, ref.version_stage, ref.version_id)
+
+    if cache_key in _secret_cache:
+        return _secret_cache[cache_key]
+
+    logger.debug(f"Fetching secret {ref}")
+
+    response = client.get_secret_value(  # type: ignore[union-attr]
+        SecretId=ref.arn,
+        VersionStage=ref.version_stage,
+        VersionId=ref.version_id,
+    )
+
+    try:
+        secret = json.loads(response["SecretString"])
+    except json.JSONDecodeError as error:
+        raise ValueError(f"Could not decode secret as json for {ref}")
+
+    _secret_cache[cache_key] = secret
+
+    return secret
+
+
+def resolve_secrets_into_env() -> None:
+    """Scan env vars for ``AWS_SECRETS_MANAGER_*`` references and resolve them.
+
+    For each matching env var the resolved value is written back into
+    ``os.environ`` under the unprefixed name.
+
+    Raises ``ValueError`` at cold start if:
+    - A reference is malformed (wrong segment count, any empty field)
+    - The target env var is already set — use ``AWS_SECRETS_MANAGER_FOO`` or
+      ``FOO``, not both
+
+    This function is idempotent — calling it multiple times is safe and cheap
+    because resolved secrets are cached after the first fetch.
+    """
+    references: dict[str, _SecretReference] = {}  # target_name → _SecretReference
+
+    for key, value in os.environ.items():
+        if not key.startswith(_PREFIX):
+            continue
+        target = key[len(_PREFIX) :]
+        references[target] = _parse_reference(env_var=key, value=value)
+
+    if not references:
+        return
+
+    # Fail fast if any target is already set in the environment.
+    conflicts = [target for target in references if target in os.environ]
+    if conflicts:
+        raise ValueError(
+            "The following environment variables are set both directly and via "
+            f"AWS_SECRETS_MANAGER_*: {', '.join(sorted(conflicts))}. "
+            "Use one or the other, not both."
+        )
+
+    # Unique (arn, version_stage, version_id) combinations to minimise API calls.
+    unique_cache_keys = {
+        (ref.arn, ref.version_stage, ref.version_id) for ref in references.values()
+    }
+    uncached = unique_cache_keys - _secret_cache.keys()
+    logger.info(
+        f"Resolving {len(references)} secret reference(s) from"
+        f" {len(unique_cache_keys)} unique secret version(s)"
+        f" ({len(unique_cache_keys) - len(uncached)} already cached)"
+    )
+
+    client = boto3.client("secretsmanager") if uncached else None
+
+    for target, ref in references.items():
+        secret = _fetch_secret(client=client, ref=ref)
+        secret_value = secret[ref.json_key]
+
+        os.environ[target] = secret_value
+
+        logger.info(f"Resolved {target} from secret {ref}")

lambda_tasks/settings.py

@@ -0,0 +1,76 @@
+"""
+Lazy settings object for lambda_tasks.
+Reads from django.conf.settings on first attribute access.
+"""
+
+from django.conf import settings as django_settings
+from django.core.exceptions import ImproperlyConfigured
+
+_UNSET = object()
+
+MAX_TIMEOUT = 900
+
+
+class LambdaTasksSettings:
+    """Lazy wrapper around Django settings for the lambda_tasks library."""
+
+    def _get(self, *, name: str, default: object = _UNSET) -> object:
+        val = getattr(django_settings, name, _UNSET)
+        if val is _UNSET:
+            if default is _UNSET:
+                return None
+            return default
+        if val is None:
+            if default is _UNSET:
+                return None
+            return None
+        return val
+
+    def _resolve_queues(self) -> dict[str, str]:
+        queues = self._get(name="LAMBDA_TASKS_QUEUES")
+
+        if queues:
+            if "default" not in queues:  # type: ignore[operator]
+                raise ImproperlyConfigured(
+                    "LAMBDA_TASKS_QUEUES must contain a 'default' key."
+                )
+            return queues  # type: ignore[return-value]
+
+        raise ImproperlyConfigured(
+            "LAMBDA_TASKS_QUEUES must be defined in Django settings."
+        )
+
+    @property
+    def QUEUES(self) -> dict[str, str]:
+        return self._resolve_queues()
+
+    @property
+    def EAGER(self) -> bool:
+        val = self._get(name="LAMBDA_TASKS_EAGER", default=False)
+        return bool(val)
+
+    @property
+    def DEFAULT_SOFT_TIMEOUT(self) -> int:
+        default = 270
+        val = self._get(name="LAMBDA_TASKS_DEFAULT_SOFT_TIMEOUT", default=default)
+        if val is None:
+            return default
+        else:
+            return int(val)  # type: ignore[return-value]
+
+    @property
+    def DEFAULT_HARD_TIMEOUT(self) -> int:
+        default = 300
+        val = self._get(name="LAMBDA_TASKS_DEFAULT_HARD_TIMEOUT", default=default)
+        if val is None:
+            return default
+        else:
+            return int(val)  # type: ignore[return-value]
+
+    @property
+    def MAX_RETRIES(self) -> int:
+        default = 2880  # 60 * 24 * 2
+        val = self._get(name="LAMBDA_TASKS_MAX_RETRIES", default=default)
+        if val is None:
+            return default
+        return int(val)  # type: ignore[return-value]

lambda_tasks/timeouts.py

@@ -0,0 +1,78 @@
+"""Timeout enforcement for lambda tasks using Unix SIGALRM."""
+
+import signal
+
+from lambda_tasks.settings import LambdaTasksSettings
+
+
+class SoftTimeLimitExceeded(Exception):
+    """Raised inside a running task when the soft timeout is exceeded.
+
+    The task may catch this exception to perform cleanup before the hard
+    timeout forcibly terminates execution.
+    """
+
+
+class HardTimeLimitExceeded(Exception):
+    """Raised by the executor when the hard timeout is exceeded.
+
+    The task is terminated immediately; this exception is not intended to
+    be caught by task code.
+    """
+
+
+class TimeoutContext:
+    """Context manager that enforces soft and hard timeouts via SIGALRM.
+
+    Two-phase approach:
+    1. Arm soft_timeout seconds. On SIGALRM, raise SoftTimeLimitExceeded
+       inside the task and re-arm for (hard_timeout - soft_timeout) seconds.
+    2. Second SIGALRM raises HardTimeLimitExceeded.
+
+    Any pre-existing alarm is saved on enter and restored on exit.
+
+    Usage::
+
+        with TimeoutContext(soft_timeout=270, hard_timeout=300):
+            run_task()
+    """
+
+    def __init__(self, soft_timeout: int, hard_timeout: int) -> None:
+        self.soft_timeout = soft_timeout
+        self.hard_timeout = hard_timeout
+        self._previous_alarm: int = 0
+        self._previous_handler = None
+        self._in_hard_phase: bool = False
+
+    def _handler(self, signum: int, frame) -> None:
+        if not self._in_hard_phase:
+            # Phase 1 → Phase 2: raise soft, re-arm for remaining hard window
+            self._in_hard_phase = True
+            remaining = self.hard_timeout - self.soft_timeout
+            signal.alarm(max(1, remaining))
+            raise SoftTimeLimitExceeded()
+        else:
+            # Phase 2: hard timeout expired
+            raise HardTimeLimitExceeded()
+
+    def __enter__(self) -> "TimeoutContext":
+        conf = LambdaTasksSettings()
+
+        if not conf.EAGER:
+            self._in_hard_phase = False
+            self._previous_handler = signal.signal(signal.SIGALRM, self._handler)
+            self._previous_alarm = signal.alarm(self.soft_timeout)
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        conf = LambdaTasksSettings()
+
+        if not conf.EAGER:
+            # Cancel any pending alarm we armed
+            signal.alarm(0)
+            # Restore the original signal handler
+            signal.signal(signal.SIGALRM, self._previous_handler)
+            # Restore any pre-existing alarm
+            if self._previous_alarm > 0:
+                signal.alarm(self._previous_alarm)