Merge pull request #508 from github/nsone-country-filter-chain

Add country based filter to NS1 filter chain

octodns/provider/ns1.py

@@ -238,29 +238,102 @@ class Ns1Provider(BaseProvider):
 
     ZONE_NOT_FOUND_MESSAGE = 'server error: zone not found'
 
-    _DYNAMIC_FILTERS = [{
-        'config': {},
-        'filter': 'up'
-    }, {
-        'config': {},
-        'filter': u'geofence_regional'
-    }, {
-        'config': {},
-        'filter': u'select_first_region'
-    }, {
-        'config': {
-            'eliminate': u'1'
-        },
-        'filter': 'priority'
-    }, {
-        'config': {},
-        'filter': u'weighted_shuffle'
-    }, {
-        'config': {
-            'N': u'1'
-        },
-        'filter': u'select_first_n'
-    }]
+    def _update_filter(self, filter, with_disabled):
+        if with_disabled:
+            filter['disabled'] = False
+            return (dict(sorted(filter.items(), key=lambda t: t[0])))
+        return filter
+
+    def _UP_FILTER(self, with_disabled):
+        return self._update_filter({
+            'config': {},
+            'filter': 'up'
+        }, with_disabled)
+
+    def _REGION_FILTER(self, with_disabled):
+        return self._update_filter({
+            'config': {},
+            'filter': u'geofence_regional'
+        }, with_disabled)
+
+    def _COUNTRY_FILTER(self, with_disabled):
+        return self._update_filter({
+            'config': {
+                'remove_no_location': True
+            },
+            'filter': u'geofence_country'
+        }, with_disabled)
+
+    # In the NS1 UI/portal, this filter is called "SELECT FIRST GROUP" though
+    # the filter name in the NS1 api is 'select_first_region'
+    def _SELECT_FIRST_REGION_FILTER(self, with_disabled):
+        return self._update_filter({
+            'config': {},
+            'filter': u'select_first_region'
+        }, with_disabled)
+
+    def _PRIORITY_FILTER(self, with_disabled):
+        return self._update_filter({
+            'config': {
+                'eliminate': u'1'
+            },
+            'filter': 'priority'
+        }, with_disabled)
+
+    def _WEIGHTED_SHUFFLE_FILTER(self, with_disabled):
+        return self._update_filter({
+            'config': {},
+            'filter': u'weighted_shuffle'
+        }, with_disabled)
+
+    def _SELECT_FIRST_N_FILTER(self, with_disabled):
+        return self._update_filter({
+            'config': {
+                'N': u'1'
+            },
+            'filter': u'select_first_n'
+        }, with_disabled)
+
+    def _BASIC_FILTER_CHAIN(self, with_disabled):
+        return [
+            self._UP_FILTER(with_disabled),
+            self._SELECT_FIRST_REGION_FILTER(with_disabled),
+            self._PRIORITY_FILTER(with_disabled),
+            self._WEIGHTED_SHUFFLE_FILTER(with_disabled),
+            self._SELECT_FIRST_N_FILTER(with_disabled)
+        ]
+
+    def _FILTER_CHAIN_WITH_REGION(self, with_disabled):
+        return [
+            self._UP_FILTER(with_disabled),
+            self._REGION_FILTER(with_disabled),
+            self._SELECT_FIRST_REGION_FILTER(with_disabled),
+            self._PRIORITY_FILTER(with_disabled),
+            self._WEIGHTED_SHUFFLE_FILTER(with_disabled),
+            self._SELECT_FIRST_N_FILTER(with_disabled)
+        ]
+
+    def _FILTER_CHAIN_WITH_COUNTRY(self, with_disabled):
+        return [
+            self._UP_FILTER(with_disabled),
+            self._COUNTRY_FILTER(with_disabled),
+            self._SELECT_FIRST_REGION_FILTER(with_disabled),
+            self._PRIORITY_FILTER(with_disabled),
+            self._WEIGHTED_SHUFFLE_FILTER(with_disabled),
+            self._SELECT_FIRST_N_FILTER(with_disabled)
+        ]
+
+    def _FILTER_CHAIN_WITH_REGION_AND_COUNTRY(self, with_disabled):
+        return [
+            self._UP_FILTER(with_disabled),
+            self._REGION_FILTER(with_disabled),
+            self._COUNTRY_FILTER(with_disabled),
+            self._SELECT_FIRST_REGION_FILTER(with_disabled),
+            self._PRIORITY_FILTER(with_disabled),
+            self._WEIGHTED_SHUFFLE_FILTER(with_disabled),
+            self._SELECT_FIRST_N_FILTER(with_disabled)
+        ]
+
     _REGION_TO_CONTINENT = {
         'AFRICA': 'AF',
         'ASIAPAC': 'AS',
@@ -298,6 +371,29 @@ class Ns1Provider(BaseProvider):
         self._client = Ns1Client(api_key, parallelism, retry_count,
                                  client_config)
 
+    def _valid_filter_config(self, filter_cfg, domain):
+        with_disabled = self._disabled_flag_in_filters(filter_cfg, domain)
+        has_region = self._REGION_FILTER(with_disabled) in filter_cfg
+        has_country = self._COUNTRY_FILTER(with_disabled) in filter_cfg
+        expected_filter_cfg = self._get_updated_filter_chain(has_region,
+                                                             has_country,
+                                                             with_disabled)
+        return filter_cfg == expected_filter_cfg
+
+    def _get_updated_filter_chain(self, has_region, has_country,
+                                  with_disabled=True):
+        if has_region and has_country:
+            filter_chain = self._FILTER_CHAIN_WITH_REGION_AND_COUNTRY(
+                with_disabled)
+        elif has_region:
+            filter_chain = self._FILTER_CHAIN_WITH_REGION(with_disabled)
+        elif has_country:
+            filter_chain = self._FILTER_CHAIN_WITH_COUNTRY(with_disabled)
+        else:
+            filter_chain = self._BASIC_FILTER_CHAIN(with_disabled)
+
+        return filter_chain
+
     def _encode_notes(self, data):
         return ' '.join(['{}:{}'.format(k, v)
                          for k, v in sorted(data.items())])
@@ -358,7 +454,7 @@ class Ns1Provider(BaseProvider):
 
     def _data_for_dynamic_A(self, _type, record):
         # First make sure we have the expected filters config
-        if self._DYNAMIC_FILTERS != record['filters']:
+        if not self._valid_filter_config(record['filters'], record['domain']):
             self.log.error('_data_for_dynamic_A: %s %s has unsupported '
                            'filters', record['domain'], _type)
             raise Ns1Exception('Unrecognized advanced record')
@@ -843,6 +939,8 @@ class Ns1Provider(BaseProvider):
         pools = record.dynamic.pools
 
         # Convert rules to regions
+        has_country = False
+        has_region = False
         regions = {}
         for i, rule in enumerate(record.dynamic.rules):
             pool_name = rule.data['pool']
@@ -860,17 +958,23 @@ class Ns1Provider(BaseProvider):
             us_state = set()
 
             for geo in rule.data.get('geos', []):
+
                 n = len(geo)
                 if n == 8:
                     # US state, e.g. NA-US-KY
                     us_state.add(geo[-2:])
+                    # For filtering. State filtering is done by the country
+                    # filter
+                    has_country = True
                 elif n == 5:
                     # Country, e.g. EU-FR
                     country.add(geo[-2:])
+                    has_country = True
                 else:
                     # Continent, e.g. AS
                     if geo in self._CONTINENT_TO_REGIONS:
                         georegion.update(self._CONTINENT_TO_REGIONS[geo])
+                        has_region = True
                     else:
                         # No maps for geo in _CONTINENT_TO_REGIONS.
                         # Use the country list
@@ -878,6 +982,7 @@ class Ns1Provider(BaseProvider):
                                        format(geo))
                         for c in self._CONTINENT_TO_LIST_OF_COUNTRIES[geo]:
                             country.add(c)
