jerry 8 godzin temu
rodzic
commit
87f9d7f95e

+ 3 - 2
app/api/router.py

@@ -1520,10 +1520,11 @@ async def vas_task_manual_confirm(
 
 @admin_required_router.get("/vas/task/pop", summary="任务出队(pop)", tags=["Visafly签证系统"], response_model=ApiResponse[VasTaskOut])
 async def vas_task_pop_task(
-    queue_name: str,
+    queue_name: str = Query("", description="队列名字"),
+    cooldown_seconds: int = Query(30, description="冷却时间 默认30秒"),
     db: AsyncSession = Depends(get_db),
 ):
-    task = await VasTaskService.pop_vas_task(db, queue_name, 180)
+    task = await VasTaskService.pop_vas_task(db, queue_name, 30)
     return success(data=task)
 
 @protected_router.post("/vas/ticket/create", summary="创建工单", tags=["Visafly签证系统"], response_model=ApiResponse[VasTicketOut])

+ 1 - 1
app/models/proxy_pool.py

@@ -15,6 +15,6 @@ class ProxyPool(Base):
     password = Column(String(64), default=None, comment="代理密码")
     next_use_time = Column(DateTime, nullable=False, default=func.now(), comment="下次允许使用的时间")
     status = Column(String(16), nullable=False, default="active", comment="active=可用, disable=禁用")
-    
+    time_zone = Column(String(100), nullable=True, comment="时区")
     created_at = Column(DateTime, default=datetime.utcnow)
     updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

+ 2 - 0
app/schemas/proxy_pool.py

@@ -12,6 +12,7 @@ class ProxyBase(BaseModel):
     username: Optional[str] = None
     password: Optional[str] = None
     next_use_time: Optional[datetime] = None
+    time_zone: Optional[str] = None
     status: Optional[str] = None
 
 # ================= Request Schemas =================
@@ -31,6 +32,7 @@ class ProxyUpdate(ProxyBase):
     username: Optional[str] = None
     password: Optional[str] = None
     next_use_time: Optional[datetime] = None
+    time_zone: Optional[str] = None
     status: Optional[str] = None
 
 class GetNextIpPayload(BaseModel):

+ 4 - 4
app/services/account_service.py

@@ -1,7 +1,7 @@
 import time
 from typing import Optional
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy import select, delete
+from sqlalchemy import select, func, text, delete
 from datetime import datetime, timedelta
 from app.utils.search import apply_keyword_search_stmt
 from app.utils.pagination import paginate
@@ -68,13 +68,12 @@ class AccountService:
         pool_name: str, 
         account_cd: int
     ) -> Account:
