ACM-6082: Support cluster notifications (#327)

Co-authored-by: Vasant M <[email protected]> Co-authored-by: Joy Lal Chattaraj <[email protected]> (cherry picked from commit ed05e0a)
qubole · Oct 6, 2020 · 5f1ab81 · 5f1ab81
1 parent f00a76e
commit 5f1ab81
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 2 deletions.
diff --git 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
@@ -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):