Upgrade from psycopg2 to psycopg (version 3)

.pre-commit-config.yaml

@@ -17,4 +17,4 @@ repos:
   hooks:
   - id: mypy
     files: ^blitzortung/
-    additional_dependencies: [pyproj, shapely, psycopg2-binary, fasteners, txpostgres]
+    additional_dependencies: [pyproj, shapely, psycopg, fasteners, txpostgres]

blitzortung/db/__init__.py

@@ -20,12 +20,9 @@
 
 import atexit
 
-from . import compat  # Register psycopg2cffi compatibility
-
-import psycopg2
-import psycopg2.pool
-import psycopg2.extras
-import psycopg2.extensions
+import psycopg
+import psycopg.rows
+import psycopg_pool
 from injector import Module, singleton, inject, provider
 
 from .. import config
@@ -41,8 +38,12 @@ def cleanup(connection_pool):
     @singleton
     @provider
     @inject
-    def provide_psycopg2_connection_pool(self, config: config.Config) -> psycopg2.pool.ThreadedConnectionPool:
-        connection_pool = psycopg2.pool.ThreadedConnectionPool(4, 50, config.get_db_connection_string())
+    def provide_psycopg_connection_pool(self, config: config.Config) -> psycopg_pool.ConnectionPool:
+        connection_pool = psycopg_pool.ConnectionPool(
+            config.get_db_connection_string(),
+            min_connections=4,
+            max_connections=50
+        )
         atexit.register(self.cleanup, connection_pool)
         return connection_pool
 

blitzortung/db/query.py

@@ -23,7 +23,7 @@
 import shapely.geometry.base
 import shapely.wkb
 
-import psycopg2
+import psycopg
 
 class BaseInterval:
     """
@@ -200,13 +200,13 @@ def add_id_interval(self, id_interval):
     def add_geometry(self, geometry):
         if geometry.is_valid:
             self.add_condition('ST_GeomFromWKB(%(envelope)s, %(srid)s) && geog',
-                               envelope=psycopg2.Binary(shapely.wkb.dumps(geometry.envelope)))
+                               envelope=psycopg.Binary(shapely.wkb.dumps(geometry.envelope)))
 
             if not geometry.equals(geometry.envelope):
                 self.add_condition(
                     'ST_Intersects(ST_GeomFromWKB(%(geometry)s, %(srid)s), ' +
                     'ST_Transform(geog::geometry, %(srid)s))',
-                    geometry=psycopg2.Binary(shapely.wkb.dumps(geometry)))
+                    geometry=psycopg.Binary(shapely.wkb.dumps(geometry)))
 
         else:
             raise ValueError("invalid geometry in db.Strike.select()")
@@ -272,7 +272,7 @@ def __init__(self, grid, count_threshold=0):
 
         if env.is_valid:
             self.add_condition('ST_GeomFromWKB(%(envelope)s, %(envelope_srid)s) && geog',
-                               envelope=psycopg2.Binary(shapely.wkb.dumps(env)),
+                               envelope=psycopg.Binary(shapely.wkb.dumps(env)),
                                envelope_srid=grid.srid)
         else:
             raise ValueError("invalid Raster geometry in db.query.GridQuery.__init__()")

blitzortung/db/query_builder.py

@@ -19,7 +19,7 @@
 """
 from typing import Optional
 
-import psycopg2
+import psycopg
 
 import shapely.wkb
 