-        now = datetime.utcnow()
         stmt = (
             select(Account)
             .where(
                 Account.pool_name == pool_name,
                 Account.status == 'active',
-                Account.next_use_time <= now
+                Account.next_use_time <= func.utc_timestamp()
             )
             .order_by(Account.next_use_time.asc())
             .limit(1)
@@ -84,7 +83,8 @@ class AccountService:
         obj = result.scalar_one_or_none()
         if not obj:
             raise NotFoundError('Account not found')
-        obj.next_use_time = now + timedelta(seconds=account_cd)
+        
+        obj.next_use_time = func.utc_timestamp() + text(f"INTERVAL {account_cd} SECOND")     
         await db.commit()
         await db.refresh(obj)
         return obj

+ 3 - 4
app/services/proxy_service.py

@@ -2,7 +2,7 @@ import time
 from datetime import datetime, timedelta
 from typing import Optional
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy import select
+from sqlalchemy import select, func, text
 from typing import List, Dict
 from app.utils.search import apply_keyword_search_stmt
 from app.utils.pagination import paginate
@@ -58,13 +58,12 @@ class ProxyService:
 
     @staticmethod
     async def get_next_ip(db: AsyncSession, pools: list[str], proxy_cd: int):
-        now = datetime.utcnow()
         stmt = (
             select(ProxyPool)
             .where(
                 ProxyPool.status == 'active',
                 ProxyPool.pool_name.in_(pools),
-                ProxyPool.next_use_time <= now
+                ProxyPool.next_use_time <= func.utc_timestamp()
             )
             .order_by(ProxyPool.next_use_time.asc())
             .limit(1)
@@ -74,7 +73,7 @@ class ProxyService:
         obj = result.scalar_one_or_none()
         if not obj:
             raise NotFoundError('Proxy not found')
-        obj.next_use_time = now + timedelta(seconds=proxy_cd)
+        obj.next_use_time = func.utc_timestamp() + text(f"INTERVAL {proxy_cd} SECOND")
         await db.commit()
         await db.refresh(obj)
         return obj

+ 1 - 1
app/services/ticket_service.py

@@ -142,7 +142,7 @@ class TicketService:
             task_res = await db.execute(
                 select(VasTask).where(
                     VasTask.order_id == order.id,
-                    VasTask.status.in_(["pending", "grabbed", "running", "completed"]),
+                    VasTask.status.in_(["pending", "grabbed", "running", "pause", "completed"]),
                 )
             )
             for task in task_res.scalars().all():

+ 25 - 56
app/services/vas_task_service.py

@@ -4,7 +4,7 @@ from datetime import datetime, date, timedelta
 from typing import List, Optional
 
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy import select, or_, and_
+from sqlalchemy import func, text, select, or_, and_
 from redis.asyncio import Redis  # 引入 Redis 类型
 
 
@@ -32,62 +32,34 @@ class VasTaskService:
         return rec
     
     @staticmethod
-    async def pop_vas_task(session: AsyncSession, routing_key: str, cooldown_seconds: int = 60):
-        """
-        异步获取任务,支持冷却期机制。
-        
-        :param session: 数据库异步会话
-        :param routing_key: 队列键值
-        :param cooldown_seconds: 失败后的冷却时间(秒),默认60秒
-        :return: VasTask 对象 or None
-        """
-        # 计算冷却截止时间:当前时间 - 冷却秒数
-        # 只有 updated_at 早于这个时间的重试任务,才会被提取
-        cutoff_time = datetime.utcnow() - timedelta(seconds=cooldown_seconds)
-
-        try:
-            # --- 构造查询语句 ---
-            stmt = (
-                select(VasTask)
-                .where(VasTask.routing_key == routing_key)
-                .where(VasTask.status == 'pending')
-                # === 核心逻辑:冷却期筛选 ===
-                .where(
-                    or_(
-                        # 情况1:这是一个全新任务 (从未尝试过)
-                        VasTask.attempt_count == 0,
-                        
-                        # 情况2:这是一个重试任务,且距离上次更新(失败)已经过了冷却期
-                        and_(
-                            VasTask.attempt_count > 0,
-                            VasTask.updated_at < cutoff_time
-                        )
+    async def pop_vas_task(db: AsyncSession, routing_key: str, cooldown_seconds: int = 60):
+        stmt = (
+            select(VasTask)
+            .where(VasTask.routing_key == routing_key)
+            .where(VasTask.status == 'pending')
+            .where(
+                or_(
+                    VasTask.attempt_count == 0,
+                    and_(
+                        VasTask.attempt_count > 0,
+                        VasTask.updated_at < (func.utc_timestamp() - text(f"INTERVAL {cooldown_seconds} SECOND"))
                     )
                 )
-                # 排序:优先级优先(0假设是最高优先级?),其次是先创建的
-                # 注意:根据你的业务,priority 可能 desc 才是高优先级,这里按 asc 写
-                .order_by(VasTask.priority.desc(), VasTask.id.asc())
-                .limit(1)
-                # MySQL 8.0+ 必加,跳过被锁定的行
-                .with_for_update(skip_locked=True)
             )
-
-            result = await session.execute(stmt)
-            task = result.scalar_one_or_none()
-
-            # --- 更新状态 ---
-            if task:
-                task.status = 'running'     # 标记为已被抓取
-                task.attempt_count += 1     # 增加尝试次数
-                task.updated_at = datetime.utcnow() # 更新时间(重置冷却计时起点)
-                
-                await session.commit()
-                # session.begin() 结束时自动 commit
-                return task
+            .order_by(VasTask.priority.desc(), VasTask.id.asc())
+            .limit(1)
+            .with_for_update(skip_locked=True)
+        )
+        result = await db.execute(stmt)
+        task = result.scalar_one_or_none()
+        if not task:
             raise NotFoundError(message="Task not found")
-        except Exception as e:
-            # 记录日志
-            raise e
+        task.status = 'running'
+        task.attempt_count += 1
+        task.updated_at = func.utc_timestamp() 
+        await db.commit()
+        await db.refresh(task)
+        return task
         
     @staticmethod
     async def get_expiring_tasks(db: AsyncSession, threshold_days: int = 3):
@@ -237,9 +209,6 @@ class VasTaskService:
         if not rec:
             raise NotFoundError("Task not exist")
         
-        if rec.status == "pending":
-            raise BizLogicError("Task is in queue already")
-        
         rec.status = "pending"
         await db.commit()
         await db.refresh(rec)