Add commands transforming alpha3 country codes to alpha2

Countries are stored in ping_v1 and activation_v1 tables. As many third-party
tools use 2-letter ISO codes, we use these commands to store alpha2 instead of
alpha3 codes.

azafea/event_processors/endless/activation/tests/integration/test_activation_cli.py

@@ -216,3 +216,61 @@ def test_parse_old_unknown_images(self, capfd):
 
         capture = capfd.readouterr()
         assert 'No activation record with unparsed image ids' in capture.out
+
+    def test_transform_countries_alpha_3_to_2(self, capfd):
+        from azafea.event_processors.endless.activation.v1.handler import Activation
+
+        # Create the table
+        self.run_subcommand('initdb')
+        self.ensure_tables(Activation)
+
+        created_at = datetime.utcnow().replace(tzinfo=timezone.utc)
+
+        with self.db as dbsession:
+            activation = Activation(
+                image='unknown', product='product', release='release', country='FR',
+                created_at=created_at, vendor='vendor')
+            dbsession.add(activation)
+
+        with self.db as dbsession:
+            # Bypass ORM country validation
+            dbsession.execute("UPDATE activation_v1 SET country='FRA'")
+
+        with self.db as dbsession:
+            activation = dbsession.query(Activation).one()
+            assert activation.country == 'FRA'
+
+        # Parse the image for old activation records
+        self.run_subcommand(
+            'test_transform_countries_alpha_3_to_2',
+            'transform-countries-alpha-3-to-2')
+
+        with self.db as dbsession:
+            activation = dbsession.query(Activation).one()
+            assert activation.country == 'FR'
+
+    def test_transform_countries_alpha_3_to_2_only_2(self, capfd):
+        from azafea.event_processors.endless.activation.v1.handler import Activation
+
+        # Create the table
+        self.run_subcommand('initdb')
+        self.ensure_tables(Activation)
+
+        created_at = datetime.utcnow().replace(tzinfo=timezone.utc)
+
+        with self.db as dbsession:
+            activation = Activation(
+                image='unknown', product='product', release='release', country='FR',
+                created_at=created_at, vendor='vendor')
+            dbsession.add(activation)
+
+        self.run_subcommand(
+            'test_transform_countries_alpha_3_to_2_only_2',
+            'transform-countries-alpha-3-to-2')
+
+        with self.db as dbsession:
+            activation = dbsession.query(Activation).one()
+            assert activation.country == 'FR'
+
+        capture = capfd.readouterr()
+        assert 'No activation with alpha-3 country code found' in capture.out

azafea/event_processors/endless/activation/v1/cli.py

@@ -10,12 +10,14 @@
 import logging
 
 from sqlalchemy.orm.query import Query
+from sqlalchemy.sql.expression import func
 
 from azafea.config import Config
 from azafea.model import Db
 from azafea.utils import progress
 from azafea.vendors import normalize_vendor
 
+from ...country import tranform_alpha_3_into_2
 from ...image import parse_endless_os_image
 from .handler import Activation
 
@@ -38,6 +40,14 @@ def register_commands(subs: argparse._SubParsersAction) -> None:
                               help='The size of the chunks to operate on')
     parse_images.set_defaults(subcommand=do_parse_images)
 