@@ -71,7 +71,7 @@ def histogram_query(table_name: str, time_interval: TimeInterval, binsize:int, r
 
         if envelope and envelope.env.is_valid:
             query.add_condition('ST_SetSRID(CAST(%(envelope)s AS geometry), %(envelope_srid)s) && geog',
-                                envelope=psycopg2.Binary(shapely.wkb.dumps(envelope.env)),
+                                envelope=psycopg.Binary(shapely.wkb.dumps(envelope.env)),
                                 envelope_srid=envelope.srid)
 
         return query

blitzortung/db/table.py

@@ -22,10 +22,8 @@
 import logging
 from typing import Optional
 
-import psycopg2
-import psycopg2.extensions
-import psycopg2.extras
-import psycopg2.pool
+import psycopg
+import psycopg_pool
 from injector import inject
 
 from blitzortung.db.grid_result import build_grid_result
@@ -78,24 +76,20 @@ def __init__(self, db_connection_pool):
             self.conn.cancel()
             try:
                 self.conn.reset()
-            except psycopg2.OperationalError:
+            except psycopg.OperationalError:
                 print("reconnect to db")
                 self.db_connection_pool.putconn(self.conn, close=True)
                 continue
             break
 
-        psycopg2.extensions.register_type(psycopg2.extensions.UNICODE, self.conn)
-        psycopg2.extensions.register_type(psycopg2.extensions.UNICODEARRAY, self.conn)
-        self.conn.set_client_encoding('UTF8')
-
         self.srid = geom.Geometry.default_srid
         self.tz = None
         self.set_timezone(Base.default_timezone)
 
         cur = None
         try:
-            cur = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
-        except psycopg2.DatabaseError as error:
+            cur = self.conn.cursor(row_factory=psycopg.rows.DictRow)
+        except psycopg.DatabaseError as error:
             self.logger.error(error)
 
             if self.conn:
@@ -164,7 +158,7 @@ def select(self, **kwargs):
         pass
 
     def execute(self, sql_statement, parameters=None, factory_method=None, **factory_method_args):
-        with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cursor:
+        with self.conn.cursor(row_factory=psycopg.rows.DictRow) as cursor:
             cursor.execute(sql_statement, parameters)
             if factory_method:
                 method = factory_method(cursor, **factory_method_args)
@@ -179,7 +173,7 @@ def single_cursor_factory(cursor):
 
     def execute_many(self, sql_statement, parameters=None, factory_method=None, **factory_method_args):
         factory_method = factory_method or (lambda values, **_: values)
-        with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cursor:
+        with self.conn.cursor(row_factory=psycopg.rows.DictRow) as cursor:
             cursor.execute(sql_statement, parameters)
             for value in cursor:
                 yield factory_method(value, **factory_method_args)
@@ -219,7 +213,7 @@ class Strike(Base):
     table_name = 'strikes'
 
     @inject
-    def __init__(self, db_connection_pool: psycopg2.pool.ThreadedConnectionPool, query_builder_: query_builder.Strike,
+    def __init__(self, db_connection_pool: psycopg_pool.ConnectionPool, query_builder_: query_builder.Strike,
                  strike_mapper: mapper.Strike):
         super().__init__(db_connection_pool)
 

blitzortung/service/db.py

@@ -16,8 +16,8 @@
 
 """
 
-import psycopg2
-import psycopg2.extras
+import psycopg
+import psycopg.rows
 from twisted.internet.defer import Deferred
 from twisted.python import log
 from txpostgres import reconnection
@@ -28,9 +28,9 @@
 
 
 def connection_factory(*args, **kwargs):
-    """Create a psycopg2 connection with DictConnection factory."""
-    kwargs['connection_factory'] = psycopg2.extras.DictConnection
-    return psycopg2.connect(*args, **kwargs)
+    """Create a psycopg connection with DictRow factory."""
+    kwargs['row_factory'] = psycopg.rows.DictRow
+    return psycopg.connect(*args, **kwargs)
 
 
 class LoggingDetector(reconnection.DeadConnectionDetector):

mypy.ini

@@ -16,10 +16,10 @@ ignore_missing_imports = True
 [mypy-shapely.*]
 ignore_missing_imports = True
 
-[mypy-psycopg2.*]
+[mypy-psycopg.*]
 ignore_missing_imports = True
 
-[mypy-psycopg2cffi.*]
+[mypy-psycopg_pool.*]
 ignore_missing_imports = True
 
 [mypy-fasteners.*]