3
\$\begingroup\$

I want to implement a persistent queue for file system operations (copy, delete, move) that needs to be process and thread safe. The module will be used by a daemon (called qcp) which can accept Tasks from multiple clients. It must also be possible to process several queued elements in parallel.

I am new to python and looking for general advice (even the nitpicky kind) on python best practices and on object-oriented programming in general. I am especially curious if my implementation of Task.from_json() as @staticmethod is appropriate. I wasn't sure about this because it conditionally instantiates subclasses of Task rather than a plain Task.

tasks.py

from pathlib import Path
from typing import Union, Optional
import shutil
import os
import json
import sqlite3    
import logging

Pathish = Union[Path, str]


class Task:
    """Abstract class for qcp Tasks. Should not be instantiated directly."""

    def __init__(self) -> None:
        self.type = 0

    def run(self) -> None:
        """Run the Task"""
        pass

    @staticmethod
    def from_dict(x, validate: bool = False) -> "Task":
        """Create a Task of the appropriate subclass from a python dict"""
        task_type = x["type"]

        if task_type == -1:
            return KillTask()
        elif task_type == 0:
            return Task()
        elif task_type == 1:
            return EchoTask(x["msg"])
        elif task_type == 2:
            return FileTask(x["src"], validate=validate)
        elif task_type == 3:
            return DeleteTask(x["src"], validate=validate)
        elif task_type == 4:
            return CopyTask(x["src"], x["dst"], validate=validate)
        elif task_type == 5:
            return MoveTask(x["src"], x["dst"], validate=validate)
        elif task_type == 6:
            raise NotImplementedError
        else:
            raise ValueError

    def __repr__(self) -> str:
        return 'NULL'

    def __eq__(self, other) -> bool:
        return self.__dict__ == other.__dict__

    def __ne__(self, other) -> bool:
        return self.__dict__ != other.__dict__


class KillTask(Task):
    """Kill the qcp server"""
    def __init__(self) -> None:
        self.type = -1
        super().__init__()

    def run(self) -> None:
        raise NotImplementedError

    def __repr__(self) -> str:
        return 'KILL'


class EchoTask(Task):
    """Log a message"""
    def __init__(self,  msg: str) -> None:
        super().__init__()
        self.msg = msg
        self.type = 1

    def run(self) -> None:
        print(self.msg)

    def __repr__(self) -> str:
        return f'Echo: "{self.msg}"'


class FileTask(Task):
    """Abstract class for all file-based tasks"""
    def __init__(self, src: Pathish, validate: bool = True) -> None:
        super().__init__()
        self.validate = validate
        self.src = Path(src).as_posix()
        self.type = 2
        if validate:
            self.__validate__()

    def __validate__(self) -> None:
        if not Path(self.src).exists():
            raise FileNotFoundError(f'{self.src} does not exist')
        elif not (Path(self.src).is_dir() or Path(self.src).is_file()):
            raise TypeError(f'{self.src} is neither a file nor directory')


class DeleteTask(FileTask):
    """Delete a file"""
    def __init__(self, src: Pathish, validate: bool = True) -> None:
        super().__init__(src=src, validate=validate)
        self.type = 3

    def run(self) -> None:
        os.unlink(self.src)

    def __repr__(self) -> str:
        return f'DEL {self.src}'


class CopyTask(FileTask):
    """Copy a file"""
    def __init__(self, src: Pathish, dst: Pathish, validate: bool = True) -> None:
        super().__init__(src=src, validate=False)
        self.dst = Path(dst).as_posix()
        self.type = 4
        self.validate = validate
        if validate:
            self.__validate__()

    def __repr__(self) -> str:
        return f'COPY {self.src} -> {self.dst}'

    def __validate__(self) -> None:
        super().__validate__()
        if Path(self.dst).exists():
            raise FileExistsError

    def run(self) -> None:
        self.__validate__()
        shutil.copy(self.src, self.dst)


class MoveTask(CopyTask):
    """Move a file"""
    def __init__(self, src: Pathish, dst: Pathish, validate: bool = True) -> None:
        super().__init__(src=src, dst=dst, validate=validate)
        self.type = 5

    def run(self) -> None:
        super().__validate__()
        shutil.move(self.src, self.dst)

    def __repr__(self) -> str:
        return f'MOVE {self.src} -> {self.dst}'


