Merge pull request #8986 from OpenMined/bschell/update-migrations-for…

…-enclave-pr Add Migrations for Enclave Changes to UserCode and CodeHistory
OpenMined · Jul 2, 2024 · a3b5de3 · a3b5de3
2 parents 46207a5 + 5a3f5a5
commit a3b5de3
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/packages/syft/src/syft/service/code/user_code.py b/packages/syft/src/syft/service/code/user_code.py
@@ -1952,27 +1952,41 @@ 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),
+        make_set_default("runtime_policy_type", EmptyRuntimePolicy),
+        make_set_default("runtime_policy_init_kwargs", None),
+        make_set_default("runtime_policy_state", b""),
         drop("enclave_metadata"),
     ]
 
 
 @migrate(UserCode, UserCodeV4)
 def migrate_usercode_v5_to_v4() -> list[Callable]:
     return [
-        drop(["origin_node_side_type", "l0_deny_reason"]),
         make_set_default("enclave_metadata", None),
+        drop(
+            [
+                "origin_node_side_type",
+                "l0_deny_reason",
+                "runtime_policy_type",
+                "runtime_policy_init_kwargs",
+                "runtime_policy_state",
+            ]
+        ),
     ]
 
 
 @migrate(SubmitUserCodeV4, SubmitUserCode)
 def upgrade_submitusercode() -> list[Callable]:
     return [
+        make_set_default("runtime_policy_type", EmptyRuntimePolicy),
+        make_set_default("runtime_policy_init_kwargs", None),
         drop("enclave_metadata"),
     ]
 
 
 @migrate(SubmitUserCode, SubmitUserCodeV4)
 def downgrade_submitusercode() -> list[Callable]:
     return [
+        drop(["runtime_policy_type", "runtime_policy_init_kwargs"]),
         make_set_default("enclave_metadata", None),
     ]