Add postgresql+asyncpg scheme prefix (closes #395)

.github/workflows/test-suite.yml

@@ -47,5 +47,5 @@ jobs:
         run: "scripts/install"
       - name: "Run tests"
         env:
-          TEST_DATABASE_URLS: "sqlite:///testsuite, mysql://username:password@localhost:3306/testsuite, postgresql://username:password@localhost:5432/testsuite, postgresql+aiopg://username:[email protected]:5432/testsuite"
+          TEST_DATABASE_URLS: "sqlite:///testsuite, sqlite+aiosqlite:///testsuite, mysql://username:password@localhost:3306/testsuite, mysql+aiomysql://username:password@localhost:3306/testsuite, postgresql://username:password@localhost:5432/testsuite, postgresql+aiopg://username:[email protected]:5432/testsuite, postgresql+asyncpg://username:password@localhost:5432/testsuite"
         run: "scripts/test"

databases/core.py

@@ -62,7 +62,7 @@ def __init__(
 
         self._force_rollback = force_rollback
 
-        backend_str = self.SUPPORTED_BACKENDS[self.url.scheme]
+        backend_str = self._get_backend()
         backend_cls = import_from_string(backend_str)
         assert issubclass(backend_cls, DatabaseBackend)
         self._backend = backend_cls(self.url, **self.options)
@@ -220,6 +220,12 @@ def force_rollback(self) -> typing.Iterator[None]:
         finally:
             self._force_rollback = initial
 
+    def _get_backend(self) -> str:
+        try:
+            return self.SUPPORTED_BACKENDS[self.url.scheme]
+        except KeyError:
+            return self.SUPPORTED_BACKENDS[self.url.dialect]
+
 
 class Connection:
     def __init__(self, backend: DatabaseBackend) -> None:

tests/test_databases.py

@@ -83,9 +83,11 @@ def create_test_database():
     # Create test databases with tables creation
     for url in DATABASE_URLS:
         database_url = DatabaseURL(url)
-        if database_url.scheme == "mysql":
+        if database_url.scheme in ["mysql", "mysql+aiomysql"]:
             url = str(database_url.replace(driver="pymysql"))
-        elif database_url.scheme == "postgresql+aiopg":
+        elif database_url.scheme in ["postgresql+aiopg", "postgresql+asyncpg"]:
+            url = str(database_url.replace(driver=None))
+        elif database_url.scheme in ["sqlite", "sqlite+aiosqlite"]:
             url = str(database_url.replace(driver=None))
         engine = sqlalchemy.create_engine(url)
         metadata.create_all(engine)
@@ -96,9 +98,11 @@ def create_test_database():
     # Drop test databases
     for url in DATABASE_URLS:
         database_url = DatabaseURL(url)
-        if database_url.scheme == "mysql":
+        if database_url.scheme in ["mysql", "mysql+aiomysql"]:
             url = str(database_url.replace(driver="pymysql"))
-        elif database_url.scheme == "postgresql+aiopg":
+        elif database_url.scheme in ["postgresql+aiopg", "postgresql+asyncpg"]:
+            url = str(database_url.replace(driver=None))
+        elif database_url.scheme in ["sqlite", "sqlite+aiosqlite"]:
             url = str(database_url.replace(driver=None))
         engine = sqlalchemy.create_engine(url)
         metadata.drop_all(engine)
@@ -855,7 +859,7 @@ async def test_queries_with_expose_backend_connection(database_url):
                 if database.url.scheme in ["mysql", "postgresql+aiopg"]:
                     cursor = await raw_connection.cursor()
                     await cursor.execute(insert_query, values)
-                elif database.url.scheme == "postgresql":
+                elif database.url.scheme in ["postgresql", "postgresql+asyncpg"]:
                     await raw_connection.execute(insert_query, *values)
                 elif database.url.scheme == "sqlite":
                     await raw_connection.execute(insert_query, values)
@@ -882,7 +886,7 @@ async def test_queries_with_expose_backend_connection(database_url):
                     cursor = await raw_connection.cursor()
                     await cursor.execute(select_query)
                     results = await cursor.fetchall()
-                elif database.url.scheme == "postgresql":
+                elif database.url.scheme in ["postgresql", "postgresql+asyncpg"]:
                     results = await raw_connection.fetch(select_query)
                 elif database.url.scheme == "sqlite":
                     results = await raw_connection.execute_fetchall(select_query)
@@ -897,7 +901,7 @@ async def test_queries_with_expose_backend_connection(database_url):
                 assert results[2][2] == True
 
                 # fetch_one()
-                if database.url.scheme == "postgresql":
+                if database.url.scheme in ["postgresql", "postgresql+asyncpg"]:
                     result = await raw_connection.fetchrow(select_query)
                 else:
                     cursor = await raw_connection.cursor()
@@ -1065,7 +1069,7 @@ async def test_posgres_interface(database_url):
     """
     database_url = DatabaseURL(database_url)
 
-    if database_url.scheme != "postgresql":
+    if database_url.scheme not in ["postgresql", "postgresql+asyncpg"]:
         pytest.skip("Test is only for postgresql")
 
     async with Database(database_url) as database:

tests/test_integration.py

@@ -28,9 +28,11 @@ def create_test_database():
     # Create test databases
     for url in DATABASE_URLS:
         database_url = DatabaseURL(url)
-        if database_url.scheme == "mysql":
+        if database_url.scheme in ["mysql", "mysql+aiomysql"]:
             url = str(database_url.replace(driver="pymysql"))
-        elif database_url.scheme == "postgresql+aiopg":
+        elif database_url.scheme in ["postgresql+aiopg", "postgresql+asyncpg"]:
+            url = str(database_url.replace(driver=None))
+        elif database_url.scheme in ["sqlite", "sqlite+aiosqlite"]:
             url = str(database_url.replace(driver=None))
         engine = sqlalchemy.create_engine(url)
         metadata.create_all(engine)
@@ -41,9 +43,11 @@ def create_test_database():
     # Drop test databases
     for url in DATABASE_URLS:
         database_url = DatabaseURL(url)
-        if database_url.scheme == "mysql":
+        if database_url.scheme in ["mysql", "mysql+aiomysql"]:
             url = str(database_url.replace(driver="pymysql"))
-        elif database_url.scheme == "postgresql+aiopg":
+        elif database_url.scheme in ["postgresql+aiopg", "postgresql+asyncpg"]:
+            url = str(database_url.replace(driver=None))
+        elif database_url.scheme in ["sqlite", "sqlite+aiosqlite"]:
             url = str(database_url.replace(driver=None))
         engine = sqlalchemy.create_engine(url)
         metadata.drop_all(engine)