|
|
@@ -93,6 +93,34 @@ class ArticlesDecodeTaskMapper(DecodeArticleConst):
|
|
|
),
|
|
|
)
|
|
|
|
|
|
+ async def count_pending_tasks(self, source: int = None) -> int:
|
|
|
+ """统计队列中正在进行(INIT + PROCESSING)的任务数量,用于控制消费端压力"""
|
|
|
+ if source is not None:
|
|
|
+ query = f"""
|
|
|
+ SELECT COUNT(*) AS cnt
|
|
|
+ FROM {TABLE}
|
|
|
+ WHERE status IN (%s, %s) AND source = %s AND config_id = %s
|
|
|
+ """
|
|
|
+ params = (
|
|
|
+ self.TaskStatus.INIT,
|
|
|
+ self.TaskStatus.PROCESSING,
|
|
|
+ source,
|
|
|
+ self.CONFIG_ID,
|
|
|
+ )
|
|
|
+ else:
|
|
|
+ query = f"""
|
|
|
+ SELECT COUNT(*) AS cnt
|
|
|
+ FROM {TABLE}
|
|
|
+ WHERE status IN (%s, %s) AND config_id = %s
|
|
|
+ """
|
|
|
+ params = (
|
|
|
+ self.TaskStatus.INIT,
|
|
|
+ self.TaskStatus.PROCESSING,
|
|
|
+ self.CONFIG_ID,
|
|
|
+ )
|
|
|
+ rows = await self.pool.async_fetch(query=query, params=params)
|
|
|
+ return rows[0]["cnt"] if rows else 0
|
|
|
+
|
|
|
async def fetch_pending_tasks(self, source: int = None) -> List[Dict]:
|
|
|
if source is not None:
|
|
|
query = f"""
|
|
|
@@ -202,7 +230,7 @@ class AdPlatformArticlesDecodeTaskMapper(ArticlesDecodeTaskMapper):
|
|
|
query=query, params=(new_status, id_, ori_status)
|
|
|
)
|
|
|
|
|
|
- async def fetch_decode_articles(self) -> List[Dict]:
|
|
|
+ async def fetch_decode_articles(self, limit: int = None) -> List[Dict]:
|
|
|
query = """
|
|
|
SELECT id, account_name, gh_id, article_title, article_cover,
|
|
|
article_text, article_images, wx_sn
|
|
|
@@ -212,7 +240,7 @@ class AdPlatformArticlesDecodeTaskMapper(ArticlesDecodeTaskMapper):
|
|
|
"""
|
|
|
return await self.pool.async_fetch(
|
|
|
query=query,
|
|
|
- params=(self.TaskStatus.SUCCESS, self.TaskStatus.INIT, self.TASK_BATCH),
|
|
|
+ params=(self.TaskStatus.SUCCESS, self.TaskStatus.INIT, limit if limit is not None else self.TASK_BATCH),
|
|
|
)
|
|
|
|
|
|
|
|
|
@@ -222,7 +250,7 @@ class InnerArticlesDecodeTaskMapper(ArticlesDecodeTaskMapper):
|
|
|
def __init__(self, pool: DatabaseManager):
|
|
|
super().__init__(pool)
|
|
|
|
|
|
- async def fetch_inner_articles(self) -> List[Dict]:
|
|
|
+ async def fetch_inner_articles(self, limit: int = None) -> List[Dict]:
|
|
|
query = f"""
|
|
|
SELECT id, title, source_id, coverimgurl, article_text, summary, card_title
|
|
|
FROM {self.TABLE_INNER}
|
|
|
@@ -230,7 +258,7 @@ class InnerArticlesDecodeTaskMapper(ArticlesDecodeTaskMapper):
|
|
|
LIMIT %s
|
|
|
"""
|
|
|
return await self.pool.async_fetch(
|
|
|
- query=query, params=(self.TaskStatus.INIT, self.TASK_BATCH)
|
|
|
+ query=query, params=(self.TaskStatus.INIT, limit if limit is not None else self.TASK_BATCH),
|
|
|
)
|
|
|
|
|
|
async def update_inner_article_status(
|