+                            has_country = True
 
             meta = {
                 'note': self._encode_notes(notes),
@@ -965,9 +1070,12 @@ class Ns1Provider(BaseProvider):
                 }
                 answers.append(answer)
 
+        # Update filters as necessary
+        filters = self._get_updated_filter_chain(has_region, has_country)
+
         return {
             'answers': answers,
-            'filters': self._DYNAMIC_FILTERS,
+            'filters': filters,
             'regions': regions,
             'ttl': record.ttl,
         }, active_monitors
@@ -1020,17 +1128,64 @@ class Ns1Provider(BaseProvider):
                   for v in record.values]
         return {'answers': values, 'ttl': record.ttl}, None
 
+    def _get_ns1_filters(self, ns1_zone_name):
+        ns1_filters = {}
+        ns1_zone = {}
+
+        try:
+            ns1_zone = self._client.zones_retrieve(ns1_zone_name)
+        except ResourceException as e:
+            if e.message != self.ZONE_NOT_FOUND_MESSAGE:
+                raise
+
+        if 'records' in ns1_zone:
+            for ns1_record in ns1_zone['records']:
+                if ns1_record.get('tier', 1) > 1:
+                    # Need to get the full record data for geo records
+                    full_rec = self._client.records_retrieve(
+                        ns1_zone_name,
+                        ns1_record['domain'],
+                        ns1_record['type'])
+                    if 'filters' in full_rec:
+                        filter_key = '{}.'.format(ns1_record['domain'])
+                        ns1_filters[filter_key] = full_rec['filters']
+
+        return ns1_filters
+
+    def _disabled_flag_in_filters(self, filters, domain):
+        disabled_count = ['disabled' in f for f in filters].count(True)
+        if disabled_count and disabled_count != len(filters):
+            # Some filters have the disabled flag, and some don't. Disallow
+            exception_msg = 'Mixed disabled flag in filters for {}'.format(
+                            domain)
+            raise Ns1Exception(exception_msg)
+        return disabled_count == len(filters)
+
     def _extra_changes(self, desired, changes, **kwargs):
         self.log.debug('_extra_changes: desired=%s', desired.name)