+    transform_countries_alpha_3_to_2 = subs.add_parser(
+        'transform-countries-alpha-3-to-2',
+        help='Transform 3-letter country codes to 2-letter country codes',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    transform_countries_alpha_3_to_2.add_argument('--chunk-size', type=int, default=5000,
+                                                  help='The size of the chunks to operate on')
+    transform_countries_alpha_3_to_2.set_defaults(subcommand=do_transform_countries_alpha_3_to_2)
+
 
 def _normalize_chunk(chunk: Query) -> None:
     for record in chunk:
@@ -100,3 +110,27 @@ def do_parse_images(config: Config, args: argparse.Namespace) -> None:
     progress(num_records, num_records, end='\n')
 
     log.info('All done!')
+
+
+def do_transform_countries_alpha_3_to_2(config: Config, args: argparse.Namespace) -> None:
+    db = Db(config.postgresql)
+    log.info('Tranform alpha-3 country codes into alpha-2 for activations')
+
+    with db as dbsession:
+        query = dbsession.chunked_query(Activation, chunk_size=args.chunk_size)
+        query = query.filter(func.length(Activation.country) == 3)
+        num_records = query.count()
+
+        if num_records == 0:
+            log.info('-> No activation with alpha-3 country code found')
+            return None
+
+        for chunk_number, chunk in enumerate(query, start=1):
+            for activation in chunk:
+                activation.country = tranform_alpha_3_into_2(activation.country)
+            dbsession.commit()
+            progress(chunk_number * args.chunk_size, num_records)
+
+    progress(num_records, num_records, end='\n')
+
+    log.info('All done!')

azafea/event_processors/endless/country.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2020 - Endless
+#
+# This file is part of Azafea
+#
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+COUNTRIES = {
+    'ABW': 'AW', 'AGO': 'AO', 'ALA': 'AX', 'AND': 'AD', 'ARE': 'AE', 'ARM': 'AM', 'ATA': 'AQ',
+    'ATF': 'TF', 'ATG': 'AG', 'AUT': 'AT', 'BDI': 'BI', 'BEN': 'BJ', 'BES': 'BQ', 'BGD': 'BD',
+    'BHS': 'BS', 'BIH': 'BA', 'BLR': 'BY', 'BLZ': 'BZ', 'BRB': 'BB', 'BRN': 'BN', 'CAF': 'CF',
+    'CHL': 'CL', 'CHN': 'CN', 'COD': 'CD', 'COG': 'CG', 'COK': 'CK', 'COM': 'KM', 'CPV': 'CV',
+    'CUW': 'CW', 'CYM': 'KY', 'DNK': 'DK', 'ESH': 'EH', 'EST': 'EE', 'FLK': 'FK', 'FRO': 'FO',
+    'FSM': 'FM', 'GIN': 'GN', 'GLP': 'GP', 'GNB': 'GW', 'GNQ': 'GQ', 'GRD': 'GD', 'GRL': 'GL',
+    'GUF': 'GF', 'GUY': 'GY', 'IRL': 'IE', 'IRQ': 'IQ', 'ISR': 'IL', 'JAM': 'JM', 'KAZ': 'KZ',
+    'KOR': 'KR', 'LBR': 'LR', 'LBY': 'LY', 'MAC': 'MO', 'MAF': 'MF', 'MDG': 'MG', 'MDV': 'MV',
+    'MEX': 'MX', 'MLT': 'MT', 'MNE': 'ME', 'MNP': 'MP', 'MOZ': 'MZ', 'MTQ': 'MQ', 'MYT': 'YT',
+    'NIU': 'NU', 'PAK': 'PK', 'PCN': 'PN', 'PLW': 'PW', 'PNG': 'PG', 'POL': 'PL', 'PRK': 'KP',
+    'PRT': 'PT', 'PRY': 'PY', 'PYF': 'PF', 'SEN': 'SN', 'SGS': 'GS', 'SLB': 'SB', 'SLV': 'SV',
+    'SPM': 'PM', 'SRB': 'RS', 'SUR': 'SR', 'SVK': 'SK', 'SVN': 'SI', 'SWE': 'SE', 'SWZ': 'SZ',
+    'SYC': 'SC', 'TCD': 'TD', 'TKM': 'TM', 'TUN': 'TN', 'TUR': 'TR', 'TUV': 'TV', 'UKR': 'UA',
+    'URY': 'UY', 'WLF': 'WF'}
+
+
+def tranform_alpha_3_into_2(alpha_3):
+    return COUNTRIES.get(alpha_3, alpha_3[:2])

azafea/event_processors/endless/ping/tests/integration/test_ping_cli.py

@@ -346,3 +346,64 @@ def test_parse_old_unknown_images(self, capfd):
 
         capture = capfd.readouterr()
         assert 'No ping record with unparsed image ids' in capture.out
+
+    def test_transform_countries_alpha_3_to_2(self, capfd):
+        from azafea.event_processors.endless.ping.v1.handler import Ping, PingConfiguration
+
+        # Create the table
+        self.run_subcommand('initdb')
+        self.ensure_tables(Ping)
+
+        created_at = datetime.utcnow().replace(tzinfo=timezone.utc)
+
+        with self.db as dbsession:
+            ping_config = PingConfiguration(image='eos-eos3.7-amd64-amd64.190419-225606.base',
+                                            product='product', dualboot=True, created_at=created_at,
+                                            vendor='EnDlEsS')
+            dbsession.add(ping_config)
+            dbsession.add(Ping(config=ping_config, release='release', count=0,
+                               created_at=created_at, country='FR'))
+
+        with self.db as dbsession:
+            # Bypass ORM country validation
+            dbsession.execute("UPDATE ping_v1 SET country='FRA'")
+
+        with self.db as dbsession:
+            ping = dbsession.query(Ping).one()
+            assert ping.country == 'FRA'
+
+        self.run_subcommand(
+            'test_transform_countries_alpha_3_to_2',
+            'transform-countries-alpha-3-to-2')
+
+        with self.db as dbsession:
+            ping = dbsession.query(Ping).one()
+            assert ping.country == 'FR'
+
+    def test_transform_countries_alpha_3_to_2_only_2(self, capfd):
+        from azafea.event_processors.endless.ping.v1.handler import Ping, PingConfiguration
+
+        # Create the table
+        self.run_subcommand('initdb')
+        self.ensure_tables(Ping)
+
+        created_at = datetime.utcnow().replace(tzinfo=timezone.utc)
+
+        with self.db as dbsession:
+            ping_config = PingConfiguration(image='eos-eos3.7-amd64-amd64.190419-225606.base',
+                                            product='product', dualboot=True, created_at=created_at,
+                                            vendor='EnDlEsS')
+            dbsession.add(ping_config)
+            dbsession.add(Ping(config=ping_config, release='release', count=0,
+                               created_at=created_at, country='FR'))
+
+        self.run_subcommand(
+            'test_transform_countries_alpha_3_to_2_only_2',
+            'transform-countries-alpha-3-to-2')
+
+        with self.db as dbsession:
+            ping = dbsession.query(Ping).one()
+            assert ping.country == 'FR'
+
+        capture = capfd.readouterr()
+        assert 'No ping with alpha-3 country code found' in capture.out

azafea/event_processors/endless/ping/v1/cli.py

@@ -11,12 +11,14 @@
 
 from sqlalchemy.orm.exc import NoResultFound
 from sqlalchemy.orm.query import Query
+from sqlalchemy.sql.expression import func
 
 from azafea.config import Config
 from azafea.model import Db
 from azafea.utils import progress
 from azafea.vendors import normalize_vendor
 
+from ...country import tranform_alpha_3_into_2
 from ...image import parse_endless_os_image
 from .handler import Ping, PingConfiguration
 
@@ -39,6 +41,14 @@ def register_commands(subs: argparse._SubParsersAction) -> None:
                               help='The size of the chunks to operate on')
     parse_images.set_defaults(subcommand=do_parse_images)
 
