Extend trailbase-sqlites execution model to allow for parallel reads. This reduces the latency long-tail for slow reads.

Note that this complicates the APIs, since it pushes the responsibility
of declaring a query a read or write to the user to then be scheduled
appropriately.

Add `.(read_|)query_row_f` APIs similar to rusqlites
`conn.query_row` accepting a `|row| -> Result<T>` to reduce the use of
`Row` and `Rows`.

Make benchmarks more isolated by not sharing a DB across runs
accumulating writes.

Author: ignatz

trailbase-cli/src/bin/trail.rs

@@ -55,8 +55,8 @@ async fn get_user_by_email(
   email: &str,
 ) -> Result<DbUser, BoxError> {
   if let Some(user) = conn
-    .query_value::<DbUser>(
-      &format!("SELECT * FROM {USER_TABLE} WHERE email = $1"),
+    .read_query_value::<DbUser>(
+      format!("SELECT * FROM {USER_TABLE} WHERE email = $1"),
       (email.to_string(),),
     )
     .await?
@@ -165,7 +165,7 @@ async fn async_main() -> Result<(), BoxError> {
       match cmd {
         Some(AdminSubCommands::List) => {
           let users = conn
-            .query_values::<DbUser>(&format!("SELECT * FROM {USER_TABLE} WHERE admin > 0"), ())
+            .read_query_values::<DbUser>(format!("SELECT * FROM {USER_TABLE} WHERE admin > 0"), ())
             .await?;
 
           println!("{: >36}\temail\tcreated\tupdated", "id");
@@ -183,7 +183,7 @@ async fn async_main() -> Result<(), BoxError> {
         Some(AdminSubCommands::Demote { email }) => {
           conn
             .execute(
-              &format!("UPDATE {USER_TABLE} SET admin = FALSE WHERE email = $1"),
+              format!("UPDATE {USER_TABLE} SET admin = FALSE WHERE email = $1"),
               (email.clone(),),
             )
             .await?;
@@ -193,7 +193,7 @@ async fn async_main() -> Result<(), BoxError> {
         Some(AdminSubCommands::Promote { email }) => {
           conn
             .execute(
-              &format!("UPDATE {USER_TABLE} SET admin = TRUE WHERE email = $1"),
+              format!("UPDATE {USER_TABLE} SET admin = TRUE WHERE email = $1"),
               (email.clone(),),
             )
             .await?;

trailbase-core/benches/benchmark.rs

@@ -60,13 +60,13 @@ async fn add_room(
   name: &str,
 ) -> Result<[u8; 16], anyhow::Error> {
   let room: [u8; 16] = conn
-    .query_row(
+    .query_row_f(
       "INSERT INTO room (name) VALUES ($1) RETURNING id",
       params!(name.to_string()),
+      |row| row.get(0),
     )
     .await?
-    .unwrap()
-    .get(0)?;
+    .unwrap();
 
   return Ok(room);
 }
@@ -111,11 +111,10 @@ async fn setup_app() -> Result<Setup, anyhow::Error> {
   })
   .await?;
 
-  let state = app.state();
-  let conn = state.conn();
+  let conn = app.state.conn();
 
   create_chat_message_app_tables(conn).await?;
-  state.refresh_table_cache().await?;
+  app.state.refresh_table_cache().await?;
 
   let room = add_room(conn, "room0").await?;
   let password = "Secret!1!!";
@@ -124,7 +123,7 @@ async fn setup_app() -> Result<Setup, anyhow::Error> {
     r#"(SELECT 1 FROM room_members WHERE user = _USER_.id AND room = _REQ_.room)"#;
 
   add_record_api_config(
-    &state,
+    &app.state,
     RecordApiConfig {
       name: Some("messages_api".to_string()),
       table_name: Some("message".to_string()),
@@ -137,7 +136,7 @@ async fn setup_app() -> Result<Setup, anyhow::Error> {
 
   let email = "[email protected]";
   let user_x = create_user_handler(
-    State(state.clone()),
+    State(app.state.clone()),
     Json(CreateUserRequest {
       email: email.to_string(),
       password: password.to_string(),
@@ -149,7 +148,7 @@ async fn setup_app() -> Result<Setup, anyhow::Error> {
   .id
   .into_bytes();
 
-  let user_x_token = login_with_password(&state, email, password)
+  let user_x_token = login_with_password(&app.state, email, password)
     .await?
     .auth_token;
 
@@ -208,7 +207,7 @@ fn create_message_benchmark(b: &mut Bencher, runtime: &tokio::runtime::Runtime,
     let tasks = (0..iters).map(|_i| {
       let body = body.clone();
       let auth = format!("Bearer {user_x_token}");
-      let mut router = setup.app.router().clone();
+      let mut router = setup.app.main_router.1.clone();
 
       return runtime.spawn(async move {
         let response = router
@@ -244,7 +243,7 @@ fn benchmark_group(c: &mut Criterion) {
 
   let setup = runtime.block_on(async {
     let setup = setup_app().await.unwrap();
-    let mut router = setup.app.router().clone();
+    let mut router = setup.app.main_router.1.clone();
 
     ServiceExt::<Request<Body>>::ready(&mut router)
       .await

trailbase-core/src/admin/list_logs.rs

@@ -126,12 +126,13 @@ pub async fn list_logs_handler(
     build_filter_where_clause(&table_metadata.schema.columns, filter_params)?;
 
   let total_row_count: i64 = conn
-    .query_value(
-      &format!(
+    .read_query_row_f(
+      format!(
         "SELECT COUNT(*) FROM {LOGS_TABLE_NAME} WHERE {where_clause}",
         where_clause = filter_where_clause.clause
       ),
       filter_where_clause.params.clone(),
+      |row| row.get(0),
     )
     .await?
     .unwrap_or(-1);
@@ -243,7 +244,11 @@ async fn fetch_logs(
     "#,
   );
 
-  return Ok(conn.query_values::<LogEntry>(&sql_query, params).await?);
+  return Ok(
+    conn
+      .read_query_values::<LogEntry>(sql_query, params)
+      .await?,
+  );
 }
 
 #[derive(Debug, Serialize, TS)]
@@ -315,7 +320,7 @@ async fn fetch_aggregate_stats(
     params.extend(filter.params.clone())
   }
 
-  let rows = conn.query_values::<AggRow>(&qps_query, params).await?;
+  let rows = conn.read_query_values::<AggRow>(qps_query, params).await?;
 
   let mut rate: Vec<(i64, f64)> = vec![];
   for r in rows.iter() {
@@ -348,7 +353,7 @@ async fn fetch_aggregate_stats(
   "#
     );
 
-    let rows = conn.query(&cc_query, ()).await?;
+    let rows = conn.read_query_rows(cc_query, ()).await?;
 
     let mut country_codes = HashMap::<String, usize>::new();
     for row in rows.iter() {
@@ -409,7 +414,7 @@ mod tests {
       let smack_in_there1 = (from + Duration::seconds(12 * 3600 + 1)).timestamp();
 
       conn
-        .execute_batch(&format!(
+        .execute_batch(format!(
           r#"
             INSERT INTO {LOGS_TABLE_NAME} (created) VALUES({before});
             INSERT INTO {LOGS_TABLE_NAME} (created) VALUES({after});

trailbase-core/src/admin/query.rs

@@ -72,7 +72,7 @@ pub async fn query_handler(
     ));
   }
 
-  let batched_rows_result = state.conn().execute_batch(&request.query).await;
+  let batched_rows_result = state.conn().execute_batch(request.query).await;
 
   // In the fallback case we always need to invalidate the cache.
   if must_invalidate_table_cache {

trailbase-core/src/admin/rows/delete_rows.rs

@@ -182,8 +182,8 @@ mod tests {
 
       return state
         .conn()
-        .query_value::<TestTable>(
-          &format!("SELECT * FROM {table_name} WHERE _rowid_ = ?1"),
+        .read_query_value::<TestTable>(
+          format!("SELECT * FROM {table_name} WHERE _rowid_ = ?1"),
           trailbase_sqlite::params!(row_id),
         )
         .await
@@ -203,12 +203,12 @@ mod tests {
 
     let count = || async {
       conn
-        .query_row(&format!("SELECT COUNT(*) FROM '{table_name}'"), ())
+        .read_query_row_f(format!("SELECT COUNT(*) FROM '{table_name}'"), (), |row| {
+          row.get::<_, i64>(0)
+        })
         .await
         .unwrap()
         .unwrap()
-        .get::<i64>(0)
-        .unwrap()
     };
 
     assert_eq!(count().await, 2);

trailbase-core/src/admin/rows/list_rows.rs

@@ -72,12 +72,14 @@ pub async fn list_rows_handler(
     }
   };
 
-  let total_row_count = {
+  let total_row_count: i64 = {
     let where_clause = &filter_where_clause.clause;
     let count_query = format!("SELECT COUNT(*) FROM '{table_name}' WHERE {where_clause}");
     state
       .conn()
-      .query_value::<i64>(&count_query, filter_where_clause.params.clone())
+      .read_query_row_f(count_query, filter_where_clause.params.clone(), |row| {
+        row.get(0)
+      })
       .await?
       .unwrap_or(-1)
   };
@@ -207,12 +209,15 @@ async fn fetch_rows(
     "#,
   );
 
-  let result_rows = conn.query(&query, params).await.map_err(|err| {
-    #[cfg(debug_assertions)]
-    error!("QUERY: {query}\n\t=> {err}");
+  let result_rows = conn
+    .read_query_rows(query.clone(), params)
+    .await
+    .map_err(|err| {
+      #[cfg(debug_assertions)]
+      error!("QUERY: {query}\n\t=> {err}");
 
-    return err;
-  })?;
+      return err;
+    })?;
 
   return Ok(rows_to_json_arrays(result_rows, 1024)?);
 }
@@ -259,7 +264,7 @@ mod tests {
       .unwrap();
 
     let cnt: i64 = conn
-      .query_value("SELECT COUNT(*) FROM test_table", ())
+      .read_query_row_f("SELECT COUNT(*) FROM test_table", (), |row| row.get(0))
       .await
       .unwrap()
       .unwrap();

trailbase-core/src/admin/table/alter_table.rs

@@ -174,7 +174,9 @@ mod tests {
     );
     let _ = create_table_handler(State(state.clone()), Json(create_table_request.clone())).await?;
 
-    conn.query(&format!("SELECT {pk_col} FROM foo"), ()).await?;
+    conn
+      .read_query_rows(format!("SELECT {pk_col} FROM foo"), ())
+      .await?;
 
     {
       // Noop: source and target identical.
@@ -187,7 +189,9 @@ mod tests {
         .await
         .unwrap();
 
-      conn.query(&format!("SELECT {pk_col} FROM foo"), ()).await?;
+      conn
+        .read_query_rows(format!("SELECT {pk_col} FROM foo"), ())
+        .await?;
     }
 
     {
@@ -215,7 +219,7 @@ mod tests {
         .unwrap();
 
       conn
-        .query(&format!("SELECT {pk_col}, new FROM foo"), ())
+        .read_query_rows(format!("SELECT {pk_col}, new FROM foo"), ())
         .await?;
     }
 
@@ -236,8 +240,10 @@ mod tests {
         .await
         .unwrap();
 
-      assert!(conn.query("SELECT * FROM foo", ()).await.is_err());
-      conn.query(&format!("SELECT {pk_col} FROM bar"), ()).await?;
+      assert!(conn.read_query_rows("SELECT * FROM foo", ()).await.is_err());
+      conn
+        .read_query_rows(format!("SELECT {pk_col} FROM bar"), ())
+        .await?;
     }
 
     return Ok(());

trailbase-core/src/admin/table/list_tables.rs

@@ -41,8 +41,8 @@ pub async fn list_tables_handler(
   // "view"s.
   let rows = state
     .conn()
-    .query_values::<SqliteSchema>(
-      &format!("SELECT type, name, tbl_name, sql FROM {SQLITE_SCHEMA_TABLE} ORDER BY type"),
+    .read_query_values::<SqliteSchema>(
+      format!("SELECT type, name, tbl_name, sql FROM {SQLITE_SCHEMA_TABLE} ORDER BY type"),
       (),
     )
     .await?;

trailbase-core/src/admin/user/create_user.rs

@@ -65,8 +65,8 @@ pub async fn create_user_handler(
 
   let Some(user) = state
     .user_conn()
-    .query_value::<DbUser>(
-      &INSERT_USER_QUERY,
+    .write_query_value::<DbUser>(
+      &*INSERT_USER_QUERY,
       named_params! {
         ":email": normalized_email,
         ":password_hash": hashed_password,

trailbase-core/src/admin/user/list_users.rs

@@ -81,12 +81,13 @@ pub async fn list_users_handler(
     build_filter_where_clause(&table_metadata.schema.columns, filter_params)?;
 
   let total_row_count: i64 = conn
-    .query_value(
-      &format!(
+    .read_query_row_f(
+      format!(
         "SELECT COUNT(*) FROM {USER_TABLE} WHERE {where_clause}",
         where_clause = filter_where_clause.clause
       ),
       filter_where_clause.params.clone(),
+      |row| row.get(0),
     )
     .await?
     .unwrap_or(-1);
@@ -161,6 +162,6 @@ async fn fetch_users(
     "#,
   );
 
-  let users = conn.query_values::<DbUser>(&sql_query, params).await?;
+  let users = conn.read_query_values::<DbUser>(sql_query, params).await?;
   return Ok(users);
 }