Skip to content

Commit b42c27d

Browse files
author
Dylan Huang
authored
increase sql retries (#413)
* increase sql retries * make connection more robust * Implement robust database connection with retry logic to handle pragma execution failures
1 parent 8e2686b commit b42c27d

2 files changed

Lines changed: 52 additions & 6 deletions

File tree

eval_protocol/dataset_logger/sqlite_evaluation_row_store.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from eval_protocol.event_bus.sqlite_event_bus_database import (
88
SQLITE_HARDENED_PRAGMAS,
99
check_and_repair_database,
10+
connect_with_retry,
1011
execute_with_sqlite_retry,
1112
)
1213
from eval_protocol.models import EvaluationRow
@@ -42,9 +43,10 @@ class EvaluationRow(BaseModel): # type: ignore
4243

4344
self._EvaluationRow = EvaluationRow
4445

45-
self._db.connect()
46+
# Connect with retry logic that properly handles pragma execution failures
47+
connect_with_retry(self._db)
4648
# Use safe=True to avoid errors when tables/indexes already exist
47-
self._db.create_tables([EvaluationRow], safe=True)
49+
execute_with_sqlite_retry(lambda: self._db.create_tables([EvaluationRow], safe=True))
4850

4951
@property
5052
def db_path(self) -> str:

eval_protocol/event_bus/sqlite_event_bus_database.py

Lines changed: 48 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@
1111

1212

1313
# Retry configuration for database operations
14-
SQLITE_RETRY_MAX_TRIES = 5
15-
SQLITE_RETRY_MAX_TIME = 30 # seconds
14+
SQLITE_RETRY_MAX_TRIES = 10
15+
SQLITE_RETRY_MAX_TIME = 60 # seconds
1616

1717

1818
def _is_database_locked_error(e: Exception) -> bool:
@@ -55,6 +55,49 @@ def _execute() -> T:
5555
return _execute()
5656

5757

58+
def connect_with_retry(db: SqliteDatabase) -> None:
59+
"""
60+
Connect to the database with retry logic, ensuring pragmas are always applied.
61+
62+
Peewee's connect() method sets the connection state *before* executing pragmas
63+
(in _initialize_connection). If pragma execution fails with "database is locked",
64+
the connection is marked as open but pragmas are not applied. Subsequent calls
65+
to connect(reuse_if_open=True) would see the connection as already open and
66+
skip pragma execution entirely.
67+
68+
This function handles this edge case by:
69+
1. Closing the connection if a lock error occurs during connect
70+
2. Retrying with exponential backoff until pragmas are successfully applied
71+
72+
Args:
73+
db: The SqliteDatabase instance to connect
74+
"""
75+
76+
@backoff.on_exception(
77+
backoff.expo,
78+
OperationalError,
79+
max_tries=SQLITE_RETRY_MAX_TRIES,
80+
max_time=SQLITE_RETRY_MAX_TIME,
81+
giveup=lambda e: not _is_database_locked_error(e),
82+
jitter=backoff.full_jitter,
83+
)
84+
def _connect() -> None:
85+
try:
86+
# Close any partially-open connection before retrying to ensure
87+
# a fresh connection is opened and pragmas are executed
88+
if not db.is_closed():
89+
db.close()
90+
db.connect()
91+
except OperationalError:
92+
# If connect fails (e.g., during pragma execution), ensure the
93+
# connection is closed so the next retry starts fresh
94+
if not db.is_closed():
95+
db.close()
96+
raise
97+
98+
_connect()
99+
100+
58101
# SQLite pragmas for hardened concurrency safety
59102
SQLITE_HARDENED_PRAGMAS = {
60103
"journal_mode": "wal", # Write-Ahead Logging for concurrent reads/writes
@@ -181,9 +224,10 @@ class Event(BaseModel): # type: ignore
181224
processed = BooleanField(default=False) # Track if event has been processed
182225

183226
self._Event = Event
184-
self._db.connect()
227+
# Connect with retry logic that properly handles pragma execution failures
228+
connect_with_retry(self._db)
185229
# Use safe=True to avoid errors when tables already exist
186-
self._db.create_tables([Event], safe=True)
230+
execute_with_sqlite_retry(lambda: self._db.create_tables([Event], safe=True))
187231

188232
def publish_event(self, event_type: str, data: Any, process_id: str) -> None:
189233
"""Publish an event to the database."""

0 commit comments

Comments
 (0)