Merge pull request #164 from nsone/master

Add geo support for NS1 provider

octodns/provider/ns1.py

@@ -6,8 +6,11 @@ from __future__ import absolute_import, division, print_function, \
     unicode_literals
 
 from logging import getLogger
+from itertools import chain
+from collections import OrderedDict, defaultdict
 from nsone import NSONE
 from nsone.rest.errors import RateLimitException, ResourceException
+from incf.countryutils import transformations
 from time import sleep
 
 from ..record import Record
@@ -22,9 +25,9 @@ class Ns1Provider(BaseProvider):
         class: octodns.provider.ns1.Ns1Provider
         api_key: env/NS1_API_KEY
     '''
-    SUPPORTS_GEO = False
-    SUPPORTS = set(('A', 'AAAA', 'ALIAS', 'CAA', 'CNAME', 'MX', 'NAPTR', 'NS',
-                    'PTR', 'SPF', 'SRV', 'TXT'))
+    SUPPORTS_GEO = True
+    SUPPORTS = set(('A', 'AAAA', 'ALIAS', 'CAA', 'CNAME', 'MX', 'NAPTR',
+                    'NS', 'PTR', 'SPF', 'SRV', 'TXT'))
 
     ZONE_NOT_FOUND_MESSAGE = 'server error: zone not found'
 
@@ -35,11 +38,50 @@ class Ns1Provider(BaseProvider):
         self._client = NSONE(apiKey=api_key)
 
     def _data_for_A(self, _type, record):
-        return {
+        # record meta (which would include geo information is only
+        # returned when getting a record's detail, not from zone detail
+        geo = defaultdict(list)
+        data = {
             'ttl': record['ttl'],
             'type': _type,
-            'values': record['short_answers'],
         }
+        values, codes = [], []
+        if 'answers' not in record:
+            values = record['short_answers']
+        for answer in record.get('answers', []):
+            meta = answer.get('meta', {})
+            if meta:
+                # country + state and country + province are allowed
+                # in that case though, supplying a state/province would
+                # be redundant since the country would supercede in when
+                # resolving the record.  it is syntactically valid, however.
+                country = meta.get('country', [])
+                us_state = meta.get('us_state', [])
+                ca_province = meta.get('ca_province', [])
+                for cntry in country:
+                    cn = transformations.cc_to_cn(cntry)
+                    con = transformations.cn_to_ctca2(cn)
+                    key = '{}-{}'.format(con, cntry)
+                    geo[key].extend(answer['answer'])
+                for state in us_state:
+                    key = 'NA-US-{}'.format(state)
+                    geo[key].extend(answer['answer'])
+                for province in ca_province:
+                    key = 'NA-CA-{}'.format(province)
+                    geo[key].extend(answer['answer'])
+                for code in meta.get('iso_region_code', []):
+                    key = code
+                    geo[key].extend(answer['answer'])
+            else:
+                values.extend(answer['answer'])
+                codes.append([])
+        values = [str(x) for x in values]
+        geo = OrderedDict(
+            {str(k): [str(x) for x in v] for k, v in geo.items()}
+        )
+        data['values'] = values
+        data['geo'] = geo
+        return data
 
     _data_for_AAAA = _data_for_A
 
@@ -140,39 +182,79 @@ class Ns1Provider(BaseProvider):
         }
 
     def populate(self, zone, target=False, lenient=False):
-        self.log.debug('populate: name=%s, target=%s, lenient=%s', zone.name,
+        self.log.debug('populate: name=%s, target=%s, lenient=%s',
+                       zone.name,
                        target, lenient)
 
         try:
             nsone_zone = self._client.loadZone(zone.name[:-1])
             records = nsone_zone.data['records']
+            geo_records = nsone_zone.search(has_geo=True)
         except ResourceException as e:
             if e.message != self.ZONE_NOT_FOUND_MESSAGE:
                 raise
             records = []
+            geo_records = []
 
         before = len(zone.records)
-        for record in records:
+        # geo information isn't returned from the main endpoint, so we need
+        # to query for all records with geo information
+        zone_hash = {}
+        for record in chain(records, geo_records):
             _type = record['type']
             data_for = getattr(self, '_data_for_{}'.format(_type))
             name = zone.hostname_from_fqdn(record['domain'])
             record = Record.new(zone, name, data_for(_type, record),
                                 source=self, lenient=lenient)
-            zone.add_record(record)
-
+            zone_hash[(_type, name)] = record
+        [zone.add_record(r) for r in zone_hash.values()]
         self.log.info('populate:   found %s records',
                       len(zone.records) - before)
 
     def _params_for_A(self, record):
-        return {'answers': record.values, 'ttl': record.ttl}
+        params = {'answers': record.values, 'ttl': record.ttl}
+        if hasattr(record, 'geo'):
+            # purposefully set non-geo answers to have an empty meta,
+            # so that we know we did this on purpose if/when troubleshooting
+            params['answers'] = [{"answer": [x], "meta": {}}
+                                 for x in record.values]
+            has_country = False
+            for iso_region, target in record.geo.items():
+                key = 'iso_region_code'
+                value = iso_region
+                if not has_country and \
+                   len(value.split('-')) > 1:  # pragma: nocover
+                    has_country = True
+                for answer in target.values:
+                    params['answers'].append(
+                        {
+                            'answer': [answer],
+                            'meta': {key: [value]},
+                        },
+                    )
+            params['filters'] = []
+            if len(params['answers']) > 1:
+                params['filters'].append(
+                    {"filter": "shuffle", "config": {}}
+                )
+            if has_country:
+                params['filters'].append(
+                    {"filter": "geotarget_country", "config": {}}
+                )
+                params['filters'].append(
+                    {"filter": "select_first_n",
+                     "config": {"N": 1}}
+                )
+        self.log.debug("params for A: %s", params)
+        return params
 
     _params_for_AAAA = _params_for_A
     _params_for_NS = _params_for_A
 
     def _params_for_SPF(self, record):
-        # NS1 seems to be the only provider that doesn't want things escaped in
-        # values so we have to strip them here and add them when going the
-        # other way
+        # NS1 seems to be the only provider that doesn't want things
+        # escaped in values so we have to strip them here and add
+        # them when going the other way
         values = [v.replace('\;', ';') for v in record.values]
         return {'answers': values, 'ttl': record.ttl}
 
@@ -264,4 +346,5 @@ class Ns1Provider(BaseProvider):
 
         for change in changes:
             class_name = change.__class__.__name__
-            getattr(self, '_apply_{}'.format(class_name))(nsone_zone, change)
+            getattr(self, '_apply_{}'.format(class_name))(nsone_zone,
+                                                          change)

setup.cfg

@@ -19,7 +19,7 @@ classifiers =
     Programming Language :: Python :: 3.6
 
 [options]
-install_requires = 
+install_requires =
     PyYaml>=3.12
     dnspython>=1.15.0
     futures>=3.1.1
@@ -32,7 +32,7 @@ packages = find:
 include_package_data = True
 
 [options.entry_points]
-console_scripts = 
+console_scripts =
     octodns-compare = octodns.cmds.compare:main
     octodns-dump = octodns.cmds.dump:main
     octodns-report = octodns.cmds.report:main
@@ -44,7 +44,7 @@ exclude =
     tests
 
 [options.extras_require]
-dev = 
+dev =
     azure-mgmt-dns==1.0.1
     azure-common==1.1.6
     boto3>=1.4.6
@@ -54,7 +54,7 @@ dev =
     google-cloud>=0.27.0
     jmespath>=0.9.3
     msrestazure==0.4.10
-    nsone>=0.9.14
+    nsone>=0.9.17
     ovh>=0.4.7
     s3transfer>=0.1.10
     six>=1.10.0

tests/test_octodns_provider_ns1.py

@@ -30,11 +30,20 @@ class TestNs1Provider(TestCase):
         'ttl': 32,
         'type': 'A',
         'value': '1.2.3.4',
+        'meta': {},
     }))
     expected.add(Record.new(zone, 'foo', {
         'ttl': 33,
         'type': 'A',
         'values': ['1.2.3.4', '1.2.3.5'],
+        'meta': {},
+    }))
+    expected.add(Record.new(zone, 'geo', {
+        'ttl': 34,
+        'type': 'A',
+        'values': ['101.102.103.104', '101.102.103.105'],
+        'geo': {'NA-US-NY': ['201.202.203.204']},
+        'meta': {},
     }))
     expected.add(Record.new(zone, 'cname', {
         'ttl': 34,
@@ -116,6 +125,11 @@ class TestNs1Provider(TestCase):
         'ttl': 33,
         'short_answers': ['1.2.3.4', '1.2.3.5'],
         'domain': 'foo.unit.tests.',
+    }, {
+        'type': 'A',
+        'ttl': 34,
+        'short_answers': ['101.102.103.104', '101.102.103.105'],
+        'domain': 'geo.unit.tests',
     }, {
         'type': 'CNAME',
         'ttl': 34,
@@ -190,15 +204,53 @@ class TestNs1Provider(TestCase):
         load_mock.reset_mock()
         nsone_zone = DummyZone([])
         load_mock.side_effect = [nsone_zone]
+        zone_search = Mock()
+        zone_search.return_value = [
+            {
+                "domain": "geo.unit.tests",
+                "zone": "unit.tests",
+                "type": "A",
+                "answers": [
+                    {'answer': ['1.1.1.1'], 'meta': {}},
+                    {'answer': ['1.2.3.4'],
+                     'meta': {'ca_province': ['ON']}},
+                    {'answer': ['2.3.4.5'], 'meta': {'us_state': ['NY']}},
+                    {'answer': ['3.4.5.6'], 'meta': {'country': ['US']}},
+                    {'answer': ['4.5.6.7'],
+                     'meta': {'iso_region_code': ['NA-US-WA']}},
+                ],
+                'ttl': 34,
+            },
+        ]
+        nsone_zone.search = zone_search
         zone = Zone('unit.tests.', [])
         provider.populate(zone)
-        self.assertEquals(set(), zone.records)
+        self.assertEquals(1, len(zone.records))
         self.assertEquals(('unit.tests',), load_mock.call_args[0])
 
         # Existing zone w/records
         load_mock.reset_mock()
         nsone_zone = DummyZone(self.nsone_records)
         load_mock.side_effect = [nsone_zone]
+        zone_search = Mock()
+        zone_search.return_value = [
+            {
+                "domain": "geo.unit.tests",
+                "zone": "unit.tests",
+                "type": "A",
+                "answers": [
+                    {'answer': ['1.1.1.1'], 'meta': {}},
+                    {'answer': ['1.2.3.4'],
+                     'meta': {'ca_province': ['ON']}},
+                    {'answer': ['2.3.4.5'], 'meta': {'us_state': ['NY']}},
+                    {'answer': ['3.4.5.6'], 'meta': {'country': ['US']}},
+                    {'answer': ['4.5.6.7'],
+                     'meta': {'iso_region_code': ['NA-US-WA']}},
+                ],
+                'ttl': 34,
+            },
+        ]
+        nsone_zone.search = zone_search
         zone = Zone('unit.tests.', [])
         provider.populate(zone)
         self.assertEquals(self.expected, zone.records)
@@ -264,11 +316,30 @@ class TestNs1Provider(TestCase):
         }])
         nsone_zone.data['records'][0]['short_answers'][0] = '2.2.2.2'
         nsone_zone.loadRecord = Mock()
+        zone_search = Mock()
+        zone_search.return_value = [
+            {
+                "domain": "geo.unit.tests",
+                "zone": "unit.tests",
+                "type": "A",
+                "answers": [
+                    {'answer': ['1.1.1.1'], 'meta': {}},
+                    {'answer': ['1.2.3.4'],
+                     'meta': {'ca_province': ['ON']}},
+                    {'answer': ['2.3.4.5'], 'meta': {'us_state': ['NY']}},
+                    {'answer': ['3.4.5.6'], 'meta': {'country': ['US']}},
+                    {'answer': ['4.5.6.7'],
+                     'meta': {'iso_region_code': ['NA-US-WA']}},
+                ],
+                'ttl': 34,
+            },
+        ]
+        nsone_zone.search = zone_search
         load_mock.side_effect = [nsone_zone, nsone_zone]
         plan = provider.plan(desired)
-        self.assertEquals(2, len(plan.changes))
+        self.assertEquals(3, len(plan.changes))
         self.assertIsInstance(plan.changes[0], Update)
-        self.assertIsInstance(plan.changes[1], Delete)
+        self.assertIsInstance(plan.changes[2], Delete)
         # ugh, we need a mock record that can be returned from loadRecord for
         # the update and delete targets, we can add our side effects to that to
         # trigger rate limit handling
@@ -276,26 +347,52 @@ class TestNs1Provider(TestCase):
         mock_record.update.side_effect = [
             RateLimitException('one', period=0),
             None,
+            None,
         ]
         mock_record.delete.side_effect = [
             RateLimitException('two', period=0),
             None,
+            None,
         ]
-        nsone_zone.loadRecord.side_effect = [mock_record, mock_record]
+        nsone_zone.loadRecord.side_effect = [mock_record, mock_record,
+                                             mock_record]
         got_n = provider.apply(plan)
-        self.assertEquals(2, got_n)
+        self.assertEquals(3, got_n)
         nsone_zone.loadRecord.assert_has_calls([
             call('unit.tests', u'A'),
+            call('geo', u'A'),
             call('delete-me', u'A'),
         ])
         mock_record.assert_has_calls([
-            call.update(answers=[u'1.2.3.4'], ttl=32),
+            call.update(answers=[{'answer': [u'1.2.3.4'], 'meta': {}}],
+                        filters=[],
+                        ttl=32),
+            call.update(answers=[{u'answer': [u'1.2.3.4'], u'meta': {}}],
+                        filters=[],
+                        ttl=32),
+            call.update(
+                answers=[
+                    {u'answer': [u'101.102.103.104'], u'meta': {}},
+                    {u'answer': [u'101.102.103.105'], u'meta': {}},
+                    {
+                        u'answer': [u'201.202.203.204'],
+                        u'meta': {
+                            u'iso_region_code': [u'NA-US-NY']
+                        },
+                    },
+                ],
+                filters=[
+                    {u'filter': u'shuffle', u'config': {}},
+                    {u'filter': u'geotarget_country', u'config': {}},
+                    {u'filter': u'select_first_n', u'config': {u'N': 1}},
+                ],
+                ttl=34),
+            call.delete(),
             call.delete()
         ])
 
     def test_escaping(self):
         provider = Ns1Provider('test', 'api-key')
-
         record = {
             'ttl': 31,
             'short_answers': ['foo; bar baz; blip']