diff options
Diffstat (limited to 'data-collection')
| -rw-r--r-- | data-collection/components/subscriptions/main.py | 30 |
1 files changed, 18 insertions, 12 deletions
diff --git a/data-collection/components/subscriptions/main.py b/data-collection/components/subscriptions/main.py index 6a3f5cd..26b0c64 100644 --- a/data-collection/components/subscriptions/main.py +++ b/data-collection/components/subscriptions/main.py @@ -1,11 +1,11 @@ from dataclasses import dataclass, field, asdict from datetime import datetime, UTC from sys import stderr -from typing import TypedDict, List +from typing import TypedDict, List, cast from bson.objectid import ObjectId from feedparser import parse # type: ignore from pymongo.collection import Collection -from pymongo.results import UpdateResult +from pymongo.results import InsertOneResult, UpdateResult import schedule from components.database import subscriptions from components.subscriptions.typing import SubsDict @@ -44,15 +44,21 @@ class Subscription: if last_video_update > self.last_video_update: print("Updating", self._id) self.last_video_update = last_video_update - self.update_fields(["videos", "last_video_update"]) + self.update_videos() self.last_fetch = datetime.now(tz=UTC) - def update_fields(self, fields: List[str]) -> UpdateResult: - sub = asdict(self) - if self._in_db: - return self._collection.update_one( - {"_id": self._id}, - {"$set": {key: sub[key] for key in fields}}, - ) - self._in_db = True - return self._collection.replace_one({"_id": self._id}, sub, upsert=True) + def asdict(self) -> SubsDict: + return cast(SubsDict, asdict(self)) + + def insert(self) -> InsertOneResult: + return self._collection.insert_one(self.asdict()) + + def update_videos(self) -> UpdateResult: + return self._collection.update_one( + {"_id": self._id}, + {"$set": { + "videos": self.videos, + "last_video_update": self.last_video_update, + "last_fetch": self.last_fetch, + }}, + ) |