✨ support task services

alembic/versions/1df05b897d3f_tasks.py

@@ -0,0 +1,134 @@
+"""tasks
+
+Revision ID: 1df05b897d3f
+Revises: a1c10da5704b
+Create Date: 2023-07-23 14:44:59.592519
+
+"""
+import logging
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy import text
+from sqlalchemy.dialects import mysql
+from sqlalchemy.exc import NoSuchTableError
+
+# revision identifiers, used by Alembic.
+revision = "1df05b897d3f"
+down_revision = "a1c10da5704b"
+branch_labels = None
+depends_on = None
+
+logger = logging.getLogger(__name__)
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    connection = op.get_bind()
+    task_table = op.create_table(
+        "task",
+        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
+        sa.Column("user_id", sa.BigInteger(), nullable=True),
+        sa.Column("chat_id", sa.BigInteger(), nullable=True),
+        sa.Column(
+            "time_created",
+            sa.DateTime(),
+            server_default=sa.text("now()"),
+            nullable=True,
+        ),
+        sa.Column("time_updated", sa.DateTime(), nullable=True),
+        sa.Column(
+            "type",
+            sa.Enum(
+                "SIGN",
+                "RESIN",
+                "REALM",
+                "EXPEDITION",
+                "TRANSFORMER",
+                "CARD",
+                name="tasktypeenum",
+            ),
+            nullable=True,
+        ),
+        sa.Column(
+            "status",
+            sa.Enum(
+                "STATUS_SUCCESS",
+                "INVALID_COOKIES",
+                "ALREADY_CLAIMED",
+                "NEED_CHALLENGE",
+                "GENSHIN_EXCEPTION",
+                "TIMEOUT_ERROR",
+                "BAD_REQUEST",
+                "FORBIDDEN",
+                name="taskstatusenum",
+            ),
+            nullable=True,
+        ),
+        sa.Column("data", sa.JSON(), nullable=True),
+        sa.PrimaryKeyConstraint("id"),
+        mysql_charset="utf8mb4",
+        mysql_collate="utf8mb4_general_ci",
+    )
+    op.create_index("task_1", "task", ["user_id"], unique=False)
+    try:
+        statement = "SELECT * FROM sign;"
+        old_sign_table_data = connection.execute(text(statement))
+    except NoSuchTableError:
+        logger.warning("Table 'sign' doesn't exist")
+        return  # should not happen
+    if old_sign_table_data is not None:
+        for row in old_sign_table_data:
+            try:
+                user_id = row["user_id"]
+                chat_id = row["chat_id"]
+                time_created = row["time_created"]
+                time_updated = row["time_updated"]
+                status = row["status"]
+                task_type = "SIGN"
+                insert = task_table.insert().values(
+                    user_id=int(user_id),
+                    chat_id=int(chat_id),
+                    time_created=time_created,
+                    time_updated=time_updated,
+                    type=task_type,
+                    status=status,
+                )
+                with op.get_context().autocommit_block():
+                    connection.execute(insert)
+            except Exception as exc:  # pylint: disable=W0703
+                logger.error("Process sign->task Exception", exc_info=exc)  # pylint: disable=W0703
+    op.drop_table("sign")
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "sign",
+        sa.Column("id", mysql.INTEGER(), autoincrement=False, nullable=False),
+        sa.Column("user_id", mysql.BIGINT(), autoincrement=False, nullable=False),
+        sa.Column("chat_id", mysql.BIGINT(), autoincrement=False, nullable=True),
+        sa.Column("time_created", mysql.DATETIME(), nullable=True),
+        sa.Column("time_updated", mysql.DATETIME(), nullable=True),
+        sa.Column(
+            "status",
+            mysql.ENUM(
+                "STATUS_SUCCESS",
+                "INVALID_COOKIES",
+                "ALREADY_CLAIMED",
+                "GENSHIN_EXCEPTION",
+                "TIMEOUT_ERROR",
+                "BAD_REQUEST",
+                "FORBIDDEN",
+            ),
+            nullable=True,
+        ),
+        sa.PrimaryKeyConstraint("id", "user_id"),
+        mysql_collate="utf8mb4_general_ci",
+        mysql_default_charset="utf8mb4",
+        mysql_engine="InnoDB",
+    )
+    op.drop_index("task_1", table_name="task")
+    op.drop_table("task")
+    # ### end Alembic commands ###

core/services/task/__init__.py

@@ -0,0 +1 @@
+"""TaskService"""

core/services/sign/models.py → core/services/task/models.py

@@ -1,15 +1,15 @@
 import enum
 from datetime import datetime
-from typing import Optional
+from typing import Optional, Dict, Any
 
-from sqlalchemy import func, BigInteger
+from sqlalchemy import func, BigInteger, JSON
 from sqlmodel import Column, DateTime, Enum, Field, SQLModel, Integer
 
