providers/ec2: support VPC

juju/machine/unit.py

 import os
 import shutil
 import logging
 
 import juju
 
 from twisted.internet.defer import inlineCallbacks, returnValue
 
 from juju.charm.bundle import CharmBundle
 from juju.errors import ServiceError
-from juju.lib.lxc import LXCContainer, get_containers, LXCError
+from juju.lib.lxc import LXCContainer, get_containers
 from juju.lib.twistutils import get_module_directory
 from juju.lib.upstart import UpstartService
 from juju.providers.common.cloudinit import CloudInit
 
 from .errors import UnitDeploymentError
 
 log = logging.getLogger("unit.deploy")
 
 
 def get_deploy_factory(provider_type):
     if provider_type == "local":
         return UnitContainerDeployment
     elif provider_type == "subordinate":
         return SubordinateContainerDeployment
 
     return UnitMachineDeployment
 
 
-def _get_environment(unit_name, juju_home, machine_id, zookeeper_hosts):
-    environ = dict()
-    environ["JUJU_MACHINE_ID"] = str(machine_id)
-    environ["JUJU_UNIT_NAME"] = unit_name
-    environ["JUJU_HOME"] = juju_home
-    environ["JUJU_ZOOKEEPER"] = zookeeper_hosts
-    environ["PYTHONPATH"] = ":".join(
+def _get_environment(unit_name, juju_home, machine_id, zookeeper_hosts, base=None):
+    if not base:
+        base = {}
+    base["JUJU_MACHINE_ID"] = str(machine_id)
+    base["JUJU_UNIT_NAME"] = unit_name
+    base["JUJU_HOME"] = juju_home
+    base["JUJU_ZOOKEEPER"] = zookeeper_hosts
+    base["PYTHONPATH"] = ":".join(
         filter(None, [
             os.path.dirname(get_module_directory(juju)),
             os.environ.get("PYTHONPATH")]))
-    return environ
-
-
-class UnitMachineDeployment(object):
+    return base
+
+
+class ExtraEnvironmentDeployment(object):
+
+    def __init__(self):
+        self._extra_environment = {}
+
+    def set_env(self, name, value):
+        self._extra_environment[name] = value
+
+    def unset_env(self, name):
+        del self._extra_environment[name]
+
+
+class UnitMachineDeployment(ExtraEnvironmentDeployment):
     """ Deploy a unit directly onto a machine.
 
     A service unit deployed directly to a machine has full access
     to the machine resources.
 
     Units deployed in such a manner have no isolation from other units
     on the machine, and may leave artifacts on the machine even upon
     service destruction.
     """
 
     unit_agent_module = "juju.agents.unit"
 
     def __init__(self, unit_name, juju_home):
+        super(UnitMachineDeployment, self).__init__()
         assert ".." not in unit_name, "Invalid Unit Name"
         self.unit_name = unit_name
         self.juju_home = juju_home
         self.unit_path_name = unit_name.replace("/", "-")
         self.directory = os.path.join(
             self.juju_home, "units", self.unit_path_name)
         self.service = UpstartService(
             # NOTE: we need use_sudo to work correctly during tests that
             # launch actual processes (rather than just mocking/trusting).
             "juju-%s" % self.unit_path_name, use_sudo=True)
 
     @inlineCallbacks
     def start(self, machine_id, zookeeper_hosts, bundle):
         """Start a service unit agent."""
         self.unpack_charm(bundle)
         self.service.set_description(
             "Juju unit agent for %s" % self.unit_name)
         self.service.set_environ(_get_environment(
-            self.unit_name, self.juju_home, machine_id, zookeeper_hosts))
+            self.unit_name, self.juju_home, machine_id, zookeeper_hosts, self._extra_environment))
         self.service.set_command(" ".join((
             "/usr/bin/python", "-m", self.unit_agent_module,
             "--nodaemon",
             "--logfile", os.path.join(self.directory, "charm.log"),
             "--session-file",
             "/var/run/juju/unit-%s-agent.zksession" % self.unit_path_name)))
         try:
             yield self.service.start()
         except ServiceError as e:
             raise UnitDeploymentError(str(e))
@@ skipping 37 matching lines @@
         self.juju_home = juju_home
         self.unit_path_name = unit_name.replace("/", "-")
         self.directory = os.path.join(
             self.juju_home, "units", self.unit_path_name)
         self.service = UpstartService(
             # NOTE: we need use_sudo to work correctly during tests that
             # launch actual processes (rather than just mocking/trusting).
             "juju-%s" % self.unit_path_name, use_sudo=True)
 
 
-class UnitContainerDeployment(object):
+class UnitContainerDeployment(ExtraEnvironmentDeployment):
     """Deploy a service unit in a container.
 
     Units deployed in a container have strong isolation between
     others units deployed in a container on the same machine.
 
     From the perspective of the service unit, the container deployment
     should be indistinguishable from a machine deployment.
 
     Note, strong isolation properties are still fairly trivial
     to escape for a user with a root account within the container.
     This is an ongoing development topic for LXC.
     """
 
     def __init__(self, unit_name, juju_home):
+        super(UnitContainerDeployment, self).__init__()
         self.unit_name = unit_name
         self.juju_home = juju_home
         self.unit_path_name = unit_name.replace("/", "-")
 
         self._juju_origin = os.environ.get("JUJU_ORIGIN")
         self._juju_series = os.environ.get("JUJU_SERIES")
         assert self._juju_series is not None, "Required juju series not found"
         self._unit_namespace = os.environ.get("JUJU_UNIT_NS")
         assert self._unit_namespace is not None, "Required unit ns not found"
         self.container_name = "%s-%s" % (
@@ skipping 24 matching lines @@
         authorized_keys = authorized_keys.strip("'\"")
 
         cloud_init.add_ssh_key(authorized_keys)
 
         zks = []
         for zk in zookeepers.split(','):
             if ':' in zk:
                 (zk, port) = zk.split(':')
             else:
                 port = 2181
-            zks.append((zk,port))
+            zks.append((zk, port))
 
         cloud_init.set_zookeeper_hosts(zks)
         # XXX Very hard to access the provider's notion of network
         # or even env configs, so just assume the first ZK is running
         # the apt-cacher-ng since this is meant for local provider.
         cloud_init.set_apt_proxy('http://%s:3142' % zks[0][0])
+
+        for name, value in self._extra_environment.iteritems():
+            cloud_init.set_env(name, value)
 
         if self._juju_origin:
             if self._juju_origin.startswith("lp:"):
                 cloud_init.set_juju_source(branch=self._juju_origin)
             elif self._juju_origin == "ppa":
                 cloud_init.set_juju_source(ppa=True)
             elif self._juju_origin == "proposed":
                 cloud_init.set_juju_source(proposed=True)
+            elif self._juju_origin == "image":
+                cloud_init.set_juju_source(image=True)
             else:
                 # Ignore other values, just use the distro for sanity
                 cloud_init.set_juju_source(distro=True)
         cloud_init.set_unit_name(self.unit_name)
         cloud_init.set_juju_home(self._container_juju_home)
         return cloud_init
 
     @inlineCallbacks
     def _get_container(self, machine_id, cloud_init):
         container = LXCContainer(
@@ skipping 70 matching lines @@
     @inlineCallbacks
     def is_running(self):
         """Is the unit container running?"""
         # TODO: container running may not imply agent running.
         # query zookeeper for the unit agent presence node?
         if not self.container:
             returnValue(False)
         container_map = yield get_containers(
             prefix=self.container.container_name)
         returnValue(container_map.get(self.container.container_name, False))