Merge pull request #701 from LerikP/master

Add psycopg3 support for database functions

setup.py

@@ -27,6 +27,7 @@ def get_version():
         'Jinja2>=2.3',
         'docutils>=0.10',
         'flexmock>=0.9.7',
+        'psycopg>=3.1.8',
         'psycopg2>=2.5.1',
         'psycopg2cffi>=2.8.1',
         'pg8000>=1.12.4',

sqlalchemy_utils/functions/database.py

@@ -556,7 +556,7 @@ def create_database(url, encoding='utf8', template=None):
 
     if (dialect_name == 'mssql' and dialect_driver in {'pymssql', 'pyodbc'}) \
             or (dialect_name == 'postgresql' and dialect_driver in {
-            'asyncpg', 'pg8000', 'psycopg2', 'psycopg2cffi'}):
+            'asyncpg', 'pg8000', 'psycopg', 'psycopg2', 'psycopg2cffi'}):
         engine = sa.create_engine(url, isolation_level='AUTOCOMMIT')
     else:
         engine = sa.create_engine(url)
@@ -625,7 +625,7 @@ def drop_database(url):
     if dialect_name == 'mssql' and dialect_driver in {'pymssql', 'pyodbc'}:
         engine = sa.create_engine(url, connect_args={'autocommit': True})
     elif dialect_name == 'postgresql' and dialect_driver in {
-            'asyncpg', 'pg8000', 'psycopg2', 'psycopg2cffi'}:
+            'asyncpg', 'pg8000', 'psycopg', 'psycopg2', 'psycopg2cffi'}:
         engine = sa.create_engine(url, isolation_level='AUTOCOMMIT')
     else:
         engine = sa.create_engine(url)

tests/functions/test_database.py

@@ -2,13 +2,16 @@
 import sqlalchemy as sa
 
 from sqlalchemy_utils import create_database, database_exists, drop_database
+from sqlalchemy_utils.compat import get_sqlalchemy_version
 
 pymysql = None
 try:
     import pymysql  # noqa
 except ImportError:
     pass
 
+sqlalchemy_version = get_sqlalchemy_version()
+
 
 class DatabaseTest:
     def test_create_and_drop(self, dsn):
@@ -98,6 +101,18 @@ def dsn(self, postgresql_db_user, postgresql_db_password):
         )
 
 
+@pytest.mark.skipif('sqlalchemy_version < (2, 0, 0)')
+class TestDatabasePostgresPsycoPG3(DatabaseTest):
+
+    @pytest.fixture
+    def dsn(self, postgresql_db_user, postgresql_db_password):
+        return 'postgresql+psycopg://{}:{}@localhost/{}'.format(
+            postgresql_db_user,
+            postgresql_db_password,
+            'db_to_test_create_and_drop_via_psycopg3_driver'
+        )
+
+
 @pytest.mark.usefixtures('postgresql_dsn')
 class TestDatabasePostgresWithQuotedName(DatabaseTest):