diff --git a/src/aws_durable_execution_sdk_python_testing/stores/sqlite.py b/src/aws_durable_execution_sdk_python_testing/stores/sqlite.py index 4eb4222..eeb0a95 100644 --- a/src/aws_durable_execution_sdk_python_testing/stores/sqlite.py +++ b/src/aws_durable_execution_sdk_python_testing/stores/sqlite.py @@ -17,7 +17,10 @@ from aws_durable_execution_sdk_python_testing.stores.base import ( ExecutionStore, ) -from aws_durable_execution_sdk_python_testing.stores.filesystem import DateTimeEncoder +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + DateTimeEncoder, + datetime_object_hook, +) class SQLiteExecutionStore(ExecutionStore): @@ -122,7 +125,9 @@ def load(self, execution_arn: str) -> Execution: if not row: raise ResourceNotFoundException(f"Execution {execution_arn} not found") - return Execution.from_dict(json.loads(row[0])) + return Execution.from_dict( + json.loads(row[0], object_hook=datetime_object_hook) + ) except sqlite3.Error as e: raise RuntimeError(f"Failed to load execution {execution_arn}: {e}") from e except json.JSONDecodeError as e: @@ -217,7 +222,11 @@ def query( executions: list[Execution] = [] for durable_execution_arn, data in rows: try: - executions.append(Execution.from_dict(json.loads(data))) + executions.append( + Execution.from_dict( + json.loads(data, object_hook=datetime_object_hook) + ) + ) except (json.JSONDecodeError, ValueError) as e: # Log corrupted data but continue with other records print(