LazyLoad database engines (piccolo-orm#147)

aminalaee · dantownsend · web-flow · commit 9ae5cf2164cc · 2021-08-09T07:39:45.000+01:00
* LazyLoader draft

* modify extra requirements

* fix github workflow

* Update the installation docs

* throw specific errors if driver is not installed

* use `str` on exception for backwards compatibility

* increase `lazy_laoder.py` coverage

* mock module instead of object

* set `postgres_only` and `sqlite_only` for tests

* increase coverage of `postgres.py` and `sqlite.py`

* update the installation message for the playground

* change LazyLoader super method

* Update README

* Remove Travis badge and add other badges

* add codecov badge

* update asgi template requirements.txt file to install asyncpg

Co-authored-by: Daniel Townsend &lt;dan@dantownsend.co.uk&gt;
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
@@ -64,6 +64,7 @@ jobs:
                   python -m pip install --upgrade pip
                   pip install -r requirements/requirements.txt
                   pip install -r requirements/test-requirements.txt
+                  pip install -r requirements/extras/postgres.txt
             - name: Setup postgres
               run: |
                   export PGPASSWORD=postgres
@@ -99,6 +100,7 @@ jobs:
                   python -m pip install --upgrade pip
                   pip install -r requirements/requirements.txt
                   pip install -r requirements/test-requirements.txt
+                  pip install -r requirements/extras/sqlite.txt
             - name: Test with pytest, SQLite
               run: ./scripts/test-sqlite.sh
             - name: Upload coverage
diff --git a/README.md b/README.md
@@ -1,9 +1,12 @@
 # Piccolo
 
-[![Build Status](https://travis-ci.com/piccolo-orm/piccolo.svg?branch=master)](https://travis-ci.com/piccolo-orm/piccolo)
+![Tests](https://github.com/piccolo-orm/piccolo/actions/workflows/tests.yaml/badge.svg)
+![Release](https://github.com/piccolo-orm/piccolo/actions/workflows/release.yaml/badge.svg)
 [![Documentation Status](https://readthedocs.org/projects/piccolo-orm/badge/?version=latest)](https://piccolo-orm.readthedocs.io/en/latest/?badge=latest)
+[![PyPI](https://img.shields.io/pypi/v/piccolo?color=%2334D058&label=pypi)](https://pypi.org/project/piccolo/)
 [![Language grade: Python](https://img.shields.io/lgtm/grade/python/g/piccolo-orm/piccolo.svg?logo=lgtm&logoWidth=18)](https://lgtm.com/projects/g/piccolo-orm/piccolo/context:python)
 [![Total alerts](https://img.shields.io/lgtm/alerts/g/piccolo-orm/piccolo.svg?logo=lgtm&logoWidth=18)](https://lgtm.com/projects/g/piccolo-orm/piccolo/alerts/)
+[![codecov](https://codecov.io/gh/piccolo-orm/piccolo/branch/master/graph/badge.svg?token=V19CWH7MXX)](https://codecov.io/gh/piccolo-orm/piccolo)
 
 A fast, user friendly ORM and query builder which supports asyncio. [Read the docs](https://piccolo-orm.readthedocs.io/en/latest/).
 
@@ -66,8 +69,16 @@ await b.remove().run()
 
 ## Installation
 
+Installing with PostgreSQL driver:
+
+```
+pip install piccolo[postgres]
+```
+
+Installing with SQLite driver:
+
 ```
-pip install piccolo
+pip install piccolo[sqlite]
 ```
 
 ## Building a web app?
diff --git a/docs/src/piccolo/getting_started/installing_piccolo.rst b/docs/src/piccolo/getting_started/installing_piccolo.rst
@@ -21,6 +21,12 @@ Now install piccolo, ideally inside a `virtualenv <https://docs.python-guide.org
     # The important bit:
     pip install piccolo
 
+    # Install Piccolo with PostgreSQL driver:
+    pip install piccolo[postgres]
+
+    # Install Piccolo with SQLite driver:
+    pip install piccolo[sqlite]
+
     # Optional: orjson for improved JSON serialisation performance
     pip install piccolo[orjson]
 
diff --git a/piccolo/apps/asgi/commands/templates/app/requirements.txt.jinja b/piccolo/apps/asgi/commands/templates/app/requirements.txt.jinja
@@ -1,5 +1,5 @@
 {{ router }}
 {{ server }}
 jinja2
-piccolo
+piccolo[postgres]
 piccolo_admin
diff --git a/piccolo/apps/playground/commands/run.py b/piccolo/apps/playground/commands/run.py
@@ -145,7 +145,8 @@ def run(
         import IPython  # type: ignore
     except ImportError:
         sys.exit(
-            "Install iPython using `pip install ipython` to use this feature."
+            "Install iPython using `pip install piccolo[playground,sqlite]` "
+            "to use this feature."
         )
 
     if engine.upper() == "POSTGRES":
diff --git a/piccolo/conf/apps.py b/piccolo/conf/apps.py
@@ -316,7 +316,7 @@ def get_piccolo_conf_module(
 
         try:
             module = t.cast(PiccoloConfModule, import_module(module_name))
-        except ModuleNotFoundError:
+        except ModuleNotFoundError as exc:
             if self.diagnose:
                 colored_warning(
                     (
@@ -326,7 +326,19 @@ def get_piccolo_conf_module(
                     level=Level.high,
                 )
                 print(traceback.format_exc())
-            return None
+
+            if str(exc) == "No module named 'asyncpg'":
+                raise ModuleNotFoundError(
+                    "PostgreSQL driver not found. "
+                    "Try running 'pip install piccolo[postgres]'"
+                )
+            elif str(exc) == "No module named 'aiosqlite'":
+                raise ModuleNotFoundError(
+                    "SQLite driver not found. "
+                    "Try running 'pip install piccolo[sqlite]'"
+                )
+            else:
+                raise exc
         else:
             return module
 
diff --git a/piccolo/engine/postgres.py b/piccolo/engine/postgres.py
@@ -4,19 +4,21 @@
 import typing as t
 from dataclasses import dataclass
 
-import asyncpg  # type: ignore
-from asyncpg.connection import Connection  # type: ignore
-from asyncpg.cursor import Cursor  # type: ignore
-from asyncpg.exceptions import InsufficientPrivilegeError  # type: ignore
-from asyncpg.pool import Pool  # type: ignore
-
 from piccolo.engine.base import Batch, Engine
 from piccolo.engine.exceptions import TransactionError
 from piccolo.query.base import Query
 from piccolo.querystring import QueryString
+from piccolo.utils.lazy_loader import LazyLoader
 from piccolo.utils.sync import run_sync
 from piccolo.utils.warnings import Level, colored_string, colored_warning
 
+asyncpg = LazyLoader("asyncpg", globals(), "asyncpg")
+
+if t.TYPE_CHECKING:  # pragma: no cover
+    from asyncpg.connection import Connection  # type: ignore
+    from asyncpg.cursor import Cursor  # type: ignore
+    from asyncpg.pool import Pool  # type: ignore
+
 
 @dataclass
 class AsyncBatch(Batch):
@@ -283,7 +285,7 @@ async def prep_database(self):
                 await self._run_in_new_connection(
                     f'CREATE EXTENSION IF NOT EXISTS "{extension}"',
                 )
-            except InsufficientPrivilegeError:
+            except asyncpg.exceptions.InsufficientPrivilegeError:
                 print(
                     colored_string(
                         f"=> Unable to create {extension} extension - some "
diff --git a/piccolo/engine/sqlite.py b/piccolo/engine/sqlite.py
@@ -9,16 +9,20 @@
 from dataclasses import dataclass
 from decimal import Decimal
 
-from aiosqlite import Connection, Cursor, connect
-
 from piccolo.engine.base import Batch, Engine
 from piccolo.engine.exceptions import TransactionError
 from piccolo.query.base import Query
 from piccolo.querystring import QueryString
 from piccolo.utils.encoding import dump_json, load_json
+from piccolo.utils.lazy_loader import LazyLoader
 from piccolo.utils.sync import run_sync
 
-if t.TYPE_CHECKING:
+aiosqlite = LazyLoader("aiosqlite", globals(), "aiosqlite")
+
+
+if t.TYPE_CHECKING:  # pragma: no cover
+    from aiosqlite import Connection, Cursor  # type: ignore
+
     from piccolo.table import Table
 
 ###############################################################################
@@ -414,7 +418,7 @@ async def batch(self, query: Query, batch_size: int = 100) -> AsyncBatch:
     ###########################################################################
 
     async def get_connection(self) -> Connection:
-        connection = await connect(**self.connection_kwargs)
+        connection = await aiosqlite.connect(**self.connection_kwargs)
         connection.row_factory = dict_factory  # type: ignore
         await connection.execute("PRAGMA foreign_keys = 1")
         return connection
@@ -442,7 +446,7 @@ async def _run_in_new_connection(
         query_type: str = "generic",
         table: t.Optional[t.Type[Table]] = None,
     ):
-        async with connect(**self.connection_kwargs) as connection:
+        async with aiosqlite.connect(**self.connection_kwargs) as connection:
             await connection.execute("PRAGMA foreign_keys = 1")
 
             connection.row_factory = dict_factory  # type: ignore
diff --git a/piccolo/utils/lazy_loader.py b/piccolo/utils/lazy_loader.py
@@ -0,0 +1,57 @@
+# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/util/lazy_loader.py
+from __future__ import absolute_import, division, print_function
+
+import importlib
+import types
+import typing as t
+
+
+class LazyLoader(types.ModuleType):
+    """
+    Lazily import a module.
+
+    `PostgresEngine` and `SQLiteEngine` are example use cases.
+    """
+
+    def __init__(self, local_name, parent_module_globals, name):
+        self._local_name = local_name
+        self._parent_module_globals = parent_module_globals
+
+        super().__init__(name)
+
+    def _load(self) -> types.ModuleType:
+        try:
+            # Import the target module and
+            # insert it into the parent's namespace
+            module = importlib.import_module(self.__name__)
+            self._parent_module_globals[self._local_name] = module
+
+            # Update this object's dict so that
+            # if someone keeps a reference to the
+            #   LazyLoader, lookups are efficient
+            #  (__getattr__ is only called on lookups that fail).
+            self.__dict__.update(module.__dict__)
+
+            return module
+
+        except ModuleNotFoundError as exc:
+            if str(exc) == "No module named 'asyncpg'":
+                raise ModuleNotFoundError(
+                    "PostgreSQL driver not found. "
+                    "Try running 'pip install piccolo[postgres]'"
+                )
+            elif str(exc) == "No module named 'aiosqlite'":
+                raise ModuleNotFoundError(
+                    "SQLite driver not found. "
+                    "Try running 'pip install piccolo[sqlite]'"
+                )
+            else:
+                raise exc
+
+    def __getattr__(self, item) -> t.Any:
+        module = self._load()
+        return getattr(module, item)
+
+    def __dir__(self) -> t.List[str]:
+        module = self._load()
+        return dir(module)
diff --git a/requirements/dev-requirements.txt b/requirements/dev-requirements.txt
@@ -1,6 +1,7 @@
 black>=21.7b0
 ipdb==0.12.2
 ipython==7.8.0
+flake8==3.8.4
 isort==5.9.2
 twine==3.1.1
 mypy==0.782
diff --git a/requirements/extras/postgres.txt b/requirements/extras/postgres.txt
@@ -0,0 +1 @@
+asyncpg>=0.21.0
diff --git a/requirements/extras/sqlite.txt b/requirements/extras/sqlite.txt
@@ -0,0 +1 @@
+aiosqlite>=0.16.0
diff --git a/requirements/requirements.txt b/requirements/requirements.txt
@@ -1,5 +1,3 @@
-aiosqlite>=0.16.0
-asyncpg>=0.21.0
 black
 colorama>=0.4.0
 Jinja2>=2.11.0
diff --git a/requirements/test-requirements.txt b/requirements/test-requirements.txt
@@ -1,5 +1,4 @@
 coveralls==2.2.0
-flake8==3.8.4
 pytest-cov==2.10.1
 pytest==6.2.1
 python-dateutil==2.8.1
diff --git a/setup.py b/setup.py
@@ -10,7 +10,7 @@
 
 directory = os.path.abspath(os.path.dirname(__file__))
 
-extras = ["orjson", "playground", "uvloop"]
+extras = ["orjson", "playground", "postgres", "sqlite", "uvloop"]
 
 
 with open(os.path.join(directory, "README.md")) as f:
diff --git a/tests/apps/migrations/auto/test_migration_manager.py b/tests/apps/migrations/auto/test_migration_manager.py
@@ -1,17 +1,18 @@
 import asyncio
 from unittest.mock import MagicMock, patch
 
-from asyncpg.exceptions import UniqueViolationError  # type: ignore
-
 from piccolo.apps.migrations.auto import MigrationManager
 from piccolo.apps.migrations.commands.base import BaseMigrationManager
 from piccolo.columns import Text, Varchar
 from piccolo.columns.base import OnDelete, OnUpdate
 from piccolo.columns.column_types import ForeignKey
 from piccolo.conf.apps import AppConfig
+from piccolo.utils.lazy_loader import LazyLoader
 from tests.base import DBTestCase, postgres_only, set_mock_return_value
 from tests.example_app.tables import Manager
 
+asyncpg = LazyLoader("asyncpg", globals(), "asyncpg")
+
 
 class TestMigrationManager(DBTestCase):
     @postgres_only
@@ -346,7 +347,7 @@ def test_alter_column_unique(self):
 
         asyncio.run(manager.run())
 
-        with self.assertRaises(UniqueViolationError):
+        with self.assertRaises(asyncpg.exceptions.UniqueViolationError):
             self.run_sync(
                 "INSERT INTO manager VALUES "
                 "(default, 'Dave'), (default, 'Dave');"
diff --git a/tests/utils/test_lazy_loader.py b/tests/utils/test_lazy_loader.py
@@ -0,0 +1,27 @@
+from unittest import TestCase, mock
+
+from piccolo.utils.lazy_loader import LazyLoader
+from tests.base import postgres_only, sqlite_only
+
+
+class TestLazyLoader(TestCase):
+    def test_lazy_loading_database_driver(self):
+        _ = LazyLoader("asyncpg", globals(), "asyncpg")
+
+    @postgres_only
+    def test_lazy_loader_asyncpg_exception(self):
+        lazy_loader = LazyLoader("asyncpg", globals(), "asyncpg.connect")
+
+        with mock.patch("asyncpg.connect") as module:
+            module.side_effect = ModuleNotFoundError()
+            with self.assertRaises(ModuleNotFoundError):
+                lazy_loader._load()
+
+    @sqlite_only
+    def test_lazy_loader_aiosqlite_exception(self):
+        lazy_loader = LazyLoader("aiosqlite", globals(), "aiosqlite.connect")
+
+        with mock.patch("aiosqlite.connect") as module:
+            module.side_effect = ModuleNotFoundError()
+            with self.assertRaises(ModuleNotFoundError):
+                lazy_loader._load()

Original file line number	Diff line number	Diff line change
`@@ -145,7 +145,8 @@ def run(`
`145`	`145`	`import IPython # type: ignore`
`146`	`146`	`except ImportError:`
`147`	`147`	`sys.exit(`
`148`		- "Install iPython using `pip install ipython` to use this feature."
	`148`	+ "Install iPython using `pip install piccolo[playground,sqlite]` "
	`149`	`+ "to use this feature."`
`149`	`150`	`)`
`150`	`151`
`151`	`152`	`if engine.upper() == "POSTGRES":`