Perf: retryable mysql connection. (#8858)

### What problem does this PR solve?

### Type of change

- [x] Performance Improvement
This commit is contained in:
Kevin Hu
2025-07-15 19:05:48 +08:00
committed by GitHub
parent 9371d7b19c
commit f2909ea0c4
2 changed files with 45 additions and 2 deletions

View File

@ -243,8 +243,51 @@ class JsonSerializedField(SerializedField):
super(JsonSerializedField, self).__init__(serialized_type=SerializedType.JSON, object_hook=object_hook, object_pairs_hook=object_pairs_hook, **kwargs)
class RetryingPooledMySQLDatabase(PooledMySQLDatabase):
def __init__(self, *args, **kwargs):
self.max_retries = kwargs.pop('max_retries', 5)
self.retry_delay = kwargs.pop('retry_delay', 1)
super().__init__(*args, **kwargs)
def execute_sql(self, sql, params=None, commit=True):
from peewee import OperationalError
for attempt in range(self.max_retries + 1):
try:
return super().execute_sql(sql, params, commit)
except OperationalError as e:
if e.args[0] in (2013, 2006) and attempt < self.max_retries:
logging.warning(
f"Lost connection (attempt {attempt+1}/{self.max_retries}): {e}"
)
self._handle_connection_loss()
time.sleep(self.retry_delay * (2 ** attempt))
else:
logging.error(f"DB execution failure: {e}")
raise
return None
def _handle_connection_loss(self):
self.close_all()
self.connect()
def begin(self):
from peewee import OperationalError
for attempt in range(self.max_retries + 1):
try:
return super().begin()
except OperationalError as e:
if e.args[0] in (2013, 2006) and attempt < self.max_retries:
logging.warning(
f"Lost connection during transaction (attempt {attempt+1}/{self.max_retries})"
)
self._handle_connection_loss()
time.sleep(self.retry_delay * (2 ** attempt))
else:
raise
class PooledDatabase(Enum):
MYSQL = PooledMySQLDatabase
MYSQL = RetryingPooledMySQLDatabase
POSTGRES = PooledPostgresqlDatabase