traceroute_map/app/db.py

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()