Increase timeout on model disconnect

Fix typehint on column_type / sql_type
typehint Model.columns
optional `cascade` kwarg on `drop_table_sql()`

 + Applied **`black`** formatting.

docs/conf.py

@@ -31,7 +31,9 @@
 
 release = ""
 with open("../postDB/__init__.py") as f:
-    release = re.search(r'^__version__\s*=\s*[\'"]([^\'"]*)[\'"]', f.read(), re.MULTILINE).group(1)
+    release = re.search(
+        r'^__version__\s*=\s*[\'"]([^\'"]*)[\'"]', f.read(), re.MULTILINE
+    ).group(1)
 
 # The short X.Y version.
 version = ".".join(release.split(".")[:2])
@@ -48,7 +50,7 @@
 extensions = [
     "sphinx.ext.autodoc",
     "sphinx.ext.coverage",
-    'sphinx.ext.extlinks',
+    "sphinx.ext.extlinks",
     "sphinx.ext.intersphinx",
     "sphinx.ext.napoleon",
 ]

examples/internal_pool.py

@@ -54,7 +54,7 @@ async def main():
     print(await Model.pool.fetch("SELECT * FROM users"))
 
     try:
-        await asyncio.wait_for(Model.pool.close(), timeout=3.0)
+        await asyncio.wait_for(Model.pool.close(), timeout=10.0)
     except asyncio.TimeoutError:
         print("Exiting by timeout...")
 

postDB/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.2.0"
+__version__ = "0.2.1"
 __author__ = "Sylte"
 __licence__ = "MIT"
 __copyright__ = "Copyright (c) 2020 Sylte"
@@ -18,7 +18,7 @@ def get_version(version: tuple):
     return "%s %s.%s.%s" % (version[3], version[0], version[1], version[2])
 
 
-VERSION = (0, 2, 0, "final", 0)
+VERSION = (0, 2, 1, "final", 0)
 version_info = VersionInfo(*VERSION)
 
 logging.getLogger(__name__).addHandler(logging.NullHandler())

postDB/exceptions.py

@@ -1,8 +1,10 @@
 class SchemaError(Exception):
     """Base error for the module."""
+
     pass
 
 
 class UniqueViolationError(SchemaError):
     """Raised when a unique constraint is violated."""
+
     pass

postDB/model/column.py

@@ -1,9 +1,9 @@
-import inspect
-import typing
-
 from postDB.exceptions import SchemaError
 from postDB.types import SQLType, String
 
+import inspect
+import typing
+
 
 class Column:
     """Class to define a column in a :class:`Model`."""
@@ -21,7 +21,7 @@ class Column:
 
     def __init__(
         self,
-        column_type: SQLType,
+        column_type: typing.Union[typing.Type[SQLType], SQLType],
         *,
         index: bool = False,
         primary_key: bool = False,

postDB/model/meta.py

@@ -1,5 +1,7 @@
 from postDB import Column
 
+from typing import List
+
 
 class ModelMeta(type):
     """Metaclass for Model class."""
@@ -12,7 +14,7 @@ def __new__(mcs, name, parents, data, **kwargs):
 
         tablename = kwargs.get("tablename", name.lower() + "s")
 
-        columns = []
+        columns: List[Column] = []
         for key, col in data.items():
             if isinstance(col, Column):
                 if col.name is None:

postDB/model/model.py

@@ -1,9 +1,9 @@
-import json
-from asyncio import BaseEventLoop
 from typing import Optional, List
+from asyncio import BaseEventLoop
+import json
 
-from asyncpg import create_pool
 from asyncpg.connection import Connection
+from asyncpg import create_pool
 from asyncpg.pool import Pool
 
 from postDB.model.meta import ModelMeta
@@ -95,14 +95,15 @@ def create_table_sql(cls, *, exists_ok: bool = True) -> str:
         return "\n".join(statements)
 
     @classmethod
-    def drop_table_sql(cls, exists_ok: bool = True) -> str:
+    def drop_table_sql(cls, *, exists_ok: bool = True, cascade: bool = False) -> str:
         """Generates the ``DROP TABLE`` SQL statement."""
         builder = ["DROP TABLE"]
 
         if exists_ok:
             builder.append("IF EXISTS")
 
-        builder.append("%s CASCADE;" % cls.__tablename__)
+        to_cascade = "CASCADE" if cascade else "RESTRICT"
+        builder.append("%s %s;" % (cls.__tablename__, to_cascade))
         return " ".join(builder)
 
     @classmethod
@@ -145,7 +146,7 @@ async def ensure_con(cls, con: Optional[Connection] = None) -> Connection:
             return await cls.pool.acquire()
 
         raise RuntimeWarning(
-            "Unable to get Connection, either call `Model.create_pool` or provide a con` argument."
+            "Unable to get Connection, either call `Model.create_pool` or provide a `con` argument."
         )
 
     @classmethod
@@ -172,12 +173,13 @@ async def drop_table(
         con: Optional[Connection] = None,
         *,
         verbose: bool = False,
+        cascade: bool = True,
         exists_ok: bool = True,
     ):
         """Drop the PostgreSQL Table for this Model."""
         con = await cls.ensure_con(con=con)
 
-        sql = cls.drop_table_sql(exists_ok=exists_ok)
+        sql = cls.drop_table_sql(exists_ok=exists_ok, cascade=cascade)
 
         if verbose:
             print(sql)

postDB/types.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Union, Type
 import datetime
 import inspect
 import decimal
@@ -285,7 +285,7 @@ def __init__(
         model: str,
         column: str,
         *,
-        sql_type: Optional[SQLType] = None,
+        sql_type: Optional[Union[Type[SQLType], SQLType]] = None,
         on_delete: str = "CASCADE",
         on_update: str = "NO ACTION"
     ):
@@ -344,7 +344,7 @@ class Array(SQLType):
 
     python = list
 
-    def __init__(self, sql_type: SQLType):
+    def __init__(self, sql_type: Union[Type[SQLType], SQLType]):
         if inspect.isclass(sql_type):
             sql_type = sql_type()