Skip to main content

mini_apm/
db.rs

1use crate::config::Config;
2use r2d2::Pool;
3use r2d2_sqlite::SqliteConnectionManager;
4use std::fs;
5use std::path::Path;
6
7pub type DbPool = Pool<SqliteConnectionManager>;
8
9const SCHEMA: &str = r#"
10PRAGMA journal_mode = WAL;
11PRAGMA busy_timeout = 100;
12PRAGMA synchronous = NORMAL;
13PRAGMA foreign_keys = ON;
14
15CREATE TABLE IF NOT EXISTS api_keys (
16    id INTEGER PRIMARY KEY,
17    name TEXT NOT NULL,
18    key_hash TEXT NOT NULL UNIQUE,
19    created_at TEXT NOT NULL,
20    last_used_at TEXT
21);
22
23CREATE TABLE IF NOT EXISTS requests (
24    id INTEGER PRIMARY KEY,
25    project_id INTEGER REFERENCES projects(id) ON DELETE CASCADE,
26    request_id TEXT NOT NULL,
27    method TEXT NOT NULL,
28    path TEXT NOT NULL,
29    controller TEXT,
30    action TEXT,
31    status INTEGER NOT NULL,
32    total_ms REAL NOT NULL,
33    db_ms REAL DEFAULT 0,
34    db_count INTEGER DEFAULT 0,
35    view_ms REAL DEFAULT 0,
36    host TEXT,
37    env TEXT,
38    git_sha TEXT,
39    happened_at TEXT NOT NULL
40);
41
42CREATE INDEX IF NOT EXISTS idx_requests_project_id ON requests(project_id);
43CREATE INDEX IF NOT EXISTS idx_requests_happened_at ON requests(happened_at);
44CREATE INDEX IF NOT EXISTS idx_requests_path_method ON requests(path, method);
45CREATE INDEX IF NOT EXISTS idx_requests_total_ms ON requests(total_ms DESC);
46
47CREATE TABLE IF NOT EXISTS errors (
48    id INTEGER PRIMARY KEY,
49    project_id INTEGER REFERENCES projects(id) ON DELETE CASCADE,
50    fingerprint TEXT NOT NULL,
51    exception_class TEXT NOT NULL,
52    message TEXT NOT NULL,
53    first_seen_at TEXT NOT NULL,
54    last_seen_at TEXT NOT NULL,
55    occurrence_count INTEGER DEFAULT 1,
56    status TEXT DEFAULT 'open',
57    UNIQUE(project_id, fingerprint)
58);
59
60CREATE INDEX IF NOT EXISTS idx_errors_project_id ON errors(project_id);
61CREATE INDEX IF NOT EXISTS idx_errors_status ON errors(status);
62CREATE INDEX IF NOT EXISTS idx_errors_last_seen ON errors(last_seen_at DESC);
63
64CREATE TABLE IF NOT EXISTS error_occurrences (
65    id INTEGER PRIMARY KEY,
66    error_id INTEGER NOT NULL REFERENCES errors(id) ON DELETE CASCADE,
67    request_id TEXT,
68    user_id TEXT,
69    backtrace TEXT NOT NULL,
70    params TEXT,
71    happened_at TEXT NOT NULL
72);
73
74CREATE INDEX IF NOT EXISTS idx_error_occurrences_error_id ON error_occurrences(error_id);
75CREATE INDEX IF NOT EXISTS idx_error_occurrences_happened_at ON error_occurrences(happened_at);
76
77CREATE TABLE IF NOT EXISTS rollups_hourly (
78    id INTEGER PRIMARY KEY,
79    hour TEXT NOT NULL,
80    path TEXT NOT NULL,
81    method TEXT NOT NULL,
82    request_count INTEGER NOT NULL,
83    error_count INTEGER DEFAULT 0,
84    total_ms_sum REAL NOT NULL,
85    total_ms_p50 REAL,
86    total_ms_p95 REAL,
87    total_ms_p99 REAL,
88    db_ms_sum REAL DEFAULT 0,
89    db_count_sum INTEGER DEFAULT 0,
90    UNIQUE(hour, path, method)
91);
92
93CREATE INDEX IF NOT EXISTS idx_rollups_hourly_hour ON rollups_hourly(hour);
94
95CREATE TABLE IF NOT EXISTS rollups_daily (
96    id INTEGER PRIMARY KEY,
97    date TEXT NOT NULL,
98    path TEXT NOT NULL,
99    method TEXT NOT NULL,
100    request_count INTEGER NOT NULL,
101    error_count INTEGER DEFAULT 0,
102    total_ms_p50 REAL,
103    total_ms_p95 REAL,
104    total_ms_p99 REAL,
105    avg_db_ms REAL,
106    avg_db_count REAL,
107    UNIQUE(date, path, method)
108);
109
110CREATE INDEX IF NOT EXISTS idx_rollups_daily_date ON rollups_daily(date);
111
112CREATE TABLE IF NOT EXISTS deploys (
113    id INTEGER PRIMARY KEY,
114    project_id INTEGER REFERENCES projects(id) ON DELETE CASCADE,
115    git_sha TEXT NOT NULL,
116    version TEXT,
117    env TEXT,
118    deployed_at TEXT NOT NULL,
119    description TEXT,
120    deployer TEXT
121);
122
123CREATE INDEX IF NOT EXISTS idx_deploys_project_id ON deploys(project_id);
124CREATE INDEX IF NOT EXISTS idx_deploys_deployed_at ON deploys(deployed_at);
125
126CREATE TABLE IF NOT EXISTS projects (
127    id INTEGER PRIMARY KEY,
128    name TEXT NOT NULL UNIQUE,
129    slug TEXT NOT NULL UNIQUE,
130    api_key TEXT NOT NULL UNIQUE,
131    created_at TEXT NOT NULL
132);
133
134CREATE INDEX IF NOT EXISTS idx_projects_slug ON projects(slug);
135CREATE INDEX IF NOT EXISTS idx_projects_api_key ON projects(api_key);
136
137CREATE TABLE IF NOT EXISTS users (
138    id INTEGER PRIMARY KEY,
139    username TEXT NOT NULL UNIQUE,
140    password_hash TEXT,
141    is_admin INTEGER NOT NULL DEFAULT 0,
142    must_change_password INTEGER NOT NULL DEFAULT 0,
143    invite_token TEXT UNIQUE,
144    invite_expires_at TEXT,
145    created_at TEXT NOT NULL,
146    last_login_at TEXT
147);
148
149CREATE TABLE IF NOT EXISTS sessions (
150    id INTEGER PRIMARY KEY,
151    token TEXT NOT NULL UNIQUE,
152    user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
153    created_at TEXT NOT NULL,
154    expires_at TEXT NOT NULL
155);
156
157CREATE INDEX IF NOT EXISTS idx_sessions_token ON sessions(token);
158CREATE INDEX IF NOT EXISTS idx_sessions_expires ON sessions(expires_at);
159
160CREATE TABLE IF NOT EXISTS settings (
161    key TEXT PRIMARY KEY,
162    value TEXT NOT NULL,
163    updated_at TEXT NOT NULL
164);
165
166CREATE TABLE IF NOT EXISTS spans (
167    id INTEGER PRIMARY KEY,
168    project_id INTEGER REFERENCES projects(id) ON DELETE CASCADE,
169    trace_id TEXT NOT NULL,
170    span_id TEXT NOT NULL,
171    parent_span_id TEXT,
172    start_time_unix_nano INTEGER NOT NULL,
173    end_time_unix_nano INTEGER NOT NULL,
174    duration_ms REAL,
175    name TEXT NOT NULL,
176    kind INTEGER NOT NULL DEFAULT 0,
177    status_code INTEGER DEFAULT 0,
178    status_message TEXT,
179    span_category TEXT NOT NULL,
180    root_span_type TEXT,
181    service_name TEXT,
182    http_method TEXT,
183    http_url TEXT,
184    http_status_code INTEGER,
185    db_system TEXT,
186    db_statement TEXT,
187    db_operation TEXT,
188    messaging_system TEXT,
189    messaging_operation TEXT,
190    request_id TEXT,
191    attributes_json TEXT,
192    events_json TEXT,
193    resource_attributes_json TEXT,
194    happened_at TEXT NOT NULL,
195    UNIQUE(trace_id, span_id)
196);
197
198CREATE INDEX IF NOT EXISTS idx_spans_project_id ON spans(project_id);
199CREATE INDEX IF NOT EXISTS idx_spans_trace_id ON spans(trace_id);
200CREATE INDEX IF NOT EXISTS idx_spans_happened_at ON spans(happened_at);
201CREATE INDEX IF NOT EXISTS idx_spans_root_type ON spans(root_span_type) WHERE root_span_type IS NOT NULL;
202CREATE INDEX IF NOT EXISTS idx_spans_category ON spans(span_category);
203"#;
204
205pub fn init(config: &Config) -> anyhow::Result<DbPool> {
206    // Ensure data directory exists
207    if let Some(parent) = Path::new(&config.sqlite_path).parent() {
208        fs::create_dir_all(parent)?;
209    }
210
211    let manager = SqliteConnectionManager::file(&config.sqlite_path);
212    let pool = Pool::builder().max_size(10).build(manager)?;
213
214    // Run migrations
215    migrate(&pool)?;
216
217    Ok(pool)
218}
219
220fn migrate(pool: &DbPool) -> anyhow::Result<()> {
221    let conn = pool.get()?;
222
223    // Run migrations for existing tables BEFORE executing schema
224    // This ensures columns exist before indexes are created
225
226    // Check if deploys table exists and add project_id if missing
227    let deploys_exists: bool = conn
228        .query_row(
229            "SELECT 1 FROM sqlite_master WHERE type='table' AND name='deploys'",
230            [],
231            |_| Ok(true),
232        )
233        .unwrap_or(false);
234
235    if deploys_exists {
236        let _ = conn.execute(
237            "ALTER TABLE deploys ADD COLUMN project_id INTEGER REFERENCES projects(id) ON DELETE CASCADE",
238            [],
239        );
240    }
241
242    // Execute schema (creates tables if not exist, creates indexes)
243    conn.execute_batch(SCHEMA)?;
244
245    // Add invite columns if they don't exist (for existing databases)
246    let _ = conn.execute("ALTER TABLE users ADD COLUMN invite_token TEXT UNIQUE", []);
247    let _ = conn.execute("ALTER TABLE users ADD COLUMN invite_expires_at TEXT", []);
248
249    // Add source_context column if it doesn't exist
250    let _ = conn.execute(
251        "ALTER TABLE error_occurrences ADD COLUMN source_context TEXT",
252        [],
253    );
254
255    tracing::debug!("Database schema initialized");
256    Ok(())
257}
258
259pub fn get_db_size(pool: &DbPool) -> anyhow::Result<f64> {
260    let conn = pool.get()?;
261    let size: i64 = conn.query_row(
262        "SELECT page_count * page_size FROM pragma_page_count(), pragma_page_size()",
263        [],
264        |row| row.get(0),
265    )?;
266    Ok(size as f64 / 1_048_576.0) // Convert to MB
267}