actions: add test for database migration

[nominatim.git] / src / nominatim_db / data / country_info.py

diff --git a/src/nominatim_db/data/country_info.py b/src/nominatim_db/data/country_info.py
index c8002ee7023ca27a9e680003e92a66e3d8f0b428..9b71405993f35f855e57465564a2b1eebd29ee9e 100644 (file)
--- a/src/nominatim_db/data/country_info.py
+++ b/src/nominatim_db/data/country_info.py
@@ -9,10 +9,9 @@ Functions for importing and managing static country information.
 """
 from typing import Dict, Any, Iterable, Tuple, Optional, Container, overload
 from pathlib import Path
 """
 from typing import Dict, Any, Iterable, Tuple, Optional, Container, overload
 from pathlib import Path

-import psycopg2.extras

 
 from ..db import utils as db_utils
 
 from ..db import utils as db_utils

-from ..db.connection import connect, Connection
+from ..db.connection import connect, Connection, register_hstore

 from ..errors import UsageError
 from ..config import Configuration
 from ..tokenizer.base import AbstractTokenizer
 from ..errors import UsageError
 from ..config import Configuration
 from ..tokenizer.base import AbstractTokenizer


@@ -129,8 +128,8 @@ def setup_country_tables(dsn: str, sql_dir: Path, ignore_partitions: bool = Fals
 
             params.append((ccode, props['names'], lang, partition))
     with connect(dsn) as conn:
 
             params.append((ccode, props['names'], lang, partition))
     with connect(dsn) as conn:

+        register_hstore(conn)

         with conn.cursor() as cur:
         with conn.cursor() as cur:

-            psycopg2.extras.register_hstore(cur)

             cur.execute(
                 """ CREATE TABLE public.country_name (
                         country_code character varying(2),
             cur.execute(
                 """ CREATE TABLE public.country_name (
                         country_code character varying(2),


@@ -139,9 +138,10 @@ def setup_country_tables(dsn: str, sql_dir: Path, ignore_partitions: bool = Fals
                         country_default_language_code text,
                         partition integer
                     ); """)
                         country_default_language_code text,
                         partition integer
                     ); """)

-            cur.execute_values(
+            cur.executemany(

                 """ INSERT INTO public.country_name
                 """ INSERT INTO public.country_name

-                    (country_code, name, country_default_language_code, partition) VALUES %s
+                    (country_code, name, country_default_language_code, partition)
+                    VALUES (%s, %s, %s, %s)

                 """, params)
         conn.commit()
 
                 """, params)
         conn.commit()
 


@@ -157,8 +157,8 @@ def create_country_names(conn: Connection, tokenizer: AbstractTokenizer,
         return ':' not in key or not languages or \
                key[key.index(':') + 1:] in languages
 
         return ':' not in key or not languages or \
                key[key.index(':') + 1:] in languages
 

+    register_hstore(conn)

     with conn.cursor() as cur:
     with conn.cursor() as cur:

-        psycopg2.extras.register_hstore(cur)

         cur.execute("""SELECT country_code, name FROM country_name
                        WHERE country_code is not null""")
 
         cur.execute("""SELECT country_code, name FROM country_name
                        WHERE country_code is not null""")