Merge pull request #49 from lonvia/upgrade-to-psycopg3

Upgrade to psycopg3

pyproject.toml

@@ -20,7 +20,7 @@ classifiers = [
 dependencies = [
     "pyYAML",
     "geojson",
-    "psycopg2"
+    "psycopg"
 ]
 version = "0.1.0"
 
@@ -34,7 +34,6 @@ tests = [
 
 [tool.setuptools]
 packages = ["nominatim_data_analyser",
-            "nominatim_data_analyser.database",
             "nominatim_data_analyser.config",
             "nominatim_data_analyser.logger",
             "nominatim_data_analyser.core",

src/nominatim_data_analyser/core/pipes/data_fetching/sql_processor.py

@@ -1,11 +1,10 @@
 from typing import Any
 
-import psycopg2.extras
+import psycopg
+
 from ....config.config import Config
 from ... import Pipe
-from ....database.connection import connect
 from ....logger.timer import Timer
-from psycopg2._psycopg import connection
 
 
 class SQLProcessor(Pipe):
@@ -16,24 +15,19 @@ class SQLProcessor(Pipe):
     def on_created(self) -> None:
         self.query = self.extract_data('query', required=True)
 
-    def process(self, data: Any) -> list[dict[str, Any]]:
+    def process(self, _: Any) -> list[dict[str, Any]]:
         """
             Executes the query and returns the results.
         """
-        with connect(Config.values['Dsn']) as conn:
-            return self.execute_query(conn)
-
-    def execute_query(self, conn: connection) -> list[dict[str, Any]]:
-        """
-            Executes the query and returns the results.
-            Takes a database connection as input.
-        """
         results: list[dict[str, Any]]
-        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
-            timer = Timer().start_timer()
-            cur.execute(self.query)
-            results = list(cur)
-            elapsed_mins, elapsed_secs = timer.get_elapsed()
-            self.log(f'Query {self.id} executed in {elapsed_mins} mins {elapsed_secs} secs.')
-        self.log(f'Query {self.id} returned {len(results)} results.')
+
+        with psycopg.connect(Config.values['Dsn']) as conn:
+            with conn.cursor(row_factory=psycopg.rows.dict_row) as cur:
+                timer = Timer().start_timer()
+                cur.execute(self.query)
+                results = cur.fetchall()
+                elapsed_mins, elapsed_secs = timer.get_elapsed()
+                self.log(f'Query {self.id} executed in {elapsed_mins} mins {elapsed_secs} secs.')
+            self.log(f'Query {self.id} returned {len(results)} results.')
+
         return results

src/nominatim_data_analyser/core/pipes/output_formatters/osmoscope_layer_formatter.py

@@ -1,10 +1,11 @@
-from __future__ import annotations
-from ....database.connection import connect
-from ....config import Config
-from ... import Pipe
 from pathlib import Path
 import json
 
+import psycopg
+
+from ....config import Config
+from ... import Pipe
+
 class OsmoscopeLayerFormatter(Pipe):
     """
         Handles the creation of the layer JSON file.
@@ -39,7 +40,7 @@ def add_last_update_date_layer_info(self) -> None:
             This field contains the date of the last database update.
             The date is extracted from the lastimportdate table of the database.
         """
-        with connect(Config.values['Dsn']) as conn:
+        with psycopg.connect(Config.values['Dsn']) as conn:
             with conn.cursor() as cur:
                 cur.execute("SELECT to_char(lastimportdate at time zone 'UTC', "
                             "               'YYYY-MM-DD HH24:MI:SS UTC') FROM import_status")

tests/conftest.py

@@ -1,7 +1,9 @@
 import sys
 import sysconfig
 from pathlib import Path
-import psycopg2
+
+import psycopg
+from psycopg import sql as pysql
 import pytest
 
 SRC_DIR = Path(__file__, '..', '..').resolve()
@@ -25,8 +27,6 @@
                                                    OsmoscopeLayerFormatter,
                                                    VectorTileFormatter)
 from nominatim_data_analyser.core.qa_rule import ExecutionContext
-from nominatim_data_analyser.database.connection import connect
-from psycopg2._psycopg import connection, cursor
 
 
 @pytest.fixture
@@ -35,45 +35,38 @@ def temp_db() -> str:
         Create an empty database for the test.
     """
     name = 'test_qa_tool_python_unittest'
-    conn = psycopg2.connect(database='postgres')
-
-    conn.set_isolation_level(0)
-    with conn.cursor() as cur:
-        cur.execute('DROP DATABASE IF EXISTS {}'.format(name))
-        cur.execute('CREATE DATABASE {}'.format(name))
-    conn.close()
+    with psycopg.connect(dbname='postgres', autocommit=True) as conn:
+        with conn.cursor() as cur:
+            cur.execute(pysql.SQL('DROP DATABASE IF EXISTS') + pysql.Identifier(name))
+            cur.execute(pysql.SQL('CREATE DATABASE') + pysql.Identifier(name))
 
     yield name
 
-    conn = psycopg2.connect(database='postgres')
-    conn.set_isolation_level(0)
-    with conn.cursor() as cur:
-        cur.execute('DROP DATABASE IF EXISTS {}'.format(name))
-    conn.close()
+    with psycopg.connect(dbname='postgres', autocommit=True) as conn:
+        with conn.cursor() as cur:
+            cur.execute(pysql.SQL('DROP DATABASE IF EXISTS') + pysql.Identifier(name))
 
 @pytest.fixture
 def dsn(temp_db: str) -> str:
     return 'dbname=' + temp_db
 
 @pytest.fixture
-def temp_db_conn(dsn: str) -> connection:
+def temp_db_conn(dsn: str):
     """
         Connection to the test database.
     """
-    with connect(dsn) as conn:
+    with psycopg.connect(dsn) as conn:
         yield conn
 
 @pytest.fixture
-def temp_db_cursor(dsn: str) -> cursor:
-    """     
+def temp_db_cursor(dsn: str):
+    """ 
         Connection and cursor towards the test database. 
         The connection will be in auto-commit mode.
     """
-    conn = psycopg2.connect(dsn)
-    conn.set_isolation_level(0)
-    with conn.cursor() as cur:
-        yield cur
-    conn.close()
+    with psycopg.connect(dsn, autocommit=True) as conn:
+        with conn.cursor() as cur:
+            yield cur
 
 @pytest.fixture
 def config() -> Config:

tests/pipes/test_sql_processor.py

@@ -1,8 +1,8 @@
 from nominatim_data_analyser.core.pipes.data_fetching import SQLProcessor
-from psycopg2._psycopg import connection, cursor
+from nominatim_data_analyser.config import Config
 
 
-def test_on_created_sql_processor(sql_processor: SQLProcessor) -> None:
+def test_on_created_sql_processor(sql_processor: SQLProcessor):
     """
         Test the on_created() method of the SQLProcessor.
     """
@@ -11,9 +11,7 @@ def test_on_created_sql_processor(sql_processor: SQLProcessor) -> None:
     sql_processor.on_created()
     assert sql_processor.query == 'QUERY'
 
-def test_execute_query(sql_processor: SQLProcessor, 
-                       temp_db_conn: connection,
-                       temp_db_cursor: cursor) -> None:
+def test_execute_query(sql_processor: SQLProcessor, dsn, temp_db_cursor):
     """
         Test the execute_query() method of the SQLProcessor.
     """
@@ -23,5 +21,6 @@ def test_execute_query(sql_processor: SQLProcessor,
     """)
     sql_processor.query = 'SELECT * FROM test_table'
 
-    results = sql_processor.execute_query(temp_db_conn)
+    Config.values['Dsn'] = dsn
+    results = sql_processor.process(None)
     assert len(results) == 3 and results[0]['val'] == 'test1' and results[1]['val'] == 'test2' and results[2]['val'] == 'test3'