diff --git a/src/app.py b/src/app.py index 2d283051..eb42bc4f 100644 --- a/src/app.py +++ b/src/app.py @@ -530,12 +530,10 @@ def job_statuses(user: User = Depends(current_user)) -> JobsSchema: """Returns statuses of all the jobs in the system. May take some time (> 1s) when there are a lot of them. """ - jobs = [db.get_job(job) for job in db.get_job_ids()] - jobs = sorted(jobs, key=lambda j: j.submitted, reverse=True) - jobs = [j for j in jobs if j.status != "removed"] - if "can_list_all_queries" not in get_user_roles(user): - jobs = [j for j in jobs if j.rule_author == user.name] - + username_filter: Optional[str] = user.name + if "can_list_all_queries" in get_user_roles(user): + username_filter = None + jobs = db.get_valid_jobs(username_filter) return JobsSchema(jobs=jobs) diff --git a/src/db.py b/src/db.py index a609dcb5..abea95f2 100644 --- a/src/db.py +++ b/src/db.py @@ -7,7 +7,15 @@ from redis import StrictRedis from enum import Enum from rq import Queue # type: ignore -from sqlmodel import Session, SQLModel, create_engine, select, and_, update +from sqlmodel import ( + Session, + SQLModel, + create_engine, + select, + and_, + update, + col, +) from .models.agentgroup import AgentGroup from .models.configentry import ConfigEntry @@ -51,12 +59,6 @@ def session(self): with Session(self.engine) as session: yield session - def get_job_ids(self) -> List[JobId]: - """Gets IDs of all jobs in the database.""" - with self.session() as session: - jobs = session.exec(select(Job)).all() - return [j.id for j in jobs] - def cancel_job(self, job: JobId, error=None) -> None: """Sets the job status to cancelled, with optional error message.""" with self.session() as session: @@ -80,6 +82,18 @@ def get_job(self, job: JobId) -> Job: with self.session() as session: return self.__get_job(session, job) + def get_valid_jobs(self, username_filter: Optional[str]) -> List[Job]: + """Retrieves valid (accessible and not removed) jobs from the database.""" + with self.session() as session: + query = ( + select(Job) + .where(Job.status != "removed") + .order_by(col(Job.submitted).desc()) + ) + if username_filter: + query = query.where(Job.rule_author == username_filter) + return session.exec(query).all() + def remove_query(self, job: JobId) -> None: """Sets the job status to removed.""" with self.session() as session: @@ -139,7 +153,8 @@ def agent_finish_job(self, job: Job) -> None: def init_jobagent(self, job: Job, agent_id: int, tasks: int) -> None: """Creates a new JobAgent object. - If tasks==0 then finishes job immediately""" + If tasks==0 then finishes job immediately. + """ with self.session() as session: obj = JobAgent( task_in_progress=tasks, @@ -293,7 +308,8 @@ def register_active_agent( active_plugins: List[str], ) -> None: """Update or create a Agent information row in the database. - Returns the new or existing agent ID.""" + Returns the new or existing agent ID. + """ # Currently this is done by workers when starting. In the future, # this should be configured by the admin, and workers should just read # their configuration from the database.