Add psycopg3 support for database functions

kvesteri · Apr 7, 2023 · 0d0d02e · 0d0d02e
1 parent d81ea06
commit 0d0d02e
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/sqlalchemy_utils/functions/database.py b/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)

diff --git a/tests/functions/test_database.py b/tests/functions/test_database.py
@@ -98,6 +98,17 @@ def dsn(self, postgresql_db_user, postgresql_db_password):
         )
 
 
+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):