Merge pull request piccolo-orm#65 from piccolo-orm/improved_engine_setup

modified postgres engine - can specify extensions to install on start

Author: dantownsend

docs/src/piccolo/engines/index.rst

@@ -10,6 +10,8 @@ It's important that each ``Table`` class knows which engine to use. There are
 two ways of doing this - setting it explicitly via the ``db`` argument, or
 letting Piccolo find it using ``engine_finder``.
 
+-------------------------------------------------------------------------------
+
 Explicit
 --------
 
@@ -30,6 +32,7 @@ connect to a database.
     class MyTable(Table, db=DB):
         name = Varchar()
 
+-------------------------------------------------------------------------------
 
 engine_finder
 -------------
@@ -94,6 +97,8 @@ variable accordingly.
 
 .. hint:: You can also specify sub modules, like `my_module.piccolo_conf`.
 
+-------------------------------------------------------------------------------
+
 .. _EngineTypes:
 
 Engine types
@@ -103,83 +108,8 @@ Engine types
  production. It is the most feature complete.
 
 
-SQLiteEngine
-~~~~~~~~~~~~
-
-.. code-block:: python
-
-    # piccolo_conf.py
-    from piccolo.engine.sqlite import SQLiteEngine
-
-
-    DB = SQLiteEngine(path='my_app.sqlite')
-
-
-PostgresEngine
-~~~~~~~~~~~~~~
-
-.. code-block:: python
-
-    # piccolo_conf.py
-    from piccolo.engine.postgres import PostgresEngine
-
-
-    DB = PostgresEngine({
-        'host': 'localhost',
-        'database': 'my_app',
-        'user': 'postgres',
-        'password': ''
-    })
-
-Connection pool
----------------
-
-.. warning:: This is currently only available for Postgres.
-
-
-To use a connection pool, you need to first initialise it. The best place to do
-this is in the startup event handler of whichever web framework you are using.
-
-Here's an example using Starlette. Notice that we also close the connection
-pool in the shutdown event handler.
-
-.. code-block:: python
-
-    from piccolo.engine import from starlette.applications import Starlette
-    from starlette.applications import Starlette
-
-
-    app = Starlette()
-
-
-    @app.on_event('startup')
-    async def open_database_connection_pool():
-        engine = engine_finder()
-        await engine.start_connnection_pool()
-
-
-    @app.on_event('shutdown')
-    async def close_database_connection_pool():
-        engine = engine_finder()
-        await engine.close_connnection_pool()
-
-.. hint:: Using a connection pool helps with performance, since connections
-    are reused instead of being created for each query.
-
-Once a connection pool has been started, the engine will use it for making
-queries.
-
-.. hint:: If you're running several instances of an app on the same server,
-    you may prefer an external connection pooler - like pgbouncer.
-
-Configuration
-~~~~~~~~~~~~~
-
-The connection pool uses the same configuration as your engine. You can also
-pass in additional parameters, which are passed to the underlying database
-adapter. Here's an example:
-
-.. code-block:: python
+.. toctree::
+    :maxdepth: 1
 
-    # To increase the number of connections available:
-    await engine.start_connnection_pool(max_size=20)
+    ./sqlite_engine
+    ./postgres_engine

docs/src/piccolo/engines/postgres_engine.rst

@@ -0,0 +1,86 @@
+PostgresEngine
+==============
+
+Configuration
+-------------
+
+.. code-block:: python
+
+    # piccolo_conf.py
+    from piccolo.engine.postgres import PostgresEngine
+
+
+    DB = PostgresEngine(config={
+        'host': 'localhost',
+        'database': 'my_app',
+        'user': 'postgres',
+        'password': ''
+    })
+
+config
+~~~~~~
+
+The config dictionary is passed directly to the underlying database adapter,
+asyncpg. See the `asyncpg docs <https://magicstack.github.io/asyncpg/current/api/index.html#connection>`_
+to learn more.
+
+-------------------------------------------------------------------------------
+
+Connection pool
+---------------
+
+To use a connection pool, you need to first initialise it. The best place to do
+this is in the startup event handler of whichever web framework you are using.
+
+Here's an example using Starlette. Notice that we also close the connection
+pool in the shutdown event handler.
+
+.. code-block:: python
+
+    from piccolo.engine import from starlette.applications import Starlette
+    from starlette.applications import Starlette
+
+
+    app = Starlette()
+
+
+    @app.on_event('startup')
+    async def open_database_connection_pool():
+        engine = engine_finder()
+        await engine.start_connnection_pool()
+
+
+    @app.on_event('shutdown')
+    async def close_database_connection_pool():
+        engine = engine_finder()
+        await engine.close_connnection_pool()
+
+.. hint:: Using a connection pool helps with performance, since connections
+    are reused instead of being created for each query.
+
+Once a connection pool has been started, the engine will use it for making
+queries.
+
+.. hint:: If you're running several instances of an app on the same server,
+    you may prefer an external connection pooler - like pgbouncer.
+
+Configuration
+~~~~~~~~~~~~~
+
+The connection pool uses the same configuration as your engine. You can also
+pass in additional parameters, which are passed to the underlying database
+adapter. Here's an example:
+
+.. code-block:: python
+
+    # To increase the number of connections available:
+    await engine.start_connnection_pool(max_size=20)
+
+-------------------------------------------------------------------------------
+
+Source
+------
+
+.. currentmodule:: piccolo.engine.postgres
+
+.. autoclass:: PostgresEngine

docs/src/piccolo/engines/sqlite_engine.rst

