Sticking with postgresql+aiopg and some clean up

Author: gvbgduh

databases/backends/aiopg.py

@@ -1,9 +1,9 @@
+import getpass
 import logging
 import typing
 import uuid
 
 import aiopg
-
 from sqlalchemy.dialects.postgresql import pypostgresql
 from sqlalchemy.engine.interfaces import Dialect, ExecutionContext
 from sqlalchemy.engine.result import ResultMetaData, RowProxy
@@ -38,7 +38,7 @@ def _get_dialect(self) -> Dialect:
 
         return dialect
 
-    def _get_connection_kwargs(self) -> dict:  # TODO move to `core.py`
+    def _get_connection_kwargs(self) -> dict:
         url_options = self._database_url.options
 
         kwargs = {}

databases/core.py

@@ -18,17 +18,9 @@
 
 
 class Database:
-    # TODO Nested schema?
-    # {
-    #     "postgresql": {
-    #         "asyncpg": "...",  # Default
-    #         "aiopg": "..."
-    #     }
-    # }
     SUPPORTED_BACKENDS = {
-        # TODO `postgresql+asyncpg`?
         "postgresql": "databases.backends.postgres:PostgresBackend",
-        "postgresql+psycopg2": "databases.backends.aiopg:AiopgBackend",
+        "postgresql+aiopg": "databases.backends.aiopg:AiopgBackend",
         "mysql": "databases.backends.mysql:MySQLBackend",
         "sqlite": "databases.backends.sqlite:SQLiteBackend",
     }

tests/conftest.py

