1#[cfg(test)]
2#[path = "mod_test.rs"]
3mod tests;
4
5pub(crate) mod migration;
6
7use crate::storage::Storage;
8use crate::{
9 config::resolve_path,
10 models::{
11 Context as ConvoContext, Conversation, Message, message::Issuer,
12 storage::FilterConversation,
13 },
14};
15use async_trait::async_trait;
16use eyre::{Context, Result, bail};
17use migration::MIGRATION;
18use std::collections::HashMap;
19use tokio_rusqlite::{Connection, OpenFlags, ToSql, named_params, params};
20
21pub struct Sqlite {
22 conn: Connection,
23}
24
25impl Sqlite {
26 pub async fn new(path: Option<&str>) -> Result<Self> {
27 let conn = match path {
28 Some(path) => Connection::open_with_flags(
29 resolve_path(path).wrap_err("resolving path")?,
30 OpenFlags::SQLITE_OPEN_READ_WRITE | OpenFlags::SQLITE_OPEN_CREATE,
31 )
32 .await
33 .wrap_err(format!("opening database path: {}", path))?,
34 None => Connection::open_in_memory()
35 .await
36 .wrap_err("opening in-memory database")?,
37 };
38
39 let ret = Self { conn };
40 ret.run_migration().await.wrap_err("running migration")?;
41 Ok(ret)
42 }
43
44 async fn run_migration(&self) -> Result<()> {
45 self.conn
46 .call(|conn| Ok(conn.execute_batch(MIGRATION)?))
47 .await
48 .wrap_err("executing migration")?;
49 Ok(())
50 }
51}
52
53#[async_trait]
54impl Storage for Sqlite {
55 async fn get_conversation(&self, id: &str) -> Result<Option<Conversation>> {
56 let conversations = self
57 .get_conversations(FilterConversation::default().with_id(id))
58 .await
59 .wrap_err("getting conversation")?;
60
61 let conversation = match conversations.get(id) {
62 Some(conversation) => conversation.clone(),
63 None => return Ok(None),
64 };
65 let messages = self.get_messages(conversation.id()).await?;
66 Ok(Some(conversation.with_messages(messages)))
67 }
68
69 async fn get_messages(&self, conversation_id: &str) -> Result<Vec<Message>> {
70 let conversation_id = conversation_id.to_string();
71 let messages = self.conn.call(move |conn| {
72 let mut stmt = conn.prepare(
73 "SELECT id, conversation_id, text, issuer, system, token_count, created_at FROM messages WHERE conversation_id = ?",
74 )?;
75
76 let mut rows = stmt.query(params![conversation_id])?;
77 let mut messages = vec![];
78 while let Some(row) = rows.next()? {
79 let id: String = row.get(0)?;
80 let text: String = row.get(2)?;
81 let issuer: String = row.get(3)?;
82 let system: i32 = row.get(4)?;
83 let token_count: usize = row.get(5)?;
84 let created_at: i64 = row.get(6)?;
85
86 let issuer = if system == 1 {
87 Issuer::System(issuer)
88 } else {
89 Issuer::User(issuer)
90 };
91
92 let created_at = chrono::DateTime::from_timestamp_millis(created_at).ok_or(tokio_rusqlite::Error::Other(eyre::eyre!("invalid timestamp").into()))?;
93
94 messages.push(Message::new(issuer, text).with_id(id).with_created_at(created_at).with_token_count(token_count));
95 }
96 messages.sort_by(|a, b| {
97 a.created_at()
98 .timestamp_millis()
99 .cmp(&b.created_at().timestamp_millis())
100 });
101 Ok(messages)
102 }).await?;
103 Ok(messages)
104 }
105
106 async fn get_conversations(
107 &self,
108 filter: FilterConversation,
109 ) -> Result<HashMap<String, Conversation>> {
110 let mut conversations = self
111 .conn
112 .call(move |conn| {
113 let (query, params) = filter_to_query(&filter);
114 let mut stmt = conn.prepare(&query)?;
115 let params: Vec<(&str, &dyn ToSql)> =
116 params.iter().map(|(n, v)| (*n, v.as_ref())).collect();
117 let mut rows = stmt.query(params.as_slice())?;
118
119 let mut conversations: HashMap<String, Conversation> = HashMap::new();
120
121 while let Some(row) = rows.next()? {
122 let id: String = row.get(0)?;
123 let title: String = row.get(1)?;
124 let created_at: i64 = row.get(2)?;
125 let created_at = chrono::DateTime::from_timestamp_millis(created_at).ok_or(
126 tokio_rusqlite::Error::Other(eyre::eyre!("invalid created_at").into()),
127 )?;
128
129 let updated_at: i64 = row.get(3)?;
130 let updated_at = chrono::DateTime::from_timestamp_millis(updated_at).ok_or(
131 tokio_rusqlite::Error::Other(eyre::eyre!("invalid updated_at").into()),
132 )?;
133
134 let mut con = Conversation::default()
135 .with_id(&id)
136 .with_title(title)
137 .with_created_at(created_at);
138
139 if updated_at.timestamp_millis() > 0 {
140 con = con.with_updated_at(updated_at);
141 }
142
143 conversations.insert(id, con);
144 }
145 Ok(conversations)
146 })
147 .await?;
148
149 for conversation in conversations.values_mut() {
150 let messages = self.get_messages(conversation.id()).await?;
151 conversation.messages_mut().extend(messages);
152 let ctx = self.get_contexts(conversation.id()).await?;
153 conversation.contexts_mut().extend(ctx);
154 }
155
156 Ok(conversations)
157 }
158
159 async fn upsert_conversation(&self, conversation: Conversation) -> Result<()> {
160 if conversation.id().is_empty() {
161 bail!("conversation id is empty");
162 }
163
164 self.conn
165 .call(move |conn| {
166 let tx = conn.transaction()?;
167 tx.execute(
168 r#"INSERT INTO conversations (id, title, created_at, updated_at)
169 VALUES (:id, :title, :created_at, :updated_at)
170 ON CONFLICT(id) DO UPDATE SET
171 title = excluded.title,
172 created_at = excluded.created_at,
173 updated_at = excluded.updated_at
174 "#,
175 named_params! {
176 ":id": conversation.id(),
177 ":title": conversation.title(),
178 ":created_at": conversation.created_at().timestamp_millis(),
179 ":updated_at": conversation.updated_at().timestamp_millis(),
180 },
181 )?;
182 tx.commit()?;
183 Ok(())
184 })
185 .await?;
186 Ok(())
187 }
188
189 async fn delete_conversation(&self, id: &str) -> Result<()> {
190 let id = id.to_string();
191 self.conn
192 .call(move |conn| {
193 let tx = conn.transaction()?;
194 tx.execute("DELETE FROM conversations WHERE id = ?", params![id])?;
195 Ok(tx.commit()?)
196 })
197 .await?;
198 Ok(())
199 }
200
201 async fn add_messages(&self, conversation_id: &str, messages: &[Message]) -> Result<()> {
202 let conversation_id = conversation_id.to_string();
203 let messages = messages.to_vec();
204 self.conn
205 .call(move |conn| {
206 let tx = conn.transaction()?;
207 for message in messages {
208 tx.execute(
209 r#"INSERT INTO messages (id, conversation_id, text, issuer, system, token_count, created_at)
210 VALUES (:id, :conversation_id, :text, :issuer, :system, :token_count, :created_at)
211 ON CONFLICT(id, conversation_id) DO UPDATE SET
212 text = excluded.text,
213 issuer = excluded.issuer,
214 system = excluded.system,
215 token_count = excluded.token_count,
216 created_at = excluded.created_at
217 "#,
218 named_params! {
219 ":id": message.id(),
220 ":conversation_id": conversation_id,
221 ":text": message.text(),
222 ":issuer": message.issuer_str(),
223 ":system": message.is_system() as i32,
224 ":token_count": message.token_count() as i32,
225 ":created_at": message.created_at().timestamp_millis()
226 },
227 )?;
228 }
229 Ok(tx.commit()?)
230 })
231 .await?;
232 Ok(())
233 }
234
235 async fn upsert_message(&self, conversation_id: &str, message: Message) -> Result<()> {
236 let conversation_id = conversation_id.to_string();
237 let id = message.id().to_string();
238 let text = message.text().to_string();
239 let issuer = message.issuer_str().to_string();
240 let system = message.is_system() as i32;
241 let token_count = message.token_count() as i32;
242 let timestamp = message.created_at().timestamp_millis();
243 let affected_rows = self
244 .conn
245 .call(move |conn| {
246 Ok(conn.execute(
247 r#"INSERT INTO messages (id, conversation_id, text, issuer, system, token_count, created_at)
248 VALUES (:id, :conversation_id, :text, :issuer, :system, :token_count, :created_at)
249 ON CONFLICT(id, conversation_id) DO UPDATE SET
250 text = excluded.text,
251 issuer = excluded.issuer,
252 system = excluded.system,
253 token_count = excluded.token_count,
254 created_at = excluded.created_at
255 "#,
256 named_params! {
257 ":id": id,
258 ":conversation_id": conversation_id,
259 ":text": text,
260 ":issuer": issuer,
261 ":system": system,
262 ":token_count":token_count,
263 ":created_at": timestamp
264 },
265 )?)
266 })
267 .await?;
268
269 if affected_rows == 0 {
270 bail!("no rows updated for message with id {}", message.id());
271 }
272 Ok(())
273 }
274
275 async fn delete_messsage(&self, id: &str) -> Result<()> {
276 let id = id.to_string();
277 self.conn
278 .call(move |conn| {
279 let tx = conn.transaction()?;
280 tx.execute("DELETE FROM messages WHERE id = ?", params![id])?;
281 Ok(tx.commit()?)
282 })
283 .await?;
284 Ok(())
285 }
286
287 async fn upsert_context(&self, conversation_id: &str, ctx: ConvoContext) -> Result<()> {
288 let conversation_id = conversation_id.to_string();
289 let ctx_id = ctx.id().to_string();
290 let affected_rows = self
291 .conn
292 .call(move |conn| {
293 Ok(conn.execute(
294 r#"INSERT INTO contexts (id, conversation_id, last_message_id, content, token_count, created_at)
295 VALUES (:id, :conversation_id, :last_message_id, :content, :token_count, :created_at)
296 ON CONFLICT(id, conversation_id, last_message_id) DO UPDATE SET
297 content = excluded.content,
298 token_count = excluded.token_count,
299 created_at = excluded.created_at
300 "#,
301 named_params! {
302 ":id": ctx.id(),
303 ":conversation_id": conversation_id,
304 ":last_message_id": ctx.last_message_id(),
305 ":content": ctx.content(),
306 ":token_count":ctx.token_count() as i32,
307 ":created_at": ctx.created_at().timestamp_millis(),
308 },
309 )?)
310 })
311 .await?;
312
313 if affected_rows == 0 {
314 bail!("no rows updated for context with id {}", ctx_id);
315 }
316 Ok(())
317 }
318}
319
320impl Sqlite {
321 async fn get_contexts(&self, conversation_id: &str) -> Result<Vec<ConvoContext>> {
322 let conversation_id = conversation_id.to_string();
323 let contexts = self.conn.call(move |conn| {
324 let mut stmt = conn.prepare(
325 "SELECT id, conversation_id, last_message_id, content, token_count, created_at FROM contexts WHERE conversation_id = ?",
326 )?;
327
328 let mut rows = stmt.query(params![conversation_id])?;
329 let mut contexts = vec![];
330 while let Some(row) = rows.next()? {
331 let id: String = row.get(0)?;
332 let last_message_id: String = row.get(2)?;
333 let content: String = row.get(3)?;
334 let token_count: usize = row.get(4)?;
335 let created_at: i64 = row.get(5)?;
336 let created_at =
337 chrono::DateTime::from_timestamp_millis(created_at).ok_or(tokio_rusqlite::Error::Other(eyre::eyre!("invalid timestamp").into()))?;
338
339 contexts.push(
340 ConvoContext::new(&last_message_id)
341 .with_id(id)
342 .with_content(content)
343 .with_token_count(token_count)
344 .with_created_at(created_at),
345 );
346 }
347
348 contexts.sort_by(|a, b| {
349 a.created_at()
350 .timestamp_millis()
351 .cmp(&b.created_at().timestamp_millis())
352 });
353 Ok(contexts)
354 }).await?;
355 Ok(contexts)
356 }
357}
358
359type Param = (&'static str, Box<dyn ToSql>);
360
361fn filter_to_query(filter: &FilterConversation) -> (String, Vec<Param>) {
362 let mut query =
363 String::from("SELECT id, title, created_at, updated_at FROM conversations WHERE 1=1");
364 let mut params: Vec<(&str, Box<dyn ToSql>)> = vec![];
365
366 if let Some(id) = filter.id() {
367 query.push_str(" AND id = :id");
368 params.push((":id", Box::new(id.to_string())));
369 }
370
371 if let Some(title) = filter.title() {
372 query.push_str(" AND title LIKE :title");
373 params.push((":title", Box::new(format!("%{}%", title))));
374 }
375
376 if let Some(message_contains) = filter.message_contains() {
377 query.push_str(" AND EXISTS (SELECT 1 FROM messages WHERE conversation_id = conversations.id AND text LIKE :message_contains)");
378 params.push((
379 ":message_contains",
380 Box::new(format!("%{}%", message_contains)),
381 ));
382 }
383
384 if let Some(from) = filter.updated_at_from() {
385 query.push_str(" AND updated_at >= :updated_at_from");
386 params.push((":updated_at_from", Box::new(from.timestamp_millis())));
387 }
388
389 if let Some(to) = filter.updated_at_to() {
390 query.push_str(" AND updated_at <= :updated_at_to");
391 params.push((":updated_at_to", Box::new(to.timestamp_millis())));
392 }
393
394 if let Some(from) = filter.created_at_from() {
395 query.push_str(" AND created_at >= :created_at_from");
396 params.push((":created_at_from", Box::new(from.timestamp_millis())));
397 }
398
399 if let Some(to) = filter.created_at_to() {
400 query.push_str(" AND created_at <= :created_at_to");
401 params.push((":created_at_to", Box::new(to.timestamp_millis())));
402 }
403
404 (query, params)
405}