Replace manual exclude handling in ORM

MolSSI · Jan 12, 2024 · e9c8742 · e9c8742
1 parent 369cd9f
commit e9c8742
Show file tree

Hide file tree

Showing 23 changed files with 97 additions and 310 deletions.
diff --git a/qcfractal/qcfractal/components/auth/db_models.py b/qcfractal/qcfractal/components/auth/db_models.py
@@ -51,10 +51,7 @@ class RoleORM(BaseORM):
 
     __table_args__ = (UniqueConstraint("rolename", name="ux_role_rolename"),)
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Remove fields not present in the model
-        exclude = self.append_exclude(exclude, "id")
-        return BaseORM.model_dict(self, exclude)
+    _qcportal_model_excludes = ["id"]
 
 
 class UserORM(BaseORM):
@@ -82,9 +79,9 @@ class UserORM(BaseORM):
 
     __table_args__ = (UniqueConstraint("username", name="ux_user_username"),)
 
+    _qcportal_model_excludes = ["role_orm", "groups_orm", "role_id", "password"]
+
     def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Removes some sensitive or useless fields
-        exclude = self.append_exclude(exclude, "role_orm", "groups_orm", "role_id", "password")
         d = BaseORM.model_dict(self, exclude)
 
         d["role"] = self.role_orm.rolename

diff --git a/qcfractal/qcfractal/components/dataset_db_models.py b/qcfractal/qcfractal/components/dataset_db_models.py
@@ -91,11 +91,9 @@ class BaseDatasetORM(BaseORM):
 
     __mapper_args__ = {"polymorphic_on": "dataset_type"}
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # lname is only for the server
-        # strip user/group ids
-        exclude = self.append_exclude(exclude, "lname", "owner_user_id", "owner_group_id")
+    _qcportal_model_excludes = ["lname", "owner_user_id", "owner_group_id"]
 
+    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
         d = BaseORM.model_dict(self, exclude)
 
         # meta -> metadata
@@ -130,6 +128,4 @@ class ContributedValuesORM(BaseORM):
 
     __table_args__ = (Index("ix_contributed_values_dataset_id", "dataset_id"),)
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        exclude = self.append_exclude(exclude, "dataset_id")
-        return BaseORM.model_dict(self, exclude)
+    _qcportal_model_excludes = ["dataset_id"]
diff --git a/qcfractal/qcfractal/components/gridoptimization/dataset_db_models.py b/qcfractal/qcfractal/components/gridoptimization/dataset_db_models.py
@@ -1,8 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
-from sqlalchemy import Column, Integer, ForeignKey, String, ForeignKeyConstraint, UniqueConstraint, Index
+from sqlalchemy import Column, Integer, ForeignKey, String, ForeignKeyConstraint, Index
 from sqlalchemy.dialects.postgresql import JSONB
 from sqlalchemy.orm import relationship
 
@@ -14,9 +12,6 @@
 from qcfractal.components.molecules.db_models import MoleculeORM
 from qcfractal.db_socket import BaseORM
 
-if TYPE_CHECKING:
-    from typing import Dict, Any, Optional, Iterable
-
 
 class GridoptimizationDatasetEntryORM(BaseORM):
     __tablename__ = "gridoptimization_dataset_entry"
@@ -39,10 +34,7 @@ class GridoptimizationDatasetEntryORM(BaseORM):
         Index("ix_gridoptimization_dataset_entry_initial_molecule_id", "initial_molecule_id"),
     )
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Remove fields not present in the model
-        exclude = self.append_exclude(exclude, "dataset_id", "initial_molecule_id")
-        return BaseORM.model_dict(self, exclude)
+    _qcportal_model_excludes = ["dataset_id", "initial_molecule_id"]
 
 
 class GridoptimizationDatasetSpecificationORM(BaseORM):
