fix: state size was not checked for dill (#4431)

This commit is contained in:
benedikt-bartscher 2024-11-25 19:38:52 +01:00 committed by GitHub
parent 697e26c25b
commit d7d46e431b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -2175,11 +2175,9 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
Returns:
The serialized state.
"""
payload = b""
try:
pickle_state = pickle.dumps((self._to_schema(), self))
if environment.REFLEX_PERF_MODE.get() != PerformanceMode.OFF:
self._check_state_size(len(pickle_state))
return pickle_state
payload = pickle.dumps((self._to_schema(), self))
except HANDLED_PICKLE_ERRORS as og_pickle_error:
error = (
f"Failed to serialize state {self.get_full_name()} due to unpicklable object. "
@ -2188,7 +2186,7 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
try:
import dill
return dill.dumps((self._to_schema(), self))
payload = dill.dumps((self._to_schema(), self))
except ImportError:
error += (
f"Pickle error: {og_pickle_error}. "
@ -2196,8 +2194,10 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
)
except HANDLED_PICKLE_ERRORS as ex:
error += f"Dill was also unable to pickle the state: {ex}"
console.warn(error)
return b""
console.warn(error)
if environment.REFLEX_PERF_MODE.get() != PerformanceMode.OFF:
self._check_state_size(len(payload))
return payload
@classmethod
def _deserialize(