fm: FMVirtualNetworkCreator

Author: Thibaud Baas

dataikuapi/fm/virtualnetworks.py

@@ -1,17 +1,123 @@
 from dataikuapi.fm.future import FMFuture
 
+
+class FMVirtualNetworkCreator(object):
+    def __init__(self, client, label):
+        """
+        A builder class to create a Virtual Network
+
+        :param str label: The label of the Virtual Network
+        """
+        self.client = client
+        self.data = {}
+        self.data["label"] = label
+        self.data["mode"] = "EXISTING_MONOTENANT"
+
+    def with_internet_access_mode(self, internet_access_mode):
+        """
+        :param str internet_access_mode: The internet access mode of the instances created in this virtual network. Accepts "YES", "NO", "EGRESS_ONLY". Defaults to "YES"
+        """
+        if internet_access_mode not in ["YES", "NO", "EGRESS_ONLY"]:
+            raise ValueError(
+                'internet_access_mode should be either "YES", "NO", or "EGRESS_ONLY"'
+            )
+
+        self.data["internetAccessMode"] = internet_access_mode
+        return self
+
+
+class FMAWSVirtualNetworkCreator(FMVirtualNetworkCreator):
+    def with_vpc(self, aws_vpc_id, aws_subnet_id):
+        """
+        Setup the VPC and Subnet to with the VirtualNetwork
+
+        :param str aws_vpc_id: ID of the VPC to use
+        :param str aws_subnet_id: ID of the subnet to use
+        """
+        self.data["awsVpcId"] = aws_vpc_id
+        self.data["awsSubnetId"] = aws_subnet_id
+        return self
+
+    def with_auto_create_security_groups(self):
+        """
+        Automatically create the AWS Security Groups when creating this VirtualNetwork
+        """
+        self.data["awsAutoCreateSecurityGroups"] = True
+        return self
+
+    def with_aws_security_groups(self, aws_security_groups):
+        """
+        Use pre-created AWS Security Groups
+
+        :param list aws_security_groups: A list of up to 5 security group ids to assign to the instances created in this virtual network.
+        """
+        self.data["awsAutoCreateSecurityGroups"] = False
+        self.data["awsSecurityGroups"] = aws_security_groups
+        return self
+
+    def create(self):
+        """
+        Create the VirtualNetwork
+
+        :return: Created VirtualNetwork
+        :rtype: :class:`dataikuapi.fm.virtualnetworks.FMAWSVirtualNetwork`
+        """
+        vn = self.client._perform_tenant_json(
+            "POST", "/virtual-networks", body=self.data
+        )
+        return FMAWSVirtualNetwork(self, vn)
+
+
+class FMAzureVirtualNetworkCreator(FMVirtualNetworkCreator):
+    def with_virtual_network(self, azure_vn_id, azure_subnet_id):
+        """
+        Setup the Azure Virtual Network and Subnet to with the VirtualNetwork
+
+        :param str azure_vn_id: Resource ID of the Azure Virtual Network to use
+        :param str azure_subnet_id: Resource ID of the subnet to use
+        """
+        self.data["azureVnId"] = azure_vn_id
+        self.data["azureSubnetId"] = azure_subnet_id
+        return self
+
+    def with_auto_update_security_groups(self, auto_update_security_groups=True):
+        """
+        Auto update the security groups of the Azure Virtual Network
+
+        :param boolean auto_update_security_groups: Optional, Auto update the subnet security group. Defaults to True
+        """
+        self.data["azureAutoUpdateSecurityGroups"] = auto_update_security_groups
+        return self
+
+    def create(self):
+        """
+        Create the VirtualNetwork
+
+        :return: Created VirtualNetwork
+        :rtype: :class:`dataikuapi.fm.virtualnetworks.FMAzureVirtualNetwork`
+        """
+        vn = self.client._perform_tenant_json(
+            "POST", "/virtual-networks", body=self.data
+        )
+        return FMAzureVirtualNetwork(self, vn)
+
+
 class FMVirtualNetwork(object):
     def __init__(self, client, vn_data):
-        self.client  = client
+        self.client = client
         self.vn_data = vn_data
-        self.id = self.vn_data['id']
+        self.id = self.vn_data["id"]
 
     def save(self):
         """
         Update the Virtual Network.
         """
-        self.client._perform_tenant_empty("PUT", "/virtual-networks/%s" % self.id, body=self.vn_data)
-        self.vn_data = self.client._perform_tenant_json("GET", "/virtual-networks/%s" % self.id)
+        self.client._perform_tenant_empty(
+            "PUT", "/virtual-networks/%s" % self.id, body=self.vn_data
+        )
+        self.vn_data = self.client._perform_tenant_json(
+            "GET", "/virtual-networks/%s" % self.id
+        )
 
     def delete(self):
         """
@@ -20,10 +126,14 @@ def delete(self):
         :return: A :class:`~dataikuapi.fm.future.FMFuture` representing the deletion process
         :rtype: :class:`~dataikuapi.fm.future.FMFuture`
         """