class TaskQueueElement:
    """An enqueued Task"""

    task = None  #: A Task
    status = None  #: Status of the queued Task
    priority = None  #: Priority of the queued Task

    def __init__(self, task: Task, priority: 1) -> None:
        self.task = task
        self.priority = priority

    def __lt__(self, other) -> bool:
        return self.priority < other.priority

    def __gt__(self, other) -> bool:
        return self.priority > other.priority

    def __eq__(self, other) -> bool:
        return self.__dict__ == other.__dict__

    def __ne__(self, other) -> bool:
        return self.__dict__ != other.__dict__


class TaskQueue:
    """A prioritzed queue for tasks"""
    def __init__(self, path: Pathish = 'qcp.db') -> None:
        """
        Instantiate a TaskQueue

        :param path: Path to store the persistent queue
        :type path: Path or str
        """

        self.con = sqlite3.connect(path, isolation_level="EXCLUSIVE")
        self.path = Path(path)

        cur = self.con.cursor()
        cur.execute("""
           CREATE TABLE IF NOT EXISTS tasks (
              priority INTEGER,
              task TEXT,
              status INTEGER,
              owner INTEGER              
            )              
        """)
        self.con.commit()

    @property
    def n_total(self) -> int:
        """Count of all tasks in queue (including failed and completed)"""
        cur = self.con.cursor()
        return cur.execute("SELECT COUNT(1) from tasks").fetchall()[0][0]

    @property
    def n_pending(self) -> int:
        """Number of pending tasks"""
        cur = self.con.cursor()
        return cur.execute("SELECT COUNT(1) FROM tasks WHERE status = 0").fetchall()[0][0]

    @property
    def n_running(self) -> int:
        """Count of currently running tasks"""
        cur = self.con.cursor()
        return cur.execute("SELECT COUNT(1) FROM tasks WHERE status = 1").fetchall()[0][0]

    @property
    def n_done(self) -> int:
        """count of completed tasks"""
        cur = self.con.cursor()
        return cur.execute("SELECT COUNT(1) from tasks WHERE status = 2").fetchall()[0][0]

    @property
    def n_failed(self) -> int:
        """count of completed tasks"""
        cur = self.con.cursor()
        return cur.execute("SELECT COUNT(1) from tasks WHERE status = -1").fetchall()[0][0]

    def put(self, task: "Task", priority: Optional[int] = None) -> None:
        """
        Enqueue a task

        :param task: Task to be added to the queue
        :type task: Task
        :param priority: (optional) priority for executing `task` (tasks with lower priority will be executed earlier)
        :type priority: int
        """

        cur = self.con.cursor()
        cur.execute(
            "INSERT INTO tasks (priority, task, status) VALUES (?, ?, ?)", (priority, json.dumps(task.__dict__), 0)
        )
        self.con.commit()

    def pop(self) -> "Task":
        """
        Retrieves Task object and sets status of Task in database to "in progress" (1)

        :raises AlreadyUnderEvaluationError: If trying to pop a tasks that is already being processed  (i.e. if a race
        condition occurs if the queue is processed in parallel)
        """
        cur = self.con.cursor()
        cur.execute("SELECT _ROWID_ from tasks WHERE status = 0 ORDER BY priority LIMIT 1")
        oid = cur.fetchall()[0][0].__str__()
        self.mark_running(oid, id(self))

        cur.execute("SELECT owner, task FROM tasks WHERE _ROWID_ = ?", oid)
        record = cur.fetchall()[0]
        if record[0] != id(self):
            raise AlreadyUnderEvaluationError

        task = Task.from_dict(json.loads(record[1]))
        task.oid = oid
        return task

    def peek(self) -> "Task":
        """
        Retrieves Task object without changing its status in the queue
        """
        cur = self.con.cursor()
        cur.execute("SELECT * from tasks ORDER BY priority LIMIT 1")
        record = cur.fetchall()[0]
        oid = record[0].__str__()
        task = Task.from_dict(json.loads(record[1]), validate=False)
        task.oid = oid
        return task

    def print(self, n: int = 10) -> None:
        """
        Print an overview of the queue

        :param n: number of tasks to preview
        :type n: int
        """
        assert isinstance(n, int) and n > 0
        cur = self.con.cursor()
        cur.execute("SELECT status, task from tasks ORDER BY priority LIMIT ?", (str(n), ))
        records = cur.fetchall()
        for record in records:
            print(f"[{record[0]}] {Task.from_dict(json.loads(record[1]))}")

    def mark_pending(self, oid: int) -> None:
        """
        Mark the operation with the _ROWID_ `oid` as "pending" (0)

        :param oid: ID of the task to mark
        :type oid: int
        """
        cur = self.con.cursor()
        cur.execute("UPDATE tasks SET status = 0, owner = NULL where _ROWID_ = ?", (oid, ))
        self.con.commit()

    def mark_running(self, oid: int, owner: int) -> None:
        """Mark the operation with the _ROWID_ `oid` as "running" (1). The "owner" Id is to ensure no two processes
        are trying to execute the same operation

        :param oid: ID of the task to mark
        :type oid: int
        :param owner: Id of the process that is handling the operation
        :type owner: int
        """
        cur = self.con.cursor()
        cur.execute("UPDATE tasks SET status = 1, owner = ? where _ROWID_ = ?", (owner, oid))
        self.con.commit()

    def mark_done(self, oid: int) -> None:
        """
        Mark the operation with the _ROWID_ `oid` as "done" (2)
        :param oid: ID of the task to mark
        :type oid: int
        """
        cur = self.con.cursor()
        cur.execute("UPDATE tasks SET status = 2, owner = NULL where _ROWID_ = ?", (oid, ))
        self.con.commit()

    def mark_failed(self, oid: int) -> None:
        """
        Mark the operation with the _ROWID_ `oid` as "failed" (-1)

        :param oid: ID of the task to mark
        :type oid: int
        """
        cur = self.con.cursor()
        cur.execute("UPDATE tasks SET status = -1, owner = NULL where _ROWID_ = ?", (oid, ))
        self.con.commit()

    def run(self) -> None:
        """Execute all pending tasks"""
        if self.n_pending < 1:
            logging.getLogger().warn("Queue is empty")

        while self.n_pending > 0:
            op = self.pop()
            op.run()
            self.mark_done(op.oid)