-
+        ns1_filters = self._get_ns1_filters(desired.name[:-1])
         changed = set([c.record for c in changes])
-
         extra = []
         for record in desired.records:
             if record in changed or not getattr(record, 'dynamic', False):
                 # Already changed, or no dynamic , no need to check it
                 continue
 
+            # Filter normalization
+            # Check if filters for existing domains need an update
+            # Needs an explicit check since there might be no change in the
+            # config at all. Filters however might still need an update
+            domain = '{}.{}'.format(record.name, record.zone.name)
+            if domain in ns1_filters:
+                domain_filters = ns1_filters[domain]
+                if not self._disabled_flag_in_filters(domain_filters, domain):
+                    # 'disabled' entry absent in filter config. Need to update
+                    # filters. Update record
+                    self.log.info('_extra_changes: change in filters for %s',
+                                  domain)
+                    extra.append(Update(record, record))
+                    continue
+
             for have in self._monitors_for(record).values():
                 value = have['config']['host']
                 expected = self._monitor_gen(record, value)

tests/test_octodns_provider_ns1.py

@@ -316,6 +316,7 @@ class TestNs1Provider(TestCase):
 
         # Fails, general error
         zone_retrieve_mock.reset_mock()
+        record_retrieve_mock.reset_mock()
         zone_create_mock.reset_mock()
         zone_retrieve_mock.side_effect = ResourceException('boom')
         with self.assertRaises(ResourceException) as ctx:
@@ -324,6 +325,7 @@ class TestNs1Provider(TestCase):
 
         # Fails, bad auth
         zone_retrieve_mock.reset_mock()
+        record_retrieve_mock.reset_mock()
         zone_create_mock.reset_mock()
         zone_retrieve_mock.side_effect = \
             ResourceException('server error: zone not found')
@@ -334,6 +336,7 @@ class TestNs1Provider(TestCase):
 
         # non-existent zone, create
         zone_retrieve_mock.reset_mock()
+        record_retrieve_mock.reset_mock()
         zone_create_mock.reset_mock()
         zone_retrieve_mock.side_effect = \
             ResourceException('server error: zone not found')
@@ -364,6 +367,7 @@ class TestNs1Provider(TestCase):
 
         # Update & delete
         zone_retrieve_mock.reset_mock()
+        record_retrieve_mock.reset_mock()
         zone_create_mock.reset_mock()
 
         ns1_zone = {
@@ -389,7 +393,7 @@ class TestNs1Provider(TestCase):
         }
         ns1_zone['records'][0]['short_answers'][0] = '2.2.2.2'
 
-        record_retrieve_mock.side_effect = [{
+        ns1_record = {
             "domain": "geo.unit.tests",
             "zone": "unit.tests",
             "type": "A",
@@ -404,8 +408,9 @@ class TestNs1Provider(TestCase):
             ],
             'tier': 3,
             'ttl': 34,
-        }]
+        }
 
+        record_retrieve_mock.side_effect = [ns1_record, ns1_record]
         zone_retrieve_mock.side_effect = [ns1_zone, ns1_zone]
         plan = provider.plan(desired)
         self.assertEquals(3, len(plan.changes))
@@ -427,6 +432,8 @@ class TestNs1Provider(TestCase):
             None,
         ]
 
+        record_retrieve_mock.side_effect = [ns1_record, ns1_record]
+        zone_retrieve_mock.side_effect = [ns1_zone, ns1_zone]
         got_n = provider.apply(plan)
         self.assertEquals(3, got_n)
 
@@ -926,6 +933,42 @@ class TestNs1ProviderDynamic(TestCase):
         monitors_delete_mock.assert_has_calls([call('mon-id2')])
         notifylists_delete_mock.assert_has_calls([call('nl-id2')])
 
+    @patch('octodns.provider.ns1.Ns1Provider._monitor_sync')
+    @patch('octodns.provider.ns1.Ns1Provider._monitors_for')
+    def test_params_for_dynamic_region_only(self, monitors_for_mock,
+                                            monitor_sync_mock):
+        provider = Ns1Provider('test', 'api-key')
+
+        # pre-fill caches to avoid extranious calls (things we're testing
+        # elsewhere)
+        provider._client._datasource_id = 'foo'
+        provider._client._feeds_for_monitors = {
+            'mon-id': 'feed-id',
+        }
+
+        # provider._params_for_A() calls provider._monitors_for() and
+        # provider._monitor_sync(). Mock their return values so that we don't
+        # make NS1 API calls during tests
+        monitors_for_mock.reset_mock()
+        monitor_sync_mock.reset_mock()
+        monitors_for_mock.side_effect = [{
+            '3.4.5.6': 'mid-3',
+        }]
+        monitor_sync_mock.side_effect = [
+            ('mid-1', 'fid-1'),
+            ('mid-2', 'fid-2'),
+            ('mid-3', 'fid-3'),
+        ]
+
+        rule0 = self.record.data['dynamic']['rules'][0]
+        saved_geos = rule0['geos']
+        rule0['geos'] = ['AF', 'EU']
+        ret, _ = provider._params_for_A(self.record)
+        self.assertEquals(ret['filters'],
+                          Ns1Provider._FILTER_CHAIN_WITH_REGION(provider,
+                                                                True))
+        rule0['geos'] = saved_geos
+
     @patch('octodns.provider.ns1.Ns1Provider._monitor_sync')
     @patch('octodns.provider.ns1.Ns1Provider._monitors_for')
     def test_params_for_dynamic_oceania(self, monitors_for_mock,
@@ -961,6 +1004,11 @@ class TestNs1ProviderDynamic(TestCase):
         ret, _ = provider._params_for_A(self.record)
         self.assertEquals(set(ret['regions']['lhr']['meta']['country']),
                           Ns1Provider._CONTINENT_TO_LIST_OF_COUNTRIES['OC'])
+        # When rules has 'OC', it is converted to list of countries in the
+        # params. Look if the returned filters is the filter chain with country
+        self.assertEquals(ret['filters'],
+                          Ns1Provider._FILTER_CHAIN_WITH_COUNTRY(provider,
+                                                                 True))
         rule0['geos'] = saved_geos
 
     @patch('octodns.provider.ns1.Ns1Provider._monitor_sync')
@@ -987,7 +1035,14 @@ class TestNs1ProviderDynamic(TestCase):
         ]
         # This indirectly calls into _params_for_dynamic_A and tests the
         # handling to get there
-        provider._params_for_A(self.record)
+        ret, _ = provider._params_for_A(self.record)
+
+        # Given that self.record has both country and region in the rules,
+        # the returned filter chain should be one with region and country
+        self.assertEquals(ret['filters'],
+                          Ns1Provider._FILTER_CHAIN_WITH_REGION_AND_COUNTRY(
+                          provider, True))
+
         monitors_for_mock.assert_has_calls([call(self.record)])
         monitors_sync_mock.assert_has_calls([
             call(self.record, '1.2.3.4', None),
@@ -1012,7 +1067,7 @@ class TestNs1ProviderDynamic(TestCase):
         ns1_record = {
             'answers': [],
             'domain': 'unit.tests',
-            'filters': Ns1Provider._DYNAMIC_FILTERS,
+            'filters': Ns1Provider._BASIC_FILTER_CHAIN(provider, True),
             'regions': {},
             'ttl': 42,
         }
@@ -1028,6 +1083,8 @@ class TestNs1ProviderDynamic(TestCase):
         }, data)
 
         # Test out a small, but realistic setup that covers all the options
+        # We have country and region in the test config
+        filters = provider._get_updated_filter_chain(True, True)
         ns1_record = {
             'answers': [{
                 'answer': ['3.4.5.6'],
@@ -1068,7 +1125,7 @@ class TestNs1ProviderDynamic(TestCase):
                 'region': 'iad',
             }],
             'domain': 'unit.tests',
-            'filters': Ns1Provider._DYNAMIC_FILTERS,
+            'filters': filters,
             'regions': {
                 'lhr': {
                     'meta': {
@@ -1144,20 +1201,46 @@ class TestNs1ProviderDynamic(TestCase):
             self.assertTrue(
                 'OC-{}'.format(c) in data4['dynamic']['rules'][0]['geos'])
 
+    @patch('ns1.rest.records.Records.retrieve')
+    @patch('ns1.rest.zones.Zones.retrieve')
     @patch('octodns.provider.ns1.Ns1Provider._monitors_for')
-    def test_extra_changes(self, monitors_for_mock):
+    def test_extra_changes(self, monitors_for_mock, zones_retrieve_mock,
+                           records_retrieve_mock):
         provider = Ns1Provider('test', 'api-key')
 
         desired = Zone('unit.tests.', [])
 
         # Empty zone and no changes
         monitors_for_mock.reset_mock()
+        zones_retrieve_mock.reset_mock()
+        records_retrieve_mock.reset_mock()
+
         extra = provider._extra_changes(desired, [])
         self.assertFalse(extra)
         monitors_for_mock.assert_not_called()
 
-        # Simple record, ignored
+        # Non-existent zone. No changes
         monitors_for_mock.reset_mock()
+        zones_retrieve_mock.side_effect = \
+            ResourceException('server error: zone not found')
+        records_retrieve_mock.reset_mock()
+        extra = provider._extra_changes(desired, [])
+        self.assertFalse(extra)
+
+        # Unexpected exception message
+        zones_retrieve_mock.reset_mock()
+        zones_retrieve_mock.side_effect = ResourceException('boom')
+        with self.assertRaises(ResourceException) as ctx:
+            extra = provider._extra_changes(desired, [])
+        self.assertEquals(zones_retrieve_mock.side_effect, ctx.exception)
+
+        # Simple record, ignored, filter update lookups ignored
+        monitors_for_mock.reset_mock()
+        zones_retrieve_mock.reset_mock()
+        records_retrieve_mock.reset_mock()
+        zones_retrieve_mock.side_effect = \
+            ResourceException('server error: zone not found')
+
         simple = Record.new(desired, '', {
             'ttl': 32,
             'type': 'A',
@@ -1200,6 +1283,8 @@ class TestNs1ProviderDynamic(TestCase):
 
         # untouched, but everything in sync so no change needed
         monitors_for_mock.reset_mock()
+        zones_retrieve_mock.reset_mock()
+        records_retrieve_mock.reset_mock()
         # Generate what we expect to have
         gend = provider._monitor_gen(dynamic, '1.2.3.4')
         gend.update({
@@ -1218,6 +1303,8 @@ class TestNs1ProviderDynamic(TestCase):
         # If we don't have a notify list we're broken and we'll expect to see
         # an Update
         monitors_for_mock.reset_mock()
+        zones_retrieve_mock.reset_mock()
+        records_retrieve_mock.reset_mock()
         del gend['notify_list']
         monitors_for_mock.side_effect = [{
             '1.2.3.4': gend,
@@ -1232,6 +1319,8 @@ class TestNs1ProviderDynamic(TestCase):
         # Add notify_list back and change the healthcheck protocol, we'll still
         # expect to see an update
         monitors_for_mock.reset_mock()
+        zones_retrieve_mock.reset_mock()
+        records_retrieve_mock.reset_mock()
         gend['notify_list'] = 'xyz'
         dynamic._octodns['healthcheck']['protocol'] = 'HTTPS'
         del gend['notify_list']
@@ -1247,10 +1336,74 @@ class TestNs1ProviderDynamic(TestCase):
 
         # If it's in the changed list, it'll be ignored
         monitors_for_mock.reset_mock()
+        zones_retrieve_mock.reset_mock()
+        records_retrieve_mock.reset_mock()
         extra = provider._extra_changes(desired, [update])
         self.assertFalse(extra)
         monitors_for_mock.assert_not_called()
 
+        # Test changes in filters
+
+        # No change in filters
+        monitors_for_mock.reset_mock()
+        zones_retrieve_mock.reset_mock()
+        records_retrieve_mock.reset_mock()
+        ns1_zone = {
+            'records': [{
+                "domain": "dyn.unit.tests",
+                "zone": "unit.tests",
+                "type": "A",
+                "tier": 3,
+                "filters": Ns1Provider._BASIC_FILTER_CHAIN(provider, True)
+            }],
+        }
+        monitors_for_mock.side_effect = [{}]
+        zones_retrieve_mock.side_effect = [ns1_zone]
+        records_retrieve_mock.side_effect = ns1_zone['records']
+        extra = provider._extra_changes(desired, [])
+        self.assertFalse(extra)
+
+        # filters need an update
+        monitors_for_mock.reset_mock()
+        zones_retrieve_mock.reset_mock()
+        records_retrieve_mock.reset_mock()
+        ns1_zone = {
+            'records': [{
+                "domain": "dyn.unit.tests",
+                "zone": "unit.tests",
+                "type": "A",
+                "tier": 3,
+                "filters": Ns1Provider._BASIC_FILTER_CHAIN(provider, False)
+            }],
+        }
+        monitors_for_mock.side_effect = [{}]
+        zones_retrieve_mock.side_effect = [ns1_zone]
+        records_retrieve_mock.side_effect = ns1_zone['records']
+        extra = provider._extra_changes(desired, [])
+        self.assertTrue(extra)
+
+        # Mixed disabled in filters. Raise Ns1Exception
+        monitors_for_mock.reset_mock()
+        zones_retrieve_mock.reset_mock()
+        records_retrieve_mock.reset_mock()
+        ns1_zone = {
+            'records': [{
+                "domain": "dyn.unit.tests",
+                "zone": "unit.tests",
+                "type": "A",
+                "tier": 3,
+                "filters": Ns1Provider._BASIC_FILTER_CHAIN(provider, True)
+            }],
+        }
+        del ns1_zone['records'][0]['filters'][0]['disabled']
+        monitors_for_mock.side_effect = [{}]
+        zones_retrieve_mock.side_effect = [ns1_zone]
+        records_retrieve_mock.side_effect = ns1_zone['records']
+        with self.assertRaises(Ns1Exception) as ctx:
+            extra = provider._extra_changes(desired, [])
+        self.assertTrue('Mixed disabled flag in filters' in
+                        text_type(ctx.exception))
+
     DESIRED = Zone('unit.tests.', [])
 
     SIMPLE = Record.new(DESIRED, 'sim', {