-        future = self.client._perform_tenant_json("DELETE", "/virtual-networks/%s" % self.id)
+        future = self.client._perform_tenant_json(
+            "DELETE", "/virtual-networks/%s" % self.id
+        )
         return FMFuture.from_resp(self.client, future)
 
-    def set_fleet_management(self, enable, event_server=None, deployer_management="NO_MANAGED_DEPLOYER"):
+    def set_fleet_management(
+        self, enable, event_server=None, deployer_management="NO_MANAGED_DEPLOYER"
+    ):
         """
         When enabled, all instances in this virtual network know each other and can centrally manage deployer and logs centralization
 
@@ -36,40 +146,64 @@ def set_fleet_management(self, enable, event_server=None, deployer_management="N
             - "EACH_DESIGN_NODE": Deployer from design. Recommanded if you have a single design node and want a simpler setup.
         """
 
-        self.vn_data['managedNodesDirectory'] = enable
-        self.vn_data['eventServerNodeLabel'] = event_server
-        self.vn_data['nodesDirectoryDeployerMode'] = deployer_management
+        self.vn_data["managedNodesDirectory"] = enable
+        self.vn_data["eventServerNodeLabel"] = event_server
+        self.vn_data["nodesDirectoryDeployerMode"] = deployer_management
+        self.save()
+
+    def set_https_strategy(self, https_strategy):
+        """
+        Set the HTTPS strategy for this virtual network
+
+        :param object: a :class:`dataikuapi.fm.virtualnetworks.FMHTTPSStrategy`
+        """
+        self.vn_data.update(https_strategy)
         self.save()
 
-    def set_dns_strategy(self, assign_domain_name, aws_private_ip_zone53_id=None, aws_public_ip_zone53_id=None, azure_dns_zone_id=None):
+
+class FMAWSVirtualNetwork(FMVirtualNetwork):
+    def set_dns_strategy(
+        self,
+        assign_domain_name,
+        aws_private_ip_zone53_id=None,
+        aws_public_ip_zone53_id=None,
+    ):
         """
         Set the DNS strategy for this virtual network
 
         :param boolean assign_domain_name: If false, don't assign domain names, use ip_only
         :param str aws_private_ip_zone53_id: Optional, AWS Only, the ID of the AWS Route53 Zone to use for private ip
         :param str aws_public_ip_zone53_id: Optional, AWS Only, the ID of the AWS Route53 Zone to use for public ip
-        :param str azure_dns_zone_id: Optional, Azure Only, the ID of the Azure DNS zone to use
         """
 
         if assign_domain_name:
-            self.vn_data['dnsStrategy'] = "VN_SPECIFIC_CLOUD_DNS_SERVICE"
-            self.vn_data['awsRoute53PrivateIPZoneId'] = aws_private_ip_zone53_id
-            self.vn_data['awsRoute53PublicIPZoneId'] = aws_public_ip_zone53_id
-            self.vn_data['azureDnsZoneId'] = azure_dns_zone_id
-        else :
-            self.vn_data['dnsStrategy'] = "NONE"
+            self.vn_data["dnsStrategy"] = "VN_SPECIFIC_CLOUD_DNS_SERVICE"
+            self.vn_data["awsRoute53PrivateIPZoneId"] = aws_private_ip_zone53_id
+            self.vn_data["awsRoute53PublicIPZoneId"] = aws_public_ip_zone53_id
+        else:
+            self.vn_data["dnsStrategy"] = "NONE"
 
         self.save()
 
-    def set_https_strategy(self, https_strategy):
+
+class FMAzureVirtualNetwork(FMVirtualNetwork):
+    def set_dns_strategy(self, assign_domain_name, azure_dns_zone_id=None):
         """
-        Set the HTTPS strategy for this virtual network
+        Set the DNS strategy for this virtual network
 
-        :param object: a :class:`dataikuapi.fm.virtualnetworks.FMHTTPSStrategy`
+        :param boolean assign_domain_name: If false, don't assign domain names, use ip_only
+        :param str azure_dns_zone_id: Optional, Azure Only, the ID of the Azure DNS zone to use
         """
-        self.vn_data.update(https_strategy)
+
+        if assign_domain_name:
+            self.vn_data["dnsStrategy"] = "VN_SPECIFIC_CLOUD_DNS_SERVICE"
+            self.vn_data["azureDnsZoneId"] = azure_dns_zone_id
+        else:
+            self.vn_data["dnsStrategy"] = "NONE"
+
         self.save()
 
