Add checkpointing functionality to safely copy SQLite databases

This update introduces a new function, _checkpoint_and_copy_database, which ensures that all data in the Write-Ahead Logging (WAL) file is flushed to the main database file before copying. This function is now utilized in the SQLResource class for database forking, checkpointing, and restoring operations, enhancing data integrity during these processes.

Author: Dylan Huang

eval_protocol/agent/resources/sql_resource.py

@@ -27,6 +27,35 @@ def _apply_hardened_pragmas(conn: sqlite3.Connection) -> None:
     conn.execute("PRAGMA temp_store=MEMORY")  # Store temp tables in memory
 
 
+def _checkpoint_and_copy_database(
+    source_path: Path, dest_path: Path, timeout: int = SQLITE_CONNECTION_TIMEOUT
+) -> None:
+    """
+    Safely copy a SQLite database by checkpointing WAL first.
+
+    In WAL mode, data may exist in the -wal file that hasn't been written
+    to the main database file. This function performs a TRUNCATE checkpoint
+    to flush all WAL data to the main file before copying, ensuring a
+    complete and consistent copy.
+
+    Args:
+        source_path: Path to the source database file.
+        dest_path: Path where the copy should be created.
+        timeout: Connection timeout in seconds.
+    """
+    # First, checkpoint the WAL to ensure all data is in the main file
+    conn = sqlite3.connect(str(source_path), timeout=timeout)
+    try:
+        # TRUNCATE mode: checkpoint and truncate the WAL file to zero bytes
+        # This ensures all data is flushed to the main database file
+        conn.execute("PRAGMA wal_checkpoint(TRUNCATE)")
+    finally:
+        conn.close()
+
+    # Now safely copy just the main database file
+    shutil.copyfile(str(source_path), str(dest_path))
+
+
 class SQLResource(ForkableResource):
     """
     A ForkableResource for managing SQL database states, primarily SQLite.
@@ -134,7 +163,8 @@ async def fork(self) -> "SQLResource":
         forked_db_name = f"fork_{uuid.uuid4().hex}.sqlite"
         forked_resource._db_path = self._temp_dir / forked_db_name
 
-        shutil.copyfile(str(self._db_path), str(forked_resource._db_path))
+        # Use checkpoint-and-copy to ensure WAL data is flushed before copying
+        _checkpoint_and_copy_database(self._db_path, forked_resource._db_path)
         return forked_resource
 
     async def checkpoint(self) -> Dict[str, Any]:
@@ -148,7 +178,8 @@ async def checkpoint(self) -> Dict[str, Any]:
 
         checkpoint_name = f"checkpoint_{self._db_path.stem}_{uuid.uuid4().hex}.sqlite"
         checkpoint_path = self._temp_dir / checkpoint_name
-        shutil.copyfile(str(self._db_path), str(checkpoint_path))
+        # Use checkpoint-and-copy to ensure WAL data is flushed before copying
+        _checkpoint_and_copy_database(self._db_path, checkpoint_path)
         return {"db_type": "sqlite", "checkpoint_path": str(checkpoint_path)}
 
     async def restore(self, state_data: Dict[str, Any]) -> None:
@@ -170,7 +201,8 @@ async def restore(self, state_data: Dict[str, Any]) -> None:
         if not self._db_path:
             self._db_path = self._temp_dir / f"restored_{uuid.uuid4().hex}.sqlite"
 
-        shutil.copyfile(str(checkpoint_path), str(self._db_path))
+        # Use checkpoint-and-copy to ensure WAL data is flushed before copying
+        _checkpoint_and_copy_database(checkpoint_path, self._db_path)
         self._base_db_path = self._db_path  # The restored state becomes the new base for future forks
 
     async def step(self, action_name: str, action_params: Dict[str, Any]) -> Any: