feat: auto-bump schema versions to current version (#1193)

* feat: auto-bump schema versions to current version

* chore: lint

* chore: lint

* chore: missing docstrings

src/aind_data_schema/base.py

@@ -3,7 +3,7 @@
 import json
 import re
 from pathlib import Path
-from typing import Any, Generic, Optional, TypeVar
+from typing import Any, Generic, Optional, TypeVar, get_args
 
 from pydantic import (
     AwareDatetime,
@@ -16,6 +16,7 @@
     ValidatorFunctionWrapHandler,
     create_model,
     model_validator,
+    field_validator,
 )
 from pydantic.functional_validators import WrapValidator
 from typing_extensions import Annotated
@@ -122,6 +123,12 @@ class AindCoreModel(AindModel):
         ..., pattern=r"^\d+.\d+.\d+$", description="schema version", title="Version", frozen=True
     )
 
+    @field_validator("schema_version", mode="before")
+    @classmethod
+    def coerce_version(cls, v: str) -> str:
+        """Update the schema version to the latest version"""
+        return get_args(cls.model_fields["schema_version"].annotation)[0]
+
     @classmethod
     def default_filename(cls):
         """

tests/test_base.py

@@ -6,9 +6,10 @@
 from pathlib import Path
 from unittest.mock import MagicMock, call, mock_open, patch
 
-from pydantic import ValidationError, create_model
+from pydantic import ValidationError, create_model, SkipValidation
+from typing import Literal
 
-from aind_data_schema.base import AindGeneric, AwareDatetimeWithDefault, is_dict_corrupt, AindModel
+from aind_data_schema.base import AindGeneric, AwareDatetimeWithDefault, is_dict_corrupt, AindModel, AindCoreModel
 from aind_data_schema.core.subject import Subject
 from aind_data_schema_models.brain_atlas import CCFStructure
 
@@ -118,10 +119,36 @@ def test_ccf_validator(self):
 
         class StructureModel(AindModel):
             """Test model with a targeted_structure"""
+
             targeted_structure: CCFStructure.ONE_OF
 
         self.assertRaises(ValueError, StructureModel, targeted_structure="invalid")
 
+    def test_schema_bump(self):
+        """Test that schema version are bumped successfully
+        and that validation errors prevent bumping"""
+
+        class Modelv1(AindCoreModel):
+            """test class"""
+            describedBy: str = "modelv1"
+            schema_version: SkipValidation[Literal["1.0.0"]] = "1.0.0"
+
+        class Modelv2(AindCoreModel):
+            """test class"""
+            describedBy: str = "modelv2"
+            schema_version: SkipValidation[Literal["1.0.1"]] = "1.0.1"
+            extra_field: str = "extra_field"
+
+        v1_init = Modelv1()
+        self.assertEqual("1.0.0", v1_init.schema_version)
+
+        v2_from_v1 = Modelv2(**v1_init.model_dump())
+        self.assertEqual("1.0.1", v2_from_v1.schema_version)
+
+        # Check that adding additional fields still fails validation
+        # this is to ensure you can't get a bumped schema_version without passing validation
+        self.assertRaises(ValidationError, lambda: Modelv1(**v2_from_v1.model_dump()))
+
 
 if __name__ == "__main__":
     unittest.main()