From 42014695e3523d8cff41757ba966c968886ca2fc Mon Sep 17 00:00:00 2001 From: EstrellaXD Date: Sat, 20 May 2023 15:22:08 +0800 Subject: [PATCH] fix: Close 272 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit change: 重构种子管理器 --- src/module/api/bangumi.py | 51 ++++-------------- src/module/database/bangumi.py | 3 +- src/module/database/connector.py | 7 ++- src/module/database/user.py | 6 +-- src/module/manager/torrent.py | 89 +++++++++++++++++++++++++------- src/module/models/bangumi.py | 1 + 6 files changed, 91 insertions(+), 66 deletions(-) diff --git a/src/module/api/bangumi.py b/src/module/api/bangumi.py index 4ede6b3f..3fb79763 100644 --- a/src/module/api/bangumi.py +++ b/src/module/api/bangumi.py @@ -1,5 +1,3 @@ -import sqlite3 - from .log import router from module.models import BangumiData @@ -9,59 +7,32 @@ from module.manager import TorrentManager @router.get("/api/v1/bangumi/getAll", tags=["bangumi"], response_model=list[BangumiData]) async def get_all_data(): - try: - with BangumiDatabase() as database: - return database.search_all() - except sqlite3.OperationalError: - return [] + with TorrentManager() as torrent: + return torrent.search_all() @router.get("/api/v1/bangumi/getData/{bangumi_id}", tags=["bangumi"], response_model=BangumiData) async def get_data(bangumi_id: str): - with BangumiDatabase() as database: - data = database.search_id(int(bangumi_id)) - if data: - return data - else: - return {"": "data not exist"} + with TorrentManager() as torrent: + return torrent.search_data(bangumi_id) @router.post("/api/v1/bangumi/updateData", tags=["bangumi"]) -async def update_data(old_data: BangumiData, new_data: BangumiData): - with BangumiDatabase() as database: - updated = database.update_one(new_data) - if updated: - with TorrentManager() as torrent: - torrent.set_new_path(old_data, new_data) - return {"status": "ok"} - else: - return {"status": "data not exist"} +async def update_data(data: BangumiData): + with TorrentManager() as torrent: + return torrent.update_rule(data) @router.delete("/api/v1/bangumi/deleteData/{bangumi_id}", tags=["bangumi"]) async def delete_data(bangumi_id: str): - with BangumiDatabase() as database: - _id = int(bangumi_id) - deleted = database.delete_one(_id) - if deleted: - return {"status": "ok"} - else: - return {"status": "data not exist"} + with TorrentManager() as torrent: + return torrent.delete_data(bangumi_id) @router.delete("/api/v1/bangumi/deleteRule/{bangumi_id}", tags=["bangumi"]) async def delete_rule(bangumi_id: str, file: bool = False): - with BangumiDatabase() as database: - _id = int(bangumi_id) - data = database.search_id(_id) - if data: - with TorrentManager() as torrent: - torrent.delete_rule(data) - if file: - torrent.delete_bangumi(data) - return {"status": "ok"} - else: - return {"status": "data not exist"} + with TorrentManager() as torrent: + return torrent.delete_rule(bangumi_id, file) @router.get("/api/v1/bangumi/resetAll", tags=["bangumi"]) diff --git a/src/module/database/bangumi.py b/src/module/database/bangumi.py index 8c652c30..083f5cff 100644 --- a/src/module/database/bangumi.py +++ b/src/module/database/bangumi.py @@ -200,7 +200,8 @@ class BangumiDatabase(DataConnector): def not_added(self) -> list[BangumiData]: self._cursor.execute( """ - SELECT * FROM bangumi WHERE added = 0 + SELECT * FROM bangumi + WHERE added = 0 or save_path is null or rule_name is null """ ) return self.__fetch_data() diff --git a/src/module/database/connector.py b/src/module/database/connector.py index 03060f53..f8cea42c 100644 --- a/src/module/database/connector.py +++ b/src/module/database/connector.py @@ -23,7 +23,10 @@ class DataConnector: existing_columns = {column_info[1]: column_info for column_info in self._cursor.fetchall()} for key, value in db_data.items(): if key not in existing_columns: - add_column_sql = f"ALTER TABLE {table_name} ADD COLUMN {key} {self.__python_to_sqlite_type(value)} DEFAULT {value};" + insert_column = self.__python_to_sqlite_type(value) + if value is None: + value = "NULL" + add_column_sql = f"ALTER TABLE {table_name} ADD COLUMN {key} {insert_column} DEFAULT {value};" self._cursor.execute(add_column_sql) self._conn.commit() logger.debug(f"Create / Update table {table_name}.") @@ -74,7 +77,7 @@ class DataConnector: self._conn.commit() def _table_exists(self, table_name: str) -> bool: - self._cursor.execute(f"SELECT name FROM sqlite_master WHERE type='table' AND name='{table_name}'") + self._cursor.execute(f"SELECT name FROM sqlite_master WHERE type='table' AND name=?;", (table_name,)) return len(self._cursor.fetchall()) == 1 @staticmethod diff --git a/src/module/database/user.py b/src/module/database/user.py index ddca0cc8..5a1ef836 100644 --- a/src/module/database/user.py +++ b/src/module/database/user.py @@ -5,7 +5,7 @@ from fastapi import HTTPException from module.database.connector import DataConnector from module.security.jwt import get_password_hash, verify_password -from module.models.user import UserLogin, User +from module.models.user import User logger = logging.getLogger(__name__) @@ -33,7 +33,7 @@ class AuthDB(DataConnector): return User(**db_data) def get_user(self, username): - self._cursor.execute(f"SELECT * FROM {self.__table_name} WHERE username='{username}'") + self._cursor.execute(f"SELECT * FROM {self.__table_name} WHERE username=?", (username,)) result = self._cursor.fetchone() if not result: return None @@ -41,7 +41,7 @@ class AuthDB(DataConnector): return self.__db_to_data(db_data) def auth_user(self, username, password) -> bool: - self._cursor.execute(f"SELECT username, password FROM {self.__table_name} WHERE username='{username}'") + self._cursor.execute(f"SELECT username, password FROM {self.__table_name} WHERE username=?", (username,)) result = self._cursor.fetchone() if not result: raise HTTPException(status_code=404, detail="User not found") diff --git a/src/module/manager/torrent.py b/src/module/manager/torrent.py index 5b13f0a1..ec129038 100644 --- a/src/module/manager/torrent.py +++ b/src/module/manager/torrent.py @@ -1,37 +1,86 @@ import logging from module.downloader import DownloadClient -from module.conf import settings from module.models import BangumiData +from module.database import BangumiDatabase logger = logging.getLogger(__name__) -class TorrentManager(DownloadClient): - def __match_torrents_list(self, data: BangumiData) -> list: - torrents = self.get_torrent_info() +class TorrentManager(BangumiDatabase): + @staticmethod + def __match_torrents_list(data: BangumiData) -> list: + with DownloadClient() as client: + torrents = client.get_torrent_info() matched_list = [] for torrent in torrents: if data.save_path == torrent.save_path: matched_list.append(torrent.hash) return matched_list - def delete_bangumi(self, data: BangumiData): - hash_list = self.__match_torrents_list(data) - self.delete_torrent(hash_list) + def delete_torrents(self, _id: int | str): + data = self.search_one(int(_id)) + if isinstance(data, BangumiData): + hash_list = self.__match_torrents_list(data) + with DownloadClient() as client: + client.delete_torrent(hash_list) + return {"status": "success", "msg": f"Delete torrents for {data.official_title}"} + else: + return data - def delete_rule(self, data: BangumiData): - rule_name = f"{data.official_title}({data.year})" if data.year else data.title_raw - if settings.bangumi_manage.group_tag: - rule_name = f"[{data.group_name}] {rule_name}" - self.remove_rule(rule_name) + def delete_data(self, _id: int | str): + data = self.search_id(int(_id)) + if isinstance(data, BangumiData): + self.delete_one(int(_id)) + logger.info(f"Delete {data.official_title}") + return {"status": "success", "msg": f"Delete {data.official_title}"} + else: + return data - def set_new_path(self, old_data: BangumiData, new_data: BangumiData): - # set download rule - self.remove_rule(old_data.rule_name) - self.set_rule(new_data) - # set torrent path - match_list = self.__match_torrents_list(old_data) - path = self._gen_save_path(new_data) - self.move_torrent(match_list, path) + def delete_rule(self, _id: str | int, file: bool = False): + data = self.search_id(int(_id)) + if isinstance(data, BangumiData): + with DownloadClient() as client: + client.remove_rule(data.rule_name) + data.deleted = True + self.update_one(data) + if file: + self.delete_torrents(data.id) + logger.info(f"Delete rule and torrents for {data.official_title}") + return {"status": "success", "msg": f"Delete rule and torrents for {data.official_title}"} + logger.info(f"Delete rule for {data.official_title}") + return {"status": "success", "msg": f"Delete rule for {data.official_title}"} + else: + return data + + def update_rule(self, data: BangumiData): + old_data = self.search_id(data.id) + if not old_data: + logger.error(f"Can't find data with id {data.id}") + return {"status": "error", "msg": f"Can't find data with id {data.id}"} + else: + # Set torrent path + match_list = self.__match_torrents_list(data) + with DownloadClient() as client: + path = client._gen_save_path(data) + client.move_torrent(match_list, path) + # Set new download rule + client.remove_rule(data.rule_name) + client.set_rule(data) + self.update_one(data) + return {"status": "success", "msg": f"Set new path for {data.official_title}"} + + def search_all_bangumi(self): + datas = self.search_all() + if not datas: + return [] + return [data for data in datas if not data.deleted] + + def search_one(self, _id: int | str): + data = self.search_id(int(_id)) + if not data: + logger.error(f"Can't find data with id {_id}") + return {"status": "error", "msg": f"Can't find data with id {_id}"} + else: + return data diff --git a/src/module/models/bangumi.py b/src/module/models/bangumi.py index a2af9851..fb9c1495 100644 --- a/src/module/models/bangumi.py +++ b/src/module/models/bangumi.py @@ -21,6 +21,7 @@ class BangumiData(BaseModel): added: bool = Field(False, alias="added", title="是否已添加") rule_name: str | None = Field(None, alias="rule_name", title="番剧规则名") save_path: str | None = Field(None, alias="save_path", title="番剧保存路径") + deleted: bool = Field(False, alias="deleted", title="是否已删除") class Notification(BaseModel):