Merge pull request #2450 from mtmail/tiger-data-2021

[nominatim.git] / nominatim / tools / country_info.py

diff --git a/nominatim/tools/country_info.py b/nominatim/tools/country_info.py
index 1b61ae685cc6106ac522b00c9c08d582983e2039..635d15840a84b8197efb9f5cb358344a78a0c2b9 100644 (file)
--- a/nominatim/tools/country_info.py
+++ b/nominatim/tools/country_info.py
@@ -6,6 +6,51 @@ import psycopg2.extras
 from nominatim.db import utils as db_utils
 from nominatim.db.connection import connect
 
+class _CountryInfo:
+    """ Caches country-specific properties from the configuration file.
+    """
+
+    def __init__(self):
+        self._info = {}
+
+
+    def load(self, config):
+        """ Load the country properties from the configuration files,
+            if they are not loaded yet.
+        """
+        if not self._info:
+            self._info = config.load_sub_configuration('country_settings.yaml')
+            # Convert languages into a list for simpler handling.
+            for prop in self._info.values():
+                if 'languages' not in prop:
+                    prop['languages'] = []
+                elif not isinstance(prop['languages'], list):
+                    prop['languages'] = [x.strip()
+                                         for x in prop['languages'].split(',')]
+
+
+    def items(self):
+        """ Return tuples of (country_code, property dict) as iterable.
+        """
+        return self._info.items()
+
+
+_COUNTRY_INFO = _CountryInfo()
+
+def setup_country_config(config):
+    """ Load country properties from the configuration file.
+        Needs to be called before using any other functions in this
+        file.
+    """
+    _COUNTRY_INFO.load(config)
+
+
+def iterate():
+    """ Iterate over country code and properties.
+    """
+    return _COUNTRY_INFO.items()
+
+
 def setup_country_tables(dsn, sql_dir, ignore_partitions=False):
     """ Create and populate the tables with basic static data that provides
         the background for geocoding. Data is assumed to not yet exist.
@@ -13,11 +58,24 @@ def setup_country_tables(dsn, sql_dir, ignore_partitions=False):
     db_utils.execute_file(dsn, sql_dir / 'country_name.sql')
     db_utils.execute_file(dsn, sql_dir / 'country_osm_grid.sql.gz')
 
-    if ignore_partitions:
-        with connect(dsn) as conn:
-            with conn.cursor() as cur:
-                cur.execute('UPDATE country_name SET partition = 0')
-            conn.commit()
+    params = []
+    for ccode, props in _COUNTRY_INFO.items():
+        if ccode is not None and props is not None:
+            if ignore_partitions:
+                partition = 0
+            else:
+                partition = props.get('partition')
+            lang = props['languages'][0] if len(props['languages']) == 1 else None
+            params.append((ccode, partition, lang))
+
+    with connect(dsn) as conn:
+        with conn.cursor() as cur:
+            cur.execute_values(
+                """ UPDATE country_name
+                    SET partition = part, country_default_language_code = lang
+                    FROM (VALUES %s) AS v (cc, part, lang)
+                    WHERE country_code = v.cc""", params)
+        conn.commit()
 
 
 def create_country_names(conn, tokenizer, languages=None):