traceroute_map/app/db.py
2024-06-02 19:38:39 +03:00

188 lines
4.9 KiB
Python

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