Skip to content

Commit

Permalink
fail safely when pickling (#4085)
Browse files Browse the repository at this point in the history
* fail safely when pickling

* why did i do that
  • Loading branch information
adhami3310 authored and simon committed Oct 23, 2024
1 parent c310f6c commit db292c4
Showing 1 changed file with 18 additions and 9 deletions.
27 changes: 18 additions & 9 deletions reflex/state.py
Original file line number Diff line number Diff line change
Expand Up @@ -1999,7 +1999,14 @@ def _serialize(self) -> bytes:
Returns:
The serialized state.
"""
return pickle.dumps((self._to_schema(), self))
try:
return pickle.dumps((self._to_schema(), self))
except pickle.PicklingError:
console.warn(
f"Failed to serialize state {self.get_full_name()} due to unpicklable object. "
"This state will not be persisted."
)
return b""

@classmethod
def _deserialize(
Expand Down Expand Up @@ -2826,9 +2833,10 @@ async def set_state_for_substate(self, client_token: str, substate: BaseState):
if substate._get_was_touched():
substate._was_touched = False # Reset the touched flag after serializing.
pickle_state = substate._serialize()
if not self.states_directory.exists():
self.states_directory.mkdir(parents=True, exist_ok=True)
self.token_path(substate_token).write_bytes(pickle_state)
if pickle_state:
if not self.states_directory.exists():
self.states_directory.mkdir(parents=True, exist_ok=True)
self.token_path(substate_token).write_bytes(pickle_state)

for substate_substate in substate.substates.values():
await self.set_state_for_substate(client_token, substate_substate)
Expand Down Expand Up @@ -3121,11 +3129,12 @@ async def set_state(
if state._get_was_touched():
pickle_state = state._serialize()
self._warn_if_too_large(state, len(pickle_state))
await self.redis.set(
_substate_key(client_token, state),
pickle_state,
ex=self.token_expiration,
)
if pickle_state:
await self.redis.set(
_substate_key(client_token, state),
pickle_state,
ex=self.token_expiration,
)

# Wait for substates to be persisted.
for t in tasks:
Expand Down

0 comments on commit db292c4

Please sign in to comment.