@@ -61,10 +53,7 @@ class GridoptimizationDatasetSpecificationORM(BaseORM):
         Index("ix_gridoptimization_dataset_specification_specification_id", "specification_id"),
     )
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Remove fields not present in the model
-        exclude = self.append_exclude(exclude, "dataset_id", "specification_id")
-        return BaseORM.model_dict(self, exclude)
+    _qcportal_model_excludes = ["dataset_id", "specification_id"]
 
 
 class GridoptimizationDatasetRecordItemORM(BaseORM):
@@ -93,10 +82,7 @@ class GridoptimizationDatasetRecordItemORM(BaseORM):
         Index("ix_gridoptimization_dataset_record_record_id", "record_id"),
     )
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Remove fields not present in the model
-        exclude = self.append_exclude(exclude, "dataset_id")
-        return BaseORM.model_dict(self, exclude)
+    _qcportal_model_excludes = ["dataset_id"]
 
 
 class GridoptimizationDatasetORM(BaseDatasetORM):

diff --git a/qcfractal/qcfractal/components/gridoptimization/record_db_models.py b/qcfractal/qcfractal/components/gridoptimization/record_db_models.py
@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from sqlalchemy import (
     select,
     UniqueConstraint,
@@ -26,9 +24,6 @@
 from qcfractal.components.record_db_models import BaseRecordORM
 from qcfractal.db_socket import BaseORM
 
-if TYPE_CHECKING:
-    from typing import Dict, Any, Optional, Iterable
-
 
 class GridoptimizationOptimizationORM(BaseORM):
     """
@@ -51,10 +46,7 @@ class GridoptimizationOptimizationORM(BaseORM):
 
     __table_args__ = (Index("ix_gridoptimization_optimization_id", "optimization_id"),)
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Remove fields not present in the model
-        exclude = self.append_exclude(exclude, "gridoptimization_id")
-        return BaseORM.model_dict(self, exclude)
+    _qcportal_model_excludes = ["gridoptimization_id"]
 
 
 class GridoptimizationSpecificationORM(BaseORM):
@@ -89,10 +81,7 @@ class GridoptimizationSpecificationORM(BaseORM):
         CheckConstraint("program = LOWER(program)", name="ck_gridoptimization_specification_program_lower"),
     )
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Remove fields not present in the model
-        exclude = self.append_exclude(exclude, "id", "keywords_hash", "optimization_specification_id")
-        return BaseORM.model_dict(self, exclude)
+    _qcportal_model_excludes = ["id", "keywords_hash", "optimization_specification_id"]
 
     @property
     def short_description(self) -> str:
@@ -125,10 +114,7 @@ class GridoptimizationRecordORM(BaseRecordORM):
         "polymorphic_identity": "gridoptimization",
     }
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Remove fields not present in the model
-        exclude = self.append_exclude(exclude, "specification_id")
-        return BaseRecordORM.model_dict(self, exclude)
+    _qcportal_model_excludes = [*BaseRecordORM._qcportal_model_excludes, "specification_id"]
 
     @property
     def short_description(self) -> str:

diff --git a/qcfractal/qcfractal/components/internal_jobs/db_models.py b/qcfractal/qcfractal/components/internal_jobs/db_models.py
@@ -62,9 +62,9 @@ class InternalJobORM(BaseORM):
         UniqueConstraint("unique_name", name="ux_internal_jobs_unique_name"),
     )
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        exclude = self.append_exclude(exclude, "unique_name", "user_id")
+    _qcportal_model_excludes = ["unique_name", "user_id"]
 
+    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
         d = BaseORM.model_dict(self, exclude)
         d["user"] = self.user.username if self.user is not None else None
         return d

diff --git a/qcfractal/qcfractal/components/manybody/dataset_db_models.py b/qcfractal/qcfractal/components/manybody/dataset_db_models.py
@@ -1,8 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
-from sqlalchemy import Column, Integer, ForeignKey, String, ForeignKeyConstraint, UniqueConstraint, Index
+from sqlalchemy import Column, Integer, ForeignKey, String, ForeignKeyConstraint, Index
 from sqlalchemy.dialects.postgresql import JSONB
 from sqlalchemy.orm import relationship
 
@@ -11,9 +9,6 @@
 from qcfractal.components.molecules.db_models import MoleculeORM
 from qcfractal.db_socket import BaseORM
 
-if TYPE_CHECKING:
-    from typing import Dict, Any, Optional, Iterable
-
 
 class ManybodyDatasetEntryORM(BaseORM):
     __tablename__ = "manybody_dataset_entry"
@@ -35,10 +30,7 @@ class ManybodyDatasetEntryORM(BaseORM):
         Index("ix_manybody_dataset_entry_initial_molecule_id", "initial_molecule_id"),
     )
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Remove fields not present in the model
-        exclude = self.append_exclude(exclude, "dataset_id", "initial_molecule_id")
-        return BaseORM.model_dict(self, exclude)
+    _qcportal_model_excludes = ["dataset_id", "initial_molecule_id"]
 
 
 class ManybodyDatasetSpecificationORM(BaseORM):
@@ -57,10 +49,7 @@ class ManybodyDatasetSpecificationORM(BaseORM):
         Index("ix_manybody_dataset_specification_specification_id", "specification_id"),
     )
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Remove fields not present in the model
-        exclude = self.append_exclude(exclude, "dataset_id", "specification_id")
-        return BaseORM.model_dict(self, exclude)
+    _qcportal_model_excludes = ["dataset_id", "specification_id"]
 
 
 class ManybodyDatasetRecordItemORM(BaseORM):
@@ -89,10 +78,7 @@ class ManybodyDatasetRecordItemORM(BaseORM):
         Index("ix_manybody_dataset_record_record_id", "record_id"),
     )
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Remove fields not present in the model
-        exclude = self.append_exclude(exclude, "dataset_id")
-        return BaseORM.model_dict(self, exclude)
+    _qcportal_model_excludes = ["dataset_id"]
 
 
 class ManybodyDatasetORM(BaseDatasetORM):

diff --git a/qcfractal/qcfractal/components/manybody/record_db_models.py b/qcfractal/qcfractal/components/manybody/record_db_models.py
@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from sqlalchemy import Column, String, Integer, ForeignKey, UniqueConstraint, Index, CheckConstraint, event, DDL
 from sqlalchemy.dialects.postgresql import JSONB, ARRAY
 from sqlalchemy.orm import relationship
@@ -11,9 +9,6 @@
 from qcfractal.components.singlepoint.record_db_models import SinglepointRecordORM, QCSpecificationORM
 from qcfractal.db_socket import BaseORM
 
-if TYPE_CHECKING:
-    from typing import Dict, Any, Optional, Iterable
-
 
 class ManybodyClusterORM(BaseORM):
     """
@@ -41,10 +36,7 @@ class ManybodyClusterORM(BaseORM):
         Index("ix_manybody_cluster_singlepoint_id", "singlepoint_id"),
     )
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Remove fields not present in the model
-        exclude = self.append_exclude(exclude, "manybody_id")
-        return BaseORM.model_dict(self, exclude)
+    _qcportal_model_excludes = ["manybody_id"]
 
 
 class ManybodySpecificationORM(BaseORM):
