Merge remote-tracking branch 'upstream/master'

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

diff --git a/nominatim/tools/country_info.py b/nominatim/tools/country_info.py
index 1b61ae685cc6106ac522b00c9c08d582983e2039..e04a8693f116bccd6d7e609de0c463b74170e46a 100644 (file)
--- a/nominatim/tools/country_info.py
+++ b/nominatim/tools/country_info.py
@@ -6,6 +6,36 @@ import psycopg2.extras
 from nominatim.db import utils as db_utils
 from nominatim.db.connection import connect
 
 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')
+
+    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 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.
 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 +43,27 @@ 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')
 
     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')
+            if ',' in (props.get('languages', ',') or ','):
+                lang = None
+            else:
+                lang = props['languages']
+            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):
 
 
 def create_country_names(conn, tokenizer, languages=None):