feat: add async image task API

2026-07-08 22:08:48 +08:00
parent c98f713762
commit 25a4080177
26 changed files with 1531 additions and 48 deletions
@@ -1,3 +1,98 @@
-from django.db import models
+from __future__ import annotations

-# Create your models here.
+import uuid
+
+from django.conf import settings
+from django.db import models
+from django.db.models import Q
+
+
+class ImageGenerationTask(models.Model):
+    class Status(models.TextChoices):
+        QUEUED = "queued", "排队中"
+        RUNNING = "running", "处理中"
+        SUCCEEDED = "succeeded", "成功"
+        FAILED = "failed", "失败"
+        EXPIRED = "expired", "已过期"
+
+    task_id = models.UUIDField("任务 ID", default=uuid.uuid4, unique=True, editable=False)
+    user = models.ForeignKey(
+        settings.AUTH_USER_MODEL,
+        verbose_name="用户",
+        on_delete=models.PROTECT,
+        related_name="image_generation_tasks",
+    )
+    api_key = models.ForeignKey(
+        "users.ApiKey",
+        verbose_name="API 密钥",
+        on_delete=models.PROTECT,
+        related_name="image_generation_tasks",
+    )
+    call_record = models.OneToOneField(
+        "billing.CallRecord",
+        verbose_name="调用记录",
+        on_delete=models.PROTECT,
+        related_name="image_generation_task",
+    )
+    status = models.CharField(
+        "状态",
+        max_length=20,
+        choices=Status.choices,
+        default=Status.QUEUED,
+    )
+    idempotency_key = models.CharField("幂等键", max_length=128, blank=True)
+    idempotency_key_hash = models.CharField(
+        "幂等键哈希",
+        max_length=64,
+        null=True,
+        blank=True,
+        editable=False,
+    )
+    request_hash = models.CharField("请求哈希", max_length=64)
+    request_payload = models.JSONField("请求快照", default=dict, blank=True)
+    input_image = models.FileField(
+        "输入图片",
+        upload_to="generated/task_inputs/%Y/%m/%d/",
+        blank=True,
+    )
+    result_url = models.TextField("结果 URL", blank=True)
+    error_code = models.CharField("错误码", max_length=64, blank=True)
+    error_message = models.TextField("错误信息", blank=True)
+    points_balance_after_charge = models.BigIntegerField("扣费后余额", default=0)
+    started_at = models.DateTimeField("开始时间", null=True, blank=True)
+    finished_at = models.DateTimeField("完成时间", null=True, blank=True)
+    expires_at = models.DateTimeField("任务元数据过期时间", null=True, blank=True)
+    locked_at = models.DateTimeField("锁定时间", null=True, blank=True)
+    lease_expires_at = models.DateTimeField("租约过期时间", null=True, blank=True)
+    heartbeat_at = models.DateTimeField("心跳时间", null=True, blank=True)
+    worker_id = models.CharField("Worker ID", max_length=128, blank=True)
+    attempt_count = models.PositiveIntegerField("尝试次数", default=0)
+    created_at = models.DateTimeField("创建时间", auto_now_add=True)
+    updated_at = models.DateTimeField("更新时间", auto_now=True)
+
+    class Meta:
+        db_table = "image_generation_task"
+        verbose_name = "图片生成任务"
+        verbose_name_plural = "图片生成任务"
+        ordering = ("-created_at", "-id")
+        constraints = [
+            models.UniqueConstraint(
+                fields=("api_key", "idempotency_key_hash"),
+                name="unique_image_task_idempotency_hash_per_api_key",
+            ),
+            models.CheckConstraint(
+                condition=Q(points_balance_after_charge__gte=0),
+                name="image_task_balance_after_charge_non_negative",
+            ),
+        ]
+        indexes = [
+            models.Index(fields=("user", "created_at")),
+            models.Index(fields=("api_key", "created_at")),
+            models.Index(fields=("status", "created_at")),
+            models.Index(fields=("status", "lease_expires_at")),
+            models.Index(fields=("worker_id", "status")),
+            models.Index(fields=("expires_at",)),
+        ]
+
+    def __str__(self) -> str:
+        return f"{self.task_id} {self.status}"