laipeng/reflex
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fail safely when pickling (#4085)

Browse Source 
* fail safely when pickling

* why did i do that
This commit is contained in:
Khaleel Al-Adhami 2024-10-07 11:59:33 -07:00 committed by GitHub
parent 59dd54c049
commit c7c830de80
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 18 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
reflex/state.py
View File

@ -1999,7 +1999,14 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
Returns: Returns:
The serialized state. 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 @classmethod
def _deserialize( def _deserialize(
@ -2826,9 +2833,10 @@ class StateManagerDisk(StateManager):
if substate._get_was_touched(): if substate._get_was_touched():
substate._was_touched = False # Reset the touched flag after serializing. substate._was_touched = False # Reset the touched flag after serializing.
pickle_state = substate._serialize() pickle_state = substate._serialize()
if not self.states_directory.exists(): if pickle_state:
self.states_directory.mkdir(parents=True, exist_ok=True) if not self.states_directory.exists():
self.token_path(substate_token).write_bytes(pickle_state) 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(): for substate_substate in substate.substates.values():
await self.set_state_for_substate(client_token, substate_substate) await self.set_state_for_substate(client_token, substate_substate)
@ -3121,11 +3129,12 @@ class StateManagerRedis(StateManager):
if state._get_was_touched(): if state._get_was_touched():
pickle_state = state._serialize() pickle_state = state._serialize()
self._warn_if_too_large(state, len(pickle_state)) self._warn_if_too_large(state, len(pickle_state))
await self.redis.set( if pickle_state:
_substate_key(client_token, state), await self.redis.set(
pickle_state, _substate_key(client_token, state),
ex=self.token_expiration, pickle_state,
) ex=self.token_expiration,
)
# Wait for substates to be persisted. # Wait for substates to be persisted.
for t in tasks: for t in tasks: