Add OverridingYamlProvider and tests

octodns/provider/yaml.py

@@ -47,7 +47,7 @@ class YamlProvider(BaseProvider):
         self.default_ttl = default_ttl
         self.enforce_order = enforce_order
 
-    def _populate_from_file(self, filename, zone, lenient):
+    def _populate_from_file(self, filename, zone, lenient, replace=False):
         with open(filename, 'r') as fh:
             yaml_data = safe_load(fh, enforce_order=self.enforce_order)
             if yaml_data:
@@ -59,9 +59,10 @@ class YamlProvider(BaseProvider):
                             d['ttl'] = self.default_ttl
                         record = Record.new(zone, name, d, source=self,
                                             lenient=lenient)
-                        zone.add_record(record, lenient=lenient)
-            self.log.debug(
-                '_populate_from_file: successfully loaded "%s"', filename)
+                        zone.add_record(record, lenient=lenient,
+                                        replace=replace)
+            self.log.debug('_populate_from_file: successfully loaded "%s"',
+                           filename)
 
     def populate(self, zone, target=False, lenient=False):
         self.log.debug('populate: name=%s, target=%s, lenient=%s', zone.name,
@@ -211,3 +212,54 @@ class SplitYamlProvider(YamlProvider):
             self.log.debug('_apply:   writing catchall filename=%s', filename)
             with open(filename, 'w') as fh:
                 safe_dump(catchall, fh)
+
+
+class OverridingYamlProvider(YamlProvider):
+    '''
+    Provider that builds on YamlProvider to allow overriding specific records.
+
+    Works identically to YamlProvider with the additional behavior of loading
+    data from a second zonefile in override_directory if it exists. Records in
+    this second file will override (replace) those previously seen in the
+    primary. Records that do not exist in the primary will just be added. There
+    is currently no mechinism to remove records from the primary zone.
+
+    config:
+        class: octodns.provider.yaml.OverridingYamlProvider
+        # The location of yaml config files (required)
+        directory: ./config
+        # The location of overriding yaml config files (required)
+        override_directory: ./config
+        # The ttl to use for records when not specified in the data
+        # (optional, default 3600)
+        default_ttl: 3600
+        # Whether or not to enforce sorting order on the yaml config
+        # (optional, default True)
+        enforce_order: True
+    '''
+
+    def __init__(self, id, directory, override_directory, *args, **kwargs):
+        super(OverridingYamlProvider, self).__init__(id, directory, *args,
+                                                     **kwargs)
+        self.override_directory = override_directory
+
+    def populate(self, zone, target=False, lenient=False):
+        self.log.debug('populate: name=%s, target=%s, lenient=%s', zone.name,
+                       target, lenient)
+
+        if target:
+            # When acting as a target we ignore any existing records so that we
+            # create a completely new copy
+            return False
+
+        before = len(zone.records)
+        filename = join(self.directory, '{}yaml'.format(zone.name))
+        self._populate_from_file(filename, zone, lenient)
+
+        filename = join(self.override_directory, '{}yaml'.format(zone.name))
+        if isfile(filename):
+            self._populate_from_file(filename, zone, lenient, replace=True)
+
+        self.log.info('populate:   found %s records, exists=False',
+                      len(zone.records) - before)
+        return False

tests/config/override/dynamic.tests.yaml

@@ -0,0 +1,13 @@
+---
+# Replace 'a' with a generic record
+a:
+  type: A
+  values:
+    - 4.4.4.4
+    - 5.5.5.5
+# Add another record
+added:
+  type: A
+  values:
+    - 6.6.6.6
+    - 7.7.7.7

tests/test_octodns_provider_yaml.py

@@ -14,7 +14,7 @@ from yaml.constructor import ConstructorError
 from octodns.record import Create
 from octodns.provider.base import Plan
 from octodns.provider.yaml import _list_all_yaml_files, \
-    SplitYamlProvider, YamlProvider
+    OverridingYamlProvider, SplitYamlProvider, YamlProvider
 from octodns.zone import SubzoneRecordException, Zone
 
 from helpers import TemporaryDirectory
@@ -372,3 +372,38 @@ class TestSplitYamlProvider(TestCase):
             source.populate(zone)
         self.assertEquals('Record www.sub.unit.tests. is under a managed '
                           'subzone', ctx.exception.message)
+
+
+class TestOverridingYamlProvider(TestCase):
+
+    def test_provider(self):
+        config = join(dirname(__file__), 'config')
+        override_config = join(dirname(__file__), 'config', 'override')
+        source = OverridingYamlProvider('test', config, override_config)
+
+        zone = Zone('unit.tests.', [])
+        dynamic_zone = Zone('dynamic.tests.', [])
+
+        # With target we don't add anything (same as base)
+        source.populate(zone, target=source)
+        self.assertEquals(0, len(zone.records))
+
+        # without it we see everything
+        source.populate(zone)
+        self.assertEquals(18, len(zone.records))
+
+        # Load the dynamic records
+        source.populate(dynamic_zone)
+
+        got = {r.name: r for r in dynamic_zone.records}
+        # We see both the base and override files, 1 extra record
+        self.assertEquals(6, len(got))
+
+        # 'a' was replaced with a generic record
+        self.assertEquals({
+            'ttl': 3600,
+            'values': ['4.4.4.4', '5.5.5.5']
+        }, got['a'].data)
+
+        # And we have a new override
+        self.assertTrue('added' in got)