Add support for Psycopg 3

7 files changed, 19 insertions, 1 deletions

diff --git a/doc/index.rst b/doc/index.rst
index ca82512..f3874a5 100644
--- a/doc/index.rst
+++ b/doc/index.rst
@@ -120,6 +120,9 @@ Database connections are specified using a URL. Examples:
   # PostgreSQL: unix socket connection
   database = postgresql://scott:tiger@/mydatabase
 
+  # PostgreSQL with the newer psycopg 3 driver
+  database = postgresql+psycopg://scott:tiger@localhost/mydatabase
+
   # PostgreSQL: changing the schema (via set search_path)
   database = postgresql://scott:tiger@/mydatabase?schema=some_schema
 
@@ -498,7 +501,7 @@ The following example shows how to apply migrations from inside python code:
     from yoyo import read_migrations
     from yoyo import get_backend
 
-    backend = get_backend('postgres://myuser@localhost/mydatabase')
+    backend = get_backend('postgresql://myuser@localhost/mydatabase')
     migrations = read_migrations('path/to/migrations')
 
     with backend.lock():
diff --git a/setup.cfg b/setup.cfg
index edda49f..c45c98c 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -52,6 +52,7 @@ yoyo.backends =
     postgres = yoyo.backends.core.postgresql:PostgresqlBackend
     postgresql = yoyo.backends.core.postgresql:PostgresqlBackend
     psql = yoyo.backends.core.postgresql:PostgresqlBackend
+    postgresql+psycopg = yoyo.backends.core.postgresql:PostgresqlPsycopgBackend
     mysql = yoyo.backends.core.mysql:MySQLBackend
     mysql+mysqldb = yoyo.backends.core.mysql:MySQLdbBackend
     sqlite = yoyo.backends.core.sqlite3:SQLiteBackend
diff --git a/test_databases.ini b/test_databases.ini
index 9050875..e072814 100644
--- a/test_databases.ini
+++ b/test_databases.ini
@@ -1,5 +1,6 @@
 [DEFAULT]
 sqlite = sqlite:///:memory:
 postgresql = postgresql://postgres@/yoyo_test
+postgresqlpsycopg3 = postgresql+psycopg://postgres@/yoyo_test
 mysql = mysql://root@127.0.0.1/yoyo_test
 #redshift = redshift://<user>:<password>@<subdomain>.redshift.amazonaws.com:5439/<dbname>
diff --git a/tox.ini b/tox.ini
index 1bfd45e..2c80fed 100755
--- a/tox.ini
+++ b/tox.ini
@@ -5,6 +5,7 @@ envlist = py37,py38,py39,py310,py310-sphinx,py310-linters
 deps=
     PyMySQL
     psycopg2
+    psycopg
     pytest
     freezegun
     tms
diff --git a/yoyo/backends/__init__.py b/yoyo/backends/__init__.py
index c4d6f6e..2be7421 100644
--- a/yoyo/backends/__init__.py
+++ b/yoyo/backends/__init__.py
@@ -3,6 +3,7 @@ from yoyo.backends.base import get_backend_class
 from yoyo.backends.core import MySQLBackend
 from yoyo.backends.core import SQLiteBackend
 from yoyo.backends.core import PostgresqlBackend
+from yoyo.backends.core import PostgresqlPsycopgBackend
 
 __all__ = [
     "DatabaseBackend",
@@ -10,4 +11,5 @@ __all__ = [
     "MySQLBackend",
     "SQLiteBackend",
     "PostgresqlBackend",
+    "PostgresqlPsycopgBackend",
 ]
diff --git a/yoyo/backends/core/__init__.py b/yoyo/backends/core/__init__.py
index 183e24c..c7162a2 100644
--- a/yoyo/backends/core/__init__.py
+++ b/yoyo/backends/core/__init__.py
@@ -1,9 +1,11 @@
 from yoyo.backends.core.mysql import MySQLBackend
 from yoyo.backends.core.sqlite3 import SQLiteBackend
 from yoyo.backends.core.postgresql import PostgresqlBackend
+from yoyo.backends.core.postgresql import PostgresqlPsycopgBackend
 
 __all__ = [
     "MySQLBackend",
     "SQLiteBackend",
     "PostgresqlBackend",
+    "PostgresqlPsycopgBackend",
 ]
diff --git a/yoyo/backends/core/postgresql.py b/yoyo/backends/core/postgresql.py
index b5bf3c3..0400ce9 100644
--- a/yoyo/backends/core/postgresql.py
+++ b/yoyo/backends/core/postgresql.py
@@ -56,3 +56,11 @@ class PostgresqlBackend(DatabaseBackend):
     def list_tables(self):
         current_schema = self.execute("SELECT current_schema").fetchone()[0]
         return super(PostgresqlBackend, self).list_tables(schema=current_schema)
+
+
+class PostgresqlPsycopgBackend(PostgresqlBackend):
+    """
+    Like PostgresqlBackend, but using the newer Psycopg 3.
+    """
+
+    driver_module = "psycopg"