@@ -0,0 +1,25 @@
+SQLiteEngine
+============
+
+Configuration
+-------------
+
+The ``SQLiteEngine`` is very simple - just specify a file path. The database
+file will be created automatically if it doesn't exist.
+
+.. code-block:: python
+
+    # piccolo_conf.py
+    from piccolo.engine.sqlite import SQLiteEngine
+
+
+    DB = SQLiteEngine(path='my_app.sqlite')
+
+-------------------------------------------------------------------------------
+
+Source
+------
+
+.. currentmodule:: piccolo.engine.sqlite
+
+.. autoclass:: SQLiteEngine

piccolo/engine/base.py

@@ -2,6 +2,7 @@
 from abc import abstractmethod, ABCMeta
 import typing as t
 
+from piccolo.utils.sync import run_sync
 from piccolo.utils.warnings import colored_warning, Level
 
 if t.TYPE_CHECKING:  # pragma: no cover
@@ -14,8 +15,8 @@ class Batch:
 
 class Engine(metaclass=ABCMeta):
     def __init__(self):
-        self.check_version()
-        self.prep_database()
+        run_sync(self.check_version())
+        run_sync(self.prep_database())
 
     @property
     @abstractmethod
@@ -28,23 +29,23 @@ def min_version_number(self) -> float:
         pass
 
     @abstractmethod
-    def get_version(self) -> float:
+    async def get_version(self) -> float:
         pass
 
     @abstractmethod
-    def prep_database(self):
+    async def prep_database(self):
         pass
 
     @abstractmethod
     async def batch(self, query: Query, batch_size: int = 100) -> Batch:
         pass
 
-    def check_version(self):
+    async def check_version(self):
         """
         Warn if the database version isn't supported.
         """
         try:
-            version_number = self.get_version()
+            version_number = await self.get_version()
         except Exception as exception:
             colored_warning(
                 f"Unable to fetch server version: {exception}",

piccolo/engine/postgres.py

@@ -1,6 +1,4 @@
 from __future__ import annotations
-import asyncio
-from concurrent.futures import ThreadPoolExecutor
 import contextvars
 from dataclasses import dataclass
 import typing as t
@@ -195,14 +193,10 @@ async def __aexit__(self, exception_type, exception, traceback):
 
 
 class PostgresEngine(Engine):
+    """
+    Used to connect to Postgresql.
 
-    __slots__ = ("config", "pool")
-
-    engine_type = "postgres"
-    min_version_number = 9.6
-
-    def __init__(self, config: t.Dict[str, t.Any]) -> None:
-        """
+    :param config:
         The config dictionary is passed to the underlying database adapter,
         asyncpg. Common arguments you're likely to need are:
 
@@ -218,8 +212,24 @@ def __init__(self, config: t.Dict[str, t.Any]) -> None:
 
             * https://magicstack.github.io/asyncpg/current/api/index.html#connection
 
-        """  # noqa: E501
+    :param extensions:
+        When the engine starts, it will try and create these extensions
+        in Postgres.
+
+    """  # noqa: E501
+
+    __slots__ = ("config", "extensions", "pool", "transaction_connection")
+
+    engine_type = "postgres"
+    min_version_number = 9.6
+
+    def __init__(
+        self,
+        config: t.Dict[str, t.Any],
+        extensions: t.Sequence[str] = ["uuid-ossp"],
+    ) -> None:
         self.config = config
+        self.extensions = extensions
         self.pool: t.Optional[Pool] = None
         database_name = config.get("database", "Unknown")
         self.transaction_connection = contextvars.ContextVar(
@@ -241,20 +251,14 @@ def _parse_raw_version_string(version_string: str) -> float:
         version = float(f"{major}.{minor}")
         return version
 
-    def get_version(self) -> float:
+    async def get_version(self) -> float:
         """
         Returns the version of Postgres being run.
         """
-        loop = asyncio.new_event_loop()
-
-        with ThreadPoolExecutor(max_workers=1) as executor:
-            future = executor.submit(
-                loop.run_until_complete,
-                self._run_in_new_connection("SHOW server_version"),
-            )
-
         try:
-            response: t.Sequence[t.Dict] = future.result()  # type: ignore
+            response: t.Sequence[t.Dict] = await self._run_in_new_connection(
+                "SHOW server_version"
+            )
         except ConnectionRefusedError as exception:
             # Suppressing the exception, otherwise importing piccolo_conf.py
             # containing an engine will raise an ImportError.
@@ -267,26 +271,20 @@ def get_version(self) -> float:
                 version_string=version_string
             )
 
-    def prep_database(self):
-        loop = asyncio.new_event_loop()
-
-        with ThreadPoolExecutor(max_workers=1) as executor:
-            future = executor.submit(
-                loop.run_until_complete,
-                self._run_in_new_connection(
-                    'CREATE EXTENSION IF NOT EXISTS "uuid-ossp"'
-                ),
-            )
-
-        try:
-            future.result()
-        except InsufficientPrivilegeError:
-            print(
-                "Unable to create uuid-ossp extension - UUID columns might "
-                "not behave as expected. Make sure your database user has "
-                "permission to create extensions, or add it manually using "
-                '`CREATE EXTENSION "uuid-ossp";`'
-            )
+    async def prep_database(self):
+        for extension in self.extensions:
+            try:
+                await self._run_in_new_connection(
+                    f'CREATE EXTENSION IF NOT EXISTS "{extension}"',
+                )
+            except InsufficientPrivilegeError:
+                print(
+                    f"Unable to create {extension} extension - some "
+                    "functionality may not behave as expected. Make sure your "
+                    "database user has permission to create extensions, or "
+                    "add it manually using "
+                    f'`CREATE EXTENSION "{extension}";`'
+                )
 
     ###########################################################################