|
| 1 | +# -------------------------------------------------------------------------------------------- |
| 2 | +# Copyright (c) Microsoft Corporation. All rights reserved. |
| 3 | +# Licensed under the MIT License. See License.txt in the project root for license information. |
| 4 | +# -------------------------------------------------------------------------------------------- |
| 5 | + |
| 6 | +from .aaz.latest.durabletask.retention_policy import Create as _Create |
| 7 | +from azure.cli.core.aaz import AAZStrArg |
| 8 | + |
| 9 | + |
| 10 | +class CreatePolicy(_Create): |
| 11 | + """Create a retention policy for a Durabletask scheduler.""" |
| 12 | + |
| 13 | + @classmethod |
| 14 | + def _build_arguments_schema(cls, *args, **kwargs): |
| 15 | + """Define custom arguments for the command.""" |
| 16 | + # Call the parent class method to initialize the argument schema |
| 17 | + cls.args_schema = super()._build_arguments_schema(*args, **kwargs) |
| 18 | + |
| 19 | + # Remove the registration of the retention policies schema so it doesn't show up in the CLI help |
| 20 | + if hasattr(cls._args_schema.retention_policies, '_registered'): |
| 21 | + setattr(cls._args_schema.retention_policies, '_registered', False) |
| 22 | + _args_schema = cls._args_schema |
| 23 | + |
| 24 | + # Define retention policy arguments with their CLI options and descriptions |
| 25 | + retention_args = { |
| 26 | + "default_days": ("--default-days", "-d", "The number of days to retain orchestrations."), |
| 27 | + "canceled_days": ("--canceled-days", "-x", "The number of days to retain canceled orchestrations."), |
| 28 | + "completed_days": ("--completed-days", "-c", "The number of days to retain completed orchestrations."), |
| 29 | + "failed_days": ("--failed-days", "-f", "The number of days to retain failed orchestrations."), |
| 30 | + "terminated_days": ("--terminated-days", "-t", "The number of days to retain terminated orchestrations."), |
| 31 | + } |
| 32 | + |
| 33 | + # Add each retention argument to the schema |
| 34 | + for arg_name, (option, short_option, help_text) in retention_args.items(): |
| 35 | + setattr(_args_schema, arg_name, AAZStrArg( |
| 36 | + arg_group="Properties", # Group these arguments under "Properties" |
| 37 | + options=[option, short_option], # CLI options for the argument |
| 38 | + help=help_text, # Description of the argument |
| 39 | + required=False, # These arguments are optional |
| 40 | + )) |
| 41 | + |
| 42 | + return _args_schema |
| 43 | + |
| 44 | + def pre_operations(self): |
| 45 | + """Prepare retention policies before executing the operation.""" |
| 46 | + |
| 47 | + if not any([ |
| 48 | + self.ctx.args.default_days, |
| 49 | + self.ctx.args.canceled_days, |
| 50 | + self.ctx.args.completed_days, |
| 51 | + self.ctx.args.failed_days, |
| 52 | + self.ctx.args.terminated_days |
| 53 | + ]): |
| 54 | + raise ValueError("At least one retention period (e.g., --default-days, --canceled-days) must be specified.") |
| 55 | + |
| 56 | + # Build the retention policies based on the provided arguments |
| 57 | + self.ctx.args.retention_policies = _build_retention_policies({ |
| 58 | + key: value for key, value in { |
| 59 | + 'default_days': self.ctx.args.default_days, |
| 60 | + 'canceled_days': self.ctx.args.canceled_days, |
| 61 | + 'completed_days': self.ctx.args.completed_days, |
| 62 | + 'failed_days': self.ctx.args.failed_days, |
| 63 | + 'terminated_days': self.ctx.args.terminated_days |
| 64 | + }.items() if value is not None # Only include arguments that are not None |
| 65 | + }) |
| 66 | + |
| 67 | + |
| 68 | +def _build_retention_policies(args_dict): |
| 69 | + """Build a list of retention policies based on the provided arguments.""" |
| 70 | + retention_policies = [] |
| 71 | + |
| 72 | + def to_int(value): |
| 73 | + """Convert a value to an integer, handling serialization.""" |
| 74 | + try: |
| 75 | + return int(value.to_serialized_data()) # Convert serialized data to an integer |
| 76 | + except (ValueError, TypeError): |
| 77 | + return None # Return None if conversion fails |
| 78 | + |
| 79 | + # Add default retention policy if specified |
| 80 | + default_days = args_dict.get('default_days') |
| 81 | + if default_days is not None: |
| 82 | + days = to_int(default_days) |
| 83 | + if days is not None: |
| 84 | + # Add a default retention policy with the specified number of days |
| 85 | + retention_policies.append({"retentionPeriodInDays": days}) |
| 86 | + |
| 87 | + # Define a mapping of argument keys to orchestration states |
| 88 | + state_mapping = { |
| 89 | + 'canceled_days': 'Canceled', |
| 90 | + 'completed_days': 'Completed', |
| 91 | + 'failed_days': 'Failed', |
| 92 | + 'terminated_days': 'Terminated', |
| 93 | + } |
| 94 | + |
| 95 | + # Add state-specific retention policies |
| 96 | + for arg_key, state in state_mapping.items(): |
| 97 | + days_arg = args_dict.get(arg_key) |
| 98 | + if days_arg is not None: |
| 99 | + days = to_int(days_arg) |
| 100 | + if days is not None: |
| 101 | + # Add a retention policy for the specific orchestration state |
| 102 | + retention_policies.append({ |
| 103 | + "retentionPeriodInDays": days, |
| 104 | + "orchestrationState": state |
| 105 | + }) |
| 106 | + |
| 107 | + return retention_policies |
0 commit comments