Add min_size, max_size, and ssl connection options

Author: lovelydinosaur

databases/backends/mysql.py

@@ -17,20 +17,38 @@
 
 
 class MySQLBackend(DatabaseBackend):
-    def __init__(self, database_url: DatabaseURL) -> None:
-        self._database_url = database_url
+    def __init__(self, database_url: typing.Union[DatabaseURL, str]) -> None:
+        self._database_url = DatabaseURL(database_url)
         self._dialect = pymysql.dialect(paramstyle="pyformat")
         self._pool = None
 
+    def _get_connection_kwargs(self) -> dict:
+        options = self._database_url.options
+
+        kwargs = {}
+        min_size = options.get('min_size')
+        max_size = options.get('max_size')
+        ssl = options.get('ssl')
+
+        if min_size is not None:
+            kwargs['minsize'] = int(min_size)
+        if max_size is not None:
+            kwargs['maxsize'] = int(max_size)
+        if ssl is not None:
+            kwargs['ssl'] = {'true': True, 'false': False}[ssl.lower()]
+        return kwargs
+
     async def connect(self) -> None:
         assert self._pool is None, "DatabaseBackend is already running"
+        kwargs = self._get_connection_kwargs()
         self._pool = await aiomysql.create_pool(
             host=self._database_url.hostname,
             port=self._database_url.port or 3306,
             user=self._database_url.username or getpass.getuser(),
             password=self._database_url.password,
             db=self._database_url.database,
             autocommit=True,
+            **kwargs
         )
 
     async def disconnect(self) -> None:

databases/backends/postgres.py

@@ -16,8 +16,8 @@
 
 
 class PostgresBackend(DatabaseBackend):
-    def __init__(self, database_url: DatabaseURL) -> None:
-        self._database_url = database_url
+    def __init__(self, database_url: typing.Union[DatabaseURL, str]) -> None:
+        self._database_url = DatabaseURL(database_url)
         self._dialect = self._get_dialect()
         self._pool = None
 
@@ -33,9 +33,26 @@ def _get_dialect(self) -> Dialect:
 
         return dialect
 
+    def _get_connection_kwargs(self) -> dict:
+        options = self._database_url.options
+
+        kwargs = {}
+        min_size = options.get('min_size')
+        max_size = options.get('max_size')
+        ssl = options.get('ssl')
+
+        if min_size is not None:
+            kwargs['min_size'] = int(min_size)
+        if max_size is not None:
+            kwargs['max_size'] = int(max_size)
+        if ssl is not None:
+            kwargs['ssl'] = {'true': True, 'false': False}[ssl.lower()]
+        return kwargs
+
     async def connect(self) -> None:
         assert self._pool is None, "DatabaseBackend is already running"
-        self._pool = await asyncpg.create_pool(str(self._database_url))
+        kwargs = self._get_connection_kwargs()
+        self._pool = await asyncpg.create_pool(str(self._database_url), **kwargs)
 
     async def disconnect(self) -> None:
         assert self._pool is not None, "DatabaseBackend is not running"

databases/backends/sqlite.py

@@ -16,8 +16,8 @@
 
 
 class SQLiteBackend(DatabaseBackend):
-    def __init__(self, database_url: DatabaseURL) -> None:
-        self._database_url = database_url
+    def __init__(self, database_url: typing.Union[DatabaseURL, str]) -> None:
+        self._database_url = DatabaseURL(database_url)
         self._dialect = pysqlite.dialect(paramstyle="qmark")
         self._pool = SQLitePool(database_url)
 

databases/core.py

@@ -3,7 +3,7 @@
 import sys
 import typing
 from types import TracebackType
-from urllib.parse import SplitResult, urlsplit
+from urllib.parse import SplitResult, urlsplit, parse_qsl
 
 from sqlalchemy.engine import RowProxy
 from sqlalchemy.sql import ClauseElement
@@ -295,6 +295,12 @@ def netloc(self) -> typing.Optional[str]:
     def database(self) -> str:
         return self.components.path.lstrip("/")
 
+    @property
+    def options(self) -> dict:
+        if not hasattr(self, "_options"):
+            self._options = dict(parse_qsl(self.components.query))
+        return self._options
+
     def replace(self, **kwargs: typing.Any) -> "DatabaseURL":
         if (
             "username" in kwargs

tests/test_connection_options.py

@@ -0,0 +1,40 @@
+"""
+Unit tests for the backend connection arguments.
+"""
+
+from databases.backends.mysql import MySQLBackend
+from databases.backends.postgres import PostgresBackend
+
+
+def test_postgres_pool_size():
+    backend = PostgresBackend("postgres://localhost/database?min_size=1&max_size=20")
+    kwargs = backend._get_connection_kwargs()
+    assert kwargs == {
+        'min_size': 1,
+        'max_size': 20
+    }
+
+
+def test_postgres_ssl():
+    backend = PostgresBackend("postgres://localhost/database?ssl=true")
+    kwargs = backend._get_connection_kwargs()
+    assert kwargs == {
+        'ssl': True,
+    }
+
+
+def test_mysql_pool_size():
+    backend = MySQLBackend("mysql://localhost/database?min_size=1&max_size=20")
+    kwargs = backend._get_connection_kwargs()
+    assert kwargs == {
+        'minsize': 1,
+        'maxsize': 20
+    }
+
+
+def test_mysql_ssl():
+    backend = MySQLBackend("postgres://localhost/database?ssl=true")
+    kwargs = backend._get_connection_kwargs()
+    assert kwargs == {
+        'ssl': True,
+    }

tests/test_database_url.py

@@ -23,6 +23,11 @@ def test_database_url_properties():
     assert u.database == "mydatabase"
 
 
+def test_database_url_properties():
+    u = DatabaseURL("postgresql://localhost/mydatabase?pool_size=20&ssl=true")
+    assert u.options == {"pool_size": "20", "ssl": "true"}
+
+
 def test_replace_database_url_components():
     u = DatabaseURL("postgresql://localhost/mydatabase")