55from dataclasses import dataclass
66from typing import TYPE_CHECKING
77
8+ from aws_durable_execution_sdk_python .lambda_service import OperationType
89from aws_durable_execution_sdk_python .types import LoggerInterface
910
1011if TYPE_CHECKING :
1112 from collections .abc import Mapping , MutableMapping
1213
1314 from aws_durable_execution_sdk_python .identifier import OperationIdentifier
15+ from aws_durable_execution_sdk_python .state import ExecutionState
1416
1517
1618@dataclass (frozen = True )
@@ -44,13 +46,25 @@ def with_parent_id(self, parent_id: str) -> LogInfo:
4446
4547class Logger (LoggerInterface ):
4648 def __init__ (
47- self , logger : LoggerInterface , default_extra : Mapping [str , object ]
49+ self ,
50+ logger : LoggerInterface ,
51+ default_extra : Mapping [str , object ],
52+ execution_state : ExecutionState | None = None ,
53+ visited_operations : set [str ] | None = None ,
4854 ) -> None :
4955 self ._logger = logger
5056 self ._default_extra = default_extra
57+ self ._execution_state = execution_state
58+ self ._visited_operations = visited_operations or set ()
5159
5260 @classmethod
53- def from_log_info (cls , logger : LoggerInterface , info : LogInfo ) -> Logger :
61+ def from_log_info (
62+ cls ,
63+ logger : LoggerInterface ,
64+ info : LogInfo ,
65+ execution_state : ExecutionState | None = None ,
66+ visited_operations : set [str ] | None = None ,
67+ ) -> Logger :
5468 """Create a new logger with the given LogInfo."""
5569 extra : MutableMapping [str , object ] = {"execution_arn" : info .execution_arn }
5670 if info .parent_id :
@@ -59,45 +73,118 @@ def from_log_info(cls, logger: LoggerInterface, info: LogInfo) -> Logger:
5973 extra ["name" ] = info .name
6074 if info .attempt :
6175 extra ["attempt" ] = info .attempt
62- return cls (logger , extra )
76+ return cls (logger , extra , execution_state , visited_operations )
6377
6478 def with_log_info (self , info : LogInfo ) -> Logger :
6579 """Clone the existing logger with new LogInfo."""
6680 return Logger .from_log_info (
6781 logger = self ._logger ,
6882 info = info ,
83+ execution_state = self ._execution_state ,
84+ visited_operations = self ._visited_operations ,
6985 )
7086
7187 def get_logger (self ) -> LoggerInterface :
7288 """Get the underlying logger."""
7389 return self ._logger
7490
91+ def is_replay (self ) -> bool :
92+ """Check if we are currently in replay mode.
93+
94+ Returns True if there are operations in the execution state that haven't been visited yet.
95+ This indicates we are replaying previously executed operations.
96+ """
97+ if not self ._execution_state :
98+ return False
99+
100+ # If there are no operations, we're not in replay
101+ if not self ._execution_state .operations :
102+ return False
103+
104+ # Check if there are any operations in the execution state that we haven't visited
105+ # Only consider operations that are not EXECUTION type (which are system operations)
106+ for operation_id , operation in self ._execution_state .operations .items ():
107+ # Skip EXECUTION operations as they are system operations, not user operations
108+ if operation .operation_type == OperationType .EXECUTION :
109+ continue
110+ if operation_id not in self ._visited_operations :
111+ return True
112+ return False
113+
114+ def mark_operation_visited (self , operation_id : str ) -> None :
115+ """Mark an operation as visited."""
116+ self ._visited_operations .add (operation_id )
117+
118+ def _should_log (self ) -> bool :
119+ """Determine if logging should occur based on replay state."""
120+ # For the default logger, only log when not in replay
121+ return not self .is_replay ()
122+
75123 def debug (
76124 self , msg : object , * args : object , extra : Mapping [str , object ] | None = None
77125 ) -> None :
126+ if not self ._should_log ():
127+ return
78128 merged_extra = {** self ._default_extra , ** (extra or {})}
79129 self ._logger .debug (msg , * args , extra = merged_extra )
80130
81131 def info (
82132 self , msg : object , * args : object , extra : Mapping [str , object ] | None = None
83133 ) -> None :
134+ if not self ._should_log ():
135+ return
84136 merged_extra = {** self ._default_extra , ** (extra or {})}
85137 self ._logger .info (msg , * args , extra = merged_extra )
86138
87139 def warning (
88140 self , msg : object , * args : object , extra : Mapping [str , object ] | None = None
89141 ) -> None :
142+ if not self ._should_log ():
143+ return
90144 merged_extra = {** self ._default_extra , ** (extra or {})}
91145 self ._logger .warning (msg , * args , extra = merged_extra )
92146
93147 def error (
94148 self , msg : object , * args : object , extra : Mapping [str , object ] | None = None
95149 ) -> None :
150+ if not self ._should_log ():
151+ return
96152 merged_extra = {** self ._default_extra , ** (extra or {})}
97153 self ._logger .error (msg , * args , extra = merged_extra )
98154
99155 def exception (
100156 self , msg : object , * args : object , extra : Mapping [str , object ] | None = None
101157 ) -> None :
158+ if not self ._should_log ():
159+ return
102160 merged_extra = {** self ._default_extra , ** (extra or {})}
103161 self ._logger .exception (msg , * args , extra = merged_extra )
162+
163+ @property
164+ def visited_operations (self ):
165+ return self ._visited_operations
166+
167+
168+ class ReplayAwareLogger (Logger ):
169+ """A logger that provides custom replay behavior for advanced users.
170+
171+ This logger allows users to customize logging behavior during replay by overriding
172+ the _should_log method. By default, it behaves the same as the base Logger.
173+ """
174+
175+ def _should_log (self ) -> bool :
176+ """Override this method to customize replay logging behavior.
177+
178+ Returns:
179+ bool: True if logging should occur, False otherwise.
180+
181+ Example:
182+ def _should_log(self) -> bool:
183+ # Always log, even during replay
184+ return True
185+
186+ def _should_log(self) -> bool:
187+ # Only log errors during replay
188+ return not self.is_replay() or self._current_log_level == 'error'
189+ """
190+ return super ()._should_log ()
0 commit comments