Merge pull request #312 from nodestream-proj/fix/migration-failure-re…

…siliency-issues Fix Migrator Resiliency Issues
nodestream-proj · Jun 6, 2024 · 4395714 · 4395714
2 parents fed9fa1 + 64f695d
commit 4395714
Show file tree

Hide file tree

Showing 4 changed files with 75 additions and 3 deletions.
diff --git a/nodestream/schema/migrations/auto_change_detector.py b/nodestream/schema/migrations/auto_change_detector.py
@@ -377,6 +377,8 @@ def make_relationship_key_change_operations(self) -> Iterable[Operation]:
 
     def make_node_property_change_operations(self) -> Iterable[Operation]:
         for type, deleted_field in self.deleted_node_properties:
+            if (type, deleted_field) in self.added_node_keys:
+                continue
             yield DropNodeProperty(type, deleted_field)
 
         for type, old_field, new_field in self.renamed_node_properties:
@@ -387,6 +389,8 @@ def make_node_property_change_operations(self) -> Iterable[Operation]:
 
     def make_relationship_property_change_operations(self) -> Iterable[Operation]:
         for type, deleted_field in self.deleted_relationship_properties:
+            if (type, deleted_field) in self.added_relationship_keys:
+                continue
             yield DropRelationshipProperty(type, deleted_field)
 
         for type, old_field, new_field in self.renamed_relationship_properties:

diff --git a/nodestream/schema/migrations/migrator.py b/nodestream/schema/migrations/migrator.py
@@ -135,8 +135,12 @@ async def execute_migration(self, migration: Migration) -> None:
             await self.acquire_lock()
 
         async with self.transaction():
-            for operation in migration.operations:
-                await self.execute_operation(operation)
+            try:
+                for operation in migration.operations:
+                    await self.execute_operation(operation)
+            except Exception:
+                await self.release_lock()
+                raise
 
         async with self.transaction():
             await self.mark_migration_as_executed(migration)

diff --git a/tests/unit/schema/migrations/test_auto_detector.py b/tests/unit/schema/migrations/test_auto_detector.py
@@ -532,6 +532,68 @@ def get_change_operations(self) -> Iterable[Operation]:
         )
 
 
+class MoveNodePropertyToKey(Scenario):
+    # If a property is moved to a key, then the detector should detect that the
+    # key was extended only. It should not detect that the property was dropped
+    # as it was only moved to the key.
+
+    def get_intial_state_operations(self) -> Iterable[Operation]:
+        yield CreateNodeType(
+            name=self.get_name("node_type"),
+            properties={self.get_name("property")},
+            keys={self.get_name("key")},
+        )
+
+    def get_change_operations(self) -> Iterable[Operation]:
+        yield DropNodeProperty(
+            node_type=self.get_name("node_type"),
+            property_name=self.get_name("property"),
+        )
+        yield NodeKeyExtended(
+            node_type=self.get_name("node_type"),
+            added_key_property=self.get_name("property"),
+            default=None,
+        )
+
+    def get_expected_detections(self) -> Iterable[Operation]:
+        yield NodeKeyExtended(
+            node_type=self.get_name("node_type"),
+            added_key_property=self.get_name("property"),
+            default=None,
+        )
+
+
+class MoveRelationshipPropertyToKey(Scenario):
+    # If a property is moved to a key, then the detector should detect that the
+    # key was extended only. It should not detect that the property was dropped
+    # as it was only moved to the key.
+
+    def get_intial_state_operations(self) -> Iterable[Operation]:
+        yield CreateRelationshipType(
+            name=self.get_name("relationship_type"),
+            properties={self.get_name("property")},
+            keys={self.get_name("key")},
+        )
+
+    def get_change_operations(self) -> Iterable[Operation]:
+        yield DropRelationshipProperty(
+            relationship_type=self.get_name("relationship_type"),
+            property_name=self.get_name("property"),
+        )
+        yield RelationshipKeyExtended(
+            relationship_type=self.get_name("relationship_type"),
+            added_key_property=self.get_name("property"),
+            default=None,
+        )
+
+    def get_expected_detections(self) -> Iterable[Operation]:
+        yield RelationshipKeyExtended(
+            relationship_type=self.get_name("relationship_type"),
+            added_key_property=self.get_name("property"),
+            default=None,
+        )
+
+
 ALL_PERMUTABLE_SCENARIOS = [
     AddedNodeType,
     DroppedNodeType,
@@ -553,6 +615,8 @@ def get_change_operations(self) -> Iterable[Operation]:
     ExtendedRelationshipKey,
     RenamedNodeKeyPart,
     RenamedRelationshipKeyPart,
+    MoveNodePropertyToKey,
+    MoveRelationshipPropertyToKey,
 ]
 
 

diff --git a/tests/unit/schema/migrations/test_migrator.py b/tests/unit/schema/migrations/test_migrator.py
@@ -68,7 +68,7 @@ async def test_execute_migration_fail_to_execute_operation(
     assert_that(migrator.commit_transaction.await_count, equal_to(1))
     migrator.rollback_transaction.assert_awaited_once()
     migrator.acquire_lock.assert_awaited_once()
-    migrator.release_lock.assert_not_called()
+    migrator.release_lock.assert_awaited_once()
 
 
 @pytest.mark.asyncio