@@ -115,7 +115,10 @@ def _fetch_from_ssm(key: str, test_name: str | None = None) -> str:
115115 except hook .conn .exceptions .ParameterNotFound as e :
116116 log .info ("SSM does not contain any parameter for this test: %s" , e )
117117 except KeyError as e :
118- log .info ("SSM contains one parameter for this test, but not the requested value: %s" , e )
118+ log .info (
119+ "SSM contains one parameter for this test, but not the requested value: %s" ,
120+ e ,
121+ )
119122 return value
120123
121124
@@ -286,6 +289,7 @@ def prune_logs(
286289 force_delete : bool = False ,
287290 retry : bool = False ,
288291 retry_times : int = 3 ,
292+ delete_log_groups : bool = True ,
289293 ti = None ,
290294):
291295 """
@@ -300,11 +304,13 @@ def prune_logs(
300304 cases, the log group/stream is created seconds after the main resource has
301305 been created. By default, it retries for 3 times with a 5s waiting period.
302306 :param retry_times: Number of retries.
307+ :param delete_log_groups: Whether to delete the log groups if they are empty.
308+ Overridden by force_delete.
303309 :param ti: Used to check the status of the tasks. This gets pulled from the
304310 DAG's context and does not need to be passed manually.
305311 """
306312 if all_tasks_passed (ti ):
307- _purge_logs (logs , force_delete , retry , retry_times )
313+ _purge_logs (logs , force_delete , retry , retry_times , delete_log_groups )
308314 else :
309315 client : BaseClient = boto3 .client ("logs" )
310316 for group , _ in logs :
@@ -316,6 +322,7 @@ def _purge_logs(
316322 force_delete : bool = False ,
317323 retry : bool = False ,
318324 retry_times : int = 3 ,
325+ delete_log_groups : bool = True ,
319326) -> None :
320327 """
321328 Accepts a tuple in the format: ('log group name', 'log stream prefix').
@@ -332,6 +339,8 @@ def _purge_logs(
332339 is created seconds after the main resource has been created. By default, it retries for 3 times
333340 with a 5s waiting period
334341 :param retry_times: Number of retries
342+ :param delete_log_groups: Whether to delete the log groups if they are empty.
343+ Overridden by force_delete.
335344 """
336345 client : BaseClient = boto3 .client ("logs" )
337346
@@ -346,7 +355,9 @@ def _purge_logs(
346355 for stream_name in [stream ["logStreamName" ] for stream in log_streams ]:
347356 client .delete_log_stream (logGroupName = group , logStreamName = stream_name )
348357
349- if force_delete or not client .describe_log_streams (logGroupName = group )["logStreams" ]:
358+ if force_delete or (
359+ delete_log_groups and not client .describe_log_streams (logGroupName = group )["logStreams" ]
360+ ):
350361 client .delete_log_group (logGroupName = group )
351362 except ClientError as e :
352363 if not retry or retry_times == 0 or e .response ["Error" ]["Code" ] != "ResourceNotFoundException" :
0 commit comments