Merge pull request #85 from encode/connection-options

Connection options

Author: tomchristie

README.md

@@ -18,7 +18,7 @@ It allows you to make queries using the powerful [SQLAlchemy Core][sqlalchemy-co
 expression language, and provides support for PostgreSQL, MySQL, and SQLite.
 
 Databases is suitable for integrating against any async Web framework, such as [Starlette][starlette],
-[Sanic][sanic], [Responder][responder], [Quart][quart], [aiohttp][aiohttp], [Tornado][tornado], [FastAPI][fastapi], 
+[Sanic][sanic], [Responder][responder], [Quart][quart], [aiohttp][aiohttp], [Tornado][tornado], [FastAPI][fastapi],
 or [Bocadillo][bocadillo].
 
 **Requirements**: Python 3.6+
@@ -220,6 +220,13 @@ database = Database('postgresql://localhost/example?ssl=true')
 database = Database('mysql://localhost/example?min_size=5&max_size=20')
 ```
 
+You can also use keyword arguments to pass in any connection options.
+Available keyword arguments may differ between database backends.
+
+```python
+database = Database('postgresql://localhost/example', ssl=True, min_size=5, max_size=20)
+```
+
 ## Test isolation
 
 For strict test isolation you will always want to rollback the test database

databases/backends/mysql.py

@@ -17,25 +17,37 @@
 
 
 class MySQLBackend(DatabaseBackend):
-    def __init__(self, database_url: typing.Union[DatabaseURL, str]) -> None:
+    def __init__(
+        self, database_url: typing.Union[DatabaseURL, str], **options: typing.Any
+    ) -> None:
         self._database_url = DatabaseURL(database_url)
+        self._options = options
         self._dialect = pymysql.dialect(paramstyle="pyformat")
         self._pool = None
 
     def _get_connection_kwargs(self) -> dict:
-        options = self._database_url.options
+        url_options = self._database_url.options
 
         kwargs = {}
-        min_size = options.get("min_size")
-        max_size = options.get("max_size")
-        ssl = options.get("ssl")
+        min_size = url_options.get("min_size")
+        max_size = url_options.get("max_size")
+        ssl = url_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()]
+
+        for key, value in self._options.items():
+            # Coerce 'min_size' and 'max_size' for consistency.
+            if key == "min_size":
+                key = "minsize"
+            elif key == "max_size":
+                key = "maxsize"
+            kwargs[key] = value
+
         return kwargs
 
     async def connect(self) -> None:

databases/backends/postgres.py

@@ -19,8 +19,11 @@
 
 
 class PostgresBackend(DatabaseBackend):
-    def __init__(self, database_url: typing.Union[DatabaseURL, str]) -> None:
+    def __init__(
+        self, database_url: typing.Union[DatabaseURL, str], **options: typing.Any
+    ) -> None:
         self._database_url = DatabaseURL(database_url)
+        self._options = options
         self._dialect = self._get_dialect()
         self._pool = None
 
@@ -37,19 +40,22 @@ def _get_dialect(self) -> Dialect:
         return dialect
 
     def _get_connection_kwargs(self) -> dict:
-        options = self._database_url.options
+        url_options = self._database_url.options
 
         kwargs = {}
-        min_size = options.get("min_size")
-        max_size = options.get("max_size")
-        ssl = options.get("ssl")
+        min_size = url_options.get("min_size")
+        max_size = url_options.get("max_size")
+        ssl = url_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()]
+
+        kwargs.update(self._options)
+
         return kwargs
 
     async def connect(self) -> None:

databases/backends/sqlite.py

@@ -16,10 +16,13 @@
 
 
 class SQLiteBackend(DatabaseBackend):
-    def __init__(self, database_url: typing.Union[DatabaseURL, str]) -> None:
+    def __init__(
+        self, database_url: typing.Union[DatabaseURL, str], **options: typing.Any
+    ) -> None:
         self._database_url = DatabaseURL(database_url)
+        self._options = options
         self._dialect = pysqlite.dialect(paramstyle="qmark")
-        self._pool = SQLitePool(self._database_url)
+        self._pool = SQLitePool(self._database_url, **self._options)
 
     async def connect(self) -> None:
         pass
@@ -45,12 +48,13 @@ def connection(self) -> "SQLiteConnection":
 
 
 class SQLitePool:
-    def __init__(self, url: DatabaseURL) -> None:
+    def __init__(self, url: DatabaseURL, **options: typing.Any) -> None:
         self._url = url
+        self._options = options
 
     async def acquire(self) -> aiosqlite.Connection:
         connection = aiosqlite.connect(
-            database=self._url.database, isolation_level=None
+            database=self._url.database, isolation_level=None, **self._options
         )
         await connection.__aenter__()
         return connection

databases/core.py

@@ -25,16 +25,22 @@ class Database:
     }
 
     def __init__(
-        self, url: typing.Union[str, "DatabaseURL"], *, force_rollback: bool = False
+        self,
+        url: typing.Union[str, "DatabaseURL"],
+        *,
+        force_rollback: bool = False,
+        **options: typing.Any,
     ):
         self.url = DatabaseURL(url)
-        self._force_rollback = force_rollback
+        self.options = options
         self.is_connected = False
 
+        self._force_rollback = force_rollback
+
         backend_str = self.SUPPORTED_BACKENDS[self.url.dialect]
         backend_cls = import_from_string(backend_str)
         assert issubclass(backend_cls, DatabaseBackend)
-        self._backend = backend_cls(self.url)
+        self._backend = backend_cls(self.url, **self.options)
 
         # Connections are stored as task-local state.
         self._connection_context = ContextVar("connection_context")  # type: ContextVar

tests/test_connection_options.py

@@ -12,19 +12,43 @@ def test_postgres_pool_size():
     assert kwargs == {"min_size": 1, "max_size": 20}
 
 
+def test_postgres_explicit_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_postgres_explicit_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_explicit_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}
+
+
+def test_mysql_explicit_ssl():
+    backend = MySQLBackend("postgres://localhost/database", ssl=True)
+    kwargs = backend._get_connection_kwargs()
+    assert kwargs == {"ssl": True}