diff --git a/pyproject.toml b/pyproject.toml index cbeedd2..e46b01f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,10 @@ readme = "README.md" requires-python = ">=3.7" version = "0.1.0" +[project.optional-dependencies] +pyside6 = ["PySide6"] +pyqt6 = ["PyQt6"] + [[tool.uv.index]] default = true url = "https://pypi.tuna.tsinghua.edu.cn/simple" diff --git a/src/queue_sqlite/core/__init__.pyi b/src/queue_sqlite/core/__init__.pyi index 0bea59b..e3180e3 100644 --- a/src/queue_sqlite/core/__init__.pyi +++ b/src/queue_sqlite/core/__init__.pyi @@ -5,48 +5,47 @@ from ..constant import MessageStatus class core: class QueueOperation: def __init__(self, db_path: str): + """ + Initialize the queue operation. + + Args: + db_path (str): The path to the database file. + + """ ... def init_db(self): + """ + Initialize the database. + """ ... def enqueue(self, message_dict: dict) -> str: - ... - - def dequeue(self, size: int = 1) -> List[dict]: + """ + Enqueue a message. + + Args: + message_dict (dict): The message to be enqueued. + + Returns: + str: The message id. + + Raises: + ValueError: If the message is not valid. + """ ... - def get_queue_length(self) -> int: - ... + def dequeue(self, size: int = 1) -> List[dict]: ... + def get_queue_length(self) -> int: ... + def get_completed_messages(self) -> List[dict]: ... + def get_result(self, message_id: str): ... + def update_status(self, message_id: str, status: MessageStatus): ... + def update_result(self, message_id: str, result: str): ... + def delete_message(self, message_id: str): ... + def clean_old_messages(self, days: int): ... + def clean_expired_messages(self): ... - def get_completed_messages(self) -> List[dict]: - ... - - def get_result(self, message_id: str): - ... - - def update_status(self, message_id: str, status: MessageStatus): - ... - - def update_result(self, message_id: str, result: str): - ... - - def delete_message(self, message_id: str): - ... - - def clean_old_messages(self, days: int): - ... - - def clean_expired_messages(self): - ... - class TaskMounter: - def __init__(self, task_class: type[TaskMounter]): - ... - - def get_task_list(self) -> List[str]: - ... - - def get_task_function(self, name: str) -> Callable: - ... - \ No newline at end of file + def __init__(self, task_class: type[TaskMounter]): ... + def get_task_list(self) -> List[str]: ... + def get_task_function(self, name: str) -> Callable: ... diff --git a/src/queue_sqlite/model/__init__.py b/src/queue_sqlite/model/__init__.py index dd25f13..5e7bb6a 100644 --- a/src/queue_sqlite/model/__init__.py +++ b/src/queue_sqlite/model/__init__.py @@ -1,3 +1,4 @@ from .message_item import MessageItem -__all__ = ["MessageItem"] \ No newline at end of file + +__all__ = ["MessageItem"] diff --git a/src/queue_sqlite/queue_operation/listen_operation.py b/src/queue_sqlite/queue_operation/listen_operation.py index 7004dab..c54819d 100644 --- a/src/queue_sqlite/queue_operation/listen_operation.py +++ b/src/queue_sqlite/queue_operation/listen_operation.py @@ -9,6 +9,13 @@ class ListenOperation: self.listen_fields = ListenMounter.get_Listener_list() self.create_table() + def _get_connection(self): + """获取数据库连接""" + return sqlite3.connect( + self.db_dir, + check_same_thread=False, + ) + def create_table(self): if len(self.listen_fields) == 0: return @@ -37,25 +44,22 @@ class ListenOperation: VALUES ('listen_table', NEW.id, 'key', OLD.key, NEW.key); END; """ - self.conn = sqlite3.connect( - self.db_dir, - check_same_thread = False, - ) - self.conn.executescript(sql) - self.conn.execute("PRAGMA journal_mode=WAL;") - self.conn.execute("PRAGMA synchronous=NORMAL;") - self.conn.execute("PRAGMA cache_size=-20000;") - self.conn.execute("PRAGMA mmap_size=1073741824;") - self.conn.execute("PRAGMA temp_store=MEMORY;") - self.conn.execute("PRAGMA busy_timeout=5000;") - self.conn.commit() + conn = self._get_connection() + conn.executescript(sql) + conn.execute("PRAGMA journal_mode=WAL;") + conn.execute("PRAGMA synchronous=NORMAL;") + conn.execute("PRAGMA cache_size=-20000;") + conn.execute("PRAGMA mmap_size=1073741824;") + conn.execute("PRAGMA temp_store=MEMORY;") + conn.execute("PRAGMA busy_timeout=5000;") + conn.commit() # 删除原有数据 # sql = f""" # DELETE FROM listen_table # """ - # self.conn.execute(sql) - # self.conn.commit() + # conn.execute(sql) + # conn.commit() for listen_field in self.listen_fields: sql = f""" @@ -64,44 +68,48 @@ class ListenOperation: VALUES (?, ?) """ - self.conn.execute(sql, (listen_field, "null")) - self.conn.commit() - + conn.execute(sql, (listen_field, "null")) + conn.commit() + def listen_data(self) -> Tuple[bool, Union[List[Tuple], str]]: sql = f""" SELECT * FROM change_log where is_delete = 0 ORDER BY id DESC LIMIT 100 """ - result = self.conn.execute(sql).fetchall() + conn = self._get_connection() + result = conn.execute(sql).fetchall() if len(result) == 0: return False, "No data found" return True, result - - + def delete_change_log(self, delete_id): sql = f""" DELETE FROM change_log WHERE id = {delete_id} """ - self.conn.execute(sql) - + conn = self._get_connection() + conn.execute(sql) + def update_listen_data(self, key, value): sql = f""" UPDATE listen_table SET value = '{value}' WHERE key = '{key}' """ - self.conn.execute(sql) - self.conn.commit() + conn = self._get_connection() + conn.execute(sql) + conn.commit() def get_value(self, key): sql = f""" SELECT value FROM listen_table WHERE key = '{key}' """ - result = self.conn.execute(sql).fetchone() + conn = self._get_connection() + result = conn.execute(sql).fetchone() if result is None: return None return result[0] - + def get_values(self): sql = f""" SELECT key, value FROM listen_table """ - result = self.conn.execute(sql).fetchall() + conn = self._get_connection() + result = conn.execute(sql).fetchall() return result diff --git a/src/queue_sqlite/queue_operation/queue_operation.py b/src/queue_sqlite/queue_operation/queue_operation.py index 19a070d..b82e4c2 100644 --- a/src/queue_sqlite/queue_operation/queue_operation.py +++ b/src/queue_sqlite/queue_operation/queue_operation.py @@ -1,5 +1,4 @@ import os -import json import hashlib from typing import Tuple, Union, List import threading diff --git a/src/queue_sqlite/scheduler/__init__.py b/src/queue_sqlite/scheduler/__init__.py index 13bf376..641ca0e 100644 --- a/src/queue_sqlite/scheduler/__init__.py +++ b/src/queue_sqlite/scheduler/__init__.py @@ -1,49 +1,56 @@ -from .receive_scheduler import ReceiveScheduler -from .task_scheduler import TaskScheduler -from .listen_data_scheduler import ListenDataScheduler +from .standard import StandardQueueScheduler +from .base import BaseScheduler from ..model import MessageItem from typing import Callable -from ..queue_operation import QueueOperation -from ..queue_operation.listen_operation import ListenOperation -from .cleanup_scheduler import CleanupScheduler -import os -class QueueScheduler: - def __init__(self, receive_thread_num=1, task_thread_num=1, shard_num=4): - self.queue_operation = QueueOperation(shard_num) - self.listen_operation = ListenOperation(os.path.join(self.queue_operation.db_dir, "listen.db")) - self.listen_operation.create_table() - - self.listen_scheduler = ListenDataScheduler(self.listen_operation) - self.receive_scheduler = ReceiveScheduler(self.queue_operation, receive_thread_num) - self.task_scheduler = TaskScheduler(self.queue_operation, task_thread_num) - self.cleanup_scheduler = CleanupScheduler(self.queue_operation) + +class QueueScheduler(BaseScheduler): + def __init__( + self, + receive_thread_num=1, + task_thread_num=1, + shard_num=4, + scheduler_type="standard", + ): + self.scheduler = None + if scheduler_type == "standard": + self.scheduler = StandardQueueScheduler( + receive_thread_num, task_thread_num, shard_num + ) + else: + pass + if self.scheduler: + self.queue_operation = self.scheduler.queue_operation def send_message(self, message: MessageItem, callback: Callable): - self.receive_scheduler.send_message(message, callback) + if not self.scheduler: + raise Exception("Scheduler not initialized") + self.scheduler.send_message(message, callback) + + def start(self): + if not self.scheduler: + raise Exception("Scheduler not initialized") + self.scheduler.start() + + def stop(self): + if not self.scheduler: + raise Exception("Scheduler not initialized") + self.scheduler.stop() def update_listen_data(self, key, value): - self.listen_operation.update_listen_data(key, value) + if not self.scheduler: + raise Exception("Scheduler not initialized") + self.scheduler.update_listen_data(key, value) def get_listen_datas(self): - return self.listen_operation.get_values() - + if not self.scheduler: + raise Exception("Scheduler not initialized") + return self.scheduler.get_listen_datas() + def get_listen_data(self, key): - return self.listen_operation.get_value(key) + if not self.scheduler: + raise Exception("Scheduler not initialized") + return self.scheduler.get_listen_data(key) - def start_queue_scheduler(self): - self.receive_scheduler.start_receive_thread() - self.task_scheduler.start_task_thread() - self.cleanup_scheduler.start_cleanup() - self.listen_scheduler.start_listen_data() - def stop_queue_scheduler(self): - self.receive_scheduler.stop_receive_thread() - self.listen_scheduler.stop_listen_data() - self.task_scheduler.stop_task_thread() - self.cleanup_scheduler.stop_cleanup() - # self.queue_operation.close_all_connections() - -__all__ = [ - "QueueScheduler" -] \ No newline at end of file +__all__ = ["QueueScheduler"] diff --git a/src/queue_sqlite/scheduler/async/__init__.py b/src/queue_sqlite/scheduler/async/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/queue_sqlite/scheduler/base/__init__.py b/src/queue_sqlite/scheduler/base/__init__.py new file mode 100644 index 0000000..6f92426 --- /dev/null +++ b/src/queue_sqlite/scheduler/base/__init__.py @@ -0,0 +1,4 @@ +from .base_scheduler import BaseScheduler + + +__all__ = ["BaseScheduler"] diff --git a/src/queue_sqlite/scheduler/base/base_scheduler.py b/src/queue_sqlite/scheduler/base/base_scheduler.py new file mode 100644 index 0000000..b6b0231 --- /dev/null +++ b/src/queue_sqlite/scheduler/base/base_scheduler.py @@ -0,0 +1,42 @@ +from abc import ABC, abstractmethod +from typing import Callable +from ...model import MessageItem + + +class BaseScheduler(ABC): + """调度器抽象类""" + + @abstractmethod + def send_message(self, message: MessageItem, callback: Callable): + """发送消息到队列 + + Args: + message (MessageItem): 消息对象 + callback (Callable): 发送完成后的回调函数 + """ + pass + + @abstractmethod + def start(self): + """启动调度器""" + pass + + @abstractmethod + def stop(self): + """停止调度器""" + pass + + @abstractmethod + def update_listen_data(self, key, value): + """更新监听数据""" + pass + + @abstractmethod + def get_listen_datas(self) -> list: + """获取监听数据""" + pass + + @abstractmethod + def get_listen_data(self, key): + """获取单个监听数据""" + pass diff --git a/src/queue_sqlite/scheduler/pyqt6/__init__.py b/src/queue_sqlite/scheduler/pyqt6/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/queue_sqlite/scheduler/pyside6/__init__.py b/src/queue_sqlite/scheduler/pyside6/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/queue_sqlite/scheduler/standard/__init__.py b/src/queue_sqlite/scheduler/standard/__init__.py new file mode 100644 index 0000000..c036569 --- /dev/null +++ b/src/queue_sqlite/scheduler/standard/__init__.py @@ -0,0 +1,54 @@ +from .receive_scheduler import ReceiveScheduler +from .task_scheduler import TaskScheduler +from .listen_data_scheduler import ListenDataScheduler +from ...model import MessageItem +from typing import Callable +from ...queue_operation import QueueOperation +from ...queue_operation.listen_operation import ListenOperation +from .cleanup_scheduler import CleanupScheduler +import os +from ..base import BaseScheduler + + +class StandardQueueScheduler(BaseScheduler): + def __init__(self, receive_thread_num=1, task_thread_num=1, shard_num=4): + self.queue_operation = QueueOperation(shard_num) + self.listen_operation = ListenOperation( + os.path.join(self.queue_operation.db_dir, "listen.db") + ) + self.listen_operation.create_table() + + self.listen_scheduler = ListenDataScheduler(self.listen_operation) + self.receive_scheduler = ReceiveScheduler( + self.queue_operation, receive_thread_num + ) + self.task_scheduler = TaskScheduler(self.queue_operation, task_thread_num) + self.cleanup_scheduler = CleanupScheduler(self.queue_operation) + + def send_message(self, message: MessageItem, callback: Callable): + self.receive_scheduler.send_message(message, callback) + + def update_listen_data(self, key, value): + self.listen_operation.update_listen_data(key, value) + + def get_listen_datas(self): + return self.listen_operation.get_values() + + def get_listen_data(self, key): + return self.listen_operation.get_value(key) + + def start(self): + self.receive_scheduler.start_receive_thread() + self.task_scheduler.start_task_thread() + self.cleanup_scheduler.start_cleanup() + self.listen_scheduler.start_listen_data() + + def stop(self): + self.receive_scheduler.stop_receive_thread() + self.listen_scheduler.stop_listen_data() + self.task_scheduler.stop_task_thread() + self.cleanup_scheduler.stop_cleanup() + # self.queue_operation.close_all_connections() + + +__all__ = ["StandardQueueScheduler"] diff --git a/src/queue_sqlite/scheduler/cleanup_scheduler.py b/src/queue_sqlite/scheduler/standard/cleanup_scheduler.py similarity index 84% rename from src/queue_sqlite/scheduler/cleanup_scheduler.py rename to src/queue_sqlite/scheduler/standard/cleanup_scheduler.py index 0e02f1c..7f0cb56 100644 --- a/src/queue_sqlite/scheduler/cleanup_scheduler.py +++ b/src/queue_sqlite/scheduler/standard/cleanup_scheduler.py @@ -1,7 +1,7 @@ # 新建清理调度器 import threading import time -from ..queue_operation import QueueOperation +from ...queue_operation import QueueOperation class CleanupScheduler: @@ -18,14 +18,14 @@ class CleanupScheduler: # 清理7天前的已完成/失败消息 for i in range(self.queue_operation.shard_num): self.queue_operation.clean_old_messages(i, 7) - + # 清理过期但未处理的消息 for i in range(self.queue_operation.shard_num): self.queue_operation.clean_expired_messages(i) - + except Exception as e: print(f"清理消息错误: {str(e)}") - + # 休眠等待下次清理 for _ in range(self.interval): if not self.is_running: @@ -35,15 +35,17 @@ class CleanupScheduler: def start_cleanup(self): if self.is_running: return - + self.is_running = True - self.cleanup_thread = threading.Thread(target=self.cleanup_expired_messages, daemon=True) + self.cleanup_thread = threading.Thread( + target=self.cleanup_expired_messages, daemon=True + ) self.cleanup_thread.start() def stop_cleanup(self): if not self.is_running: return - + self.is_running = False if self.cleanup_thread and self.cleanup_thread.is_alive(): - self.cleanup_thread.join(timeout=2.0) \ No newline at end of file + self.cleanup_thread.join(timeout=2.0) diff --git a/src/queue_sqlite/scheduler/listen_data_scheduler.py b/src/queue_sqlite/scheduler/standard/listen_data_scheduler.py similarity index 85% rename from src/queue_sqlite/scheduler/listen_data_scheduler.py rename to src/queue_sqlite/scheduler/standard/listen_data_scheduler.py index 173a965..1f76c82 100644 --- a/src/queue_sqlite/scheduler/listen_data_scheduler.py +++ b/src/queue_sqlite/scheduler/standard/listen_data_scheduler.py @@ -1,6 +1,6 @@ -from ..queue_operation.listen_operation import ListenOperation +from ...queue_operation.listen_operation import ListenOperation from concurrent.futures import ThreadPoolExecutor -from ..mounter.listen_mounter import ListenMounter +from ...mounter.listen_mounter import ListenMounter import threading @@ -29,7 +29,9 @@ class ListenDataScheduler: key = data[6] new_value = data[7] delete_id = data[0] - self.executor.submit(self._process_listen_data, key, new_value, delete_id) + self.executor.submit( + self._process_listen_data, key, new_value, delete_id + ) def start_listen_data(self): if self.is_running: @@ -46,4 +48,3 @@ class ListenDataScheduler: self.listen_thread.join(timeout=2) self.executor.shutdown(wait=True) - \ No newline at end of file diff --git a/src/queue_sqlite/scheduler/receive_scheduler.py b/src/queue_sqlite/scheduler/standard/receive_scheduler.py similarity index 91% rename from src/queue_sqlite/scheduler/receive_scheduler.py rename to src/queue_sqlite/scheduler/standard/receive_scheduler.py index c8438a8..3d8fc59 100644 --- a/src/queue_sqlite/scheduler/receive_scheduler.py +++ b/src/queue_sqlite/scheduler/standard/receive_scheduler.py @@ -1,5 +1,5 @@ -from ..queue_operation import QueueOperation -from ..model import MessageItem +from ...queue_operation import QueueOperation +from ...model import MessageItem from typing import Callable from concurrent.futures import ThreadPoolExecutor import threading @@ -11,12 +11,14 @@ class ReceiveScheduler: self.callbacks = dict() self.receive_thread_num = receive_thread_num self.is_running = False - self.executor = ThreadPoolExecutor(max_workers=receive_thread_num) # 并行执行回调 + self.executor = ThreadPoolExecutor( + max_workers=receive_thread_num + ) # 并行执行回调 self.lock = threading.Lock() self.queue_operation = queue_operation self.receive_thread = None # 单一轮询线程 - def send_message(self, message: MessageItem, callback: Callable = None): # type: ignore + def send_message(self, message: MessageItem, callback: Callable = None): # type: ignore if callback is None: callback = lambda message: print(message) # self.queue_operation.enqueue(message.to_dict()) @@ -37,12 +39,12 @@ class ReceiveScheduler: callback = self.callbacks.pop(message.id, None) if callback is None: callback = callback_default - + # 使用线程池并行执行回调 self.executor.submit(self._safe_callback, callback, message) else: time.sleep(0.1) # 适当休眠避免CPU空转 - + def _safe_callback(self, callback, message): """安全执行回调函数""" try: @@ -60,7 +62,7 @@ class ReceiveScheduler: def start_receive_thread(self): if self.is_running: return - + self.is_running = True # 创建单一轮询线程 self.receive_thread = threading.Thread(target=self.receive_message, daemon=True) @@ -69,12 +71,11 @@ class ReceiveScheduler: def stop_receive_thread(self): if not self.is_running: return - + self.is_running = False # 等待轮询线程结束 if self.receive_thread and self.receive_thread.is_alive(): self.receive_thread.join(timeout=2.0) - + # 关闭线程池 self.executor.shutdown(wait=True) - \ No newline at end of file diff --git a/src/queue_sqlite/scheduler/task_scheduler.py b/src/queue_sqlite/scheduler/standard/task_scheduler.py similarity index 94% rename from src/queue_sqlite/scheduler/task_scheduler.py rename to src/queue_sqlite/scheduler/standard/task_scheduler.py index 12e1da9..4785a81 100644 --- a/src/queue_sqlite/scheduler/task_scheduler.py +++ b/src/queue_sqlite/scheduler/standard/task_scheduler.py @@ -1,8 +1,8 @@ -from ..queue_operation import QueueOperation -from ..model import MessageItem -from ..constant import MessageStatus -from ..cycle.task_cycle import TaskCycle -from ..mounter.task_mounter import TaskMounter +from ...queue_operation import QueueOperation +from ...model import MessageItem +from ...constant import MessageStatus +from ...cycle.task_cycle import TaskCycle +from ...mounter.task_mounter import TaskMounter from concurrent.futures import ThreadPoolExecutor from datetime import datetime import time diff --git a/tests/test_core.py b/tests/test_core.py index 1eb16a9..bd0fb60 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -10,22 +10,27 @@ class TestCore: def test_task_mounter(cls): task_mounter = core.TaskMounter(TaskMounter) task_mounter.get_task_list() - print(task_mounter.get_task_function("task")({"num": 1})(lambda x: x+1).meta["num"]) + print( + task_mounter.get_task_function("task")({"num": 1})(lambda x: x + 1).meta[ + "num" + ] + ) from tasks import example + task_mounter.get_task_list() print(task_mounter.get_task_function("")) - + @classmethod def test_queue_operation_init_db(cls): cls.queue_operation.init_db() @classmethod def test_queue_operation_enqueue(cls): - message_item = MessageItem(content={"num": 1}, destination="test") + message_item = MessageItem(content={"num": 1}, destination="test") # type: ignore cls.queue_operation.enqueue(message_item.to_dict_by_core()) @classmethod def test_queue_operation_dequeue(cls): message_item = MessageItem(content={"num": 1}, destination="test") cls.queue_operation.enqueue(message_item.to_dict_by_core()) - print(cls.queue_operation.dequeue(1)) \ No newline at end of file + print(cls.queue_operation.dequeue(1)) diff --git a/tests/test_listen.py b/tests/test_listen.py index 915c30f..9428318 100644 --- a/tests/test_listen.py +++ b/tests/test_listen.py @@ -6,7 +6,7 @@ import time class TestListen: def test_listen_data(self): scheduler = QueueScheduler() - scheduler.start_queue_scheduler() + scheduler.start() scheduler.update_listen_data("key_1", "value_1") time.sleep(0.001) scheduler.update_listen_data("key_1", "value_2") @@ -15,6 +15,6 @@ class TestListen: time.sleep(0.001) scheduler.update_listen_data("key_1", "value_4") time.sleep(0.001) - scheduler.stop_queue_scheduler() + scheduler.stop() print(scheduler.get_listen_datas()) print(scheduler.get_listen_data("key_1")) diff --git a/tests/test_stress.py b/tests/test_stress.py index a4bc901..25fcde4 100644 --- a/tests/test_stress.py +++ b/tests/test_stress.py @@ -49,12 +49,14 @@ class TestStress: for i in range(TASK_NUM): # 增加任务数量 message_item = MessageItem(content={"num": i}, destination="example") messages.append(message_item) - + cls.is_monitoring = True monitor_thread = threading.Thread(target=cls.monitor_resources, daemon=True) monitor_thread.start() - queue_scheduler = QueueScheduler(receive_thread_num=4, task_thread_num=8, shard_num=12) - + queue_scheduler = QueueScheduler( + receive_thread_num=4, task_thread_num=8, shard_num=12 + ) + # threads = [] # for message_item in messages: # t = threading.Thread(target=queue_scheduler.send_message, args=(message_item, cls.test_callback)) @@ -62,26 +64,29 @@ class TestStress: # 添加性能计时 start_time = time.perf_counter() - - queue_scheduler.start_queue_scheduler() - + + queue_scheduler.start() + for message_item in messages: # 增加任务数量 message_item: MessageItem queue_scheduler.send_message(message_item, cls._callback) - + # # 多线程发送任务 # for t in threads: # t.start() - + # # 动态等待任务完成 start_wait = time.time() - while (queue_scheduler.queue_operation.get_queue_length() > 0 or cls.callback_counter < TASK_NUM): # 30秒超时 + while ( + queue_scheduler.queue_operation.get_queue_length() > 0 + or cls.callback_counter < TASK_NUM + ): # 30秒超时 time.sleep(0.5) print(f"当前队列长度:{queue_scheduler.queue_operation.get_queue_length()}") print(f"当前回调数量:{cls.callback_counter}") - - queue_scheduler.stop_queue_scheduler() - + + queue_scheduler.stop() + print(f"当前队列长度:{queue_scheduler.queue_operation.get_queue_length()}") print(f"测试函数:斐波那契数列") print(f"次数:{TASK_NUM}") @@ -89,4 +94,3 @@ class TestStress: cls.is_monitoring = False monitor_thread.join(timeout=1) - \ No newline at end of file