Merge pull request alphagov#72 from alphagov/improve_puppet_class

samjsharpe · samjsharpe · commit 66f86b800650 · 2014-10-30T09:15:36.000Z
Improve puppet_class and deprecate puppet.loadhosts
diff --git a/fabfile.py b/fabfile.py
@@ -8,8 +8,8 @@
 import time
 
 from fabric import state
-from fabric.api import (abort, env, get, hide, local, puts, run, runs_once,
-                        serial, settings, sudo, task, warn)
+from fabric.api import (abort, env, get, hide, hosts, local, puts, run,
+                        runs_once, serial, settings, sudo, task, warn)
 from fabric.task_utils import crawl
 
 # Our command submodules
@@ -103,6 +103,15 @@ def fetch(self):
 
         self.fetched = True
 
+    def fetch_puppet_class(self, name):
+        # These cannot be prefetched because there's too many variations.
+        # But we only need to fetch once for each.
+        if self.roledefs['puppet_class-%s' % name]:
+            return
+
+        hosts = _fetch_hosts('-C %s' % name)
+        self.roledefs['puppet_class-%s' % name] = hosts
+
     def __contains__(self, key):
         return True
 
@@ -235,12 +244,11 @@ def klass(class_name):
     env.hosts.extend(env.roledefs['class-%s' % class_name]())
 
 @task
+@hosts('localhost')
 def puppet_class(class_name):
     """Select all machines which include a given puppet class"""
-    hosts = _fetch_hosts('-C %s' % class_name)
-    if (hosts == []):
-        abort("Didn't find any hosts with class %s" % class_name)
-    env.hosts.extend(sorted(hosts))
+    env.roledefs.fetch_puppet_class(class_name)
+    env.hosts.extend(env.roledefs['puppet_class-%s' % class_name]())
 
 @task
 def vdc(vdc_name):
diff --git a/puppet.py b/puppet.py
@@ -2,11 +2,9 @@
 
 @task
 def loadhosts(*classnames):
-    """Load hosts that include given Puppet class(es)"""
-    classnames = ','.join(classnames)
-    with hide('running', 'stdout'):
-          with settings(host_string=env.gateway, gateway=None):
-              env.hosts = run('govuk_node_list -C %s' % classnames).splitlines()
+    """Deprecated, use puppet_class"""
+    usage = ["puppet_class:{0}".format(name) for name in classnames]
+    abort("puppet.loadhosts is deprecated, use: {0}".format(" ".join(usage)))
 
 def puppet(*args):
     sudo('govuk_puppet %s' % ' '.join(args))
