Merge pull request #46 from encode/add-connection-options

Add connection options

Author: lovelydinosaur

README.md

@@ -173,6 +173,19 @@ async def shutdown():
     await database.disconnect()
 ```
 
+## Connection options
+
+The PostgreSQL and MySQL backends provide a few connection options for SSL
+and for configuring the connection pool.
+
+```python
+# Use an SSL connection.
+database = Database('postgresql://localhost/example?ssl=true')
+
+# Use a connection pool of between 5-20 connections.
+database = Database('mysql://localhost/example?min_size=5&max_size=20')
+```
+
 ## Test isolation
 
 For strict test isolation you will always want to rollback the test database

databases/__init__.py

@@ -1,5 +1,5 @@
 from databases.core import Database, DatabaseURL
 
 
-__version__ = "0.1.6"
+__version__ = "0.1.7"
 __all__ = ["Database", "DatabaseURL"]

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,10 +16,10 @@
 
 
 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)
+        self._pool = SQLitePool(self._database_url)
 
     async def connect(self) -> None:
         pass

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, parse_qsl, urlsplit
 
 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,30 @@
+"""
+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_options():
+    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")