traceroute_map/app/db.py
2024-06-05 22:59:20 +03:00

118 lines
3.1 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,
unparsed TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS Nodes (
id INTEGER PRIMARY KEY AUTOINCREMENT,
created TEXT NOT NULL,
name TEXT,
ip TEXT,
latency_ms REAL NOT NULL,
);
CREATE TABLE IF NOT EXISTS Links (
id INTEGER PRIMARY KEY AUTOINCREMENT,
created TEXT NOT NULL,
source TEXT NOT NULL,
target TEXT NOT NULL,
latency_ms REAL NOT NULL,
);
"""
)
def end(self):
"""Always call this after you're done with the connection / request."""
self.conn.commit()
self.conn.close()
# TODO
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
# TODO
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"], # TODO: trace.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()