Merge remote-tracking branch 'upstream/master'

[nominatim.git] / src / nominatim_db / tools / admin.py

diff --git a/src/nominatim_db/tools/admin.py b/src/nominatim_db/tools/admin.py
index cea2ad664a7e339c2c66dd23790ab517dcd7b338..b8e3cb5687161ca7e8a27e307c3b4c4e509d08ff 100644 (file)
--- a/src/nominatim_db/tools/admin.py
+++ b/src/nominatim_db/tools/admin.py
@@ -10,18 +10,19 @@ Functions for database analysis and maintenance.
 from typing import Optional, Tuple, Any, cast
 import logging
 
-from psycopg2.extras import Json, register_hstore
-from psycopg2 import DataError
+import psycopg
+from psycopg.types.json import Json
 
 from ..typing import DictCursorResult
 from ..config import Configuration
-from ..db.connection import connect, Cursor
+from ..db.connection import connect, Cursor, register_hstore
 from ..errors import UsageError
 from ..tokenizer import factory as tokenizer_factory
 from ..data.place_info import PlaceInfo
 
 LOG = logging.getLogger()
 
+
 def _get_place_info(cursor: Cursor, osm_id: Optional[str],
                     place_id: Optional[int]) -> DictCursorResult:
     sql = """SELECT place_id, extra.*
@@ -59,7 +60,7 @@ def analyse_indexing(config: Configuration, osm_id: Optional[str] = None,
     """
     with connect(config.get_libpq_dsn()) as conn:
         register_hstore(conn)
-        with conn.cursor() as cur:
+        with conn.cursor(row_factory=psycopg.rows.dict_row) as cur:
             place = _get_place_info(cur, osm_id, place_id)
 
             cur.execute("update placex set indexed_status = 2 where place_id = %s",
@@ -74,6 +75,9 @@ def analyse_indexing(config: Configuration, osm_id: Optional[str] = None,
 
             tokenizer = tokenizer_factory.get_tokenizer_for_db(config)
 
+            # Enable printing of messages.
+            conn.add_notice_handler(lambda diag: print(diag.message_primary))
+
             with tokenizer.name_analyzer() as analyzer:
                 cur.execute("""UPDATE placex
                                SET indexed_status = 0, address = %s, token_info = %s,
@@ -86,9 +90,6 @@ def analyse_indexing(config: Configuration, osm_id: Optional[str] = None,
         # we do not want to keep the results
         conn.rollback()
 
-        for msg in conn.notices:
-            print(msg)
-
 
 def clean_deleted_relations(config: Configuration, age: str) -> None:
     """ Clean deleted relations older than a given age
@@ -101,6 +102,6 @@ def clean_deleted_relations(config: Configuration, age: str) -> None:
                             WHERE p.osm_type = d.osm_type AND p.osm_id = d.osm_id
                             AND age(p.indexed_date) > %s::interval""",
                             (age, ))
-            except DataError as exc:
+            except psycopg.DataError as exc:
                 raise UsageError('Invalid PostgreSQL time interval format') from exc
         conn.commit()