od-database/crawl_server/database.py

import os
import json
import sqlite3


class Task:

    def __init__(self, url: str, priority: int = 1, callback_type: str = None, callback_args: str = None):
        self.url = url
        self.priority = priority
        self.callback_type = callback_type
        self.callback_args = json.loads(callback_args) if callback_args else {}

    def to_json(self):
        return ({
            "url": self.url,
            "priority": self.priority,
            "callback_type": self.callback_type,
            "callback_args": json.dumps(self.callback_args)
        })


class TaskManagerDatabase:

    def __init__(self, db_path):
        self.db_path = db_path

        if not os.path.exists(db_path):
            self.init_database()

    def init_database(self):

        with open("task_db_init.sql", "r") as f:
            init_script = f.read()

        with sqlite3.connect(self.db_path) as conn:
            conn.executescript(init_script)
            conn.commit()

    def pop_task(self):

        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()

            cursor.execute("SELECT id, url, priority, callback_type, callback_args"
                           " FROM Queue ORDER BY priority DESC, Queue.id ASC LIMIT 1")
            task = cursor.fetchone()

            if task:
                cursor.execute("DELETE FROM Queue WHERE id=?", (task[0],))
                conn.commit()
                return Task(task[1], task[2], task[3], task[4])
            else:
                return None

    def put_task(self, task: Task):

        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()

            cursor.execute("INSERT INTO Queue (url, priority, callback_type, callback_args) VALUES (?,?,?,?)",
                           (task.url, task.priority, task.callback_type, json.dumps(task.callback_args)))
            conn.commit()

    def get_tasks(self):

        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()

            cursor.execute("SELECT * FROM Queue")
            tasks = cursor.fetchall()

            return [Task(t[1], t[2], t[3], t[4]) for t in tasks]