fix: pull on read

Author: levydsa

libsql/src/database.rs

@@ -712,6 +712,7 @@ impl Database {
                         read_your_writes: *read_your_writes,
                         context: db.sync_ctx.clone().unwrap(),
                         state: std::sync::Arc::new(Mutex::new(State::Init)),
+                        needs_pull: std::sync::atomic::AtomicBool::new(false).into(),
                     };
 
                     let conn = std::sync::Arc::new(synced);

libsql/src/sync.rs

@@ -13,7 +13,6 @@ use uuid::Uuid;
 mod test;
 
 pub mod connection;
-pub mod statement;
 pub mod transaction;
 
 const METADATA_VERSION: u32 = 0;

libsql/src/sync/connection.rs

@@ -8,11 +8,14 @@ use crate::{
     sync::SyncContext,
     BatchRows, Error, Result, Statement, Transaction, TransactionBehavior,
 };
-use std::sync::Arc;
+use std::sync::{
+    atomic::{AtomicBool, Ordering},
+    Arc,
+};
 use std::time::Duration;
 use tokio::sync::Mutex;
 
-use super::{statement::SyncedStatement, transaction::SyncedTx};
+use super::transaction::SyncedTx;
 
 #[derive(Clone)]
 pub struct SyncedConnection {
@@ -21,6 +24,7 @@ pub struct SyncedConnection {
     pub read_your_writes: bool,
     pub context: Arc<Mutex<SyncContext>>,
     pub state: Arc<Mutex<State>>,
+    pub needs_pull: Arc<AtomicBool>,
 }
 
 impl SyncedConnection {
@@ -89,7 +93,7 @@ impl SyncedConnection {
                 _ => {
                     *state = predicted_end_state;
                     false
-                },
+                }
             };
 
             Ok(should_execute_local)
@@ -106,6 +110,10 @@ impl Conn for SyncedConnection {
 
     async fn execute_batch(&self, sql: &str) -> Result<BatchRows> {
         if self.should_execute_local(sql).await? {
+            if self.needs_pull.swap(false, Ordering::Relaxed) {
+                let mut context = self.context.lock().await;
+                crate::sync::try_pull(&mut context, &self.local).await?;
+            }
             self.local.execute_batch(sql)
         } else {
             self.remote.execute_batch(sql).await
@@ -114,6 +122,10 @@ impl Conn for SyncedConnection {
 
     async fn execute_transactional_batch(&self, sql: &str) -> Result<BatchRows> {
         if self.should_execute_local(sql).await? {
+            if self.needs_pull.swap(false, Ordering::Relaxed) {
+                let mut context = self.context.lock().await;
+                crate::sync::try_pull(&mut context, &self.local).await?;
+            }
             self.local.execute_transactional_batch(sql)?;
             Ok(BatchRows::empty())
         } else {
@@ -123,6 +135,10 @@ impl Conn for SyncedConnection {
 
     async fn prepare(&self, sql: &str) -> Result<Statement> {
         if self.should_execute_local(sql).await? {
+            if self.needs_pull.swap(false, Ordering::Relaxed) {
+                let mut context = self.context.lock().await;
+                crate::sync::try_pull(&mut context, &self.local).await?;
+            }
             Ok(Statement {
                 inner: Box::new(LibsqlStmt(self.local.prepare(sql)?)),
             })
@@ -132,16 +148,10 @@ impl Conn for SyncedConnection {
             };
 
             if self.read_your_writes {
-                Ok(Statement {
-                    inner: Box::new(SyncedStatement {
-                        conn: self.local.clone(),
-                        context: self.context.clone(),
-                        inner: stmt,
-                    }),
-                })
-            } else {
-                Ok(stmt)
+                self.needs_pull.store(true, Ordering::Relaxed);
             }
+
+            Ok(stmt)
         }
     }