@@ -2,6 +2,7 @@
 Unit tests for the backend connection arguments.
 """
 
+from databases.backends.aiopg import AiopgBackend
 from databases.backends.mysql import MySQLBackend
 from databases.backends.postgres import PostgresBackend
 
@@ -52,3 +53,29 @@ def test_mysql_explicit_ssl():
     backend = MySQLBackend("mysql://localhost/database", ssl=True)
     kwargs = backend._get_connection_kwargs()
     assert kwargs == {"ssl": True}
+
+
+def test_aiopg_pool_size():
+    backend = AiopgBackend("postgres+aiopg://localhost/database?min_size=1&max_size=20")
+    kwargs = backend._get_connection_kwargs()
+    assert kwargs == {"minsize": 1, "maxsize": 20}
+
+
+def test_aiopg_explicit_pool_size():
+    backend = AiopgBackend(
+        "postgres+aiopg://localhost/database", min_size=1, max_size=20
+    )
+    kwargs = backend._get_connection_kwargs()
+    assert kwargs == {"minsize": 1, "maxsize": 20}
+
+
+def test_aiopg_ssl():
+    backend = AiopgBackend("postgres+aiopg://localhost/database?ssl=true")
+    kwargs = backend._get_connection_kwargs()
+    assert kwargs == {"ssl": True}
+
+
+def test_aiopg_explicit_ssl():
+    backend = AiopgBackend("postgres+aiopg://localhost/database", ssl=True)
+    kwargs = backend._get_connection_kwargs()
+    assert kwargs == {"ssl": True}

tests/test_connection_options.py

@@ -72,14 +72,15 @@ def process_result_value(self, value, dialect):
 )
 
 
-# TODO Move to `conftest.py`
 @pytest.fixture(autouse=True, scope="module")
 def create_test_database():
     # Create test databases with tables creation
     for url in DATABASE_URLS:
         database_url = DatabaseURL(url)
         if database_url.dialect == "mysql":
             url = str(database_url.replace(driver="pymysql"))
+        elif database_url.driver == "aiopg":
+            url = str(database_url.replace(driver=None))
         engine = sqlalchemy.create_engine(url)
         metadata.create_all(engine)
 
@@ -91,6 +92,8 @@ def create_test_database():
         database_url = DatabaseURL(url)
         if database_url.dialect == "mysql":
             url = str(database_url.replace(driver="pymysql"))
+        elif database_url.driver == "aiopg":
+            url = str(database_url.replace(driver=None))
         engine = sqlalchemy.create_engine(url)
         metadata.drop_all(engine)
 
@@ -301,15 +304,19 @@ async def test_execute_return_val(database_url):
             query = notes.insert()
             values = {"text": "example1", "completed": True}
             pk = await database.execute(query, values)
-            # Apparently for `aiopg` it's OID that will always 0 in this case
-            # As it's only one action within this cursor life cycle
-            # Something to triple check
             assert isinstance(pk, int)
 
-            query = notes.select().where(notes.c.id == pk)
-            result = await database.fetch_one(query)
-            assert result["text"] == "example1"
-            assert result["completed"] == True
+            # Apparently for `aiopg` it's OID that will always 0 in this case
+            # As it's only one action within this cursor life cycle
+            # It's recommended to use the `RETURNING` clause
+            # For obtaining the record id
+            if database.url.scheme == "postgresql+aiopg":
+                assert pk == 0
+            else:
+                query = notes.select().where(notes.c.id == pk)
+                result = await database.fetch_one(query)
+                assert result["text"] == "example1"
+                assert result["completed"] == True
 
 
 @pytest.mark.parametrize("database_url", DATABASE_URLS)
@@ -508,14 +515,11 @@ async def test_json_field(database_url):
         async with database.transaction(force_rollback=True):
             # execute()
             query = session.insert()
-            values = {"data": {"text": "hello", "boolean": True, "int": 1}}
-            if str(database_url).startswith("postgresql+psycopg2"):
-                await database.execute(
-                    query.values(
-                        # or wrapped with `psycopg2.extras.Json`
-                        data=json.dumps({"text": "hello", "boolean": True, "int": 1})
-                    )
-                )
+            data = {"text": "hello", "boolean": True, "int": 1}
+            values = {"data": data}
+
+            if database.url.scheme == "postgresql+aiopg":
+                await database.execute(query, {"data": json.dumps(data)})
             else:
                 await database.execute(query, values)
 
@@ -679,29 +683,29 @@ async def test_queries_with_expose_backend_connection(database_url):
                 raw_connection = connection.raw_connection
 
                 # Insert query
-                if str(database_url).startswith("mysql") or str(database_url).startswith("postgresql+psycopg2"):
+                if database.url.scheme in ["mysql", "postgresql+aiopg"]:
                     insert_query = "INSERT INTO notes (text, completed) VALUES (%s, %s)"
                 else:
                     insert_query = "INSERT INTO notes (text, completed) VALUES ($1, $2)"
 
                 # execute()
                 values = ("example1", True)
 
-                if str(database_url).startswith("mysql") or str(database_url).startswith("postgresql+psycopg2"):
+                if database.url.scheme in ["mysql", "postgresql+aiopg"]:
                     cursor = await raw_connection.cursor()
                     await cursor.execute(insert_query, values)
-                elif str(database_url).startswith("postgresql"):
+                elif database.url.scheme == "postgresql":
                     await raw_connection.execute(insert_query, *values)
-                elif str(database_url).startswith("sqlite"):
+                elif database.url.scheme == "sqlite":
                     await raw_connection.execute(insert_query, values)
 
                 # execute_many()
                 values = [("example2", False), ("example3", True)]
 
-                if str(database_url).startswith("mysql"):
+                if database.url.scheme == "mysql":
                     cursor = await raw_connection.cursor()
                     await cursor.executemany(insert_query, values)
-                elif str(database_url).startswith("postgresql+psycopg2"):
+                elif database.url.scheme == "postgresql+aiopg":
                     cursor = await raw_connection.cursor()
                     # No async support for `executemany`
                     for value in values:
@@ -713,13 +717,13 @@ async def test_queries_with_expose_backend_connection(database_url):
                 select_query = "SELECT notes.id, notes.text, notes.completed FROM notes"
 
                 # fetch_all()
-                if str(database_url).startswith("mysql") or str(database_url).startswith("postgresql+psycopg2"):
+                if database.url.scheme in ["mysql", "postgresql+aiopg"]:
                     cursor = await raw_connection.cursor()
                     await cursor.execute(select_query)
                     results = await cursor.fetchall()
-                elif str(database_url).startswith("postgresql"):
+                elif database.url.scheme == "postgresql":
                     results = await raw_connection.fetch(select_query)
-                elif str(database_url).startswith("sqlite"):
+                elif database.url.scheme == "sqlite":
                     results = await raw_connection.execute_fetchall(select_query)
 
                 assert len(results) == 3
@@ -730,9 +734,9 @@ async def test_queries_with_expose_backend_connection(database_url):
                 assert results[1][2] == False
                 assert results[2][1] == "example3"
                 assert results[2][2] == True
-                
+
                 # fetch_one()
-                if str(database_url).startswith("postgresql") and not str(database_url).startswith("postgresql+psycopg2"):
+                if database.url.scheme == "postgresql":
                     result = await raw_connection.fetchrow(select_query)
                 else:
                     cursor = await raw_connection.cursor()

tests/test_databases.py

@@ -23,14 +23,15 @@
 )
 
 
-# TODO Move to `conftest.py` with tables creation
 @pytest.fixture(autouse=True, scope="module")
 def create_test_database():
     # Create test databases
     for url in DATABASE_URLS:
         database_url = DatabaseURL(url)
         if database_url.dialect == "mysql":
             url = str(database_url.replace(driver="pymysql"))
+        elif database_url.driver == "aiopg":
+            url = str(database_url.replace(driver=None))
         engine = sqlalchemy.create_engine(url)
         metadata.create_all(engine)
 
@@ -42,6 +43,8 @@ def create_test_database():
         database_url = DatabaseURL(url)
         if database_url.dialect == "mysql":
             url = str(database_url.replace(driver="pymysql"))
+        elif database_url.driver == "aiopg":
+            url = str(database_url.replace(driver=None))
         engine = sqlalchemy.create_engine(url)
         metadata.drop_all(engine)