feat: 将初始化数据库拆分为单独函数

jobs/__init__.py

@@ -20,7 +20,7 @@ def import_deployment(path: str) -> DeploymentType:
     "jobs.fetch_jianshu_lottery_win_records:deployment",
     "jobs.fetch_jpep_ftn_trade_orders:buy_deployment",
     "jobs.fetch_jpep_ftn_trade_orders:sell_deployment",
-    "jobs.fetch_lp_recommended_article_records:deployment"
+    "jobs.fetch_lp_recommended_article_records:deployment",
 }
 
 

jobs/fetch_article_earning_ranking_records.py

@@ -7,9 +7,12 @@
 from jkit.ranking.article_earning import ArticleEarningRanking, RecordField
 from prefect import flow, get_run_logger
 from prefect.states import Completed, State
-from pymongo import ASCENDING, IndexModel
+from pymongo import IndexModel
 
-from utils.config_generators import generate_deployment_config, generate_flow_config
+from utils.config_generators import (
+    generate_deployment_config,
+    generate_flow_config,
+)
 from utils.db import DB
 from utils.document_model import (
     DOCUMENT_OBJECT_CONFIG,
@@ -71,6 +74,10 @@ async def get_author_id_and_slug(
         return None, None
 
 
+async def init_db() -> None:
+    await COLLECTION.create_indexes([IndexModel([("date", "ranking")], unique=True)])
+
+
 async def process_item(
     item: RecordField, /, *, target_date: date
 ) -> ArticleEarningRankingRecordDocument:
@@ -101,9 +108,7 @@ async def process_item(
     )
 )
 async def flow_func() -> State:
-    await COLLECTION.create_indexes(
-        [IndexModel([("date", ASCENDING), ("ranking", ASCENDING)], unique=True)]
-    )
+    await init_db()
 
     target_date = datetime.now().date() - timedelta(days=1)
 

jobs/fetch_assets_ranking_records.py

@@ -8,8 +8,12 @@
 from jkit.user import User
 from prefect import flow, get_run_logger
 from prefect.states import Completed, State
+from pymongo import IndexModel
 
-from utils.config_generators import generate_deployment_config, generate_flow_config
+from utils.config_generators import (
+    generate_deployment_config,
+    generate_flow_config,
+)
 from utils.db import DB
 from utils.document_model import (
     DOCUMENT_OBJECT_CONFIG,
@@ -66,6 +70,10 @@ async def get_fp_ftn_amount(
         return None, None
 
 
+async def init_db() -> None:
+    await COLLECTION.create_indexes([IndexModel(("date", "ranking"), unique=True)])
+
+
 async def process_item(
     item: AssetsRankingRecord, /, *, target_date: date
 ) -> AssetsRankingRecordDocument:
@@ -91,6 +99,8 @@ async def process_item(
     )
 )
 async def flow_func() -> State:
+    await init_db()
+
     target_date = datetime.now().date()
 
     data: List[AssetsRankingRecordDocument] = []

jobs/fetch_daily_update_ranking_records.py

@@ -5,8 +5,12 @@
 from jkit.ranking.daily_update import DailyUpdateRanking, DailyUpdateRankingRecord
 from prefect import flow
 from prefect.states import Completed, State
+from pymongo import IndexModel
 
-from utils.config_generators import generate_deployment_config, generate_flow_config
+from utils.config_generators import (
+    generate_deployment_config,
+    generate_flow_config,
+)
 from utils.db import DB
 from utils.document_model import (
     DOCUMENT_OBJECT_CONFIG,
@@ -30,6 +34,12 @@ class DailyUpdateRankingRecordDocument(Documemt, **DOCUMENT_OBJECT_CONFIG):
     user_info: UserInfoField
 
 
+async def init_db() -> None:
+    await COLLECTION.create_indexes(
+        [IndexModel(("date", "user_info.slug"), unique=True)]
+    )
+
+
 def process_item(
     item: DailyUpdateRankingRecord, /, *, current_date: date
 ) -> DailyUpdateRankingRecordDocument:
@@ -50,6 +60,8 @@ def process_item(
     )
 )
 async def flow_func() -> State:
+    await init_db()
+
     current_date = datetime.now().date()
 
     data: List[DailyUpdateRankingRecordDocument] = []

jobs/fetch_jianshu_lottery_win_records.py

@@ -5,9 +5,12 @@
 from jkit.jianshu_lottery import JianshuLottery, JianshuLotteryWinRecord
 from prefect import flow, get_run_logger
 from prefect.states import Completed, State
-from pymongo import ASCENDING, DESCENDING, IndexModel
+from pymongo import IndexModel
 
-from utils.config_generators import generate_deployment_config, generate_flow_config
+from utils.config_generators import (
+    generate_deployment_config,
+    generate_flow_config,
+)
 from utils.db import DB
 from utils.document_model import (
     DOCUMENT_OBJECT_CONFIG,
@@ -35,14 +38,18 @@ class JianshuLotteryWinRecordDocument(Documemt, **DOCUMENT_OBJECT_CONFIG):
 async def get_latest_stored_record_id() -> int:
     try:
         latest_data = JianshuLotteryWinRecordDocument.from_dict(
-            await COLLECTION.find().sort("_id", DESCENDING).__anext__()
+            await COLLECTION.find().sort("_id", -1).__anext__()
         )
     except StopAsyncIteration:
         return 0
 
     return latest_data.record_id
 
 
+async def init_db() -> None:
+    await COLLECTION.create_indexes([IndexModel(("recordId",), unique=True)])
+
+
 def process_item(item: JianshuLotteryWinRecord, /) -> JianshuLotteryWinRecordDocument:
     return JianshuLotteryWinRecordDocument(
         record_id=item.id,
@@ -62,9 +69,7 @@ def process_item(item: JianshuLotteryWinRecord, /) -> JianshuLotteryWinRecordDoc
     )
 )
 async def flow_func() -> State:
-    await COLLECTION.create_indexes(
-        [IndexModel([("recordId", ASCENDING)], unique=True)]
-    )
+    await init_db()
 
     logger = get_run_logger()
 

jobs/fetch_jpep_ftn_trade_orders.py

@@ -9,8 +9,12 @@
 from jkit.jpep.ftn_macket import FTNMacket, FTNMacketOrderRecord
 from prefect import flow
 from prefect.states import Completed, State
+from pymongo import IndexModel
 
-from utils.config_generators import generate_deployment_config, generate_flow_config
+from utils.config_generators import (
+    generate_deployment_config,
+    generate_flow_config,
+)
 from utils.db import DB
 from utils.document_model import (
     DOCUMENT_OBJECT_CONFIG,
@@ -54,6 +58,10 @@ def get_fetch_time() -> datetime:
     return current_dt.replace(minute=current_dt.minute // 10, second=0, microsecond=0)
 
 
+async def init_db() -> None:
+    await COLLECTION.create_indexes([IndexModel(("fetchTime", "orderId"), unique=True)])
+
+
 def process_item(
     item: FTNMacketOrderRecord,
     /,
@@ -88,6 +96,8 @@ def process_item(
     ),
 )
 async def flow_func(type: Literal["buy", "sell"]) -> State:  # noqa: A002
+    await init_db()
+
     fetch_time = get_fetch_time()
 
     data: List[JPEPFTNTradeOrderDocument] = []

jobs/fetch_lp_recommended_article_records.py

@@ -12,8 +12,12 @@
 from msgspec import field
 from prefect import flow, get_run_logger
 from prefect.states import Completed, State
+from pymongo import IndexModel
 
-from utils.config_generators import generate_deployment_config, generate_flow_config
+from utils.config_generators import (
+    generate_deployment_config,
+    generate_flow_config,
+)
 from utils.db import DB
 from utils.document_model import (
     DOCUMENT_OBJECT_CONFIG,
@@ -62,6 +66,10 @@ async def is_stored(item: CollectionArticleInfo) -> bool:
     return False
 
 
+async def init_db() -> None:
+    await COLLECTION.create_indexes([IndexModel(("date", "slug"), unique=True)])
+
+
 async def process_item(
     item: CollectionArticleInfo, /, *, current_date: date
 ) -> Optional[LPRecommendedArticleRecord]:
@@ -99,6 +107,8 @@ async def process_item(
     )
 )
 async def flow_func() -> State:
+    await init_db()
+
     logger = get_run_logger()
 
     current_date = datetime.now().date()