[Enhancement] Swappable Packages (#216)

Co-authored-by: Tim Lorsbach <tim@lorsba.ch>
Reviewed-on: enviPath/enviPy#216
Reviewed-by: liambrydon <lbry121@aucklanduni.ac.nz>
Reviewed-by: Tobias O <tobias.olenyi@envipath.com>

envipath/settings.py

@@ -49,9 +49,23 @@ INSTALLED_APPS = [
     "oauth2_provider",
     # Custom
     "epdb",
-    "migration",
+    # "migration",
 ]
 
+TENANT = os.environ.get("TENANT", "public")
+
+if TENANT != "public":
+    INSTALLED_APPS.append(TENANT)
+
+EPDB_PACKAGE_MODEL = os.environ.get("EPDB_PACKAGE_MODEL", "epdb.Package")
+
+
+def GET_PACKAGE_MODEL():
+    from django.apps import apps
+
+    return apps.get_model(EPDB_PACKAGE_MODEL)
+
+
 AUTHENTICATION_BACKENDS = [
     "django.contrib.auth.backends.ModelBackend",
 ]

envipath/urls.py

@@ -23,12 +23,14 @@ from .api import api_v1, api_legacy
 
 urlpatterns = [
     path("", include("epdb.urls")),
-    path("", include("migration.urls")),
     path("admin/", admin.site.urls),
     path("api/v1/", api_v1.urls),
     path("api/legacy/", api_legacy.urls),
     path("o/", include("oauth2_provider.urls", namespace="oauth2_provider")),
 ]
 
+if "migrations" in s.INSTALLED_APPS:
+    urlpatterns.append(path("", include("migration.urls")))
+
 if s.MS_ENTRA_ENABLED:
     urlpatterns.append(path("", include("epauth.urls")))

epdb/admin.py

@@ -1,29 +1,31 @@
+from django.conf import settings as s
 from django.contrib import admin
 
 from .models import (
-    User,
-    UserPackagePermission,
-    Group,
-    GroupPackagePermission,
-    Package,
-    MLRelativeReasoning,
-    EnviFormer,
     Compound,
     CompoundStructure,
-    SimpleAmbitRule,
-    ParallelRule,
-    Reaction,
-    Pathway,
-    Node,
     Edge,
-    Scenario,
+    EnviFormer,
-    Setting,
     ExternalDatabase,
     ExternalIdentifier,
+    Group,
+    GroupPackagePermission,
     JobLog,
     License,
+    MLRelativeReasoning,
+    Node,
+    ParallelRule,
+    Pathway,
+    Reaction,
+    Scenario,
+    Setting,
+    SimpleAmbitRule,
+    User,
+    UserPackagePermission,
 )
 
+Package = s.GET_PACKAGE_MODEL()
+
 
 class UserAdmin(admin.ModelAdmin):
     list_display = ["username", "email", "is_active"]

epdb/apps.py

@@ -1,4 +1,9 @@
+import logging
+
 from django.apps import AppConfig
+from django.conf import settings
+
+logger = logging.getLogger(__name__)
 
 
 class EPDBConfig(AppConfig):
@@ -7,3 +12,6 @@ class EPDBConfig(AppConfig):
 
     def ready(self):
         import epdb.signals  # noqa: F401
+
+        model_name = getattr(settings, "EPDB_PACKAGE_MODEL", "epdb.Package")
+        logger.info(f"Using Package model: {model_name}")

epdb/context_processors.py

@@ -5,7 +5,7 @@ Context processors automatically make variables available to all templates.
 """
 
 from .logic import PackageManager
-from .models import Package
+from django.conf import settings as s
 
 
 def package_context(request):
@@ -20,7 +20,7 @@ def package_context(request):
 
     reviewed_package_qs = PackageManager.get_reviewed_packages()
 
-    unreviewed_package_qs = Package.objects.none()
+    unreviewed_package_qs = s.GET_PACKAGE_MODEL().objects.none()
 
     # Only get user-specific packages if user is authenticated
     if current_user.is_authenticated:

epdb/legacy_api.py

@@ -1,27 +1,30 @@
-from typing import List, Dict, Optional, Any
+from typing import Any, Dict, List, Optional
 
+from django.conf import settings as s
 from django.contrib.auth import get_user_model
 from django.http import HttpResponse
 from django.shortcuts import redirect
-from ninja import Router, Schema, Field, Form
+from ninja import Field, Form, Router, Schema
 
 from utilities.chem import FormatConverter
-from .logic import PackageManager, UserManager, SettingManager
+
+from .logic import PackageManager, SettingManager, UserManager
 from .models import (
     Compound,
     CompoundStructure,
-    Package,
+    Edge,
+    Node,
+    Pathway,
+    Reaction,
+    Rule,
+    Scenario,
+    SimpleAmbitRule,
     User,
     UserPackagePermission,
-    Rule,
-    Reaction,
-    Scenario,
-    Pathway,
-    Node,
-    Edge,
-    SimpleAmbitRule,
 )
 
+Package = s.GET_PACKAGE_MODEL()
+
 
 def _anonymous_or_real(request):
     if request.user.is_authenticated and not request.user.is_anonymous:
@@ -123,8 +126,7 @@ class SimpleEdge(SimpleObject):
 ################
 @router.post("/", response={200: SimpleUser, 403: Error})
 def login(request, loginusername: Form[str], loginpassword: Form[str]):
-    from django.contrib.auth import authenticate
+    from django.contrib.auth import authenticate, login
-    from django.contrib.auth import login
 
     email = User.objects.get(username=loginusername).email
     user = authenticate(username=email, password=loginpassword)

epdb/logic.py

@@ -1,39 +1,40 @@
-import re
-import logging
 import json
-from typing import Union, List, Optional, Set, Dict, Any
+import logging
+import re
+from typing import Any, Dict, List, Optional, Set, Union
 from uuid import UUID
 
 import nh3
+from django.conf import settings as s
 from django.contrib.auth import get_user_model
 from django.db import transaction
-from django.conf import settings as s
 from pydantic import ValidationError
 
 from epdb.models import (
-    User,
-    Package,
-    UserPackagePermission,
-    GroupPackagePermission,
-    Permission,
-    Group,
-    Setting,
-    EPModel,
-    UserSettingPermission,
-    Rule,
-    Pathway,
-    Node,
-    Edge,
     Compound,
-    Reaction,
     CompoundStructure,
+    Edge,
     EnzymeLink,
+    EPModel,
+    Group,
+    GroupPackagePermission,
+    Node,
+    Pathway,
+    Permission,
+    Reaction,
+    Rule,
+    Setting,
+    User,
+    UserPackagePermission,
+    UserSettingPermission,
 )
 from utilities.chem import FormatConverter
-from utilities.misc import PackageImporter, PackageExporter
+from utilities.misc import PackageExporter, PackageImporter
 
 logger = logging.getLogger(__name__)
 
+Package = s.GET_PACKAGE_MODEL()
+
 
 class EPDBURLParser:
     UUID_PATTERN = r"[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}"
@@ -583,25 +584,27 @@ class PackageManager(object):
     def import_legacy_package(
         data: dict, owner: User, keep_ids=False, add_import_timestamp=True, trust_reviewed=False
     ):
-        from uuid import UUID, uuid4
-        from datetime import datetime
         from collections import defaultdict
+        from datetime import datetime
+        from uuid import UUID, uuid4
+
+        from envipy_additional_information import AdditionalInformationConverter
+
         from .models import (
-            Package,
             Compound,
             CompoundStructure,
-            SimpleRule,
+            Edge,
-            SimpleAmbitRule,
+            Node,
+            Package,
             ParallelRule,
+            Pathway,
+            Reaction,
+            Scenario,
             SequentialRule,
             SequentialRuleOrdering,
-            Reaction,
+            SimpleAmbitRule,
-            Pathway,
+            SimpleRule,
-            Node,
-            Edge,
-            Scenario,
         )
-        from envipy_additional_information import AdditionalInformationConverter
 
         pack = Package()
         pack.uuid = UUID(data["id"].split("/")[-1]) if keep_ids else uuid4()

epdb/management/commands/create_ml_models.py

@@ -2,7 +2,9 @@ from django.conf import settings as s
 from django.core.management.base import BaseCommand
 from django.db import transaction
 
-from epdb.models import MLRelativeReasoning, EnviFormer, Package
+from epdb.models import EnviFormer, MLRelativeReasoning
+
+Package = s.GET_PACKAGE_MODEL()
 
 
 class Command(BaseCommand):
@@ -75,11 +77,13 @@ class Command(BaseCommand):
             return packages
 
         # Iteratively create models in options["model_names"]
-        print(f"Creating models: {options['model_names']}\n"
+        print(
+            f"Creating models: {options['model_names']}\n"
             f"Data packages: {options['data_packages']}\n"
             f"Rule Packages (only for MLRR): {options['rule_packages']}\n"
             f"Eval Packages: {options['eval_packages']}\n"
-              f"Threshold: {options['threshold']:.2f}")
+            f"Threshold: {options['threshold']:.2f}"
+        )
         data_packages = decode_packages(options["data_packages"])
         eval_packages = decode_packages(options["eval_packages"])
         rule_packages = decode_packages(options["rule_packages"])
@@ -90,7 +94,7 @@ class Command(BaseCommand):
                     pack,
                     data_packages=data_packages,
                     eval_packages=eval_packages,
-                    threshold=options['threshold'],
+                    threshold=options["threshold"],
                     name=f"EnviFormer - {', '.join(options['data_packages'])} - T{options['threshold']:.2f}",
                     description=f"EnviFormer transformer trained on {options['data_packages']} "
                     f"evaluated on {options['eval_packages']}.",
@@ -101,7 +105,7 @@ class Command(BaseCommand):
                     rule_packages=rule_packages,
                     data_packages=data_packages,
                     eval_packages=eval_packages,
-                    threshold=options['threshold'],
+                    threshold=options["threshold"],
                     name=f"ECC - {', '.join(options['data_packages'])} - T{options['threshold']:.2f}",
                     description=f"ML Relative Reasoning trained on {options['data_packages']} with rules from "
                     f"{options['rule_packages']} and evaluated on {options['eval_packages']}.",

epdb/management/commands/load_enviformer.py

@@ -8,7 +8,9 @@ from django.conf import settings as s
 from django.core.management.base import BaseCommand
 from django.db import transaction
 
-from epdb.models import EnviFormer, Package
+from epdb.models import EnviFormer
+
+Package = s.GET_PACKAGE_MODEL()
 
 
 class Command(BaseCommand):

epdb/management/commands/localize_urls.py

@@ -1,8 +1,8 @@
 from django.apps import apps
+from django.conf import settings as s
 from django.core.management.base import BaseCommand
-
+from django.db.models import F, JSONField, TextField, Value
-from django.db.models import F, Value, TextField, JSONField
+from django.db.models.functions import Cast, Replace
-from django.db.models.functions import Replace, Cast
 
 from epdb.models import EnviPathModel
 
@@ -23,10 +23,13 @@ class Command(BaseCommand):
         )
 
     def handle(self, *args, **options):
+        Package = s.GET_PACKAGE_MODEL()
+        print("Localizing urls for Package")
+        Package.objects.update(url=Replace(F("url"), Value(options["old"]), Value(options["new"])))
+
         MODELS = [
             "User",
             "Group",
-            "Package",
             "Compound",
             "CompoundStructure",
             "Pathway",

epdb/models.py

@@ -2,40 +2,41 @@ import abc
 import hashlib
 import json
 import logging
+import math
 import os
 import secrets
 from abc import abstractmethod
 from collections import defaultdict
 from datetime import datetime
-from typing import Union, List, Optional, Dict, Tuple, Set, Any
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
 from uuid import uuid4
-import math
+
 import joblib
 import nh3
 import numpy as np
 from django.conf import settings as s
 from django.contrib.auth.models import AbstractUser
-from django.contrib.contenttypes.fields import GenericRelation, GenericForeignKey
+from django.contrib.contenttypes.fields import GenericForeignKey, GenericRelation
 from django.contrib.contenttypes.models import ContentType
 from django.contrib.postgres.fields import ArrayField
 from django.db import models, transaction
-from django.db.models import JSONField, Count, Q, QuerySet
+from django.db.models import Count, JSONField, Q, QuerySet
 from django.utils import timezone
 from django.utils.functional import cached_property
 from envipy_additional_information import EnviPyModel
 from model_utils.models import TimeStampedModel
 from polymorphic.models import PolymorphicModel
-from sklearn.metrics import precision_score, recall_score, jaccard_score
+from sklearn.metrics import jaccard_score, precision_score, recall_score
 from sklearn.model_selection import ShuffleSplit
 
-from utilities.chem import FormatConverter, ProductSet, PredictionResult, IndigoUtils
+from utilities.chem import FormatConverter, IndigoUtils, PredictionResult, ProductSet
 from utilities.ml import (
-    RuleBasedDataset,
     ApplicabilityDomainPCA,
-    EnsembleClassifierChain,
-    RelativeReasoning,
-    EnviFormerDataset,
     Dataset,
+    EnsembleClassifierChain,
+    EnviFormerDataset,
+    RelativeReasoning,
+    RuleBasedDataset,
 )
 
 logger = logging.getLogger(__name__)
@@ -44,8 +45,6 @@ logger = logging.getLogger(__name__)
 ##########################
 # User/Groups/Permission #
 ##########################
-
-
 class User(AbstractUser):
     email = models.EmailField(unique=True)
     uuid = models.UUIDField(
@@ -53,7 +52,10 @@ class User(AbstractUser):
     )
     url = models.TextField(blank=False, null=True, verbose_name="URL", unique=True)
     default_package = models.ForeignKey(
-        "epdb.Package", verbose_name="Default Package", null=True, on_delete=models.SET_NULL
+        s.EPDB_PACKAGE_MODEL,
+        verbose_name="Default Package",
+        null=True,
+        on_delete=models.SET_NULL,
     )
     default_group = models.ForeignKey(
         "Group",
@@ -243,7 +245,7 @@ class UserPackagePermission(Permission):
     )
     user = models.ForeignKey("User", verbose_name="Permission to", on_delete=models.CASCADE)
     package = models.ForeignKey(
-        "epdb.Package", verbose_name="Permission on", on_delete=models.CASCADE
+        s.EPDB_PACKAGE_MODEL, verbose_name="Permission on", on_delete=models.CASCADE
     )
 
     class Meta:
@@ -259,7 +261,7 @@ class GroupPackagePermission(Permission):
     )
     group = models.ForeignKey("Group", verbose_name="Permission to", on_delete=models.CASCADE)
     package = models.ForeignKey(
-        "epdb.Package", verbose_name="Permission on", on_delete=models.CASCADE
+        s.EPDB_PACKAGE_MODEL, verbose_name="Permission on", on_delete=models.CASCADE
     )
 
     class Meta:
@@ -728,10 +730,13 @@ class Package(EnviPathModel):
         rules = sorted(rules, key=lambda x: x.url)
         return rules
 
+    class Meta:
+        swappable = "EPDB_PACKAGE_MODEL"
+
 
 class Compound(EnviPathModel, AliasMixin, ScenarioMixin, ChemicalIdentifierMixin):
     package = models.ForeignKey(
-        "epdb.Package", verbose_name="Package", on_delete=models.CASCADE, db_index=True
+        s.EPDB_PACKAGE_MODEL, verbose_name="Package", on_delete=models.CASCADE, db_index=True
     )
     default_structure = models.ForeignKey(
         "CompoundStructure",
@@ -781,7 +786,7 @@ class Compound(EnviPathModel, AliasMixin, ScenarioMixin, ChemicalIdentifierMixin
     @staticmethod
     @transaction.atomic
     def create(
-        package: Package, smiles: str, name: str = None, description: str = None, *args, **kwargs
+        package: "Package", smiles: str, name: str = None, description: str = None, *args, **kwargs
     ) -> "Compound":
         if smiles is None or smiles.strip() == "":
             raise ValueError("SMILES is required")
@@ -1061,7 +1066,7 @@ class EnzymeLink(EnviPathModel, KEGGIdentifierMixin):
 
 class Rule(PolymorphicModel, EnviPathModel, AliasMixin, ScenarioMixin):
     package = models.ForeignKey(
-        "epdb.Package", verbose_name="Package", on_delete=models.CASCADE, db_index=True
+        s.EPDB_PACKAGE_MODEL, verbose_name="Package", on_delete=models.CASCADE, db_index=True
     )
 
     # # https://github.com/django-polymorphic/django-polymorphic/issues/229
@@ -1167,7 +1172,7 @@ class SimpleAmbitRule(SimpleRule):
     @staticmethod
     @transaction.atomic
     def create(
-        package: Package,
+        package: "Package",
         name: str = None,
         description: str = None,
         smirks: str = None,
@@ -1241,7 +1246,7 @@ class SimpleAmbitRule(SimpleRule):
 
     @property
     def related_reactions(self):
-        qs = Package.objects.filter(reviewed=True)
+        qs = s.GET_PACKAGE_MODEL().objects.filter(reviewed=True)
         return self.reaction_rule.filter(package__in=qs).order_by("name")
 
     @property
@@ -1333,7 +1338,7 @@ class SequentialRuleOrdering(models.Model):
 
 class Reaction(EnviPathModel, AliasMixin, ScenarioMixin, ReactionIdentifierMixin):
     package = models.ForeignKey(
-        "epdb.Package", verbose_name="Package", on_delete=models.CASCADE, db_index=True
+        s.EPDB_PACKAGE_MODEL, verbose_name="Package", on_delete=models.CASCADE, db_index=True
     )
     educts = models.ManyToManyField(
         "epdb.CompoundStructure", verbose_name="Educts", related_name="reaction_educts"
@@ -1355,7 +1360,7 @@ class Reaction(EnviPathModel, AliasMixin, ScenarioMixin, ReactionIdentifierMixin
     @staticmethod
     @transaction.atomic
     def create(
-        package: Package,
+        package: "Package",
         name: str = None,
         description: str = None,
         educts: Union[List[str], List[CompoundStructure]] = None,
@@ -1514,7 +1519,7 @@ class Reaction(EnviPathModel, AliasMixin, ScenarioMixin, ReactionIdentifierMixin
 
 class Pathway(EnviPathModel, AliasMixin, ScenarioMixin):
     package = models.ForeignKey(
-        "epdb.Package", verbose_name="Package", on_delete=models.CASCADE, db_index=True
+        s.EPDB_PACKAGE_MODEL, verbose_name="Package", on_delete=models.CASCADE, db_index=True
     )
     setting = models.ForeignKey(
         "epdb.Setting", verbose_name="Setting", on_delete=models.CASCADE, null=True, blank=True
@@ -2076,7 +2081,7 @@ class Edge(EnviPathModel, AliasMixin, ScenarioMixin):
 
 class EPModel(PolymorphicModel, EnviPathModel):
     package = models.ForeignKey(
-        "epdb.Package", verbose_name="Package", on_delete=models.CASCADE, db_index=True
+        s.EPDB_PACKAGE_MODEL, verbose_name="Package", on_delete=models.CASCADE, db_index=True
     )
 
     def _url(self):
@@ -2085,17 +2090,17 @@ class EPModel(PolymorphicModel, EnviPathModel):
 
 class PackageBasedModel(EPModel):
     rule_packages = models.ManyToManyField(
-        "Package",
+        s.EPDB_PACKAGE_MODEL,
         verbose_name="Rule Packages",
         related_name="%(app_label)s_%(class)s_rule_packages",
     )
     data_packages = models.ManyToManyField(
-        "Package",
+        s.EPDB_PACKAGE_MODEL,
         verbose_name="Data Packages",
         related_name="%(app_label)s_%(class)s_data_packages",
     )
     eval_packages = models.ManyToManyField(
-        "Package",
+        s.EPDB_PACKAGE_MODEL,
         verbose_name="Evaluation Packages",
         related_name="%(app_label)s_%(class)s_eval_packages",
     )
@@ -3400,7 +3405,7 @@ class PluginModel(EPModel):
 
 class Scenario(EnviPathModel):
     package = models.ForeignKey(
-        "epdb.Package", verbose_name="Package", on_delete=models.CASCADE, db_index=True
+        s.EPDB_PACKAGE_MODEL, verbose_name="Package", on_delete=models.CASCADE, db_index=True
     )
     scenario_date = models.CharField(max_length=256, null=False, blank=False, default="No date")
     scenario_type = models.CharField(
@@ -3555,7 +3560,7 @@ class Setting(EnviPathModel):
     )
 
     rule_packages = models.ManyToManyField(
-        "Package",
+        s.EPDB_PACKAGE_MODEL,
         verbose_name="Setting Rule Packages",
         related_name="setting_rule_packages",
         blank=True,

epdb/tasks.py

@@ -6,14 +6,17 @@ from uuid import uuid4
 
 from celery import shared_task
 from celery.utils.functional import LRUCache
+from django.conf import settings as s
 from django.utils import timezone
 
 from epdb.logic import SPathway
-from epdb.models import Edge, EPModel, JobLog, Node, Package, Pathway, Rule, Setting, User
+from epdb.models import Edge, EPModel, JobLog, Node, Pathway, Rule, Setting, User
 
 logger = logging.getLogger(__name__)
 ML_CACHE = LRUCache(3)  # Cache the three most recent ML models to reduce load times.
 
+Package = s.GET_PACKAGE_MODEL()
+
 
 def get_ml_model(model_pk: int):
     if model_pk not in ML_CACHE:

epdb/views.py

@@ -1,58 +1,60 @@
 import json
 import logging
-from typing import List, Dict, Any
+from typing import Any, Dict, List
 
+import nh3
 from django.conf import settings as s
 from django.contrib.auth import get_user_model
-from django.http import JsonResponse, HttpResponse, HttpResponseNotAllowed, HttpResponseBadRequest
+from django.http import HttpResponse, HttpResponseBadRequest, HttpResponseNotAllowed, JsonResponse
-from django.shortcuts import render, redirect
+from django.shortcuts import redirect, render
 from django.urls import reverse
 from django.views.decorators.csrf import csrf_exempt
 from envipy_additional_information import NAME_MAPPING
 from oauth2_provider.decorators import protected_resource
-import nh3
 
 from utilities.chem import FormatConverter, IndigoUtils
 from utilities.decorators import package_permission_required
 from utilities.misc import HTMLGenerator
+
 from .logic import (
+    EPDBURLParser,
     GroupManager,
     PackageManager,
-    UserManager,
-    SettingManager,
     SearchManager,
-    EPDBURLParser,
+    SettingManager,
+    UserManager,
 )
 from .models import (
-    Package,
+    APIToken,
-    GroupPackagePermission,
-    Group,
-    CompoundStructure,
     Compound,
+    CompoundStructure,
+    Edge,
+    EnviFormer,
+    EnzymeLink,
+    EPModel,
+    ExternalDatabase,
+    ExternalIdentifier,
+    Group,
+    GroupPackagePermission,
+    JobLog,
+    License,
+    MLRelativeReasoning,
+    Node,
+    Pathway,
+    Permission,
     Reaction,
     Rule,
-    Pathway,
-    Node,
-    EPModel,
-    EnviFormer,
-    MLRelativeReasoning,
     RuleBasedRelativeReasoning,
     Scenario,
     SimpleAmbitRule,
-    APIToken,
-    UserPackagePermission,
-    Permission,
-    License,
     User,
-    Edge,
+    UserPackagePermission,
-    ExternalDatabase,
-    ExternalIdentifier,
-    EnzymeLink,
-    JobLog,
 )
 
 logger = logging.getLogger(__name__)
 
+Package = s.GET_PACKAGE_MODEL()
+
 
 def log_post_params(request):
     if s.DEBUG:
@@ -83,8 +85,7 @@ def login(request):
         return render(request, "static/login.html", context)
 
     elif request.method == "POST":
-        from django.contrib.auth import authenticate
+        from django.contrib.auth import authenticate, login
-        from django.contrib.auth import login
 
         username = request.POST.get("username").strip()
         if username != request.POST.get("username"):
@@ -872,7 +873,7 @@ def package_models(request, package_uuid):
                 request, "Invalid model type.", f'Model type "{model_type}" is not supported."'
             )
 
-        from .tasks import dispatch, build_model
+        from .tasks import build_model, dispatch
 
         dispatch(current_user, build_model, mod.pk)
 
@@ -2409,9 +2410,9 @@ def package_scenarios(request, package_uuid):
         context["unreviewed_objects"] = unreviewed_scenario_qs
 
         from envipy_additional_information import (
+            SEDIMENT_ADDITIONAL_INFORMATION,
             SLUDGE_ADDITIONAL_INFORMATION,
             SOIL_ADDITIONAL_INFORMATION,
-            SEDIMENT_ADDITIONAL_INFORMATION,
         )
 
         context["scenario_types"] = {

migration/views.py

@@ -1,24 +1,21 @@
-import gzip
 import json
 import logging
 import os.path
-from datetime import datetime
 
 from django.conf import settings as s
 from django.http import HttpResponseNotAllowed
 from django.shortcuts import render
-
-from epdb.logic import PackageManager
-from epdb.models import Rule, SimpleAmbitRule, Package, CompoundStructure
-from epdb.views import get_base_context, _anonymous_or_real
-from utilities.chem import FormatConverter
-
-
 from rdkit import Chem
 from rdkit.Chem.MolStandardize import rdMolStandardize
 
+from epdb.models import CompoundStructure, Rule, SimpleAmbitRule
+from epdb.views import get_base_context
+from utilities.chem import FormatConverter
+
 logger = logging.getLogger(__name__)
 
+Package = s.GET_PACKAGE_MODEL()
+
 
 def normalize_smiles(smiles):
     m1 = Chem.MolFromSmiles(smiles)
@@ -59,9 +56,7 @@ def run_both_engines(SMILES, SMIRKS):
         set(
             [
                 normalize_smiles(str(x))
-                for x in FormatConverter.sanitize_smiles(
+                for x in FormatConverter.sanitize_smiles([str(s) for s in all_rdkit_prods])[0]
-                    [str(s) for s in all_rdkit_prods]
-                )[0]
             ]
         )
     )
@@ -85,8 +80,7 @@ def migration(request):
                 url="http://localhost:8000/package/32de3cf4-e3e6-4168-956e-32fa5ddb0ce1"
             )
             ALL_SMILES = [
-                cs.smiles
+                cs.smiles for cs in CompoundStructure.objects.filter(compound__package=BBD)
-                for cs in CompoundStructure.objects.filter(compound__package=BBD)
             ]
             RULES = SimpleAmbitRule.objects.filter(package=BBD)
 
@@ -142,9 +136,7 @@ def migration(request):
             )
 
         for r in migration_status["results"]:
-            r["detail_url"] = r["detail_url"].replace(
+            r["detail_url"] = r["detail_url"].replace("http://localhost:8000", s.SERVER_URL)
-                "http://localhost:8000", s.SERVER_URL
-            )
 
         context.update(**migration_status)
 
@@ -152,8 +144,6 @@ def migration(request):
 
 
 def migration_detail(request, package_uuid, rule_uuid):
-    current_user = _anonymous_or_real(request)
-
     if request.method == "GET":
         context = get_base_context(request)
 
@@ -235,9 +225,7 @@ def compare(request):
         context["smirks"] = (
             "[#1,#6:6][#7;X3;!$(NC1CC1)!$([N][C]=O)!$([!#8]CNC=O):1]([#1,#6:7])[#6;A;X4:2][H:3]>>[#1,#6:6][#7;X3:1]([#1,#6:7])[H:3].[#6;A:2]=O"
         )
-        context["smiles"] = (
+        context["smiles"] = "C(CC(=O)N[C@@H](CS[Se-])C(=O)NCC(=O)[O-])[C@@H](C(=O)[O-])N"
-            "C(CC(=O)N[C@@H](CS[Se-])C(=O)NCC(=O)[O-])[C@@H](C(=O)[O-])N"
-        )
         return render(request, "compare.html", context)
 
     elif request.method == "POST":

tests/test_enviformer.py

@@ -1,10 +1,15 @@
 from collections import defaultdict
 from datetime import datetime
 from tempfile import TemporaryDirectory
+
+from django.conf import settings as s
 from django.test import TestCase, tag
+
 from epdb.logic import PackageManager
-from epdb.models import User, EnviFormer, Package, Setting
+from epdb.models import EnviFormer, Setting, User
-from epdb.tasks import predict_simple, predict
+from epdb.tasks import predict, predict_simple
+
+Package = s.GET_PACKAGE_MODEL()
 
 
 def measure_predict(mod, pathway_pk=None):
@@ -68,11 +73,15 @@ class EnviFormerTest(TestCase):
 
                 # Test pathway prediction
                 times = [measure_predict(mods[1], self.BBD_SUBSET.pathways[0].pk) for _ in range(5)]
-                print(f"First pathway prediction took {times[0]} seconds, subsequent ones took {times[1:]}")
+                print(
+                    f"First pathway prediction took {times[0]} seconds, subsequent ones took {times[1:]}"
+                )
 
                 # Test eviction by performing three prediction with every model, twice.
                 times = defaultdict(list)
-                for _ in range(2):  # Eviction should cause the second iteration here to have to reload the models
+                for _ in range(
+                    2
+                ):  # Eviction should cause the second iteration here to have to reload the models
                     for mod in mods:
                         for _ in range(3):
                             times[mod.pk].append(measure_predict(mod))

tests/test_model.py

@@ -1,10 +1,13 @@
 from tempfile import TemporaryDirectory
 
 import numpy as np
+from django.conf import settings as s
 from django.test import TestCase
 
 from epdb.logic import PackageManager
-from epdb.models import User, MLRelativeReasoning, Package, RuleBasedRelativeReasoning
+from epdb.models import MLRelativeReasoning, RuleBasedRelativeReasoning, User
+
+Package = s.GET_PACKAGE_MODEL()
 
 
 class ModelTest(TestCase):
@@ -85,7 +88,7 @@ class ModelTest(TestCase):
                 mod.build_model()
                 mod.evaluate_model(True, eval_packages_objs, n_splits=2)
 
-                results = mod.predict("CCN(CC)C(=O)C1=CC(=CC=C1)C")
+                _ = mod.predict("CCN(CC)C(=O)C1=CC(=CC=C1)C")
 
     def test_rbrr(self):
         with TemporaryDirectory() as tmpdir:
@@ -103,12 +106,12 @@ class ModelTest(TestCase):
                     threshold=threshold,
                     min_count=5,
                     max_count=0,
-                    name='ECC - BBD - 0.5',
+                    name="ECC - BBD - 0.5",
-                    description='Created MLRelativeReasoning in Testcase',
+                    description="Created MLRelativeReasoning in Testcase",
                 )
 
                 mod.build_dataset()
                 mod.build_model()
                 mod.evaluate_model(True, eval_packages_objs, n_splits=2)
 
-                results = mod.predict("CCN(CC)C(=O)C1=CC(=CC=C1)C")
+                _ = mod.predict("CCN(CC)C(=O)C1=CC(=CC=C1)C")

tests/test_multigen_eval.py

@@ -1,8 +1,12 @@
+from django.conf import settings as s
 from django.test import TestCase
 from networkx.utils.misc import graphs_equal
+
 from epdb.logic import PackageManager, SPathway
-from epdb.models import Pathway, User, Package
+from epdb.models import Pathway, User
-from utilities.ml import multigen_eval, pathway_edit_eval, graph_from_pathway
+from utilities.ml import graph_from_pathway, multigen_eval, pathway_edit_eval
+
+Package = s.GET_PACKAGE_MODEL()
 
 
 class MultiGenTest(TestCase):

tests/test_simpleambitrule.py

@@ -1,9 +1,10 @@
-from unittest.mock import patch, MagicMock, PropertyMock
+from unittest.mock import MagicMock, PropertyMock, patch
 
+from django.conf import settings as s
 from django.test import TestCase
 
 from epdb.logic import PackageManager
-from epdb.models import User, SimpleAmbitRule
+from epdb.models import SimpleAmbitRule, User
 
 
 class SimpleAmbitRuleTest(TestCase):
@@ -209,7 +210,7 @@ class SimpleAmbitRuleTest(TestCase):
 
         self.assertEqual(rule.products_smarts, expected_products)
 
-    @patch("epdb.models.Package.objects")
+    @patch(f"{s.EPDB_PACKAGE_MODEL.replace('.', '.models.')}.objects")
     def test_related_reactions_property(self, mock_package_objects):
         """Test related_reactions property returns correct queryset."""
         mock_qs = MagicMock()

tests/views/test_model_views.py

@@ -1,9 +1,11 @@
+from django.conf import settings as s
 from django.test import TestCase, override_settings
 from django.urls import reverse
-from django.conf import settings as s
 
 from epdb.logic import UserManager
-from epdb.models import Package, User
+from epdb.models import User
+
+Package = s.GET_PACKAGE_MODEL()
 
 
 @override_settings(MODEL_DIR=s.FIXTURE_DIRS[0] / "models", CELERY_TASK_ALWAYS_EAGER=True)

tests/views/test_package_views.py

@@ -5,14 +5,15 @@ from django.urls import reverse
 
 from epdb.logic import UserManager
 from epdb.models import (
-    Package,
-    UserPackagePermission,
-    Permission,
-    GroupPackagePermission,
     Group,
+    GroupPackagePermission,
     License,
+    Permission,
+    UserPackagePermission,
 )
 
+Package = s.GET_PACKAGE_MODEL()
+
 
 class PackageViewTest(TestCase):
     fixtures = ["test_fixtures.jsonl.gz"]

utilities/decorators.py

@@ -1,10 +1,12 @@
 # decorators.py
 from functools import wraps
 
+from django.conf import settings as s
 from django.shortcuts import get_object_or_404
 
 from epdb.logic import PackageManager
-from epdb.models import Package
+
+Package = s.GET_PACKAGE_MODEL()
 
 # Map HTTP methods to required permissions
 DEFAULT_METHOD_PERMISSIONS = {

utilities/misc.py

@@ -11,6 +11,7 @@ from enum import Enum
 from types import NoneType
 from typing import Any, Dict, List
 
+from django.conf import settings as s
 from django.db import transaction
 from envipy_additional_information import NAME_MAPPING, EnviPyModel, Interval
 from pydantic import BaseModel, HttpUrl
@@ -26,7 +27,6 @@ from epdb.models import (
     License,
     MLRelativeReasoning,
     Node,
-    Package,
     ParallelRule,
     Pathway,
     PluginModel,
@@ -42,6 +42,7 @@ from epdb.models import (
 from utilities.chem import FormatConverter
 
 logger = logging.getLogger(__name__)
+Package = s.GET_PACKAGE_MODEL()
 
 
 class HTMLGenerator: