FM: Split Instance and InstanceCreator per cloud

Author: Thibaud Baas

dataikuapi/fm/instances.py

@@ -6,9 +6,10 @@ def __init__(self, client, label, instance_settings_template_id, virtual_network
         """
         Helper to create a DSS Instance
 
+        :param object client: :class:`dataikuapi.fm.fmclient`
+        :param str label: The label of the instance
         :param str instance_settings_template: The instance settings template id this instance should be based on
         :param str virtual_network: The virtual network where the instance should be spawned
-        :param str label: The label of the instance
         :param str image_id: The ID of the DSS runtime image (ex: dss-9.0.3-default)
         """
         self.client = client
@@ -21,33 +22,29 @@ def __init__(self, client, label, instance_settings_template_id, virtual_network
         # Set the default value for dssNodeType
         self.data["dssNodeType"] = "design"
 
-    def create(self):
-        """
-        Create the DSS instance
-
-        :return: Instance
-        :rtype: :class:`dataikuapi.fm.instances.FMInstance`
-        """
-        instance = self.client._perform_tenant_json("POST", "/instances", body=self.data)
-        return FMInstance(self.client, instance)
-
-    def set_dss_node_type(self, dss_node_type):
+    def with_dss_node_type(self, dss_node_type):
         """
         Set the DSS Node type of the instance to create
 
         :param str dss_node_type: Optional , the type of the dss node to create. Supports "design", "automation ordeployer". Defaults to "design"
+        :rtype: :class:`dataikuapi.fm.instances.FMInstanceCreator`
         """
         if dss_node_type not in ["design", "automation", "deployer"]:
             raise ValueError("Only \"design\", \"automation\" or \"deployer\" dss_node_type are supported")
         self.data["dssNodeType"] = dss_node_type
+        return self
 
-    def set_cloud_instance_type(self, cloud_instance_type):
+    def with_cloud_instance_type(self, cloud_instance_type):
         """
         Set the machine type for the DSS Instance
+
+        :param str cloud_instance_type
+        :rtype: :class:`dataikuapi.fm.instances.FMInstanceCreator`
         """
         self.data["cloudInstanceType"] = cloud_instance_type
+        return self
 
-    def set_data_volume_options(self, data_volume_type=None, data_volume_size=None, data_volume_size_max=None, data_volume_IOPS=None, data_volume_encryption=None, data_volume_encryption_key=None):
+    def with_data_volume_options(self, data_volume_type=None, data_volume_size=None, data_volume_size_max=None, data_volume_IOPS=None, data_volume_encryption=None, data_volume_encryption_key=None):
         """
         Set the options of the data volume to use with the DSS Instance
 
@@ -57,6 +54,7 @@ def set_data_volume_options(self, data_volume_type=None, data_volume_size=None,
         :param int data_volume_IOPS: Optional, Data volume IOPS
         :param object data_volume_encryption: Optional, a :class:`dataikuapi.fm.instances.FMInstanceEncryptionMode` setting the encryption mode of the data volume
         :param str data_volume_encryption_key: Optional, the encryption key to use when data_volume_encryption_key is FMInstanceEncryptionMode.CUSTOM
+        :rtype: :class:`dataikuapi.fm.instances.FMInstanceCreator`
         """
         if type(data_volume_encryption) is not FMInstanceEncryptionMode:
             raise TypeError("data_volume encryption needs to be of type FMInstanceEncryptionMode")
@@ -67,42 +65,71 @@ def set_data_volume_options(self, data_volume_type=None, data_volume_size=None,
         self.data["dataVolumeIOPS"] = data_volume_IOPS
         self.data["dataVolumeEncryption"] = data_volume_encryption.value
         self.data["dataVolumeEncryptionKey"] = data_volume_encryption_key
+        return self
+
+    def with_cloud_tags(self, cloud_tags):
+        """
+        Set the tags to be applied to the cloud resources created for this DSS instance
 
-    def set_aws_root_volume_options(self, aws_root_volume_size=None, aws_root_volume_type=None, aws_root_volume_IOPS=None):
+        :param dict cloud_tags: a key value dictionary of tags to be applied on the cloud resources
+        :rtype: :class:`dataikuapi.fm.instances.FMInstanceCreator`
         """
-        AWS Only: Set the options of the root volume of the DSS Instance
+        self.data["cloudTags"] = cloud_tags
+        return self
+
+    def with_fm_tags(self, fm_tags):
+        """
+        A list of tags to add on the DSS Instance in Fleet Manager
+
+        :param list fm_tags: Optional, list of tags to be applied on the instance in the Fleet Manager
+        :rtype: :class:`dataikuapi.fm.instances.FMInstanceCreator`
+        """
+        self.data["fmTags"] = fm_tags
+        return self
+
+
+class FMAWSInstanceCreator(FMInstanceCreator):
+    def with_aws_root_volume_options(self, aws_root_volume_size=None, aws_root_volume_type=None, aws_root_volume_IOPS=None):
+        """
+        Set the options of the root volume of the DSS Instance
 
         :param int aws_root_volume_size: Optional, the root volume size
         :param str aws_root_volume_type: Optional, the root volume type
         :param int aws_root_volume_IOPS: Optional, the root volume IOPS
+        :rtype: :class:`dataikuapi.fm.instances.FMAWSInstanceCreator`
         """
-        if self.client.cloud != "AWS":
-            raise BaseException("set_aws_root_volume_options is only usable on AWS tenants")
         self.data["awsRootVolumeSizeGB"] = aws_root_volume_size
         self.data["awsRootVolumeType"] = aws_root_volume_type
         self.data["awsRootVolumeIOPS"] = aws_root_volume_IOPS
+        return self
 
-    def set_cloud_tags(self, cloud_tags):
+    def create(self):
         """
-        Set the tags to be applied to the cloud resources created for this DSS instance
+        Create the DSS instance
 
-        :param dict cloud_tags: a key value dictionary of tags to be applied on the cloud resources
+        :return: Created DSS Instance
+        :rtype: :class:`dataikuapi.fm.instances.FMAWSInstance`
         """
-        self.data["cloudTags"] = cloud_tags
+        instance = self.client._perform_tenant_json("POST", "/instances", body=self.data)
+        return FMAWSInstance(self.client, instance)
 
-    def set_fm_tags(self, fm_tags):
+
+class FMAzureInstanceCreator(FMInstanceCreator):
+    def create(self):
         """
-        A list of tags to add on the DSS Instance in Fleet Manager
+        Create the DSS instance
 
-        :param list fm_tags: Optional, list of tags to be applied on the instance in the Fleet Manager
+        :return: Created DSS Instance
+        :rtype: :class:`dataikuapi.fm.instances.FMAzureInstance`
         """
-        self.data["fmTags"] = fm_tags
+        instance = self.client._perform_tenant_json("POST", "/instances", body=self.data)
+        return FMAzureInstance(self.client, instance)
 
 
 class FMInstance(object):
     """
     A handle to interact with a DSS instance.
-    Do not create this directly, use :meth:`FMClient.get_instance` or :meth: `FMClient.create_instance`
+    Do not create this directly, use :meth:`FMClient.get_instance` or :meth: `FMClient.new_instance_creator`
     """
     def __init__(self, client, instance_data):
         self.client  = client
@@ -176,28 +203,41 @@ def set_automated_snapshots(self, enable, period, keep=0):
         self.instance_data['automatedSnapshotRetention'] = keep
         self.save()
 
+    def set_custom_certificate(self, pem_data):
+        """
+        Set the custom certificate for this instance
+
+        Only needed when Virtual Network HTTPS Strategy is set to Custom Certificate
+
+        param: str pem_data: The SSL certificate
+        """
+        self.instance_data['sslCertificatePEM'] = pem_data
+        self.save()
+
+
+class FMAWSInstance(FMInstance):
     def set_elastic_ip(self, enable, elasticip_allocation_id):
         """
         Set a public elastic ip for this instance
 
         :param boolan enable: Enable the elastic ip allocation
-        :param str elaticip_allocation_id: The AWS ElasticIP allocation ID or the Azure Public IP ID
+        :param str elaticip_allocation_id: AWS ElasticIP allocation ID
         """
         self.instance_data['awsAssignElasticIP'] = enable
         self.instance_data['awsElasticIPAllocationId'] = elasticip_allocation_id
-        self.instance_data['azureAssignElasticIP'] = enable
-        self.instance_data['azurePublicIPId'] = elasticip_allocation_id
         self.save()
 
-    def set_custom_certificate(self, pem_data):
-        """
-        Set the custom certificate for this instance
 
-        Only needed when Virtual Network HTTPS Strategy is set to Custom Certificate
+class FMAzureInstance(FMInstance):
+    def set_elastic_ip(self, enable, public_ip_id):
+        """
+        Set a public elastic ip for this instance
 
-        param: str pem_data: The SSL certificate
+        :param boolan enable: Enable the elastic ip allocation
+        :param str public_ip_id: Azure Public IP ID
         """
-        self.instance_data['sslCertificatePEM'] = pem_data
+        self.instance_data['azureAssignElasticIP'] = enable
+        self.instance_data['azurePublicIPId'] = public_ip_id
         self.save()
 
 

dataikuapi/fmclient.py

@@ -11,6 +11,7 @@
 from .fm.virtualnetworks import FMVirtualNetwork
 from .fm.instances import FMInstance, FMInstanceEncryptionMode
 from .fm.instancesettingstemplates import FMInstanceSettingsTemplate
+from dataikuapi.fm.instances import FMAWSInstanceCreator, FMAzureInstanceCreator
 
 class FMClient(object):
     """Entry point for the FM API client"""
@@ -242,58 +243,19 @@ def get_instance(self, instance_id):
         instance = self._perform_tenant_json("GET", "/instances/%s" % instance_id)
         return FMInstance(self, instance)
 
-    def create_instance(self, instance_settings_template, virtual_network, label, image_id,
-                        dss_node_type="design",
-                        cloud_instance_type=None, data_volume_type=None, data_volume_size=None,
-                        data_volume_size_max=None, data_volume_IOPS=None, data_volume_encryption=FMInstanceEncryptionMode.NONE,
-                        data_volume_encryption_key=None, aws_root_volume_size=None, aws_root_volume_type=None, aws_root_volume_IOPS=None,
-                        cloud_tags=None, fm_tags=None):
+    def new_instance_creator(self, label, instance_settings_template_id, virtual_network_id, image_id):
         """
-        Create a DSS Instance
+        Instantiate a new instance creator
 
+        :param str label: The label of the instance
         :param str instance_settings_template: The instance settings template id this instance should be based on
         :param str virtual_network: The virtual network where the instance should be spawned
-        :param str label: The label of the instance
         :param str image_id: The ID of the DSS runtime image (ex: dss-9.0.3-default)
-
-        :param str dss_node_type: Optional , the type of the dss node to create. Defaults to "design"
-        :param str cloud_instance_type: Optional, Machine type
-        :param str data_volume_type: Optional, Data volume type
-        :param int data_volume_size: Optional, Data volume initial size
-        :param int data_volume_size_max: Optional, Data volume maximum size
-        :param int data_volume_IOPS: Optional, Data volume IOPS
-        :param object data_volume_encryption: Optional, a :class:`dataikuapi.fm.instances.FMInstanceEncryptionMode` setting the encryption mode of the data volume
-        :param str data_volume_encryption_key: Optional, the encryption key to use when data_volume_encryption_key is FMInstanceEncryptionMode.CUSTOM
-        :param int aws_root_volume_size: Optional, the root volume size
-        :param str aws_root_volume_type: Optional, the root volume type
-        :param int aws_root_volume_IOPS: Optional, the root volume IOPS
-        :param dict cloud_tags: Optional, a key value dictionary of tags to be applied on the cloud resources
-        :param list fm_tags: Optional, list of tags to be applied on the instance in the Fleet Manager
-
-        :return: Instance
-        :rtype: :class:`dataikuapi.fm.instances.FMInstance`
         """
-        data = {
-            "virtualNetworkId": virtual_network.id,
-            "instanceSettingsTemplateId": instance_settings_template.id,
-            "label": label,
-            "dssNodeType": dss_node_type,
-            "imageId": image_id,
-            "cloudInstanceType": cloud_instance_type,
-            "dataVolumeType": data_volume_type,
-            "dataVolumeSizeGB": data_volume_size,
-            "dataVolumeSizeMaxGB": data_volume_size_max,
-            "dataVolumeIOPS": data_volume_IOPS,
-            "dataVolumeEncryption": data_volume_encryption.value,
-            "dataVolumeEncryptionKey": data_volume_encryption_key,
-            "awsRootVolumeSizeGB": aws_root_volume_size,
-            "awsRootVolumeType": aws_root_volume_type,
-            "awsRootVolumeIOPS": aws_root_volume_IOPS,
-            "cloudTags": cloud_tags,
-            "fmTags": fm_tags
-        }
-        instance = self._perform_tenant_json("POST", "/instances", body=data)
-        return FMInstance(self, instance)
+        if self.cloud == "AWS":
+            return FMAWSInstanceCreator(self, label, instance_settings_template_id, virtual_network_id, image_id)
+        elif self.cloud == "Azure":
+            return FMAzureInstanceCreator(self, label, instance_settings_template_id, virtual_network_id, image_id)
 
 
     ########################################################