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,11 @@ 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 +21,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 {
@@ -106,6 +107,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 +119,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 +132,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)?)),
             })
@@ -131,17 +144,9 @@ impl Conn for SyncedConnection {
                 inner: Box::new(self.remote.prepare(sql).await?),
             };
 
-            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)
         }
     }