1use std::any::Any;
10use std::sync::Arc;
11
12use async_trait::async_trait;
13
14use khive_storage::error::StorageError;
15use khive_storage::types::{SqlColumn, SqlRow, SqlStatement, SqlValue};
16use khive_storage::{AtomicUnitOp, StorageCapability};
17
18use crate::error::SqliteError;
19use crate::pool::ConnectionPool;
20
21fn row_to_sql_row(row: &rusqlite::Row<'_>, col_count: usize, col_names: &[String]) -> SqlRow {
27 let mut columns = Vec::with_capacity(col_count);
28 for i in 0..col_count {
29 let value = match row.get_ref(i) {
30 Ok(rusqlite::types::ValueRef::Null) => SqlValue::Null,
31 Ok(rusqlite::types::ValueRef::Integer(v)) => SqlValue::Integer(v),
32 Ok(rusqlite::types::ValueRef::Real(v)) => SqlValue::Float(v),
33 Ok(rusqlite::types::ValueRef::Text(bytes)) => {
34 SqlValue::Text(String::from_utf8_lossy(bytes).into_owned())
35 }
36 Ok(rusqlite::types::ValueRef::Blob(bytes)) => SqlValue::Blob(bytes.to_vec()),
37 Err(_) => SqlValue::Null,
38 };
39 columns.push(SqlColumn {
40 name: col_names.get(i).cloned().unwrap_or_default(),
41 value,
42 });
43 }
44 SqlRow { columns }
45}
46
47pub(crate) fn bind_params(
54 stmt: &mut rusqlite::Statement<'_>,
55 params: &[SqlValue],
56) -> Result<(), rusqlite::Error> {
57 for (i, param) in params.iter().enumerate() {
58 let idx = i + 1; match param {
60 SqlValue::Null => stmt.raw_bind_parameter(idx, rusqlite::types::Null)?,
61 SqlValue::Bool(v) => stmt.raw_bind_parameter(idx, *v as i64)?,
62 SqlValue::Integer(v) => stmt.raw_bind_parameter(idx, *v)?,
63 SqlValue::Float(v) => stmt.raw_bind_parameter(idx, *v)?,
64 SqlValue::Text(v) => stmt.raw_bind_parameter(idx, v.as_str())?,
65 SqlValue::Blob(v) => stmt.raw_bind_parameter(idx, v.as_slice())?,
66 SqlValue::Json(v) => {
67 let s = serde_json::to_string(v).unwrap_or_default();
68 stmt.raw_bind_parameter(idx, s.as_str())?;
69 }
70 SqlValue::Uuid(v) => stmt.raw_bind_parameter(idx, v.to_string().as_str())?,
71 SqlValue::Timestamp(v) => {
72 stmt.raw_bind_parameter(idx, v.timestamp_micros())?;
73 }
74 }
75 }
76 Ok(())
77}
78
79fn execute_query(
81 conn: &rusqlite::Connection,
82 statement: &SqlStatement,
83) -> Result<Vec<SqlRow>, rusqlite::Error> {
84 let mut stmt = conn.prepare(&statement.sql)?;
85 bind_params(&mut stmt, &statement.params)?;
86
87 let col_count = stmt.column_count();
88 let col_names: Vec<String> = (0..col_count)
89 .map(|i| stmt.column_name(i).unwrap_or("").to_string())
90 .collect();
91
92 let mut rows = Vec::new();
93 let mut raw_rows = stmt.raw_query();
94 while let Some(row) = raw_rows.next()? {
95 rows.push(row_to_sql_row(row, col_count, &col_names));
96 }
97 Ok(rows)
98}
99
100fn map_rusqlite_err(e: rusqlite::Error, op: &'static str) -> StorageError {
102 StorageError::driver(StorageCapability::Sql, op, e)
103}
104
105fn open_standalone_reader(pool: &ConnectionPool) -> Result<rusqlite::Connection, StorageError> {
110 let config = pool.config();
111 let path = config.path.as_ref().ok_or_else(|| StorageError::Pool {
112 operation: "reader".into(),
113 message: "in-memory databases do not support standalone readers; use pool-backed".into(),
114 })?;
115
116 let conn = rusqlite::Connection::open_with_flags(
117 path,
118 rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY
119 | rusqlite::OpenFlags::SQLITE_OPEN_NO_MUTEX
120 | rusqlite::OpenFlags::SQLITE_OPEN_URI,
121 )
122 .map_err(|e| map_rusqlite_err(e, "open_reader"))?;
123
124 conn.busy_timeout(config.busy_timeout)
125 .map_err(|e| map_rusqlite_err(e, "open_reader"))?;
126 conn.pragma_update(None, "cache_size", "-65536")
127 .map_err(|e| map_rusqlite_err(e, "open_reader"))?;
128 conn.pragma_update(None, "mmap_size", "1073741824")
129 .map_err(|e| map_rusqlite_err(e, "open_reader"))?;
130
131 Ok(conn)
132}
133
134fn open_standalone_writer(pool: &ConnectionPool) -> Result<rusqlite::Connection, StorageError> {
135 let config = pool.config();
136 let path = config.path.as_ref().ok_or_else(|| StorageError::Pool {
137 operation: "writer".into(),
138 message: "in-memory databases do not support standalone writer; use pool-backed".into(),
139 })?;
140
141 let conn = rusqlite::Connection::open_with_flags(
142 path,
143 rusqlite::OpenFlags::SQLITE_OPEN_READ_WRITE
144 | rusqlite::OpenFlags::SQLITE_OPEN_NO_MUTEX
145 | rusqlite::OpenFlags::SQLITE_OPEN_URI,
146 )
147 .map_err(|e| map_rusqlite_err(e, "open_writer"))?;
148
149 conn.busy_timeout(config.busy_timeout)
150 .map_err(|e| map_rusqlite_err(e, "open_writer"))?;
151 conn.pragma_update(None, "cache_size", "-65536")
152 .map_err(|e| map_rusqlite_err(e, "open_writer"))?;
153 conn.pragma_update(None, "mmap_size", "1073741824")
154 .map_err(|e| map_rusqlite_err(e, "open_writer"))?;
155
156 Ok(conn)
157}
158
159struct SqliteReader {
164 conn: Option<rusqlite::Connection>,
165}
166
167#[async_trait]
168impl khive_storage::SqlReader for SqliteReader {
169 async fn query_row(
170 &mut self,
171 statement: SqlStatement,
172 ) -> khive_storage::types::StorageResult<Option<SqlRow>> {
173 let conn = self.conn.take().ok_or_else(|| StorageError::Pool {
174 operation: "query_row".into(),
175 message: "connection already consumed".into(),
176 })?;
177 let (conn, result) = tokio::task::spawn_blocking(move || {
178 let res = execute_query(&conn, &statement);
179 (conn, res)
180 })
181 .await
182 .map_err(|e| StorageError::driver(StorageCapability::Sql, "query_row", e))?;
183 self.conn = Some(conn);
184 let rows = result.map_err(|e| map_rusqlite_err(e, "query_row"))?;
185 Ok(rows.into_iter().next())
186 }
187
188 async fn query_all(
189 &mut self,
190 statement: SqlStatement,
191 ) -> khive_storage::types::StorageResult<Vec<SqlRow>> {
192 let conn = self.conn.take().ok_or_else(|| StorageError::Pool {
193 operation: "query_all".into(),
194 message: "connection already consumed".into(),
195 })?;
196 let (conn, result) = tokio::task::spawn_blocking(move || {
197 let res = execute_query(&conn, &statement);
198 (conn, res)
199 })
200 .await
201 .map_err(|e| StorageError::driver(StorageCapability::Sql, "query_all", e))?;
202 self.conn = Some(conn);
203 result.map_err(|e| map_rusqlite_err(e, "query_all"))
204 }
205
206 async fn query_scalar(
207 &mut self,
208 statement: SqlStatement,
209 ) -> khive_storage::types::StorageResult<Option<SqlValue>> {
210 let row = self.query_row(statement).await?;
211 Ok(row.and_then(|r| r.columns.into_iter().next().map(|c| c.value)))
212 }
213
214 async fn explain(
215 &mut self,
216 statement: SqlStatement,
217 ) -> khive_storage::types::StorageResult<Vec<SqlRow>> {
218 let explain_stmt = SqlStatement {
219 sql: format!("EXPLAIN QUERY PLAN {}", statement.sql),
220 params: statement.params,
221 label: statement.label,
222 };
223 self.query_all(explain_stmt).await
224 }
225}
226
227struct SqliteWriter {
232 conn: Option<rusqlite::Connection>,
233 writer_task: Option<crate::writer_task::WriterTaskHandle>,
239}
240
241#[async_trait]
242impl khive_storage::SqlReader for SqliteWriter {
243 async fn query_row(
244 &mut self,
245 statement: SqlStatement,
246 ) -> khive_storage::types::StorageResult<Option<SqlRow>> {
247 let conn = self.conn.take().ok_or_else(|| StorageError::Pool {
248 operation: "writer.query_row".into(),
249 message: "connection already consumed".into(),
250 })?;
251 let (conn, result) = tokio::task::spawn_blocking(move || {
252 let res = execute_query(&conn, &statement);
253 (conn, res)
254 })
255 .await
256 .map_err(|e| StorageError::driver(StorageCapability::Sql, "writer.query_row", e))?;
257 self.conn = Some(conn);
258 let rows = result.map_err(|e| map_rusqlite_err(e, "writer.query_row"))?;
259 Ok(rows.into_iter().next())
260 }
261
262 async fn query_all(
263 &mut self,
264 statement: SqlStatement,
265 ) -> khive_storage::types::StorageResult<Vec<SqlRow>> {
266 let conn = self.conn.take().ok_or_else(|| StorageError::Pool {
267 operation: "writer.query_all".into(),
268 message: "connection already consumed".into(),
269 })?;
270 let (conn, result) = tokio::task::spawn_blocking(move || {
271 let res = execute_query(&conn, &statement);
272 (conn, res)
273 })
274 .await
275 .map_err(|e| StorageError::driver(StorageCapability::Sql, "writer.query_all", e))?;
276 self.conn = Some(conn);
277 result.map_err(|e| map_rusqlite_err(e, "writer.query_all"))
278 }
279
280 async fn query_scalar(
281 &mut self,
282 statement: SqlStatement,
283 ) -> khive_storage::types::StorageResult<Option<SqlValue>> {
284 let row = khive_storage::SqlReader::query_row(self, statement).await?;
285 Ok(row.and_then(|r| r.columns.into_iter().next().map(|c| c.value)))
286 }
287
288 async fn explain(
289 &mut self,
290 statement: SqlStatement,
291 ) -> khive_storage::types::StorageResult<Vec<SqlRow>> {
292 let explain_stmt = SqlStatement {
293 sql: format!("EXPLAIN QUERY PLAN {}", statement.sql),
294 params: statement.params,
295 label: statement.label,
296 };
297 khive_storage::SqlReader::query_all(self, explain_stmt).await
298 }
299}
300
301#[async_trait]
302impl khive_storage::SqlWriter for SqliteWriter {
303 async fn execute(
304 &mut self,
305 statement: SqlStatement,
306 ) -> khive_storage::types::StorageResult<u64> {
307 if let Some(writer_task) = self.writer_task.clone() {
313 return writer_task
314 .send(move |conn| {
315 let mut stmt = conn
316 .prepare(&statement.sql)
317 .map_err(|e| map_rusqlite_err(e, "execute"))?;
318 bind_params(&mut stmt, &statement.params)
319 .map_err(|e| map_rusqlite_err(e, "execute"))?;
320 let affected = stmt
321 .raw_execute()
322 .map_err(|e| map_rusqlite_err(e, "execute"))?;
323 Ok(affected as u64)
324 })
325 .await;
326 }
327
328 let conn = self.conn.take().ok_or_else(|| StorageError::Pool {
329 operation: "execute".into(),
330 message: "connection already consumed".into(),
331 })?;
332 let (conn, result) = tokio::task::spawn_blocking(move || {
333 let res = (|| -> Result<usize, rusqlite::Error> {
334 let mut stmt = conn.prepare(&statement.sql)?;
335 bind_params(&mut stmt, &statement.params)?;
336 stmt.raw_execute()
337 })();
338 (conn, res)
339 })
340 .await
341 .map_err(|e| StorageError::driver(StorageCapability::Sql, "execute", e))?;
342 self.conn = Some(conn);
343 let affected = result.map_err(|e| map_rusqlite_err(e, "execute"))?;
344 Ok(affected as u64)
345 }
346
347 async fn execute_batch(
348 &mut self,
349 statements: Vec<SqlStatement>,
350 ) -> khive_storage::types::StorageResult<u64> {
351 if let Some(writer_task) = self.writer_task.clone() {
360 return writer_task
361 .send(move |conn| {
362 let mut total: u64 = 0;
363 for statement in &statements {
364 let mut stmt = conn
365 .prepare(&statement.sql)
366 .map_err(|e| map_rusqlite_err(e, "execute_batch"))?;
367 bind_params(&mut stmt, &statement.params)
368 .map_err(|e| map_rusqlite_err(e, "execute_batch"))?;
369 total += stmt
370 .raw_execute()
371 .map_err(|e| map_rusqlite_err(e, "execute_batch"))?
372 as u64;
373 }
374 Ok(total)
375 })
376 .await;
377 }
378
379 let conn = self.conn.take().ok_or_else(|| StorageError::Pool {
380 operation: "execute_batch".into(),
381 message: "connection already consumed".into(),
382 })?;
383 let (conn, result) = tokio::task::spawn_blocking(move || {
384 if let Err(e) = conn.execute_batch("BEGIN IMMEDIATE") {
385 return (conn, Err(e));
386 }
387 let _tx_handle =
393 khive_storage::tx_registry::register(Some("execute_batch".to_string()));
394 let result = (|| -> Result<u64, rusqlite::Error> {
395 let mut total: u64 = 0;
396 for statement in &statements {
397 let mut stmt = conn.prepare(&statement.sql)?;
398 bind_params(&mut stmt, &statement.params)?;
399 total += stmt.raw_execute()? as u64;
400 }
401 conn.execute_batch("COMMIT")?;
402 Ok(total)
403 })();
404 if result.is_err() {
405 let _ = conn.execute_batch("ROLLBACK");
406 }
407 (conn, result)
409 })
410 .await
411 .map_err(|e| StorageError::driver(StorageCapability::Sql, "execute_batch", e))?;
412 self.conn = Some(conn);
413 result.map_err(|e| map_rusqlite_err(e, "execute_batch"))
414 }
415
416 async fn execute_script(&mut self, script: String) -> khive_storage::types::StorageResult<()> {
417 if let Some(writer_task) = self.writer_task.clone() {
427 return writer_task
428 .send(move |conn| {
429 conn.execute_batch(&script)
430 .map_err(|e| map_rusqlite_err(e, "execute_script"))
431 })
432 .await;
433 }
434
435 let conn = self.conn.take().ok_or_else(|| StorageError::Pool {
436 operation: "execute_script".into(),
437 message: "connection already consumed".into(),
438 })?;
439 let (conn, result) = tokio::task::spawn_blocking(move || {
440 let res = conn.execute_batch(&script);
441 (conn, res)
442 })
443 .await
444 .map_err(|e| StorageError::driver(StorageCapability::Sql, "execute_script", e))?;
445 self.conn = Some(conn);
446 result.map_err(|e| map_rusqlite_err(e, "execute_script"))
447 }
448
449 async fn execute_script_top_level(
450 &mut self,
451 script: String,
452 ) -> khive_storage::types::StorageResult<()> {
453 if let Some(writer_task) = self.writer_task.clone() {
461 return writer_task
462 .send_top_level(move |conn| {
463 conn.execute_batch(&script)
464 .map_err(|e| map_rusqlite_err(e, "execute_script_top_level"))
465 })
466 .await;
467 }
468
469 let conn = self.conn.take().ok_or_else(|| StorageError::Pool {
473 operation: "execute_script_top_level".into(),
474 message: "connection already consumed".into(),
475 })?;
476 let (conn, result) = tokio::task::spawn_blocking(move || {
477 let res = conn.execute_batch(&script);
478 (conn, res)
479 })
480 .await
481 .map_err(|e| StorageError::driver(StorageCapability::Sql, "execute_script_top_level", e))?;
482 self.conn = Some(conn);
483 result.map_err(|e| map_rusqlite_err(e, "execute_script_top_level"))
484 }
485}
486
487struct PoolBackedReader {
492 pool: Arc<ConnectionPool>,
493}
494
495#[async_trait]
496impl khive_storage::SqlReader for PoolBackedReader {
497 async fn query_row(
498 &mut self,
499 statement: SqlStatement,
500 ) -> khive_storage::types::StorageResult<Option<SqlRow>> {
501 let pool = Arc::clone(&self.pool);
502 tokio::task::spawn_blocking(move || {
503 let guard = pool
504 .reader()
505 .map_err(|e| StorageError::driver(StorageCapability::Sql, "pool_reader", e))?;
506 let rows = execute_query(&guard, &statement)
507 .map_err(|e| map_rusqlite_err(e, "pool_reader.query_row"))?;
508 Ok(rows.into_iter().next())
509 })
510 .await
511 .map_err(|e| StorageError::driver(StorageCapability::Sql, "pool_reader.query_row", e))?
512 }
513
514 async fn query_all(
515 &mut self,
516 statement: SqlStatement,
517 ) -> khive_storage::types::StorageResult<Vec<SqlRow>> {
518 let pool = Arc::clone(&self.pool);
519 tokio::task::spawn_blocking(move || {
520 let guard = pool
521 .reader()
522 .map_err(|e| StorageError::driver(StorageCapability::Sql, "pool_reader", e))?;
523 execute_query(&guard, &statement)
524 .map_err(|e| map_rusqlite_err(e, "pool_reader.query_all"))
525 })
526 .await
527 .map_err(|e| StorageError::driver(StorageCapability::Sql, "pool_reader.query_all", e))?
528 }
529
530 async fn query_scalar(
531 &mut self,
532 statement: SqlStatement,
533 ) -> khive_storage::types::StorageResult<Option<SqlValue>> {
534 let row = self.query_row(statement).await?;
535 Ok(row.and_then(|r| r.columns.into_iter().next().map(|c| c.value)))
536 }
537
538 async fn explain(
539 &mut self,
540 statement: SqlStatement,
541 ) -> khive_storage::types::StorageResult<Vec<SqlRow>> {
542 let explain_stmt = SqlStatement {
543 sql: format!("EXPLAIN QUERY PLAN {}", statement.sql),
544 params: statement.params,
545 label: statement.label,
546 };
547 self.query_all(explain_stmt).await
548 }
549}
550
551struct PoolBackedWriter {
552 pool: Arc<ConnectionPool>,
553}
554
555#[async_trait]
556impl khive_storage::SqlReader for PoolBackedWriter {
557 async fn query_row(
558 &mut self,
559 statement: SqlStatement,
560 ) -> khive_storage::types::StorageResult<Option<SqlRow>> {
561 let pool = Arc::clone(&self.pool);
562 tokio::task::spawn_blocking(move || {
563 let guard = pool.try_writer().map_err(|e: SqliteError| {
564 StorageError::driver(StorageCapability::Sql, "pool_writer.query_row", e)
565 })?;
566 let rows = execute_query(&guard, &statement)
567 .map_err(|e| map_rusqlite_err(e, "pool_writer.query_row"))?;
568 Ok(rows.into_iter().next())
569 })
570 .await
571 .map_err(|e| StorageError::driver(StorageCapability::Sql, "pool_writer.query_row", e))?
572 }
573
574 async fn query_all(
575 &mut self,
576 statement: SqlStatement,
577 ) -> khive_storage::types::StorageResult<Vec<SqlRow>> {
578 let pool = Arc::clone(&self.pool);
579 tokio::task::spawn_blocking(move || {
580 let guard = pool.try_writer().map_err(|e: SqliteError| {
581 StorageError::driver(StorageCapability::Sql, "pool_writer.query_all", e)
582 })?;
583 execute_query(&guard, &statement)
584 .map_err(|e| map_rusqlite_err(e, "pool_writer.query_all"))
585 })
586 .await
587 .map_err(|e| StorageError::driver(StorageCapability::Sql, "pool_writer.query_all", e))?
588 }
589
590 async fn query_scalar(
591 &mut self,
592 statement: SqlStatement,
593 ) -> khive_storage::types::StorageResult<Option<SqlValue>> {
594 let row = khive_storage::SqlReader::query_row(self, statement).await?;
595 Ok(row.and_then(|r| r.columns.into_iter().next().map(|c| c.value)))
596 }
597
598 async fn explain(
599 &mut self,
600 statement: SqlStatement,
601 ) -> khive_storage::types::StorageResult<Vec<SqlRow>> {
602 let explain_stmt = SqlStatement {
603 sql: format!("EXPLAIN QUERY PLAN {}", statement.sql),
604 params: statement.params,
605 label: statement.label,
606 };
607 khive_storage::SqlReader::query_all(self, explain_stmt).await
608 }
609}
610
611#[async_trait]
612impl khive_storage::SqlWriter for PoolBackedWriter {
613 async fn execute(
614 &mut self,
615 statement: SqlStatement,
616 ) -> khive_storage::types::StorageResult<u64> {
617 let pool = Arc::clone(&self.pool);
618 tokio::task::spawn_blocking(move || {
619 let guard = pool.try_writer().map_err(|e: SqliteError| {
620 StorageError::driver(StorageCapability::Sql, "pool_writer.execute", e)
621 })?;
622 let mut stmt = guard
623 .prepare(&statement.sql)
624 .map_err(|e| map_rusqlite_err(e, "pool_writer.execute"))?;
625 bind_params(&mut stmt, &statement.params)
626 .map_err(|e| map_rusqlite_err(e, "pool_writer.execute"))?;
627 let rows = stmt
628 .raw_execute()
629 .map_err(|e| map_rusqlite_err(e, "pool_writer.execute"))?;
630 Ok(rows as u64)
631 })
632 .await
633 .map_err(|e| StorageError::driver(StorageCapability::Sql, "pool_writer.execute", e))?
634 }
635
636 async fn execute_batch(
637 &mut self,
638 statements: Vec<SqlStatement>,
639 ) -> khive_storage::types::StorageResult<u64> {
640 let pool = Arc::clone(&self.pool);
641 tokio::task::spawn_blocking(move || {
642 let guard = pool.try_writer().map_err(|e: SqliteError| {
643 StorageError::driver(StorageCapability::Sql, "pool_writer.execute_batch", e)
644 })?;
645 guard
646 .execute_batch("BEGIN IMMEDIATE")
647 .map_err(|e| map_rusqlite_err(e, "pool_writer.execute_batch"))?;
648 let _tx_handle =
649 khive_storage::tx_registry::register(Some("pool_writer.execute_batch".to_string()));
650 let result = (|| -> Result<u64, StorageError> {
651 let mut total = 0u64;
652 for statement in &statements {
653 let mut stmt = guard
654 .prepare(&statement.sql)
655 .map_err(|e| map_rusqlite_err(e, "pool_writer.execute_batch"))?;
656 bind_params(&mut stmt, &statement.params)
657 .map_err(|e| map_rusqlite_err(e, "pool_writer.execute_batch"))?;
658 total += stmt
659 .raw_execute()
660 .map_err(|e| map_rusqlite_err(e, "pool_writer.execute_batch"))?
661 as u64;
662 }
663 Ok(total)
664 })();
665 match result {
666 Ok(total) => {
667 if let Err(e) = guard.execute_batch("COMMIT") {
668 let _ = guard.execute_batch("ROLLBACK");
669 Err(map_rusqlite_err(e, "pool_writer.execute_batch"))
670 } else {
671 Ok(total)
672 }
673 }
674 Err(e) => {
675 let _ = guard.execute_batch("ROLLBACK");
676 Err(e)
677 }
678 }
679 })
680 .await
681 .map_err(|e| StorageError::driver(StorageCapability::Sql, "pool_writer.execute_batch", e))?
682 }
683
684 async fn execute_script(&mut self, script: String) -> khive_storage::types::StorageResult<()> {
685 let pool = Arc::clone(&self.pool);
686 tokio::task::spawn_blocking(move || {
687 let guard = pool.try_writer().map_err(|e: SqliteError| {
688 StorageError::driver(StorageCapability::Sql, "pool_writer.execute_script", e)
689 })?;
690 guard
691 .execute_batch(&script)
692 .map_err(|e| map_rusqlite_err(e, "pool_writer.execute_script"))
693 })
694 .await
695 .map_err(|e| {
696 StorageError::driver(StorageCapability::Sql, "pool_writer.execute_script", e)
697 })?
698 }
699}
700
701struct InlineWriter {
727 conn: *const rusqlite::Connection,
728}
729
730unsafe impl Send for InlineWriter {}
737
738impl InlineWriter {
739 fn conn(&self) -> &rusqlite::Connection {
743 unsafe { &*self.conn }
744 }
745}
746
747#[async_trait]
748impl khive_storage::SqlReader for InlineWriter {
749 async fn query_row(
750 &mut self,
751 statement: SqlStatement,
752 ) -> khive_storage::types::StorageResult<Option<SqlRow>> {
753 let rows = execute_query(self.conn(), &statement)
754 .map_err(|e| map_rusqlite_err(e, "inline.query_row"))?;
755 Ok(rows.into_iter().next())
756 }
757
758 async fn query_all(
759 &mut self,
760 statement: SqlStatement,
761 ) -> khive_storage::types::StorageResult<Vec<SqlRow>> {
762 execute_query(self.conn(), &statement).map_err(|e| map_rusqlite_err(e, "inline.query_all"))
763 }
764
765 async fn query_scalar(
766 &mut self,
767 statement: SqlStatement,
768 ) -> khive_storage::types::StorageResult<Option<SqlValue>> {
769 let row = khive_storage::SqlReader::query_row(self, statement).await?;
770 Ok(row.and_then(|r| r.columns.into_iter().next().map(|c| c.value)))
771 }
772
773 async fn explain(
774 &mut self,
775 statement: SqlStatement,
776 ) -> khive_storage::types::StorageResult<Vec<SqlRow>> {
777 let explain_stmt = SqlStatement {
778 sql: format!("EXPLAIN QUERY PLAN {}", statement.sql),
779 params: statement.params,
780 label: statement.label,
781 };
782 khive_storage::SqlReader::query_all(self, explain_stmt).await
783 }
784}
785
786#[async_trait]
787impl khive_storage::SqlWriter for InlineWriter {
788 async fn execute(
789 &mut self,
790 statement: SqlStatement,
791 ) -> khive_storage::types::StorageResult<u64> {
792 let mut stmt = self
793 .conn()
794 .prepare(&statement.sql)
795 .map_err(|e| map_rusqlite_err(e, "inline.execute"))?;
796 bind_params(&mut stmt, &statement.params)
797 .map_err(|e| map_rusqlite_err(e, "inline.execute"))?;
798 let affected = stmt
799 .raw_execute()
800 .map_err(|e| map_rusqlite_err(e, "inline.execute"))?;
801 Ok(affected as u64)
802 }
803
804 async fn execute_batch(
805 &mut self,
806 statements: Vec<SqlStatement>,
807 ) -> khive_storage::types::StorageResult<u64> {
808 let mut total: u64 = 0;
809 for statement in &statements {
810 let mut stmt = self
811 .conn()
812 .prepare(&statement.sql)
813 .map_err(|e| map_rusqlite_err(e, "inline.execute_batch"))?;
814 bind_params(&mut stmt, &statement.params)
815 .map_err(|e| map_rusqlite_err(e, "inline.execute_batch"))?;
816 total += stmt
817 .raw_execute()
818 .map_err(|e| map_rusqlite_err(e, "inline.execute_batch"))?
819 as u64;
820 }
821 Ok(total)
822 }
823
824 async fn execute_script(&mut self, script: String) -> khive_storage::types::StorageResult<()> {
825 self.conn()
826 .execute_batch(&script)
827 .map_err(|e| map_rusqlite_err(e, "inline.execute_script"))
828 }
829}
830
831fn block_on_sync<F: std::future::Future>(fut: F) -> Result<F::Output, StorageError> {
854 use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
855
856 fn no_op(_: *const ()) {}
857 fn clone_waker(_: *const ()) -> RawWaker {
858 RawWaker::new(std::ptr::null(), &VTABLE)
859 }
860 static VTABLE: RawWakerVTable = RawWakerVTable::new(clone_waker, no_op, no_op, no_op);
861
862 let raw_waker = RawWaker::new(std::ptr::null(), &VTABLE);
865 let waker = unsafe { Waker::from_raw(raw_waker) };
866 let mut cx = Context::from_waker(&waker);
867
868 let mut fut = std::pin::pin!(fut);
869 match fut.as_mut().poll(&mut cx) {
870 Poll::Ready(v) => Ok(v),
871 Poll::Pending => {
872 tracing::error!(
873 "block_on_sync: atomic_unit future suspended on its first poll — \
874 the closure passed to SqlAccess::atomic_unit must be non-blocking \
875 (synchronous InlineWriter calls only, no real .await point)"
876 );
877 Err(StorageError::Internal(
878 "atomic_unit future suspended — closure must be non-blocking".to_string(),
879 ))
880 }
881 }
882}
883
884async fn run_manual_atomic_unit(
889 writer: &mut dyn khive_storage::SqlWriter,
890 op: AtomicUnitOp,
891) -> khive_storage::types::StorageResult<Box<dyn Any + Send>> {
892 fn tx_stmt(sql: &str, label: &str) -> SqlStatement {
893 SqlStatement {
894 sql: sql.to_string(),
895 params: vec![],
896 label: Some(label.to_string()),
897 }
898 }
899 khive_storage::SqlWriter::execute(writer, tx_stmt("BEGIN IMMEDIATE", "begin")).await?;
900 let _tx_handle = khive_storage::tx_registry::register(Some("atomic_unit".to_string()));
901
902 let result = op(writer).await;
903
904 match result {
905 Ok(value) => {
906 match khive_storage::SqlWriter::execute(writer, tx_stmt("COMMIT", "commit")).await {
907 Ok(_) => Ok(value),
908 Err(e) => {
909 let _ =
910 khive_storage::SqlWriter::execute(writer, tx_stmt("ROLLBACK", "rollback"))
911 .await;
912 Err(e)
913 }
914 }
915 }
916 Err(e) => {
917 let _ =
918 khive_storage::SqlWriter::execute(writer, tx_stmt("ROLLBACK", "rollback")).await;
919 Err(e)
920 }
921 }
922}
923
924pub struct SqlBridge {
936 pool: Arc<ConnectionPool>,
937 is_file_backed: bool,
938}
939
940impl SqlBridge {
941 pub fn new(pool: Arc<ConnectionPool>, is_file_backed: bool) -> Self {
943 Self {
944 pool,
945 is_file_backed,
946 }
947 }
948}
949
950#[async_trait]
951impl khive_storage::SqlAccess for SqlBridge {
952 async fn reader(
953 &self,
954 ) -> khive_storage::types::StorageResult<Box<dyn khive_storage::SqlReader>> {
955 if self.is_file_backed {
956 let conn = open_standalone_reader(&self.pool)?;
957 Ok(Box::new(SqliteReader { conn: Some(conn) }))
958 } else {
959 Ok(Box::new(PoolBackedReader {
960 pool: Arc::clone(&self.pool),
961 }))
962 }
963 }
964
965 async fn writer(
966 &self,
967 ) -> khive_storage::types::StorageResult<Box<dyn khive_storage::SqlWriter>> {
968 if self.is_file_backed {
969 if self.pool.config().read_only {
970 return Err(StorageError::Pool {
971 operation: "writer".into(),
972 message: "backend is read-only".into(),
973 });
974 }
975 let conn = open_standalone_writer(&self.pool)?;
976 let writer_task = self.pool.writer_task_handle().ok().flatten();
980 Ok(Box::new(SqliteWriter {
981 conn: Some(conn),
982 writer_task,
983 }))
984 } else {
985 Ok(Box::new(PoolBackedWriter {
986 pool: Arc::clone(&self.pool),
987 }))
988 }
989 }
990
991 async fn atomic_unit(
1002 &self,
1003 op: AtomicUnitOp,
1004 ) -> khive_storage::types::StorageResult<Box<dyn Any + Send>> {
1005 if self.is_file_backed {
1006 if self.pool.config().read_only {
1007 return Err(StorageError::Pool {
1008 operation: "atomic_unit".into(),
1009 message: "backend is read-only".into(),
1010 });
1011 }
1012 if let Some(writer_task) = self.pool.writer_task_handle()? {
1016 return writer_task
1022 .send(move |conn| {
1023 let mut inline = InlineWriter {
1024 conn: conn as *const rusqlite::Connection,
1025 };
1026 match block_on_sync(op(&mut inline)) {
1036 Ok(inner) => inner,
1037 Err(e) => Err(e),
1038 }
1039 })
1040 .await;
1041 }
1042 let conn = open_standalone_writer(&self.pool)?;
1046 let mut writer = SqliteWriter {
1047 conn: Some(conn),
1048 writer_task: None,
1049 };
1050 run_manual_atomic_unit(&mut writer, op).await
1051 } else {
1052 let mut writer = PoolBackedWriter {
1056 pool: Arc::clone(&self.pool),
1057 };
1058 run_manual_atomic_unit(&mut writer, op).await
1059 }
1060 }
1061}
1062
1063#[cfg(test)]
1064mod tests {
1065 use super::*;
1066 use crate::pool::PoolConfig;
1067 use khive_storage::types::{SqlStatement, SqlValue};
1068 use khive_storage::SqlAccess as _;
1069
1070 #[tokio::test]
1076 async fn execute_batch_routes_through_writer_task_when_flag_enabled() {
1077 let dir = tempfile::tempdir().unwrap();
1078 let path = dir.path().join("write_queue_execute_batch.db");
1079 let config = PoolConfig {
1080 path: Some(path.clone()),
1081 write_queue_enabled: true,
1082 ..PoolConfig::default()
1083 };
1084 let pool = Arc::new(ConnectionPool::new(config).unwrap());
1085 {
1086 let guard = pool.writer().unwrap();
1087 guard
1088 .conn()
1089 .execute_batch(
1090 "CREATE TABLE IF NOT EXISTS write_queue_batch_test \
1091 (id INTEGER PRIMARY KEY, val TEXT NOT NULL)",
1092 )
1093 .unwrap();
1094 }
1095
1096 let bridge = SqlBridge::new(Arc::clone(&pool), true);
1097
1098 let mut writer = bridge.writer().await.unwrap();
1099 let affected = writer
1100 .execute_batch(vec![
1101 SqlStatement {
1102 sql: "INSERT INTO write_queue_batch_test (id, val) VALUES (?1, ?2)".into(),
1103 params: vec![SqlValue::Integer(1), SqlValue::Text("a".into())],
1104 label: None,
1105 },
1106 SqlStatement {
1107 sql: "INSERT INTO write_queue_batch_test (id, val) VALUES (?1, ?2)".into(),
1108 params: vec![SqlValue::Integer(2), SqlValue::Text("b".into())],
1109 label: None,
1110 },
1111 ])
1112 .await
1113 .unwrap();
1114 assert_eq!(affected, 2);
1115
1116 let mut reader = bridge.reader().await.unwrap();
1117 let count = reader
1118 .query_scalar(SqlStatement {
1119 sql: "SELECT COUNT(*) FROM write_queue_batch_test".into(),
1120 params: vec![],
1121 label: None,
1122 })
1123 .await
1124 .unwrap();
1125 assert!(
1126 matches!(count, Some(SqlValue::Integer(2))),
1127 "expected 2 rows, got {count:?}"
1128 );
1129 assert_eq!(
1130 pool.writer_task_spawn_count(),
1131 1,
1132 "the flag-ON path must actually spawn and use the writer task"
1133 );
1134 }
1135
1136 #[tokio::test]
1143 async fn execute_batch_rolls_back_atomically_on_mid_sequence_failure() {
1144 let dir = tempfile::tempdir().unwrap();
1145 let path = dir.path().join("write_queue_execute_batch_rollback.db");
1146 let config = PoolConfig {
1147 path: Some(path.clone()),
1148 write_queue_enabled: true,
1149 ..PoolConfig::default()
1150 };
1151 let pool = Arc::new(ConnectionPool::new(config).unwrap());
1152 {
1153 let guard = pool.writer().unwrap();
1154 guard
1155 .conn()
1156 .execute_batch(
1157 "CREATE TABLE IF NOT EXISTS write_queue_rollback_test \
1158 (id INTEGER PRIMARY KEY, val TEXT NOT NULL)",
1159 )
1160 .unwrap();
1161 }
1162
1163 let bridge = SqlBridge::new(Arc::clone(&pool), true);
1164
1165 let mut writer = bridge.writer().await.unwrap();
1166 let result = writer
1167 .execute_batch(vec![
1168 SqlStatement {
1170 sql: "INSERT INTO write_queue_rollback_test (id, val) VALUES (?1, ?2)".into(),
1171 params: vec![SqlValue::Integer(1), SqlValue::Text("first".into())],
1172 label: None,
1173 },
1174 SqlStatement {
1176 sql: "INSERT INTO write_queue_rollback_test (id, val) VALUES (?1, ?2)".into(),
1177 params: vec![SqlValue::Integer(1), SqlValue::Text("duplicate".into())],
1178 label: None,
1179 },
1180 SqlStatement {
1182 sql: "INSERT INTO write_queue_rollback_test (id, val) VALUES (?1, ?2)".into(),
1183 params: vec![SqlValue::Integer(2), SqlValue::Text("third".into())],
1184 label: None,
1185 },
1186 ])
1187 .await;
1188 assert!(
1189 result.is_err(),
1190 "a batch with a mid-sequence PK conflict must return an error"
1191 );
1192
1193 let mut reader = bridge.reader().await.unwrap();
1194 let count = reader
1195 .query_scalar(SqlStatement {
1196 sql: "SELECT COUNT(*) FROM write_queue_rollback_test".into(),
1197 params: vec![],
1198 label: None,
1199 })
1200 .await
1201 .unwrap();
1202 assert!(
1203 matches!(count, Some(SqlValue::Integer(0))),
1204 "the whole request must roll back — including statement 1's \
1205 otherwise-successful INSERT — not just the failing statement; \
1206 got {count:?}"
1207 );
1208 }
1209
1210 #[tokio::test]
1227 async fn atomic_unit_pending_future_errors_without_killing_writer_task() {
1228 let dir = tempfile::tempdir().unwrap();
1229 let path = dir.path().join("atomic_unit_pending_future.db");
1230 let config = PoolConfig {
1231 path: Some(path.clone()),
1232 write_queue_enabled: true,
1233 ..PoolConfig::default()
1234 };
1235 let pool = Arc::new(ConnectionPool::new(config).unwrap());
1236 {
1237 let guard = pool.writer().unwrap();
1238 guard
1239 .conn()
1240 .execute_batch(
1241 "CREATE TABLE IF NOT EXISTS atomic_unit_pending_test \
1242 (id INTEGER PRIMARY KEY, val TEXT NOT NULL)",
1243 )
1244 .unwrap();
1245 }
1246 assert!(
1247 pool.writer_task_handle().unwrap().is_some(),
1248 "writer task must be spawned with the flag on for a file-backed pool"
1249 );
1250
1251 let bridge = SqlBridge::new(Arc::clone(&pool), true);
1252
1253 let pending_op: AtomicUnitOp = Box::new(|_writer| {
1256 Box::pin(std::future::pending::<
1257 khive_storage::types::StorageResult<Box<dyn std::any::Any + Send>>,
1258 >())
1259 });
1260
1261 let pending_result = bridge.atomic_unit(pending_op).await;
1262 assert!(
1263 pending_result.is_err(),
1264 "a Pending-on-first-poll atomic_unit closure must return Err, \
1265 not panic; got {pending_result:?}"
1266 );
1267
1268 let ok_op: AtomicUnitOp = Box::new(|writer| {
1273 Box::pin(async move {
1274 writer
1275 .execute(SqlStatement {
1276 sql: "INSERT INTO atomic_unit_pending_test (id, val) VALUES (?1, ?2)"
1277 .into(),
1278 params: vec![SqlValue::Integer(1), SqlValue::Text("survived".into())],
1279 label: None,
1280 })
1281 .await
1282 .map_err(|e| {
1283 khive_storage::StorageError::driver(
1284 StorageCapability::Sql,
1285 "atomic_unit_pending_future_test_insert",
1286 e,
1287 )
1288 })?;
1289 Ok(Box::new(()) as Box<dyn std::any::Any + Send>)
1290 })
1291 });
1292 let ok_result = bridge.atomic_unit(ok_op).await;
1293 assert!(
1294 ok_result.is_ok(),
1295 "writer task must survive a Pending misuse and keep serving \
1296 subsequent well-behaved atomic_unit requests; got {ok_result:?}"
1297 );
1298
1299 let mut reader = bridge.reader().await.unwrap();
1300 let count = reader
1301 .query_scalar(SqlStatement {
1302 sql: "SELECT COUNT(*) FROM atomic_unit_pending_test".into(),
1303 params: vec![],
1304 label: None,
1305 })
1306 .await
1307 .unwrap();
1308 assert!(
1309 matches!(count, Some(SqlValue::Integer(1))),
1310 "the well-behaved atomic_unit call after the Pending misuse must \
1311 have actually committed its write; got {count:?}"
1312 );
1313 }
1314}