Configurable base model metaclass

jpsca · Apr 10, 2022 · 20de117 · 20de117
1 parent 3de46f5
commit 20de117
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/setup.cfg b/setup.cfg
@@ -1,6 +1,6 @@
 [metadata]
 name = sqla-wrapper
-version= 5.5
+version= 5.6
 url = https://github.com/jpsca/sqla-wrapper
 project_urls =
     Documentation = https://sqla-wrapper.scaletti.dev/

diff --git a/sqla_wrapper/sqlalchemy_wrapper.py b/sqla_wrapper/sqlalchemy_wrapper.py
@@ -2,12 +2,13 @@
 
 import sqlalchemy
 import sqlalchemy.orm
+from sqlalchemy.orm.decl_api import DeclarativeMeta
 
 from .base_model import BaseModel
 from .session import PatchedScopedSession, Session
 
 
-__all__ = ("SQLAlchemy", "TestTransaction")
+__all__ = ("SQLAlchemy", "TestTransaction", "DeclarativeMeta")
 
 
 class SQLAlchemy:
@@ -60,7 +61,8 @@ def __init__(
         port: Optional[str] = None,
         engine_options: Optional[Dict[str, Any]] = None,
         session_options: Optional[Dict[str, Any]] = None,
-        base_model_class: Any = BaseModel
+        base_model_class: Any = BaseModel,
+        base_model_metaclass: Any = DeclarativeMeta,
     ) -> None:
         self.url = url or self._make_url(
             dialect=dialect,
@@ -75,7 +77,11 @@ def __init__(
         self.engine = sqlalchemy.create_engine(self.url, **engine_options)
 
         self.registry = sqlalchemy.orm.registry()
-        self.Model = self.registry.generate_base(cls=base_model_class, name="Model")
+        self.Model = self.registry.generate_base(
+            cls=base_model_class,
+            name="Model",
+            metaclass=base_model_metaclass
+        )
 
         session_options = session_options or {}
         session_options.setdefault("class_", Session)