+    transform_countries_alpha_3_to_2 = subs.add_parser(
+        'transform-countries-alpha-3-to-2',
+        help='Transform 3-letter country codes to 2-letter country codes',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    transform_countries_alpha_3_to_2.add_argument('--chunk-size', type=int, default=5000,
+                                                  help='The size of the chunks to operate on')
+    transform_countries_alpha_3_to_2.set_defaults(subcommand=do_transform_countries_alpha_3_to_2)
+
 
 def _normalize_chunk(chunk: Query) -> None:
     ping_configs = chunk.session.query(PingConfiguration)
@@ -122,3 +132,27 @@ def do_parse_images(config: Config, args: argparse.Namespace) -> None:
     progress(num_records, num_records, end='\n')
 
     log.info('All done!')
+
+
+def do_transform_countries_alpha_3_to_2(config: Config, args: argparse.Namespace) -> None:
+    db = Db(config.postgresql)
+    log.info('Tranform alpha-3 country codes into alpha-2 for pings')
+
+    with db as dbsession:
+        query = dbsession.chunked_query(Ping, chunk_size=args.chunk_size)
+        query = query.filter(func.length(Ping.country) == 3)
+        num_records = query.count()
+
+        if num_records == 0:
+            log.info('-> No ping with alpha-3 country code found')
+            return None
+
+        for chunk_number, chunk in enumerate(query, start=1):
+            for ping in chunk:
+                ping.country = tranform_alpha_3_into_2(ping.country)
+            dbsession.commit()
+            progress(chunk_number * args.chunk_size, num_records)
+
+    progress(num_records, num_records, end='\n')
+
+    log.info('All done!')