Continue working on merging prototype.

2 files changed, 38 insertions, 15 deletions

diff --git a/cloudinit/handlers/__init__.py b/cloudinit/handlers/__init__.py
index bfccfd89..566b61a7 100644
--- a/cloudinit/handlers/__init__.py
+++ b/cloudinit/handlers/__init__.py
@@ -92,14 +92,14 @@ def run_part(mod, data, filename, payload, headers, frequency):
     if not (mod_freq == PER_ALWAYS or
             (frequency == PER_INSTANCE and mod_freq == PER_INSTANCE)):
         return
-    mod_ver = mod.handler_version
     # Sanity checks on version (should be an int convertable)
     try:
+        mod_ver = mod.handler_version
         mod_ver = int(mod_ver)
-    except:
+    except (TypeError, ValueError, AttributeError):
         mod_ver = 1
-    content_type = headers['Content-Type']
     try:
+        content_type = headers['Content-Type']
         LOG.debug("Calling handler %s (%s, %s, %s) with frequency %s",
                   mod, content_type, filename, mod_ver, frequency)
         if mod_ver == 3:
@@ -110,9 +110,11 @@ def run_part(mod, data, filename, payload, headers, frequency):
             # Treat as v. 2 which does get a frequency
             mod.handle_part(data, content_type, filename,
                             payload, frequency)
-        else:
+        elif mod_ver == 1:
             # Treat as v. 1 which gets no frequency
             mod.handle_part(data, content_type, filename, payload)
+        else:
+            raise ValueError("Unknown module version %s" % (mod_ver))
     except:
         util.logexc(LOG, ("Failed calling handler %s (%s, %s, %s)"
                          " with frequency %s"),
@@ -121,11 +123,17 @@ def run_part(mod, data, filename, payload, headers, frequency):
 
 
 def call_begin(mod, data, frequency):
-    run_part(mod, data, CONTENT_START, None, None, frequency)
+    headers = {
+        'Content-Type': CONTENT_START,
+    }
+    run_part(mod, data, None, None, headers, frequency)
 
 
 def call_end(mod, data, frequency):
-    run_part(mod, data, CONTENT_END, None, None, frequency)
+    headers = {
+        'Content-Type': CONTENT_END,
+    }
+    run_part(mod, data, None, None, headers, frequency)
 
 
 def walker_handle_handler(pdata, _ctype, _filename, payload):
@@ -215,9 +223,9 @@ def walk(msg, callback, data):
         if not filename:
             filename = PART_FN_TPL % (partnum)
 
-        callback(data, ctype, filename,
-                 part.get_payload(decode=True),
-                 dict(part))
+        headers = dict(part)
+        headers['Content-Type'] = ctype
+        callback(data, filename, part.get_payload(decode=True), headers)
         partnum = partnum + 1
 
 
diff --git a/cloudinit/handlers/cloud_config.py b/cloudinit/handlers/cloud_config.py
index 86027187..22ced20d 100644
--- a/cloudinit/handlers/cloud_config.py
+++ b/cloudinit/handlers/cloud_config.py
@@ -29,6 +29,8 @@ from cloudinit.settings import (PER_ALWAYS)
 
 LOG = logging.getLogger(__name__)
 
+DEF_MERGE_TYPE = "list+dict+str"
+
 
 class CloudConfigPartHandler(handlers.Handler):
     def __init__(self, paths, **_kwargs):
@@ -44,10 +46,25 @@ class CloudConfigPartHandler(handlers.Handler):
     def _write_cloud_config(self, buf):
         if not self.cloud_fn:
             return
-        payload = util.yaml_dumps(self.cloud_buf)
-        util.write_file(self.cloud_fn, payload, 0600)
+        lines = ["#cloud-config", util.yaml_dumps(self.cloud_buf)]
+        util.write_file(self.cloud_fn, "\n".join(lines), 0600)
+
+    def _merge_part(self, payload, headers, filename):
+        merge_how = headers.get("Merge-Type")
+        try:
+            payload_y = util.load_yaml(payload)
+            if not merge_how:
+                merge_how = payload_y.pop("Merge-Type", '')
+            merge_how = merge_how.strip().lower()
+            if not merge_how:
+                merge_how = DEF_MERGE_TYPE
+            merger = mergers.construct(merge_how)
+            self.cloud_buf = merger.merge(self.cloud_buf, payload_y)
+        except:
+            util.logexc(LOG, "Failed at merging in cloud config part from %s",
+                        filename)
 
-    def handle_part(self, _data, ctype, filename, payload, _frequency, headers):
+    def handle_part(self, _data, ctype, filename, payload, _freq, headers):
         if ctype == handlers.CONTENT_START:
             self.cloud_buf = {}
             return
@@ -55,6 +72,4 @@ class CloudConfigPartHandler(handlers.Handler):
             self._write_cloud_config(self.cloud_buf)
             self.cloud_buf = {}
             return
-        merge_how = headers.get("Merge-Type", 'list+dict+str')
-        merger = mergers.construct(merge_how)
-        self.cloud_buf = merger.merge(self.cloud_buf, util.load_yaml(payload))
+        self._merge_part(payload, headers, filename)