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 35943a50d3ed08cf0a46a7658fafeffac0db2674..9b71405993f35f855e57465564a2b1eebd29ee9e 100644 (file)
--- a/src/nominatim_db/data/country_info.py
+++ b/src/nominatim_db/data/country_info.py
@@ -9,12 +9,11 @@ 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 nominatim_core.db import utils as db_utils
-from nominatim_core.db.connection import connect, Connection
-from nominatim_core.errors import UsageError
-from nominatim_core.config import Configuration
+from ..db import utils as db_utils
+from ..db.connection import connect, Connection, register_hstore
+from ..errors import UsageError
+from ..config import Configuration

 from ..tokenizer.base import AbstractTokenizer
 
 def _flatten_name_list(names: Any) -> Dict[str, str]:
 from ..tokenizer.base import AbstractTokenizer
 
 def _flatten_name_list(names: Any) -> Dict[str, str]:


@@ -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""")