feat(class): add support for Pooled database and its test cases

ruia_peewee_async/__init__.py

@@ -4,7 +4,7 @@
 from functools import wraps
 from ssl import SSLContext
 from types import MethodType
-from typing import Dict, Callable
+from typing import Callable, Dict
 from typing import Optional as TOptional
 from typing import Sequence, Tuple, Union
 
@@ -375,6 +375,13 @@ def check_config(kwargs) -> Sequence[Dict]:
                         "model": And({"table_name": And(str), str: object}),
                         Optional("port"): And(int),
                         Optional("ssl"): Use(SSLContext),
+                        Optional("pool"): And(bool),
+                        Optional("min_connections"): And(
+                            int, lambda mic: 1 <= mic <= 10
+                        ),
+                        Optional("max_connections"): And(
+                            int, lambda mac: 10 < mac <= 20
+                        ),
                     }
                 ),
             )
@@ -415,8 +422,18 @@ def create_model(spider_ins=None, create_table=False, **kwargs) -> Tuple:
         None,
     )
     if mysql:
-        mysql_db = MySQLDatabase(
-            **{key: val for key, val in mysql.items() if key != "model"}
+        mysql_db = (
+            PooledMySQLDatabase(
+                **{
+                    key: val
+                    for key, val in mysql.items()
+                    if key not in ("model", "pool")
+                }
+            )
+            if "pool" in mysql
+            else MySQLDatabase(
+                **{key: val for key, val in mysql.items() if key != "model"}
+            )
         )
         mysql_manager = Manager(mysql_db)
         meta = type("Meta", (object,), {"database": mysql_db})
@@ -432,8 +449,18 @@ def create_model(spider_ins=None, create_table=False, **kwargs) -> Tuple:
                 mysql_model.create_table(True)
         mysql_mconf["table_name"] = table_name
     if postgres:
-        postgres_db = PostgresqlDatabase(
-            **{key: val for key, val in postgres.items() if key != "model"}
+        postgres_db = (
+            PooledPostgresqlDatabase(
+                **{
+                    key: val
+                    for key, val in postgres.items()
+                    if key not in ("model", "pool")
+                }
+            )
+            if "pool" in postgres
+            else PostgresqlDatabase(
+                **{key: val for key, val in postgres.items() if key != "model"}
+            )
         )
         postgres_manager = Manager(postgres_db)
         meta = type("Meta", (object,), {"database": postgres_db})

tests/conftest.py

@@ -49,6 +49,25 @@ def mysql_config():
     }
 
 
+@pytest.fixture(scope="function")
+def pool_mysql_config():
+    return {
+        "host": "somehost",
+        "port": 1234,
+        "user": "ruiamysql",
+        "password": "abc123",
+        "database": "ruiamysql",
+        "pool": True,
+        "min_connections": 5,
+        "max_connections": 20,
+        "model": {
+            "table_name": "test",
+            "some_date": DateField(),
+            "some_char": CharField(),
+        },
+    }
+
+
 @pytest.fixture(scope="function")
 def postgres_config():
     return {
@@ -65,6 +84,25 @@ def postgres_config():
     }
 
 
+@pytest.fixture(scope="function")
+def pool_postgres_config():
+    return {
+        "host": "somehost",
+        "port": 1234,
+        "user": "ruiamysql",
+        "password": "abc123",
+        "database": "ruiamysql",
+        "pool": True,
+        "min_connections": 5,
+        "max_connections": 20,
+        "model": {
+            "table_name": "test",
+            "some_date": DateField(),
+            "some_char": CharField(),
+        },
+    }
+
+
 @pytest.mark.no_cover
 def check_postgres(postgres_conf):
     try:

tests/test_config.py

@@ -3,9 +3,11 @@
 from copy import deepcopy
 
 import pytest
+from peewee import Model
+from peewee_async import PooledMySQLDatabase, PooledPostgresqlDatabase
 from schema import SchemaError, SchemaMissingKeyError
 
-from ruia_peewee_async import after_start
+from ruia_peewee_async import after_start, create_model
 
 from .common import Insert, RuiaPeeweeInsert, RuiaPeeweeUpdate, TargetDB, Update
 
@@ -227,3 +229,26 @@ async def test_err_config(
                 loop=event_loop, after_start=after_start(postgres=postgres)
             )
         assert "Key 'model' error:\nMissing key: 'table_name'" in se4.value.args[0]
+
+    async def test_pool_config(
+        self,
+        docker_setup,
+        docker_cleanup,
+        event_loop,
+        pool_mysql_config,
+        pool_postgres_config,
+    ):  # pylint: disable=redefined-outer-name,unused-argument,unknown-option-value
+        (  # pylint: disable=unbalanced-tuple-unpacking
+            mysql_model,
+            mysql_manager,
+            postgres_model,
+            postgres_manager,
+        ) = create_model(mysql=pool_mysql_config, postgres=pool_postgres_config)
+        assert isinstance(mysql_model, Model) is True
+        assert isinstance(postgres_model, Model) is True
+        assert isinstance(mysql_manager.database, PooledMySQLDatabase) is True
+        assert isinstance(postgres_manager.database, PooledPostgresqlDatabase) is True
+        assert mysql_manager.database.min_connections == 5
+        assert mysql_manager.database.max_connections == 20
+        assert postgres_manager.database.min_connections == 5
+        assert postgres_manager.database.max_connections == 20