@@ -75,10 +67,7 @@ class ManybodySpecificationORM(BaseORM):
         CheckConstraint("program = LOWER(program)", name="ck_manybody_specification_program_lower"),
     )
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Remove fields not present in the model
-        exclude = self.append_exclude(exclude, "id", "keywords_hash", "singlepoint_specification_id")
-        return BaseORM.model_dict(self, exclude)
+    _qcportal_model_excludes = ["id", "keywords_hash", "singlepoint_specification_id"]
 
     @property
     def short_description(self) -> str:
@@ -107,10 +96,7 @@ class ManybodyRecordORM(BaseRecordORM):
         "polymorphic_identity": "manybody",
     }
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Remove fields not present in the model
-        exclude = self.append_exclude(exclude, "specification_id")
-        return BaseRecordORM.model_dict(self, exclude)
+    _qcportal_model_excludes = [*BaseRecordORM._qcportal_model_excludes, "specification_id"]
 
     @property
     def short_description(self) -> str:

diff --git a/qcfractal/qcfractal/components/molecules/db_models.py b/qcfractal/qcfractal/components/molecules/db_models.py
@@ -2,7 +2,7 @@
 
 from typing import TYPE_CHECKING
 
-from sqlalchemy import Column, Integer, String, JSON, Float, Index, CHAR, Boolean, UniqueConstraint
+from sqlalchemy import Column, Integer, String, JSON, Float, Index, CHAR, UniqueConstraint
 from sqlalchemy.dialects.postgresql import JSONB
 from sqlalchemy.orm import column_property
 