-__all__ = ("SignStatusEnum", "Sign")
+__all__ = ("Task", "TaskStatusEnum", "TaskTypeEnum")
 
 
-class SignStatusEnum(int, enum.Enum):
-    STATUS_SUCCESS = 0  # 签到成功
+class TaskStatusEnum(int, enum.Enum):
+    STATUS_SUCCESS = 0  # 任务执行成功
     INVALID_COOKIES = 1  # Cookie无效
     ALREADY_CLAIMED = 2  # 已经获取奖励
     NEED_CHALLENGE = 3  # 需要验证码
@@ -19,15 +19,26 @@ class SignStatusEnum(int, enum.Enum):
     FORBIDDEN = 7  # 这错误一般为通知失败 机器人被用户BAN
 
 
-class Sign(SQLModel, table=True):
+class TaskTypeEnum(int, enum.Enum):
+    SIGN = 0  # 签到
+    RESIN = 1  # 体力
+    REALM = 2  # 洞天宝钱
+    EXPEDITION = 3  # 委托
+    TRANSFORMER = 4  # 参量质变仪
+    CARD = 5  # 生日画片
+
+
+class Task(SQLModel, table=True):
     __table_args__ = dict(mysql_charset="utf8mb4", mysql_collate="utf8mb4_general_ci")
     id: Optional[int] = Field(
         default=None, primary_key=True, sa_column=Column(Integer(), primary_key=True, autoincrement=True)
     )
     user_id: int = Field(primary_key=True, sa_column=Column(BigInteger(), index=True))
-    chat_id: Optional[int] = Field(default=None)
+    chat_id: Optional[int] = Field(default=None, sa_column=Column(BigInteger()))
     time_created: Optional[datetime] = Field(
         sa_column=Column(DateTime, server_default=func.now())  # pylint: disable=E1102
     )
     time_updated: Optional[datetime] = Field(sa_column=Column(DateTime, onupdate=func.now()))  # pylint: disable=E1102
-    status: Optional[SignStatusEnum] = Field(sa_column=Column(Enum(SignStatusEnum)))
+    type: TaskTypeEnum = Field(primary_key=True, sa_column=Column(Enum(TaskTypeEnum)))
+    status: Optional[TaskStatusEnum] = Field(sa_column=Column(Enum(TaskStatusEnum)))
+    data: Optional[Dict[str, Any]] = Field(sa_column=Column(JSON))

core/services/sign/repositories.py → core/services/task/repositories.py

@@ -4,47 +4,47 @@
 
 from core.base_service import BaseService
 from core.dependence.database import Database
-from core.services.sign.models import Sign
+from core.services.task.models import Task, TaskTypeEnum
 from core.sqlmodel.session import AsyncSession
 
-__all__ = ("SignRepository",)
+__all__ = ("TaskRepository",)
 
 
-class SignRepository(BaseService.Component):
+class TaskRepository(BaseService.Component):
     def __init__(self, database: Database):
         self.engine = database.engine
 
-    async def add(self, sign: Sign):
+    async def add(self, task: Task):
         async with AsyncSession(self.engine) as session:
-            session.add(sign)
+            session.add(task)
             await session.commit()
 
-    async def remove(self, sign: Sign):
+    async def remove(self, task: Task):
         async with AsyncSession(self.engine) as session:
-            await session.delete(sign)
+            await session.delete(task)
             await session.commit()
 
-    async def update(self, sign: Sign) -> Sign:
+    async def update(self, task: Task) -> Task:
         async with AsyncSession(self.engine) as session:
-            session.add(sign)
+            session.add(task)
             await session.commit()
-            await session.refresh(sign)
-            return sign
+            await session.refresh(task)
+            return task
 
-    async def get_by_user_id(self, user_id: int) -> Optional[Sign]:
+    async def get_by_user_id(self, user_id: int, task_type: TaskTypeEnum) -> Optional[Task]:
         async with AsyncSession(self.engine) as session:
-            statement = select(Sign).where(Sign.user_id == user_id)
+            statement = select(Task).where(Task.user_id == user_id).where(Task.type == task_type)
             results = await session.exec(statement)
             return results.first()
 
-    async def get_by_chat_id(self, chat_id: int) -> Optional[List[Sign]]:
+    async def get_by_chat_id(self, chat_id: int, task_type: TaskTypeEnum) -> Optional[List[Task]]:
         async with AsyncSession(self.engine) as session:
-            statement = select(Sign).where(Sign.chat_id == chat_id)
+            statement = select(Task).where(Task.chat_id == chat_id).where(Task.type == task_type)
             results = await session.exec(statement)
             return results.all()
 
-    async def get_all(self) -> List[Sign]:
+    async def get_all(self, task_type: TaskTypeEnum) -> List[Task]:
         async with AsyncSession(self.engine) as session:
-            query = select(Sign)
+            query = select(Task).where(Task.type == task_type)
             results = await session.exec(query)
             return results.all()