+
 class FMHTTPSStrategy(dict):
     def __init__(self, data, https_strategy, http_redirect=False):
         """
@@ -82,11 +216,11 @@ def __init__(self, data, https_strategy, http_redirect=False):
             - :meth:`dataikuapi.fm.virtualnetwork.FMHTTPSStrategy.lets_encrypt` to use Let's Encrypt
         """
         super(FMHTTPSStrategy, self).__init__(data)
-        self['httpsStrategy'] = https_strategy
+        self["httpsStrategy"] = https_strategy
         if http_redirect:
-            self['httpStrategy'] = "REDIRECT"
+            self["httpStrategy"] = "REDIRECT"
         else:
-            self['httpStrategy'] = "DISABLE"
+            self["httpStrategy"] = "DISABLE"
 
     @staticmethod
     def disable():

dataikuapi/fmclient.py

@@ -8,7 +8,7 @@
 from .utils import DataikuException
 
 from .fm.tenant import FMCloudCredentials
-from .fm.virtualnetworks import FMVirtualNetwork
+from .fm.virtualnetworks import FMVirtualNetwork, FMAWSVirtualNetworkCreator, FMAzureVirtualNetworkCreator
 from .fm.instances import FMInstance, FMInstanceEncryptionMode, FMAWSInstanceCreator, FMAzureInstanceCreator
 from .fm.instancesettingstemplates import FMInstanceSettingsTemplate, FMAWSInstanceSettingsTemplateCreator, FMAzureInstanceSettingsTemplateCreator
 
@@ -85,51 +85,18 @@ def get_virtual_network(self, virtual_network_id):
         vn = self._perform_tenant_json("GET", "/virtual-networks/%s" % virtual_network_id)
         return FMVirtualNetwork(self, vn)
 
-    def create_virtual_network(self,
-                               label,
-                               awsVpcId=None,
-                               awsSubnetId=None,
-                               awsAutoCreateSecurityGroups=False,
-                               awsSecurityGroups=None,
-                               azureVnId=None,
-                               azureSubnetId=None,
-                               azureAutoUpdateSecurityGroups=None,
-                               internetAccessMode = "YES"):
-        """
-        Create a Virtual Network
-
-        :param str label: The label of the Virtual Network
-
-        :param str awsVpcId: AWS Only, ID of the VPC to use
-        :param str awsSubnetId: AWS Only, ID of the subnet to use
-        :param boolean awsAutoCreateSecurityGroups: Optional, AWS Only, If false, do not create security groups automatically. Defaults to false
-        :param list awsSecurityGroups: Optional, AWS Only, A list of up to 5 security group ids to assign to the instances created in this virtual network. Ignored if awsAutoCreateSecurityGroups is true
 
-        :param str azureVnId: Azure Only, ID of the Azure Virtual Network to use
-        :param str azureSubnetId: Azure Only, ID of the subnet to use
-        :param boolean azureAutoUpdateSecurityGroups: Azure Only, Auto update the subnet security group
-
-        :param str internetAccessMode: Optional, The internet access mode of the instances created in this virtual network. Accepts "YES", "NO", "EGRESS_ONLY". Defaults to "YES"
-
-        :return: requested instance settings template
-        :rtype: :class:`dataikuapi.fm.instancesettingstemplates.FMInstanceSettingsTemplate`
+    def new_virtual_network_creator(self, label):
         """
+        Instantiate a new virtual network creator
 
-        data = {
-            "label": label,
-            "awsVpcId": awsVpcId,
-            "awsSubnetId": awsSubnetId,
-            "awsAutoCreateSecurityGroups": awsAutoCreateSecurityGroups,
-            "awsSecurityGroups": awsSecurityGroups,
-            "azureVnId": azureVnId,
-            "azureSubnetId": azureSubnetId,
-            "azureAutoUpdateSecurityGroups": azureAutoUpdateSecurityGroups,
-            "internetAccessMode": internetAccessMode,
-            "mode": "EXISTING_MONOTENANT"
-        }
-
-        vn = self._perform_tenant_json("POST", "/virtual-networks", body=data)
-        return FMVirtualNetwork(self, vn)
+        :param str label: The label of the
+        :rtype: :class:`dataikuapi.fm.virtualnetworks.FMVirtualNetworkCreator`
+        """
+        if self.cloud == "AWS":
+            return FMAWSVirtualNetworkCreator(self, label)
+        elif self.cloud == "Azure":
+            return FMAzureVirtualNetworkCreator(self, label)
 
 
     ########################################################