Michael Ingvarsson
73 lines
2.4 KiB
Python
73 lines
2.4 KiB
Python
import logging
|
|
import threading
|
|
from pydantic import BaseModel
|
|
from typing import Any, Optional
|
|
from gitea.models import IssueModel, PullRequestModel
|
|
|
|
logger: logging.Logger = logging.getLogger("work-queue")
|
|
|
|
|
|
class WorkItem(BaseModel):
|
|
repo_full_name: str
|
|
task_type: str # 'issue' or 'pr'
|
|
task_number: int
|
|
task_info: IssueModel | PullRequestModel
|
|
notification_id: Optional[int] = None
|
|
priority: int = 0
|
|
|
|
|
|
class WorkQueue:
|
|
"""Thread-safe work queue grouped by repo."""
|
|
|
|
def __init__(self) -> None:
|
|
self._lock: threading.Lock = threading.Lock()
|
|
self._queue: list[WorkItem] = []
|
|
self._enqueued_repos: set[str] = set()
|
|
|
|
def enqueue(self, item: WorkItem) -> None:
|
|
with self._lock:
|
|
self._queue.append(item)
|
|
self._enqueued_repos.add(item.repo_full_name)
|
|
logger.info(f"Enqueued work item: {item.task_type} #{item.task_number} for {item.repo_full_name}")
|
|
|
|
def enqueue_batch(self, items: list[WorkItem]) -> None:
|
|
with self._lock:
|
|
for item in items:
|
|
self._queue.append(item)
|
|
self._enqueued_repos.add(item.repo_full_name)
|
|
logger.info(f"Enqueued work item: {item.task_type} #{item.task_number} for {item.repo_full_name}")
|
|
|
|
def get_repo_work(self, repo: str) -> list[WorkItem]:
|
|
"""Get all work items for a specific repo."""
|
|
with self._lock:
|
|
items: list[WorkItem] = [
|
|
item for item in self._queue if item.repo_full_name == repo
|
|
]
|
|
logger.info(f"Retrieved {len(items)} work items for repository: {repo}")
|
|
return items
|
|
|
|
def remove_repo_work(self, repo: str) -> None:
|
|
"""Remove all work items for a specific repo."""
|
|
with self._lock:
|
|
self._queue = [
|
|
item for item in self._queue if item.repo_full_name != repo
|
|
]
|
|
self._enqueued_repos.discard(repo)
|
|
logger.info(f"Removed all work items for repository: {repo}")
|
|
|
|
def get_next_repo(self) -> str | None:
|
|
"""Get the next repo with work, or None if empty."""
|
|
with self._lock:
|
|
if not self._enqueued_repos:
|
|
return None
|
|
return next(iter(self._enqueued_repos))
|
|
|
|
@property
|
|
def is_empty(self) -> bool:
|
|
with self._lock:
|
|
return len(self._queue) == 0
|
|
|
|
def __len__(self) -> int:
|
|
with self._lock:
|
|
return len(self._queue)
|