From 027dbc51f91f40a2443f0570fbc8c3163b8da884 Mon Sep 17 00:00:00 2001
From: Subhash Bhushan <subhash.bhushan@gmail.com>
Date: Tue, 29 Sep 2020 09:42:50 -0700
Subject: [PATCH] Re-Add support for restricting model to a specific database

Changes were reverted during refactoring of Domain module
---
 src/protean/adapters/repository/__init__.py | 10 +++++++++-
 src/protean/core/model.py                   |  6 ++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/protean/adapters/repository/__init__.py b/src/protean/adapters/repository/__init__.py
index 13edb479..96b87a40 100644
--- a/src/protean/adapters/repository/__init__.py
+++ b/src/protean/adapters/repository/__init__.py
@@ -33,7 +33,15 @@ def get_model(self, aggregate_cls):
         if fully_qualified_name(aggregate_cls) in self.domain._models:
             custom_model_cls = self.domain._models[fully_qualified_name(aggregate_cls)]
 
-        if custom_model_cls:
+        # FIXME This is the provide support for activating database specific models
+        #   This needs to be enhanced to allow Protean to hold multiple models per Aggregate/Entity
+        #   per database.
+        #
+        #   If no database is specified, model can be used for all databases
+        if custom_model_cls and (
+            custom_model_cls.meta_.database is None
+            or custom_model_cls.meta_.database == provider.conn_info["DATABASE"]
+        ):
             # Get the decorated model class.
             #   This is a no-op if the provider decides that the model is fully-baked
             model_cls = provider.decorate_model_class(
diff --git a/src/protean/core/model.py b/src/protean/core/model.py
index 98f9b2ae..391e271b 100644
--- a/src/protean/core/model.py
+++ b/src/protean/core/model.py
@@ -58,6 +58,12 @@ def model_factory(element_cls, **kwargs):
     if not (hasattr(element_cls.meta_, "entity_cls") and element_cls.meta_.entity_cls):
         element_cls.meta_.entity_cls = kwargs.pop("entity_cls", None)
 
+    if not (hasattr(element_cls.meta_, "schema") and element_cls.meta_.schema):
+        element_cls.meta_.schema = kwargs.pop("schema", None)
+
+    if not (hasattr(element_cls.meta_, "database") and element_cls.meta_.database):
+        element_cls.meta_.database = kwargs.pop("database", None)
+
     if not element_cls.meta_.entity_cls:
         raise IncorrectUsageError(
             "Models need to be associated with an Entity or Aggregate"