TEMP: Rewrite everything on top of FastAPI

2024-05-31 23:13:38 +03:00
parent 929a3523ba
commit f1c3cf8758
12 changed files with 1796 additions and 257 deletions
--- a/app/db.py
+++ b/app/db.py
@ -0,0 +1,182 @@
+import sqlite3
+from functools import wraps
+
+# Type alias
+Cursor = sqlite3.Cursor
+
+# Configs
+DB_FILE = "./traceroute.db"
+
+
+class Database:
+    def __init__(self):
+        self.db_file = DB_FILE
+        self.conn = sqlite3.connect(self.db_file, check_same_thread=False)
+        self.cursor = self.conn.cursor()
+
+    def create_tables(self):
+        self.cursor.executescript("""
+            CREATE TABLE IF NOT EXISTS Links (
+                id INTEGER PRIMARY KEY,
+                source_ip TEXT NOT NULL,
+                destination_ip TEXT NOT NULL,
+                UNIQUE(source_ip, destination_ip)
+            );
+
+            CREATE TABLE IF NOT EXISTS Paths (
+                id INTEGER PRIMARY KEY,
+                node TEXT NOT NULL,
+                target TEXT NOT NULL,
+                hops_json TEXT NOT NULL,
+                UNIQUE(node, target, hops_json)
+            );
+
+            CREATE TABLE IF NOT EXISTS Latency (
+                id INTEGER PRIMARY KEY,
+                link_id INTEGER NOT NULL,
+                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
+                latency_ms REAL NOT NULL,
+                FOREIGN KEY (link_id) REFERENCES Links(id)
+            );
+
+            CREATE TABLE IF NOT EXISTS HopDetails (
+                id INTEGER PRIMARY KEY,
+                hop_name TEXT,
+                hop_ip TEXT,
+                hop_latency TEXT
+            );
+        """)
+
+    def end(self):
+        """Always call this after you're done with the connection / request."""
+
+        self.conn.commit()
+        self.conn.close()
+
+    def get_traceroute(self):
+        retval = {}
+
+        self.cursor.execute('''
+            SELECT target, hops_json
+            FROM Paths
+        ''')
+        retval['path'] = self.cursor.fetchall()
+
+        self.cursor.execute('''
+            SELECT source_ip, destination_ip
+            FROM Links
+        ''')
+        retval['links'] = self.cursor.fetchall()
+
+        return retval
+
+    def create_link(self, previous_hop_ip, hop_ip):
+        self.cursor.execute(
+            "INSERT OR IGNORE INTO Links (source_ip, destination_ip) VALUES (?, ?)",
+            (previous_hop_ip, hop_ip)
+        )
+        self.cursor.execute(
+            "SELECT id FROM Links WHERE source_ip = ? AND destination_ip = ?",
+            (previous_hop_ip, hop_ip)
+        )
+
+        return self.cursor.fetchone()
+
+    def create_hop(self, name, ip, latency):
+        self.cursor.execute(
+            "INSERT INTO HopDetails (hop_name, hop_ip, hop_latency) VALUES (?, ?, ?)",
+            (name, ip, latency)
+        )
+
+    def create_latency(self, link_id, timestamp, link_latency):
+        self.cursor.execute(
+            "INSERT INTO Latency (link_id, timestamp, latency_ms) VALUES (?, NOW(), ?)",
+            (link_id, timestamp, link_latency)
+        )
+
+    def create_path(self, node, target, json):
+        self.cursor.execute(
+            "INSERT OR IGNORE INTO Paths (node, target, hops_json) VALUES (?, ?, ?)",
+            (node, target, json)
+        )
+
+
+
+
+def ensure_table_setup():
+    db = Database()
+    db.create_tables()
+    db.end()
+
+
+
+
+
+####################################################################
+####################################################################
+####################################################################
+####################################################################
+# Temp testing. Fancy decorator stuff.
+
+def with_connection(func):
+    @wraps(func)
+    def wrapped(*args, **kwargs):
+        conn = sqlite3.connect(DB_FILE)
+        cursor = conn.cursor()
+
+        result = func(cursor, *args, **kwargs)
+
+        conn.commit()
+        conn.close()
+
+        return result
+    return wrapped
+
+
+@with_connection
+def init_db(cursor: Cursor):
+    cursor.executescript("""
+        CREATE TABLE IF NOT EXISTS Links (
+            id INTEGER PRIMARY KEY,
+            source_ip TEXT NOT NULL,
+            destination_ip TEXT NOT NULL,
+            UNIQUE(source_ip, destination_ip)
+        );
+
+        CREATE TABLE IF NOT EXISTS Paths (
+            id INTEGER PRIMARY KEY,
+            node TEXT NOT NULL,
+            target TEXT NOT NULL,
+            hops_json TEXT NOT NULL,
+            UNIQUE(node, target, hops_json)
+        );
+
+        CREATE TABLE IF NOT EXISTS Latency (
+            id INTEGER PRIMARY KEY,
+            link_id INTEGER NOT NULL,
+            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
+            latency_ms REAL NOT NULL,
+            FOREIGN KEY (link_id) REFERENCES Links(id)
+        );
+
+        CREATE TABLE IF NOT EXISTS HopDetails (
+            id INTEGER PRIMARY KEY,
+            hop_name TEXT,
+            hop_ip TEXT,
+            hop_latency TEXT
+        );
+    """)
+
+@with_connection
+def create_link(cursor: Cursor, previous_hop_ip: str, hop_ip: str):
+    """Insert a new hop and return related Link id"""
+
+    cursor.execute("""
+        INSERT OR IGNORE INTO Links (source_ip, destination_ip) VALUES (?, ?)
+    """, (previous_hop_ip, hop_ip))
+
+    cursor.execute("""
+        SELECT id FROM Links WHERE source_ip = ? AND destination_ip = ?
+    """, (previous_hop_ip, hop_ip))
+
+    return cursor.fetchone()