Compatibility with psycopg 3.0 (#2088)

By default `psycopg2` is preferred. The `psycopg>=3.0` will be used only if `psycopg2` is not available or its version is too old.

features/environment.py

@@ -1,7 +1,6 @@
 import abc
 import datetime
 import os
-import psycopg2
 import json
 import shutil
 import signal
@@ -13,6 +12,8 @@ import threading
 import time
 import yaml
 
+import patroni.psycopg as psycopg
+
 from six.moves.BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
 
 
@@ -205,16 +206,16 @@ class PatroniController(AbstractController):
 
         user = config['postgresql'].get('authentication', config['postgresql']).get('superuser', {})
         self._connkwargs = {k: user[n] for n, k in [('username', 'user'), ('password', 'password')] if n in user}
-        self._connkwargs.update({'host': host, 'port': self.__PORT, 'database': 'postgres'})
+        self._connkwargs.update({'host': host, 'port': self.__PORT, 'dbname': 'postgres'})
 
         self._replication = config['postgresql'].get('authentication', config['postgresql']).get('replication', {})
-        self._replication.update({'host': host, 'port': self.__PORT, 'database': 'postgres'})
+        self._replication.update({'host': host, 'port': self.__PORT, 'dbname': 'postgres'})
 
         return patroni_config_path
 
     def _connection(self):
         if not self._conn or self._conn.closed != 0:
-            self._conn = psycopg2.connect(**self._connkwargs)
+            self._conn = psycopg.connect(**self._connkwargs)
             self._conn.autocommit = True
         return self._conn
 
@@ -228,7 +229,7 @@ class PatroniController(AbstractController):
             cursor = self._cursor()
             cursor.execute(query)
             return cursor
-        except psycopg2.Error:
+        except psycopg.Error:
             if not fail_ok:
                 raise
 
@@ -268,7 +269,7 @@ class PatroniController(AbstractController):
 
     @property
     def backup_source(self):
-        return 'postgres://{username}:{password}@{host}:{port}/{database}'.format(**self._replication)
+        return 'postgres://{username}:{password}@{host}:{port}/{dbname}'.format(**self._replication)
 
     def backup(self, dest=os.path.join('data', 'basebackup')):
         subprocess.call(PatroniPoolController.BACKUP_SCRIPT + ['--walmethod=none',

features/steps/basic_replication.py

@@ -1,4 +1,4 @@
-import psycopg2 as pg
+import patroni.psycopg as pg
 
 from behave import step, then
 from time import sleep, time

features/steps/slots.py

@@ -1,7 +1,7 @@
 import time
-import psycopg2
 
 from behave import step, then
+import patroni.psycopg as pg
 
 
 @step('I create a logical replication slot {slot_name} on {pg_name:w} with the {plugin:w} plugin')
@@ -10,7 +10,7 @@ def create_logical_replication_slot(context, slot_name, pg_name, plugin):
         output = context.pctl.query(pg_name, ("SELECT pg_create_logical_replication_slot('{0}', '{1}'),"
                                               " current_database()").format(slot_name, plugin))
         print(output.fetchone())
-    except psycopg2.Error as e:
+    except pg.Error as e:
         print(e)
         assert False, "Error creating slot {0} on {1} with plugin {2}".format(slot_name, pg_name, plugin)
 
@@ -24,7 +24,7 @@ def has_logical_replication_slot(context, pg_name, slot_name, plugin):
         assert row[0] == "logical", "Found replication slot named {0} but wasn't a logical slot".format(slot_name)
         assert row[1] == plugin, ("Found replication slot named {0} but was using plugin "
                                   "{1} rather than {2}").format(slot_name, row[1], plugin)
-    except psycopg2.Error:
+    except pg.Error:
         assert False, "Error looking for slot {0} on {1} with plugin {2}".format(slot_name, pg_name, plugin)
 
 
@@ -34,7 +34,7 @@ def does_not_have_logical_replication_slot(context, pg_name, slot_name):
         row = context.pctl.query(pg_name, ("SELECT 1 FROM pg_replication_slots"
                                            " WHERE slot_name = '{0}'").format(slot_name)).fetchone()
         assert not row, "Found unexpected replication slot named {0}".format(slot_name)
-    except psycopg2.Error:
+    except pg.Error:
         assert False, "Error looking for slot {0} on {1}".format(slot_name, pg_name)