From 1f0e1bec0ee66df7641e53c6b55e84d7eb88cec2 Mon Sep 17 00:00:00 2001
From: Sarah Hoffmann <lonvia@denofr.de>
Date: Thu, 18 May 2023 13:00:59 +0200
Subject: [PATCH] enable API use with psycopg 3

---
 nominatim/api/core.py              | 20 +++++++++-----------
 nominatim/db/async_core_library.py | 21 +++++++++++++++++++++
 2 files changed, 30 insertions(+), 11 deletions(-)
 create mode 100644 nominatim/db/async_core_library.py

diff --git a/nominatim/api/core.py b/nominatim/api/core.py
index 29325b08..3e4b175a 100644
--- a/nominatim/api/core.py
+++ b/nominatim/api/core.py
@@ -14,9 +14,10 @@ from pathlib import Path
 
 import sqlalchemy as sa
 import sqlalchemy.ext.asyncio as sa_asyncio
-import asyncpg
+
 
 from nominatim.db.sqlalchemy_schema import SearchTables
+from nominatim.db.async_core_library import PGCORE_LIB, PGCORE_ERROR
 from nominatim.config import Configuration
 from nominatim.api.connection import SearchConnection
 from nominatim.api.status import get_status, StatusResult
@@ -55,26 +56,22 @@ class NominatimAPIAsync:
 
             query = {k: v for k, v in dsn.items()
                       if k not in ('user', 'password', 'dbname', 'host', 'port')}
-            query['prepared_statement_cache_size'] = '0'
+            if PGCORE_LIB == 'asyncpg':
+                query['prepared_statement_cache_size'] = '0'
 
             dburl = sa.engine.URL.create(
-                       'postgresql+asyncpg',
+                       f'postgresql+{PGCORE_LIB}',
                        database=dsn.get('dbname'),
                        username=dsn.get('user'), password=dsn.get('password'),
                        host=dsn.get('host'), port=int(dsn['port']) if 'port' in dsn else None,
                        query=query)
-            engine = sa_asyncio.create_async_engine(
-                             dburl, future=True,
-                             connect_args={'server_settings': {
-                                'DateStyle': 'sql,european',
-                                'max_parallel_workers_per_gather': '0'
-                             }})
+            engine = sa_asyncio.create_async_engine(dburl, future=True)
 
             try:
                 async with engine.begin() as conn:
                     result = await conn.scalar(sa.text('SHOW server_version_num'))
                     server_version = int(result)
-            except asyncpg.PostgresError:
+            except (PGCORE_ERROR, sa.exc.OperationalError):
                 server_version = 0
 
             if server_version >= 110000:
@@ -82,6 +79,7 @@ class NominatimAPIAsync:
                 def _on_connect(dbapi_con: Any, _: Any) -> None:
                     cursor = dbapi_con.cursor()
                     cursor.execute("SET jit_above_cost TO '-1'")
+                    cursor.execute("SET max_parallel_workers_per_gather TO '0'")
                 # Make sure that all connections get the new settings
                 await self.close()
 
@@ -124,7 +122,7 @@ class NominatimAPIAsync:
         try:
             async with self.begin() as conn:
                 status = await get_status(conn)
-        except asyncpg.PostgresError:
+        except (PGCORE_ERROR, sa.exc.OperationalError):
             return StatusResult(700, 'Database connection failed')
 
         return status
diff --git a/nominatim/db/async_core_library.py b/nominatim/db/async_core_library.py
new file mode 100644
index 00000000..d86ce5ec
--- /dev/null
+++ b/nominatim/db/async_core_library.py
@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: GPL-3.0-or-later
+#
+# This file is part of Nominatim. (https://nominatim.org)
+#
+# Copyright (C) 2023 by the Nominatim developer community.
+# For a full list of authors see the git log.
+"""
+Import the base libary to use with asynchronous SQLAlchemy.
+"""
+# pylint: disable=invalid-name
+
+from typing import Any
+
+try:
+    import psycopg
+    PGCORE_LIB = 'psycopg'
+    PGCORE_ERROR: Any = psycopg.Error
+except ModuleNotFoundError:
+    import asyncpg
+    PGCORE_LIB = 'asyncpg'
+    PGCORE_ERROR = asyncpg.PostgresError
-- 
2.39.5