Test fixes

[nominatim.git] / nominatim / clicmd / index.py

diff --git a/nominatim/clicmd/index.py b/nominatim/clicmd/index.py
index ca3f9deedf4e0546e211679336474a1615ef524d..0225c5ed41561b35ca257eb622f6804204ac56f6 100644 (file)
--- a/nominatim/clicmd/index.py
+++ b/nominatim/clicmd/index.py
@@ -1,7 +1,7 @@
 """
 Implementation of the 'index' subcommand.
 """
-import os
+import psutil
 
 from ..db import status
 from ..db.connection import connect
@@ -11,14 +11,6 @@ from ..db.connection import connect
 # Using non-top-level imports to avoid eventually unused imports.
 # pylint: disable=E0012,C0415
 
-def _num_system_cpus():
-    try:
-        cpus = len(os.sched_getaffinity(0))
-    except NotImplementedError:
-        cpus = None
-
-    return cpus or os.cpu_count()
-
 
 class UpdateIndex:
     """\
@@ -42,7 +34,7 @@ class UpdateIndex:
         from ..indexer.indexer import Indexer
 
         indexer = Indexer(args.config.get_libpq_dsn(),
-                          args.threads or _num_system_cpus() or 1)
+                          args.threads or psutil.cpu_count() or 1)
 
         if not args.no_boundaries:
             indexer.index_boundaries(args.minrank, args.maxrank)
@@ -51,8 +43,7 @@ class UpdateIndex:
 
         if not args.no_boundaries and not args.boundaries_only \
            and args.minrank == 0 and args.maxrank == 30:
-            conn = connect(args.config.get_libpq_dsn())
-            status.set_indexed(conn, True)
-            conn.close()
+            with connect(args.config.get_libpq_dsn()) as conn:
+                status.set_indexed(conn, True)
 
         return 0