1use crate::error::Result;
2use rusqlite::Connection;
3
4pub const SCHEMA_VERSION: &str = "2";
5
6pub fn bootstrap(conn: &Connection) -> Result<()> {
7 conn.execute_batch(
8 "
9 PRAGMA journal_mode=WAL;
10 PRAGMA foreign_keys=ON;
11 PRAGMA synchronous=NORMAL;
12
13 CREATE TABLE IF NOT EXISTS _adb_meta (
14 key TEXT PRIMARY KEY,
15 value TEXT NOT NULL
16 );
17
18 CREATE TABLE IF NOT EXISTS _adb_collections (
19 id TEXT PRIMARY KEY,
20 name TEXT UNIQUE NOT NULL,
21 dim INTEGER NOT NULL,
22 metric TEXT NOT NULL DEFAULT 'cosine',
23 count INTEGER NOT NULL DEFAULT 0,
24 created_at INTEGER NOT NULL
25 );
26
27 CREATE TABLE IF NOT EXISTS _adb_vectors (
28 id TEXT NOT NULL,
29 collection_id TEXT NOT NULL,
30 vector BLOB NOT NULL,
31 metadata TEXT,
32 created_at INTEGER NOT NULL,
33 PRIMARY KEY (id, collection_id),
34 FOREIGN KEY (collection_id)
35 REFERENCES _adb_collections(id) ON DELETE CASCADE
36 );
37
38 CREATE TABLE IF NOT EXISTS _adb_hnsw_index (
39 collection_id TEXT PRIMARY KEY,
40 index_blob BLOB NOT NULL,
41 built_at INTEGER NOT NULL,
42 is_dirty INTEGER NOT NULL DEFAULT 0
43 );
44
45 CREATE TABLE IF NOT EXISTS _adb_nodes (
46 id TEXT PRIMARY KEY,
47 kind TEXT NOT NULL,
48 data TEXT,
49 created_at INTEGER NOT NULL,
50 updated_at INTEGER NOT NULL
51 );
52
53 CREATE TABLE IF NOT EXISTS _adb_edges (
54 src TEXT NOT NULL,
55 dst TEXT NOT NULL,
56 relation TEXT NOT NULL,
57 weight REAL NOT NULL DEFAULT 1.0,
58 created_at INTEGER NOT NULL,
59 PRIMARY KEY (src, dst, relation),
60 FOREIGN KEY (src) REFERENCES _adb_nodes(id) ON DELETE CASCADE,
61 FOREIGN KEY (dst) REFERENCES _adb_nodes(id) ON DELETE CASCADE
62 );
63
64 CREATE INDEX IF NOT EXISTS idx_edges_src ON _adb_edges(src);
65 CREATE INDEX IF NOT EXISTS idx_edges_dst ON _adb_edges(dst);
66 CREATE INDEX IF NOT EXISTS idx_vectors_col ON _adb_vectors(collection_id);
67 CREATE INDEX IF NOT EXISTS idx_nodes_kind ON _adb_nodes(kind);
68
69 -- Conversations / message threading
70 CREATE TABLE IF NOT EXISTS _adb_conversations (
71 id TEXT PRIMARY KEY,
72 title TEXT,
73 metadata TEXT,
74 created_at INTEGER NOT NULL DEFAULT (strftime('%s','now')),
75 updated_at INTEGER NOT NULL DEFAULT (strftime('%s','now'))
76 );
77
78 CREATE TABLE IF NOT EXISTS _adb_messages (
79 id TEXT PRIMARY KEY,
80 conversation_id TEXT NOT NULL REFERENCES _adb_conversations(id) ON DELETE CASCADE,
81 role TEXT NOT NULL,
82 content TEXT NOT NULL,
83 metadata TEXT,
84 created_at INTEGER NOT NULL DEFAULT (strftime('%s','now'))
85 );
86 CREATE INDEX IF NOT EXISTS idx_messages_conv ON _adb_messages(conversation_id, created_at);
87
88 -- Workflow persistence
89 CREATE TABLE IF NOT EXISTS _adb_workflows (
90 id TEXT PRIMARY KEY,
91 name TEXT NOT NULL,
92 status TEXT NOT NULL DEFAULT 'pending',
93 input TEXT,
94 output TEXT,
95 metadata TEXT,
96 created_at INTEGER NOT NULL DEFAULT (strftime('%s','now')),
97 updated_at INTEGER NOT NULL DEFAULT (strftime('%s','now'))
98 );
99
100 CREATE TABLE IF NOT EXISTS _adb_workflow_steps (
101 id TEXT PRIMARY KEY,
102 workflow_id TEXT NOT NULL REFERENCES _adb_workflows(id) ON DELETE CASCADE,
103 step_index INTEGER NOT NULL,
104 name TEXT NOT NULL,
105 status TEXT NOT NULL DEFAULT 'pending',
106 input TEXT,
107 output TEXT,
108 error TEXT,
109 started_at INTEGER,
110 completed_at INTEGER
111 );
112 CREATE INDEX IF NOT EXISTS idx_wf_steps ON _adb_workflow_steps(workflow_id, step_index);
113
114 -- Reasoning traces
115 CREATE TABLE IF NOT EXISTS _adb_traces (
116 id TEXT PRIMARY KEY,
117 session_id TEXT,
118 parent_id TEXT REFERENCES _adb_traces(id),
119 trace_type TEXT NOT NULL,
120 content TEXT NOT NULL,
121 metadata TEXT,
122 created_at INTEGER NOT NULL DEFAULT (strftime('%s','now'))
123 );
124 CREATE INDEX IF NOT EXISTS idx_traces_session ON _adb_traces(session_id, created_at);
125 CREATE INDEX IF NOT EXISTS idx_traces_parent ON _adb_traces(parent_id);
126 ",
127 )?;
128
129 conn.execute(
130 "INSERT OR IGNORE INTO _adb_meta (key, value) VALUES ('schema_version', ?1)",
131 rusqlite::params![SCHEMA_VERSION],
132 )?;
133 conn.execute(
134 "INSERT OR IGNORE INTO _adb_meta (key, value) VALUES ('created_at', ?1)",
135 rusqlite::params![now_ms()],
136 )?;
137 Ok(())
138}
139
140pub fn check_version(conn: &Connection) -> Result<()> {
141 let version: Option<String> = conn
142 .query_row(
143 "SELECT value FROM _adb_meta WHERE key = 'schema_version'",
144 [],
145 |row| row.get(0),
146 )
147 .ok();
148 match version.as_deref() {
149 Some(v) if v == SCHEMA_VERSION => Ok(()),
150 Some(_) => Err(crate::error::AgentDbError::SchemaMigration),
151 None => Ok(()),
152 }
153}
154
155pub fn now_ms() -> i64 {
156 std::time::SystemTime::now()
157 .duration_since(std::time::UNIX_EPOCH)
158 .unwrap_or_default()
159 .as_millis() as i64
160}