Merge pull request #3552 from lonvia/drop-some-migrations

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

diff --git a/src/nominatim_db/tools/freeze.py b/src/nominatim_db/tools/freeze.py
index bd52ba9a5a7ad194c353a721d1529ea4cefa20cb..c4eedb43262029424b120072576de0e061b32d86 100644 (file)
--- a/src/nominatim_db/tools/freeze.py
+++ b/src/nominatim_db/tools/freeze.py
@@ -10,9 +10,9 @@ Functions for removing unnecessary data from the database.
 from typing import Optional
 from pathlib import Path
 
 from typing import Optional
 from pathlib import Path
 

-from psycopg2 import sql as pysql
+from psycopg import sql as pysql

 
 

-from ..db.connection import Connection
+from ..db.connection import Connection, drop_tables, table_exists

 
 UPDATE_TABLES = [
     'address_levels',
 
 UPDATE_TABLES = [
     'address_levels',


@@ -39,9 +39,7 @@ def drop_update_tables(conn: Connection) -> None:
                     + pysql.SQL(' or ').join(parts))
         tables = [r[0] for r in cur]
 
                     + pysql.SQL(' or ').join(parts))
         tables = [r[0] for r in cur]
 

-        for table in tables:
-            cur.drop_table(table, cascade=True)
-
+    drop_tables(conn, *tables, cascade=True)

     conn.commit()
 
 
     conn.commit()
 
 


@@ -55,4 +53,4 @@ def is_frozen(conn: Connection) -> bool:
     """ Returns true if database is in a frozen state
     """
 
     """ Returns true if database is in a frozen state
     """
 

-    return conn.table_exists('place') is False
+    return table_exists(conn, 'place') is False