kodit 0.4.1__py3-none-any.whl → 0.4.3__py3-none-any.whl

kodit/infrastructure/sqlalchemy/entities.py

@@ -2,6 +2,7 @@
 
 from datetime import UTC, datetime
 from enum import Enum
+from typing import Any
 
 from git import Actor
 from sqlalchemy import (
@@ -9,6 +10,7 @@ from sqlalchemy import (
     ForeignKey,
     Integer,
     String,
+    TypeDecorator,
     UnicodeText,
     UniqueConstraint,
 )
@@ -18,6 +20,29 @@ from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
 from sqlalchemy.types import JSON
 
 
+# See <https://docs.sqlalchemy.org/en/20/core/custom_types.html#store-timezone-aware-timestamps-as-timezone-naive-utc>
+# And [this issue](https://github.com/sqlalchemy/sqlalchemy/issues/1985)
+class TZDateTime(TypeDecorator):
+    """Timezone-aware datetime type."""
+
+    impl = DateTime
+    cache_ok = True
+
+    def process_bind_param(self, value: Any, dialect: Any) -> Any:  # noqa: ARG002
+        """Process bind param."""
+        if value is not None:
+            if not value.tzinfo or value.tzinfo.utcoffset(value) is None:
+                raise TypeError("tzinfo is required")
+            value = value.astimezone(UTC).replace(tzinfo=None)
+        return value
+
+    def process_result_value(self, value: Any, dialect: Any) -> Any:  # noqa: ARG002
+        """Process result value."""
+        if value is not None:
+            value = value.replace(tzinfo=UTC)
+        return value
+
+
 class Base(AsyncAttrs, DeclarativeBase):
     """Base class for all models."""
 
@@ -27,10 +52,11 @@ class CommonMixin:
 
     id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
     created_at: Mapped[datetime] = mapped_column(
-        DateTime(timezone=True), default=lambda: datetime.now(UTC)
+        TZDateTime, nullable=False, default=lambda: datetime.now(UTC)
     )
     updated_at: Mapped[datetime] = mapped_column(
-        DateTime(timezone=True),
+        TZDateTime,
+        nullable=False,
         default=lambda: datetime.now(UTC),
         onupdate=lambda: datetime.now(UTC),
     )
@@ -236,3 +262,64 @@ class Task(Base, CommonMixin):
         self.type = type
         self.payload = payload
         self.priority = priority
+
+
+class TaskStatus(Base):
+    """Task status model."""
+
+    __tablename__ = "task_status"
+    id: Mapped[str] = mapped_column(
+        String(255), primary_key=True, index=True, nullable=False
+    )
+    created_at: Mapped[datetime] = mapped_column(
+        TZDateTime, nullable=False, default=lambda: datetime.now(UTC)
+    )
+    updated_at: Mapped[datetime] = mapped_column(
+        TZDateTime,
+        nullable=False,
+        default=lambda: datetime.now(UTC),
+        onupdate=lambda: datetime.now(UTC),
+    )
+    operation: Mapped[str] = mapped_column(String(255), index=True, nullable=False)
+    trackable_id: Mapped[int | None] = mapped_column(Integer, index=True, nullable=True)
+    trackable_type: Mapped[str | None] = mapped_column(
+        String(255), index=True, nullable=True
+    )
+    parent: Mapped[str | None] = mapped_column(
+        ForeignKey("task_status.id"), index=True, nullable=True
+    )
+    message: Mapped[str] = mapped_column(UnicodeText, default="")
+    state: Mapped[str] = mapped_column(String(255), default="")
+    error: Mapped[str] = mapped_column(UnicodeText, default="")
+    total: Mapped[int] = mapped_column(Integer, default=0)
+    current: Mapped[int] = mapped_column(Integer, default=0)
+
+    def __init__(  # noqa: PLR0913
+        self,
+        id: str,  # noqa: A002
+        operation: str,
+        created_at: datetime,
+        updated_at: datetime,
+        trackable_id: int | None,
+        trackable_type: str | None,
+        parent: str | None,
+        state: str,
+        error: str | None,
+        total: int,
+        current: int,
+        message: str,
+    ) -> None:
+        """Initialize the task status."""
+        super().__init__()
+        self.id = id
+        self.operation = operation
+        self.created_at = created_at
+        self.updated_at = updated_at
+        self.trackable_id = trackable_id
+        self.trackable_type = trackable_type
+        self.parent = parent
+        self.state = state
+        self.error = error or ""
+        self.total = total
+        self.current = current
+        self.message = message or ""