barebones crawl_server microservice

crawl_server/database.py

@@ -0,0 +1,73 @@
+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]