{Scheduled-Query} Add scheduled query rules (#3753)

* update SDK

* parameter

* Update src/scheduled-query/HISTORY.rst

Co-authored-by: kai ru <69238381+kairu-ms@users.noreply.github.com>

* Update src/scheduled-query/HISTORY.rst

Co-authored-by: kai ru <69238381+kairu-ms@users.noreply.github.com>

* Update src/scheduled-query/setup.py

Co-authored-by: kai ru <69238381+kairu-ms@users.noreply.github.com>

* update

* Update test_scheduled_query.yaml

* Update _params.py

* Update HISTORY.rst

Co-authored-by: kai ru <69238381+kairu-ms@users.noreply.github.com>

src/scheduled-query/HISTORY.rst

@@ -2,6 +2,14 @@
 
 Release History
 ===============
+0.4.0
+++++++
+* Add `--skip-query-validation` parameter
+* Add `--check-ws-alerts-storage` parameter
+* Add `--auto-mitigate` parameter
+* [Breaking Change] `--actions` are split into `--action-groups` and `--custom-properties`
+* [Breaking Change] the default value of `--mute-actions-duration` is changed to None
+
 0.3.1
 ++++++
 * Support query placeholder for `--condition` parameter.

src/scheduled-query/azext_scheduled_query/__init__.py

@@ -16,8 +16,7 @@ def __init__(self, cli_ctx=None):
         scheduled_query_custom = CliCommandType(
             operations_tmpl='azext_scheduled_query.custom#{}',
             client_factory=cf_scheduled_query)
-        super(ScheduledQueryCommandsLoader, self).__init__(cli_ctx=cli_ctx,
-                                                           custom_command_type=scheduled_query_custom)
+        super().__init__(cli_ctx=cli_ctx, custom_command_type=scheduled_query_custom)
 
     def load_command_table(self, args):
         from azext_scheduled_query.commands import load_command_table

src/scheduled-query/azext_scheduled_query/_actions.py

@@ -39,12 +39,9 @@ def __call__(self, parser, namespace, values, option_string=None):
             for item in ['time_aggregation', 'threshold', 'operator']:
                 if not getattr(scheduled_query_condition, item, None):
                     raise InvalidArgumentValueError(usage)
-        except (AttributeError, TypeError, KeyError):
-            raise InvalidArgumentValueError(usage)
-        super(ScheduleQueryConditionAction, self).__call__(parser,
-                                                           namespace,
-                                                           scheduled_query_condition,
-                                                           option_string)
+        except (AttributeError, TypeError, KeyError) as e:
+            raise InvalidArgumentValueError(usage) from e
+        super().__call__(parser, namespace, scheduled_query_condition, option_string)
 
 
 class ScheduleQueryConditionQueryAction(argparse.Action):
@@ -65,9 +62,9 @@ def __call__(self, parser, namespace, values, option_string=None):
 class ScheduleQueryAddAction(argparse._AppendAction):
 
     def __call__(self, parser, namespace, values, option_string=None):
-        from azext_scheduled_query.vendored_sdks.azure_mgmt_scheduled_query.models import Action
-        action = Action(
-            action_group_id=values[0],
-            web_hook_properties=dict(x.split('=', 1) for x in values[1:]) if len(values) > 1 else None
+        from azext_scheduled_query.vendored_sdks.azure_mgmt_scheduled_query.models import Actions
+        action = Actions(
+            action_groups=values[0],
+            custom_properties=dict(x.split('=', 1) for x in values[1:]) if len(values) > 1 else None
         )
-        super(ScheduleQueryAddAction, self).__call__(parser, namespace, action, option_string)
+        super().__call__(parser, namespace, action, option_string)

src/scheduled-query/azext_scheduled_query/_client_factory.py

@@ -6,5 +6,5 @@
 
 def cf_scheduled_query(cli_ctx, *_):
     from azure.cli.core.commands.client_factory import get_mgmt_service_client
-    from .vendored_sdks.azure_mgmt_scheduled_query._monitor_client import MonitorClient
-    return get_mgmt_service_client(cli_ctx, MonitorClient).scheduled_query_rules
+    from .vendored_sdks.azure_mgmt_scheduled_query._monitor_management_client import MonitorManagementClient
+    return get_mgmt_service_client(cli_ctx, MonitorManagementClient).scheduled_query_rules

src/scheduled-query/azext_scheduled_query/_help.py

@@ -16,12 +16,14 @@
 type: command
 short-summary: Create a scheduled query.
 parameters:
-  - name: --action -a
-    short-summary: Add an action group and optional webhook properties to fire when the alert is triggered.
+  - name: --action-groups
+    short-summary: Action Group resource Ids to invoke when the alert fires.
     long-summary: |
-        Usage:   --action ACTION_GROUP_NAME_OR_ID [KEY=VAL [KEY=VAL ...]]
-
-        Multiple action groups can be specified by using more than one `--action` argument.
+        Usage:   --action-groups ACTION_GROUP_NAME_OR_ID [NAME_OR_ID,...]
+  - name: --custom-properties
+    short-summary: The properties of an alert payload.
+    long-summary: |
+        Usage:   --custom-properties ALERT_PAYLOAD_PROPERTIES [KEY=VAL,KEY=VAL ...]
   - name: --condition
     short-summary: The condition which triggers the rule.
     long-summary: |
@@ -44,12 +46,14 @@
 type: command
 short-summary: Update a scheduled query.
 parameters:
-  - name: --action -a
-    short-summary: Add an action group and optional webhook properties to fire when the alert is triggered.
+  - name: --action-groups
+    short-summary: Action Group resource Ids to invoke when the alert fires.
     long-summary: |
-        Usage:   --action ACTION_GROUP_NAME_OR_ID [KEY=VAL [KEY=VAL ...]]
-
-        Multiple action groups can be specified by using more than one `--action` argument.
+        Usage:   --action-groups ACTION_GROUP_NAME_OR_ID [NAME_OR_ID,...]
+  - name: --custom-properties
+    short-summary: The properties of an alert payload.
+    long-summary: |
+        Usage:   --custom-properties ALERT_PAYLOAD_PROPERTIES [KEY=VAL,KEY=VAL ...]
   - name: --condition
     short-summary: The condition which triggers the rule.
     long-summary: |

src/scheduled-query/azext_scheduled_query/_params.py

@@ -6,16 +6,21 @@
 
 from azure.cli.core.commands.parameters import tags_type, get_three_state_flag
 from azure.cli.command_modules.monitor.actions import get_period_type
-from azure.cli.command_modules.monitor.validators import get_action_group_validator
+from azext_scheduled_query._validators import validate_custom_properties
 from knack.arguments import CLIArgumentType
-from ._actions import ScheduleQueryConditionAction, ScheduleQueryAddAction, ScheduleQueryConditionQueryAction
+from ._actions import ScheduleQueryConditionAction, ScheduleQueryConditionQueryAction
 
 
 def load_arguments(self, _):
 
     from azure.cli.core.commands.validators import get_default_location_from_resource_group
 
     name_arg_type = CLIArgumentType(options_list=['--name', '-n'], metavar='NAME')
+    custom_properties_arg_type = CLIArgumentType(
+        validator=validate_custom_properties,
+        options_list=['--custom-properties'],
+        nargs='*',
+        help='The properties of an alert payload.')
     with self.argument_context('monitor scheduled-query') as c:
         c.argument('rule_name', name_arg_type, id_part='name', help='Name of the scheduled query rule.')
         c.argument('location', validator=get_default_location_from_resource_group)
@@ -35,4 +40,12 @@ def load_arguments(self, _):
         c.argument('mute_actions_duration', type=get_period_type(as_timedelta=True),
                    options_list=['--mute-actions-duration', '--mad'],
                    help='Mute actions for the chosen period of time (in ISO 8601 duration format) after the alert is fired.')
-        c.argument('actions', options_list=['--action', '-a'], action=ScheduleQueryAddAction, nargs='+', validator=get_action_group_validator('actions'))
+        c.argument('action_groups', options_list=['--action-groups'], nargs='+', help='Action Group resource Ids to invoke when the alert fires.')
+        c.argument('custom_properties', custom_properties_arg_type)
+        c.argument('auto_mitigate', arg_type=get_three_state_flag(),
+                   help='The flag that indicates whether the alert should be automatically resolved or not. The default is true.')
+        c.argument('skip_query_validation', arg_type=get_three_state_flag(),
+                   help='The flag which indicates whether the provided query should be validated or not.')
+        c.argument('check_workspace_alerts_storage', options_list=['--check-ws-alerts-storage', '--cwas'],
+                   arg_type=get_three_state_flag(),
+                   help="The flag which indicates whether this scheduled query rule should be stored in the customer's storage.")

src/scheduled-query/azext_scheduled_query/_validators.py

@@ -0,0 +1,20 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+
+def validate_tag(string):
+    result = {}
+    if string:
+        comps = string.split('=', 1)
+        result = {comps[0]: comps[1]} if len(comps) > 1 else {string: ''}
+    return result
+
+
+def validate_custom_properties(ns):
+    if isinstance(ns.custom_properties, list):
+        custom_properties_dict = {}
+        for item in ns.custom_properties:
+            custom_properties_dict.update(validate_tag(item))
+        ns.custom_properties = custom_properties_dict

src/scheduled-query/azext_scheduled_query/custom.py

@@ -25,27 +25,52 @@ def _build_criteria(condition, condition_query):
     return ScheduledQueryRuleCriteria(all_of=condition)
 
 
-def create_scheduled_query(client, resource_group_name, rule_name, scopes, condition, condition_query=None,
-                           disabled=False, description=None, tags=None, location=None,
-                           actions=None, severity=2, window_size='5m', evaluation_frequency='5m',
-                           target_resource_type=None, mute_actions_duration='PT30M'):
-    from .vendored_sdks.azure_mgmt_scheduled_query.models import ScheduledQueryRuleResource
+def create_scheduled_query(client,
+                           resource_group_name,
+                           rule_name,
+                           scopes,
+                           condition,
+                           action_groups=None,
+                           custom_properties=None,
+                           condition_query=None,
+                           disabled=False,
+                           description=None,
+                           tags=None,
+                           location=None,
+                           severity=2,
+                           window_size='5m',
+                           evaluation_frequency='5m',
+                           target_resource_type=None,
+                           mute_actions_duration=None,
+                           auto_mitigate=True,
+                           skip_query_validation=False,
+                           check_workspace_alerts_storage=False):
     criteria = _build_criteria(condition, condition_query)
-    kwargs = {
-        'description': description,
-        'severity': severity,
-        'enabled': not disabled,
-        'scopes': scopes,
-        'evaluation_frequency': evaluation_frequency,
-        'window_size': window_size,
-        'criteria': criteria,
-        'target_resource_types': [target_resource_type] if target_resource_type else None,
-        'actions': actions,
-        'tags': tags,
-        'location': location,
-        'mute_actions_duration': mute_actions_duration
-    }
-    return client.create_or_update(resource_group_name, rule_name, ScheduledQueryRuleResource(**kwargs))
+    parameters = {}
+    actions = {}
+    actions['action_groups'] = action_groups if action_groups is not None else []
+    actions['custom_properties'] = custom_properties if custom_properties is not None else {}
+    parameters['actions'] = actions
+    parameters['scopes'] = scopes
+    parameters['criteria'] = criteria
+    if actions is not None:
+        parameters['actions'] = actions
+    parameters['enabled'] = not disabled
+    if description is not None:
+        parameters['description'] = description
+    if tags is not None:
+        parameters['tags'] = tags
+    if location is not None:
+        parameters['location'] = location
+    parameters['severity'] = severity
+    parameters['window_size'] = window_size
+    parameters['evaluation_frequency'] = evaluation_frequency
+    parameters['target_resource_types'] = [target_resource_type] if target_resource_type else None
+    parameters['mute_actions_duration'] = mute_actions_duration
+    parameters['auto_mitigate'] = auto_mitigate
+    parameters['skip_query_validation'] = skip_query_validation
+    parameters['check_workspace_alerts_storage_configured'] = check_workspace_alerts_storage
+    return client.create_or_update(resource_group_name=resource_group_name, rule_name=rule_name, parameters=parameters)
 
 
 def list_scheduled_query(client, resource_group_name=None):
@@ -54,19 +79,43 @@ def list_scheduled_query(client, resource_group_name=None):
     return client.list_by_subscription()
 
 
-def update_scheduled_query(cmd, instance, tags=None, disabled=False, condition=None, condition_query=None,
-                           description=None, actions=None, severity=None, window_size=None,
-                           evaluation_frequency=None, mute_actions_duration=None):
+def update_scheduled_query(cmd,
+                           instance,
+                           tags=None,
+                           disabled=None,
+                           condition=None,
+                           action_groups=None,
+                           custom_properties=None,
+                           condition_query=None,
+                           description=None,
+                           severity=None,
+                           window_size=None,
+                           evaluation_frequency=None,
+                           mute_actions_duration=None,
+                           target_resource_type=None,
+                           auto_mitigate=None,
+                           skip_query_validation=None,
+                           check_workspace_alerts_storage=None
+                           ):
     with cmd.update_context(instance) as c:
         c.set_param('tags', tags)
         c.set_param('enabled', not disabled)
         c.set_param('description', description)
-        c.set_param('actions', actions)
+        c.set_param('actions.action_groups', action_groups)
+        c.set_param('actions.custom_properties', custom_properties)
         c.set_param('severity', severity)
         c.set_param('window_size', window_size)
         c.set_param('evaluation_frequency', evaluation_frequency)
         c.set_param('mute_actions_duration', mute_actions_duration)
+        c.set_param('target_resource_type', target_resource_type)
+        c.set_param('auto_mitigate', auto_mitigate)
+        c.set_param('skip_query_validation', skip_query_validation)
+        c.set_param('check_workspace_alerts_storage_configured', check_workspace_alerts_storage)
         if condition is not None:
             criteria = _build_criteria(condition, condition_query)
             c.set_param('criteria', criteria)
+        if disabled is not None:
+            c.set_param('enabled', not disabled)
+        if auto_mitigate is not None:
+            c.set_param('auto_mitigate', auto_mitigate)
     return instance