Add migrations for latest UserCode and CodeHistory

packages/syft/src/syft/protocol/protocol_version.json

@@ -286,19 +286,9 @@
         }
       },
       "UserCode": {
-        "4": {
-          "version": 4,
-          "hash": "0a7181cd5f76800b6566175ffa7276d0cf38c4ddc5110114430147dfc8bfdb2a",
-          "action": "remove"
-        },
         "5": {
           "version": 5,
-          "hash": "128705a5fdf308055ef857b25c80966c928938a05ec03459dae9b36bd6122aa2",
-          "action": "add"
-        },
-        "6": {
-          "version": 6,
-          "hash": "c48ec3160bb34adf937e6306523c7ebc52861ff84a576a30a28cd45c224ded0f",
+          "hash": "c2409f51bf920cce557d288c40b6964ec4df3d8c23e33c5d5668addc30368632",
           "action": "add"
         }
       },

packages/syft/src/syft/service/code/user_code.py

@@ -57,7 +57,6 @@
 from ...types.syft_object import SYFT_OBJECT_VERSION_2
 from ...types.syft_object import SYFT_OBJECT_VERSION_4
 from ...types.syft_object import SYFT_OBJECT_VERSION_5
-from ...types.syft_object import SYFT_OBJECT_VERSION_6
 from ...types.syft_object import SyftObject
 from ...types.syncable_object import SyncableSyftObject
 from ...types.transforms import TransformContext
@@ -271,6 +270,7 @@ def get_sync_dependencies(self, context: AuthedServiceContext) -> list[UID]:
         return [self.user_code_link.object_uid]
 
 
+@serializable()
 class UserCodeV4(SyncableSyftObject):
     # version
     __canonical_name__ = "UserCode"
@@ -302,46 +302,11 @@ class UserCodeV4(SyncableSyftObject):
     worker_pool_name: str | None = None
 
 
-@serializable()
-class UserCodeV5(SyncableSyftObject):
-    # version
-    __canonical_name__ = "UserCode"
-    __version__ = SYFT_OBJECT_VERSION_5
-
-    id: UID
-    node_uid: UID | None = None
-    user_verify_key: SyftVerifyKey
-    raw_code: str
-    input_policy_type: type[InputPolicy] | UserPolicy
-    input_policy_init_kwargs: dict[Any, Any] | None = None
-    input_policy_state: bytes = b""
-    output_policy_type: type[OutputPolicy] | UserPolicy
-    output_policy_init_kwargs: dict[Any, Any] | None = None
-    output_policy_state: bytes = b""
-    parsed_code: str
-    service_func_name: str
-    unique_func_name: str
-    user_unique_func_name: str
-    code_hash: str
-    signature: inspect.Signature
-    status_link: LinkedObject | None = None
-    input_kwargs: list[str]
-    enclave_metadata: EnclaveMetadata | None = None
-    submit_time: DateTime | None = None
-    # tracks if the code calls domain.something, variable is set during parsing
-    uses_domain: bool = False
-
-    nested_codes: dict[str, tuple[LinkedObject, dict]] | None = {}
-    worker_pool_name: str | None = None
-    origin_node_side_type: NodeSideType
-    l0_deny_reason: str | None = None
-
-
 @serializable()
 class UserCode(SyncableSyftObject):
     # version
     __canonical_name__ = "UserCode"
-    __version__ = SYFT_OBJECT_VERSION_6
+    __version__ = SYFT_OBJECT_VERSION_5
 
     id: UID
     node_uid: UID | None = None
@@ -1924,12 +1889,27 @@ def migrate_usercode_v4_to_v5() -> list[Callable]:
     return [
         make_set_default("origin_node_side_type", NodeSideType.HIGH_SIDE),
         make_set_default("l0_deny_reason", None),
+        drop("enclave_metadata"),
     ]
 
 
 @migrate(UserCode, UserCodeV4)
 def migrate_usercode_v5_to_v4() -> list[Callable]:
     return [
-        drop("origin_node_side_type"),
-        drop("l0_deny_reason"),
+        drop(["origin_node_side_type", "l0_deny_reason"]),
+        make_set_default("enclave_metadata", None),
+    ]
+
+
+@migrate(SubmitUserCodeV4, SubmitUserCode)
+def upgrade_submitusercode() -> list[Callable]:
+    return [
+        drop("enclave_metadata"),
+    ]
+
+
+@migrate(SubmitUserCode, SubmitUserCodeV4)
+def downgrade_submitusercode() -> list[Callable]:
+    return [
+        make_set_default("enclave_metadata", None),
     ]

packages/syft/src/syft/service/code_history/code_history.py

@@ -1,4 +1,5 @@
 # stdlib
+from collections.abc import Callable
 import json
 from typing import Any
 
@@ -7,10 +8,13 @@
 from ...client.enclave_client import EnclaveMetadata
 from ...serde.serializable import serializable
 from ...service.user.user_roles import ServiceRole
+from ...types.syft_migration import migrate
 from ...types.syft_object import SYFT_OBJECT_VERSION_2
 from ...types.syft_object import SYFT_OBJECT_VERSION_3
 from ...types.syft_object import SyftObject
 from ...types.syft_object import SyftVerifyKey
+from ...types.transforms import drop
+from ...types.transforms import make_set_default
 from ...types.uid import UID
 from ...util.notebook_ui.components.tabulator_template import (
     build_tabulator_table_with_data,
@@ -171,3 +175,15 @@ def _repr_html_(self) -> str | None:
             "icon": None,
         }
         return build_tabulator_table_with_data(rows, metadata)
+
+
+@migrate(CodeHistoryV2, CodeHistory)
+def code_history_v2_to_v3() -> list[Callable]:
+    return [drop("enclave_metadata")]
+
+
+@migrate(CodeHistory, CodeHistoryV2)
+def code_history_v3_to_v2() -> list[Callable]:
+    return [
+        make_set_default("enclave_metadata", None),
+    ]