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"