modified postgres engine - can specify extensions to install on start

Author: dantownsend

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
@@ -201,25 +199,35 @@ class PostgresEngine(Engine):
     engine_type = "postgres"
     min_version_number = 9.6
 
-    def __init__(self, config: t.Dict[str, t.Any]) -> None:
+    def __init__(
+        self,
+        config: t.Dict[str, t.Any],
+        extensions: t.Sequence[str] = ["uuid-ossp"],
+    ) -> None:
         """
-        The config dictionary is passed to the underlying database adapter,
-        asyncpg. Common arguments you're likely to need are:
+        :param config:
+            The config dictionary is passed to the underlying database adapter,
+            asyncpg. Common arguments you're likely to need are:
 
-            * host
-            * port
-            * user
-            * password
-            * database
+                * host
+                * port
+                * user
+                * password
+                * database
 
-        For example, ``{'host': 'localhost', 'port': 5432}``.
+            For example, ``{'host': 'localhost', 'port': 5432}``.
 
-        To see all available options:
+            To see all available options:
 
-            * https://magicstack.github.io/asyncpg/current/api/index.html#connection
+                * https://magicstack.github.io/asyncpg/current/api/index.html#connection
+
+        :param extensions:
+            When the engine starts, it will try and create these extensions
+            in Postgres.
 
         """  # noqa: E501
         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 +249,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 +269,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}";`'
+                )
 
     ###########################################################################
 

piccolo/engine/sqlite.py

@@ -337,14 +337,14 @@ def path(self):
     def path(self, value: str):
         self.connection_kwargs["database"] = value
 
-    def get_version(self) -> float:
+    async def get_version(self) -> float:
         """
         Warn if the version of SQLite isn't supported.
         """
         major, minor, _ = sqlite3.sqlite_version_info
         return float(f"{major}.{minor}")
 
-    def prep_database(self):
+    async def prep_database(self):
         pass
 
     ###########################################################################