class AlreadyUnderEvaluationError(Exception):
    """This Task is already being processed by a different worker"""
    pass

A demo in pytest

import tasks
import pytest

def test_TaskQueue(tmp_path):
    """TaskQueue can queue and execute tasks"""
    src = tmp_path.joinpath("foo")
    src.touch()

    q = tasks.TaskQueue(tmp_path.joinpath("qcp.db"))
    q.put(tasks.CopyTask(src, tmp_path.joinpath("copied_file")))
    q.run()
    assert tmp_path.joinpath("copied_file").is_file()
    q.put(tasks.MoveTask(tmp_path.joinpath("copied_file"), tmp_path.joinpath("moved_file")))
    q.run()
    assert not tmp_path.joinpath("copied_file").is_file()
    assert tmp_path.joinpath("moved_file").is_file()
    q.put(tasks.DeleteTask(tmp_path.joinpath("moved_file")))
    q.run()
    assert not tmp_path.joinpath("moved_file").is_file()
    assert src.is_file()

p.s.: I am aware that I am importing logging for a single log call, but I plan to include more extensive logging in the future

\$\endgroup\$
  • \$\begingroup\$ -> "Task" works with quotes in place? \$\endgroup\$ – Thomas Weller Aug 18 at 16:34
  • \$\begingroup\$ Yes but i think its a fairly recent thing, see stackoverflow.com/questions/46458470/… \$\endgroup\$ – hoelk Aug 18 at 17:07
  • \$\begingroup\$ Ok, thanks. Sometimes, Python is a bit messy, I think. Not pythonic. Looks more like a workaround than a good solution. \$\endgroup\$ – Thomas Weller Aug 18 at 17:22
2
\$\begingroup\$

Overall your code is pretty clean.

Abstract class

There are a few different ways to do this. You can explicitly use the abc stuff, but the simpler thing to do is in your abstract methods (e.g. run), raise a NotImplementedError - at the Task level, not just KillTask. Then classes like KillTask that are effectively abstract children simply omit the implementation.

Enums

Python needs help when it comes to the type system. It's good that you're using type hints. You can help it out a little more by using a formal enum for the numbers in from_dict. Read here - https://docs.python.org/3/library/enum.html

Prepared statements

The series of functions beginning with n_pending should really only be one function. This is the perfect use case for prepared statements. Make that status number a parameter.

Don't conflate purpose

TaskQueue.print really does two things - a fetch and a print. You should probably separate these.

sqlite autocommit

From the brief Google I did, sqlite operates in autocommit mode by default. I don't think any of the commits you've done are necessary.

Implicit cursors

It shouldn't be necessary to explicitly get those cursors. You can iterate on the return value of fetchall.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.