diff --git a/qds_sdk/clusterv2.py b/qds_sdk/clusterv2.py index 08a7d821..cee8bebb 100755 --- a/qds_sdk/clusterv2.py +++ b/qds_sdk/clusterv2.py @@ -148,6 +148,7 @@ def set_cluster_info_from_arguments(self, arguments): enable_ganglia_monitoring=arguments.enable_ganglia_monitoring, datadog_api_token=arguments.datadog_api_token, datadog_app_token=arguments.datadog_app_token, + notification_channels=arguments.notification_channels, node_bootstrap=arguments.node_bootstrap_file, master_instance_type=arguments.master_instance_type, slave_instance_type=arguments.slave_instance_type, @@ -191,6 +192,7 @@ def set_cluster_info(self, enable_ganglia_monitoring=None, datadog_api_token=None, datadog_app_token=None, + notification_channels=None, node_bootstrap=None, master_instance_type=None, slave_instance_type=None, @@ -318,6 +320,8 @@ def set_cluster_info(self, `datadog_app_token` : Specify the Datadog APP token to use the Datadog monitoring service + `notification_channels` : Specify the list of notification channels + `image_uri_overrides` : Override the image name provided `env_name`: Name of python and R environment. (For Spark clusters) @@ -377,7 +381,7 @@ def set_cluster_info(self, stable_spot_fallback) self.set_spot_block_settings(spot_block_duration) self.set_data_disk(disk_size, disk_count, disk_type, upscaling_config, enable_encryption) - self.set_monitoring(enable_ganglia_monitoring, datadog_api_token, datadog_app_token) + self.set_monitoring(enable_ganglia_monitoring, datadog_api_token, datadog_app_token, notification_channels) self.set_internal(image_uri_overrides) self.set_env_settings(env_name, python_version, r_version) self.set_start_stop_settings(disable_cluster_pause, paused_cluster_timeout_mins, @@ -390,12 +394,18 @@ def set_datadog_setting(self, self.monitoring['datadog']['datadog_api_token'] = datadog_api_token self.monitoring['datadog']['datadog_app_token'] = datadog_app_token + def set_notification_settings(self, notification_channels): + self.monitoring["notifications"] = {} + self.monitoring["notifications"]["all"] = notification_channels + def set_monitoring(self, enable_ganglia_monitoring=None, datadog_api_token=None, - datadog_app_token=None): + datadog_app_token=None, + notification_channels=None): self.monitoring['ganglia'] = enable_ganglia_monitoring self.set_datadog_setting(datadog_api_token, datadog_app_token) + self.set_notification_settings(notification_channels) def set_spot_instance_settings(self, maximum_bid_price_percentage=None, @@ -699,6 +709,13 @@ def cluster_info_parser(argparser, action): dest="datadog_app_token", default=None, help="overrides for airflow cluster", ) + notifications_group = argparser.add_argument_group("notifications") + notifications_group.add_argument("--notification-channels", + nargs="*", + type=int, + dest="notification_channels", + default=None, + help="List of notification channel ids", ) internal_group = argparser.add_argument_group("internal settings") internal_group.add_argument("--image-overrides", diff --git a/tests/test_clusterv2.py b/tests/test_clusterv2.py index 17859456..39a4eaec 100644 --- a/tests/test_clusterv2.py +++ b/tests/test_clusterv2.py @@ -777,6 +777,15 @@ def test_start_stop_timeouts_invalid(self): with self.assertRaises(SystemExit): qds.main() + def test_notifications_given(self): + sys.argv = ['qds.py', '--version', 'v2', 'cluster', 'create', '--label', 'test_label', + '--notification-channels', '7'] + print_command() + Connection._api_call = Mock(return_value={}) + qds.main() + Connection._api_call.assert_called_with('POST', 'clusters', + {'cluster_info': {'label': ['test_label']}, + 'monitoring': {'notifications': {'all': [7]}}}) class TestClusterUpdate(QdsCliTestCase): def test_minimal(self):