add coverage for constraints

liminal/base/properties/base_field_properties.py

@@ -122,4 +122,4 @@ def __str__(self) -> str:
 
     def __repr__(self) -> str:
         """Generates a string representation of the class so that it can be executed."""
-        return f"{self.__class__.__name__}({', '.join([f'{k}={v.__repr__()}' for k, v in self.model_dump(exclude_defaults=True).items()])})"
+        return f"{self.__class__.__name__}({', '.join([f'{k}={v.__repr__()}' for k, v in self.model_dump(exclude_unset=True).items()])})"

liminal/base/properties/base_schema_properties.py

@@ -59,6 +59,8 @@ class BaseSchemaProperties(BaseModel):
         The entity type of the schema.
     mixture_schema_config : MixtureSchemaConfig | None
         The mixture schema config of the schema.
+    constraint_fields : set[str] | None
+        Set of constraints of field values for the schema. Set of column names, that specify that their values must be a unique combination in their entities.
     _archived : bool | None
         Whether the schema is archived in Benchling.
     """
@@ -69,6 +71,7 @@ class BaseSchemaProperties(BaseModel):
     entity_type: BenchlingEntityType | None = None
     naming_strategies: set[BenchlingNamingStrategy] | None = None
     mixture_schema_config: MixtureSchemaConfig | None = None
+    constraint_fields: set[str] | None = None
     _archived: bool | None = None
 
     def __init__(self, **data: Any):
@@ -109,4 +112,4 @@ def __str__(self) -> str:
 
     def __repr__(self) -> str:
         """Generates a string representation of the class so that it can be executed."""
-        return f"{self.__class__.__name__}({', '.join([f'{k}={v.__repr__()}' for k, v in self.model_dump(exclude_defaults=True).items()])})"
+        return f"{self.__class__.__name__}({', '.join([f'{k}={v.__repr__()}' for k, v in self.model_dump(exclude_unset=True).items()])})"

liminal/entity_schemas/entity_schema_models.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Any
+
 from pydantic import BaseModel
 
 from liminal.base.properties.base_field_properties import BaseFieldProperties
@@ -18,6 +20,37 @@ class FieldLinkShortModel(BaseModel):
     folderItemType: str | None = None
 
 
+class EntitySchemaConstraint(BaseModel):
+    """
+    A class to define a constraint on an entity schema.
+    """
+
+    areUniqueResiduesCaseSensitive: bool | None = None
+    fields: dict[str, Any] | None = None
+    hasUniqueCanonicalSmilers: bool | None = None
+    hasUniqueResidues: bool | None = None
+
+    @classmethod
+    def from_constraint_fields(
+        cls, constraint_fields: set[str]
+    ) -> EntitySchemaConstraint:
+        """
+        Generates a Constraint object from a set of constraint fields to create a constraint on a schema.
+        """
+        if constraint_fields is None:
+            return None
+        hasUniqueResidues = False
+        if "bases" in constraint_fields:
+            constraint_fields.discard("bases")
+            hasUniqueResidues = True
+        return cls(
+            fields=[{"name": f} for f in constraint_fields],
+            hasUniqueResidues=hasUniqueResidues,
+            hasUniqueCanonicalSmilers=False,
+            areUniqueResiduesCaseSensitive=False,
+        )
+
+
 class CreateEntitySchemaFieldModel(BaseModel):
     """A pydantic model to define a field for the create entity schema endpoint.
     This model is used as input for the benchling alpha create entity schema endpoint."""
@@ -99,6 +132,7 @@ class CreateEntitySchemaModel(BaseModel):
     type: BenchlingEntityType
     mixtureSchemaConfig: MixtureSchemaConfig | None = None
     labelingStrategies: list[str] | None = None
+    constraint: EntitySchemaConstraint | None = None
 
     @classmethod
     def from_benchling_props(
@@ -130,6 +164,11 @@ def from_benchling_props(
             type=benchling_props.entity_type,
             mixtureSchemaConfig=benchling_props.mixture_schema_config,
             labelingStrategies=[s.value for s in benchling_props.naming_strategies],
+            constraint=EntitySchemaConstraint.from_constraint_fields(
+                benchling_props.constraint_fields
+            )
+            if benchling_props.constraint_fields
+            else None,
             fields=[
                 CreateEntitySchemaFieldModel.from_benchling_props(
                     field_props, benchling_service

liminal/entity_schemas/tag_schema_models.py

@@ -32,6 +32,31 @@ class FieldRequiredLinkShortModel(BaseModel):
     storableSchema: dict[str, Any] | None = None
 
 
+class TagSchemaConstraint(BaseModel):
+    """
+    A class to define a constraint on an entity schema.
+    """
+
+    areUniqueResiduesCaseSensitive: bool | None = None
+    fields: list[TagSchemaFieldModel] | None = None
+    uniqueCanonicalSmilers: bool | None = None
+    uniqueResidues: bool | None = None
+
+    @classmethod
+    def from_constraint_fields(
+        cls, constraint_fields: list[TagSchemaFieldModel], bases: bool
+    ) -> TagSchemaConstraint:
+        """
+        Generates a Constraint object from a set of constraint fields to create a constraint on a schema.
+        """
+        return cls(
+            fields=constraint_fields,
+            uniqueResidues=bases,
+            uniqueCanonicalSmilers=False,
+            areUniqueResiduesCaseSensitive=False,
+        )
+
+
 class UpdateTagSchemaModel(BaseModel):
     """A pydantic model to define the input for the internal tag schema update endpoint."""
 
@@ -46,6 +71,7 @@ class UpdateTagSchemaModel(BaseModel):
     sequenceType: BenchlingSequenceType | None = None
     shouldCreateAsOligo: bool | None = None
     showResidues: bool | None = None
+    constraint: TagSchemaConstraint | None = None
 
 
 class CreateTagSchemaFieldModel(BaseModel):
@@ -243,7 +269,7 @@ class TagSchemaModel(BaseModel):
     authParentOption: str | None
     batchSchemaId: str | None
     childEntitySchemaSummaries: list[Any] | None
-    constraint: Any | None
+    constraint: TagSchemaConstraint | None
     containableType: str | None
     fields: list[TagSchemaFieldModel]
     folderItemType: BenchlingFolderItemType
@@ -348,6 +374,23 @@ def update_schema_props(self, update_diff: dict[str, Any]) -> TagSchemaModel:
         if "mixture_schema_config" in update_diff_names:
             self.mixtureSchemaConfig = update_props.mixture_schema_config
 
+        if "constraint_fields" in update_diff_names:
+            if update_props.constraint_fields:
+                has_bases = False
+                if "bases" in update_props.constraint_fields:
+                    has_bases = True
+                    update_props.constraint_fields.discard("bases")
+                constraint_fields = [
+                    f
+                    for f in self.fields
+                    if f.systemName in update_props.constraint_fields
+                ]
+                self.constraint = TagSchemaConstraint.from_constraint_fields(
+                    constraint_fields, has_bases
+                )
+            else:
+                self.constraint = None
+
         self.prefix = (
             update_props.prefix if "prefix" in update_diff_names else self.prefix
         )

liminal/entity_schemas/utils.py

@@ -44,6 +44,11 @@ def convert_tag_schema_to_internal_schema(
     all_fields = tag_schema.allFields
     if not include_archived_fields:
         all_fields = [f for f in all_fields if not f.archiveRecord]
+    constraint_fields: set[str] | None = None
+    if tag_schema.constraint:
+        constraint_fields = set([f.systemName for f in tag_schema.constraint.fields])
+        if tag_schema.constraint.uniqueResidues:
+            constraint_fields.add("bases")
     return (
         SchemaProperties(
             name=tag_schema.name,
@@ -63,6 +68,7 @@ def convert_tag_schema_to_internal_schema(
                 BenchlingNamingStrategy(strategy)
                 for strategy in tag_schema.labelingStrategies
             ),
+            constraint_fields=constraint_fields,
             _archived=tag_schema.archiveRecord is not None,
         ),
         {

liminal/orm/base_model.py

@@ -39,13 +39,14 @@ class BaseModel(Generic[T], Base):
 
     def __init_subclass__(cls, **kwargs: Any):
         super().__init_subclass__(**kwargs)
+        warehouse_name = cls.__schema_properties__.warehouse_name
+        cls.__tablename__ = warehouse_name + "$raw"
         if "__schema_properties__" not in cls.__dict__ or not isinstance(
             cls.__schema_properties__, SchemaProperties
         ):
             raise NotImplementedError(
                 f"{cls.__name__} must define 'schema_properties' class attribute"
             )
-        warehouse_name = cls.__schema_properties__.warehouse_name
         if warehouse_name in cls._existing_schema_warehouse_names:
             raise ValueError(
                 f"Warehouse name '{warehouse_name}' is already used by another subclass."
@@ -58,11 +59,24 @@ def __init_subclass__(cls, **kwargs: Any):
             raise ValueError(
                 f"Schema prefix '{cls.__schema_properties__.prefix}' is already used by another subclass."
             )
+        # Validate constraints
+        if cls.__schema_properties__.constraint_fields:
+            column_wh_names = [
+                c[0] for c in cls.__dict__.items() if isinstance(c[1], SqlColumn)
+            ]
+            invalid_constraints = [
+                c
+                for c in cls.__schema_properties__.constraint_fields
+                if c not in (set(column_wh_names) | {"bases"})
+            ]
+            if invalid_constraints:
+                raise ValueError(
+                    f"Constraints {', '.join(invalid_constraints)} are not fields on schema {cls.__schema_properties__.name}."
+                )
 
         cls._existing_schema_warehouse_names.add(warehouse_name)
         cls._existing_schema_names.add(cls.__schema_properties__.name)
         cls._existing_schema_prefixes.add(cls.__schema_properties__.prefix.lower())
-        cls.__tablename__ = warehouse_name + "$raw"
 
     @declared_attr
     def creator_id(cls) -> SqlColumn:

liminal/orm/schema_properties.py

@@ -24,6 +24,7 @@ class SchemaProperties(BaseSchemaProperties):
     entity_type: BenchlingEntityType
     naming_strategies: set[BenchlingNamingStrategy]
     mixture_schema_config: MixtureSchemaConfig | None = None
+    constraint_fields: set[str] | None = None
     _archived: bool = False
 
     def __init__(self, **data: Any):

liminal/tests/conftest.py

@@ -125,6 +125,7 @@ def mock_benchling_schema(
         prefix="MockEntity",
         entity_type=BenchlingEntityType.CUSTOM_ENTITY,
         naming_strategies=[BenchlingNamingStrategy.NEW_IDS],
+        constraint_fields={"enum_field", "string_field_req"},
     )
     fields = {
         "enum_field": Props(
@@ -347,6 +348,7 @@ class MockEntity(BaseModel):
             prefix="MockEntity",
             entity_type=BenchlingEntityType.CUSTOM_ENTITY,
             naming_strategies=[BenchlingNamingStrategy.NEW_IDS],
+            constraint_fields={"enum_field", "string_field_req"},
         )
         enum_field: SqlColumn = Column(
             name="Enum Field",

liminal/tests/test_entity_schema_compare.py

@@ -341,3 +341,21 @@ def test_compare_benchling_schema_fields(  # type: ignore[no-untyped-def]
                 invalid_models["mock_entity"][0].op, ArchiveEntitySchemaField
             )
             assert invalid_models["mock_entity"][0].op.wh_field_name == "archived_field"
+
+            # Test when the Benchling schema has different constraint fields
+            benchling_mismatch_constraint_fields = copy.deepcopy(mock_benchling_schema)
+            benchling_mismatch_constraint_fields[0][0].constraint_fields = {
+                "string_field_req"
+            }
+            mock_get_benchling_entity_schemas.return_value = (
+                benchling_mismatch_constraint_fields
+            )
+            invalid_models = compare_entity_schemas(mock_benchling_sdk)
+            assert len(invalid_models["mock_entity"]) == 1
+            assert isinstance(invalid_models["mock_entity"][0].op, UpdateEntitySchema)
+            assert invalid_models["mock_entity"][
+                0
+            ].op.update_props.constraint_fields == {
+                "string_field_req",
+                "enum_field",
+            }