@@ -69,10 +69,9 @@ class MoleculeORM(BaseORM):
         UniqueConstraint("molecule_hash", name="ux_molecule_molecule_hash"),
     )
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # molecule_hash is only used for indexing. It is otherwise stored in identifiers
-        exclude = self.append_exclude(exclude, "molecule_hash")
+    _qcportal_model_excludes = ["molecule_hash"]
 
+    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
         d = BaseORM.model_dict(self, exclude)
 
         # TODO - this is because the pydantic models are goofy

diff --git a/qcfractal/qcfractal/components/neb/dataset_db_models.py b/qcfractal/qcfractal/components/neb/dataset_db_models.py
@@ -1,8 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
-from sqlalchemy import Column, Integer, ForeignKey, String, ForeignKeyConstraint, UniqueConstraint, Index
+from sqlalchemy import Column, Integer, ForeignKey, String, ForeignKeyConstraint, Index
 from sqlalchemy.dialects.postgresql import JSONB
 from sqlalchemy.orm import relationship
 
@@ -11,9 +9,6 @@
 from qcfractal.components.neb.record_db_models import NEBRecordORM, NEBSpecificationORM
 from qcfractal.db_socket import BaseORM
 
-if TYPE_CHECKING:
-    from typing import Dict, Any, Optional, Iterable
-
 
 class NEBDatasetInitialMoleculeORM(BaseORM):
     """
@@ -41,10 +36,7 @@ class NEBDatasetInitialMoleculeORM(BaseORM):
         ),
     )
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Remove fields not present in the model
-        exclude = self.append_exclude(exclude, "dataset_id", "molecule_id", "entry_name", "position")
-        return BaseORM.model_dict(self, exclude)
+    _qcportal_model_excludes = ["dataset_id", "molecule_id", "entry_name", "position"]
 
 
 class NEBDatasetEntryORM(BaseORM):
@@ -81,10 +73,7 @@ class NEBDatasetEntryORM(BaseORM):
         Index("ix_neb_dataset_entry_name", "name"),
     )
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Remove fields not present in the model
-        exclude = self.append_exclude(exclude, "dataset_id")
-        return BaseORM.model_dict(self, exclude)
+    _qcportal_model_excludes = ["dataset_id"]
 
 
 class NEBDatasetSpecificationORM(BaseORM):
@@ -103,10 +92,7 @@ class NEBDatasetSpecificationORM(BaseORM):
         Index("ix_neb_dataset_specification_specification_id", "specification_id"),
     )
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Remove fields not present in the model
-        exclude = self.append_exclude(exclude, "dataset_id", "specification_id")
-        return BaseORM.model_dict(self, exclude)
+    _qcportal_model_excludes = ["dataset_id", "specification_id"]
 
 
 class NEBDatasetRecordItemORM(BaseORM):
@@ -135,10 +121,7 @@ class NEBDatasetRecordItemORM(BaseORM):
         Index("ix_neb_dataset_record_record_id", "record_id"),
     )
 
-    def model_dict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        # Remove fields not present in the model
-        exclude = self.append_exclude(exclude, "dataset_id")
-        return BaseORM.model_dict(self, exclude)
+    _qcportal_model_excludes = ["dataset_id"]
 
 
 class NEBDatasetORM(BaseDatasetORM):