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)

        # Return fetch() data as Row objects, instead of tuples.
        self.conn.row_factory = sqlite3.Row

        self.cursor = self.conn.cursor()

    def create_tables(self):
        self.cursor.executescript(
            """
            CREATE TABLE IF NOT EXISTS Traces (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                created TEXT NOT NULL,
                origin TEXT NOT NULL,
                target TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS Hops (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                trace_id INTEGER,
                created TEXT NOT NULL,
                number INTEGER NOT NULL,
                name TEXT,
                ip TEXT,
                latency TEXT,
                link_latency TEXT,

                FOREIGN KEY(trace_id) REFERENCES Traces(id)
            );
        """
        )

    def end(self):
        """Always call this after you're done with the connection / request."""

        self.conn.commit()
        self.conn.close()

    def list_traces(self):
        # TODO: time filter
        result = []

        self.cursor.execute("SELECT * FROM Traces")
        traces = self.cursor.fetchall()

        for t in traces:
            trace = dict(t)

            self.cursor.execute(
                """
                SELECT number, name, ip, latency, link_latency
                FROM Hops
                WHERE trace_id = ?
                ORDER BY number ASC
            """,
                (trace["id"],),
            )
            hops = self.cursor.fetchall()
            trace["hops"] = hops

            result.append(trace)

        return result

    def create_trace(self, trace):
        self.cursor.execute(
            "INSERT OR IGNORE INTO Traces (created, origin, target) VALUES (?, ?, ?)",
            (trace["created"], trace["origin"], trace["target"]),
        )
        trace_id = self.cursor.lastrowid

        for hop in trace["hops"]:
            self.cursor.execute(
                "INSERT OR IGNORE INTO Hops (trace_id, created, number, name, ip, latency, link_latency) VALUES (?, ?, ?, ?, ?, ?, ?)",
                (
                    trace_id,
                    hop["created"],
                    hop["number"],
                    hop["name"],
                    hop["ip"],
                    hop["latency"],
                    hop["link_latency"],
                ),
            )

        return self.cursor.fetchone()

    def create_hop(self, name, ip, latency):
        self.cursor.execute(
            "INSERT INTO Hops (name, ip, latency) VALUES (?, ?, ?)",
            (name, ip, latency),
        )


def ensure_table_setup():
    db = Database()
    db.create_tables()
    db.end()


####################################################################
####################################################################
####################################################################
####################################################################


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