Merge pull request #66 from JupiterOne/TD-8043-2.2.0-release

fix delete_relationship() method

Author: SeaBlooms

README.md

@@ -292,13 +292,11 @@ j1.update_relationship(
 ##### Delete a relationship
 
 ```python
-# Delete by relationship ID
-j1.delete_relationship(relationship_id='<id-of-relationship-to-delete>')
-
-# Delete with timestamp
+# Delete a relationship (requires relationship ID, source entity ID, and target entity ID)
 j1.delete_relationship(
     relationship_id='<id-of-relationship-to-delete>',
-    timestamp=int(time.time()) * 1000
+    from_entity_id='<id-of-source-entity>',
+    to_entity_id='<id-of-destination-entity>'
 )
 ```
 

examples/03_relationship_management.py

@@ -175,23 +175,18 @@ def update_relationship_examples(j1, relationship_id, from_entity_id, to_entity_
     )
     print(f"Updated with custom timestamp\n")
 
-def delete_relationship_examples(j1, relationship_id):
+def delete_relationship_examples(j1, relationship_id, from_entity_id, to_entity_id):
     """Demonstrate relationship deletion."""
 
     print("=== Relationship Deletion Examples ===\n")
 
-    # 1. Basic deletion
     print("1. Deleting a relationship:")
-    delete_result = j1.delete_relationship(relationship_id=relationship_id)
-    print(f"Deleted relationship: {delete_result['relationship']['_id']}\n")
-    
-    # 2. Deletion with timestamp
-    print("2. Deleting with specific timestamp:")
-    j1.delete_relationship(
+    delete_result = j1.delete_relationship(
         relationship_id=relationship_id,
-        timestamp=int(time.time()) * 1000
+        from_entity_id=from_entity_id,
+        to_entity_id=to_entity_id
     )
-    print(f"Deleted with timestamp\n")
+    print(f"Deleted relationship: {delete_result['relationship']['_id']}\n")
 
 def relationship_lifecycle_example(j1, from_entity_id, to_entity_id):
     """Demonstrate complete relationship lifecycle."""
@@ -234,7 +229,11 @@ def relationship_lifecycle_example(j1, from_entity_id, to_entity_id):
 
     # 4. Delete relationship
     print("4. Deleting relationship:")
-    j1.delete_relationship(relationship_id=relationship_id)
+    j1.delete_relationship(
+        relationship_id=relationship_id,
+        from_entity_id=from_entity_id,
+        to_entity_id=to_entity_id
+    )
     print("Deleted successfully")
 
     # 5. Verify deletion
@@ -281,14 +280,22 @@ def network_relationship_example(j1):
                 'bandwidth': '100Mbps'
             }
         )
-        relationships.append(relationship['relationship']['_id'])
+        relationships.append({
+            'id': relationship['relationship']['_id'],
+            'from': entities[i],
+            'to': entities[i+1]
+        })
         print(f"Created connection {i}: {relationship['relationship']['_id']}")
 
     print(f"Created {len(entities)} nodes with {len(relationships)} connections")
 
     # Clean up
-    for relationship_id in relationships:
-        j1.delete_relationship(relationship_id=relationship_id)
+    for rel in relationships:
+        j1.delete_relationship(
+            relationship_id=rel['id'],
+            from_entity_id=rel['from'],
+            to_entity_id=rel['to']
+        )
     for entity_id in entities:
         j1.delete_entity(entity_id=entity_id)
 
@@ -356,7 +363,11 @@ def access_control_relationship_example(j1):
     print("Updated access level to write")
 
     # Clean up
-    j1.delete_relationship(relationship_id=access_relationship['relationship']['_id'])
+    j1.delete_relationship(
+        relationship_id=access_relationship['relationship']['_id'],
+        from_entity_id=user_entity['entity']['_id'],
+        to_entity_id=resource_entity['entity']['_id']
+    )
     j1.delete_entity(entity_id=user_entity['entity']['_id'])
     j1.delete_entity(entity_id=resource_entity['entity']['_id'])
 
@@ -397,7 +408,11 @@ def main():
         relationships_to_clean = [basic_rel, props_rel, complex_rel]
         for rel in relationships_to_clean:
             try:
-                j1.delete_relationship(relationship_id=rel['relationship']['_id'])
+                j1.delete_relationship(
+                    relationship_id=rel['relationship']['_id'],
+                    from_entity_id=from_entity_id,
+                    to_entity_id=to_entity_id
+                )
                 print(f"Cleaned up relationship: {rel['relationship']['_id']}")
             except Exception:
                 # Relationship may already be deleted or not exist

examples/06_advanced_operations.py

@@ -112,7 +112,11 @@ def bulk_operations_examples(j1):
         for rel_data in relationships_to_create:
             try:
                 relationship = j1.create_relationship(**rel_data)
-                created_relationships.append(relationship['relationship']['_id'])
+                created_relationships.append({
+                    'id': relationship['relationship']['_id'],
+                    'from': rel_data['from_entity_id'],
+                    'to': rel_data['to_entity_id']
+                })
                 print(f"Created relationship: {relationship['relationship']['_id']}")
             except Exception as e:
                 print(f"Error creating relationship: {e}")
@@ -137,29 +141,35 @@ def bulk_operations_examples(j1):
 
         # 4. Bulk relationship updates
         print("4. Bulk relationship updates:")
-        for rel_id in created_relationships:
+        for rel in created_relationships:
             try:
                 j1.update_relationship(
-                    relationship_id=rel_id,
+                    relationship_id=rel['id'],
+                    from_entity_id=rel['from'],
+                    to_entity_id=rel['to'],
                     properties={
                         "lastUpdated": int(time.time()) * 1000,
                         "tag.BulkUpdated": "true"
                     }
                 )
-                print(f"Updated relationship: {rel_id}")
+                print(f"Updated relationship: {rel['id']}")
             except Exception as e:
-                print(f"Error updating relationship {rel_id}: {e}")
+                print(f"Error updating relationship {rel['id']}: {e}")
         print()
 
         # 5. Bulk deletion
         print("5. Bulk deletion:")
         # Delete relationships first
-        for rel_id in created_relationships:
+        for rel in created_relationships:
             try:
-                j1.delete_relationship(relationship_id=rel_id)
-                print(f"Deleted relationship: {rel_id}")
+                j1.delete_relationship(
+                    relationship_id=rel['id'],
+                    from_entity_id=rel['from'],
+                    to_entity_id=rel['to']
+                )
+                print(f"Deleted relationship: {rel['id']}")
             except Exception as e:
-                print(f"Error deleting relationship {rel_id}: {e}")
+                print(f"Error deleting relationship {rel['id']}: {e}")
 
         # Then delete entities
         for entity_id in created_entities:

examples/examples.py

@@ -85,7 +85,11 @@
 print(create_relationship_r)
 
 # delete_relationship
-delete_relationship_r = j1.delete_relationship(relationship_id=create_relationship_r['relationship']['_id'])
+delete_relationship_r = j1.delete_relationship(
+    relationship_id=create_relationship_r['relationship']['_id'],
+    from_entity_id=create_r['entity']['_id'],
+    to_entity_id=create_r_2['entity']['_id']
+)
 print("delete_relationship()")
 print(delete_relationship_r)
 

jupiterone/client.py

@@ -764,13 +764,37 @@ def update_relationship(self, **kwargs: Any) -> Dict[str, Any]:
         response = self._execute_query(query=UPDATE_RELATIONSHIP, variables=variables)
         return response["data"]["updateRelationship"]
 
-    def delete_relationship(self, relationship_id: Optional[str] = None) -> Dict[str, Any]:
+    def delete_relationship(
+        self,
+        relationship_id: Optional[str] = None,
+        from_entity_id: Optional[str] = None,
+        to_entity_id: Optional[str] = None,
+    ) -> Dict[str, Any]:
         """Deletes a relationship between two entities.
 
         args:
-            relationship_id (str): The ID of the relationship
+            relationship_id (str): The _id of the relationship to delete
+            from_entity_id (str): The _id of the source entity
+            to_entity_id (str): The _id of the target entity
         """
-        variables = {"relationshipId": relationship_id}
+        if not relationship_id:
+            raise JupiterOneClientError("relationship_id is required")
+        if not isinstance(relationship_id, str) or not relationship_id.strip():
+            raise JupiterOneClientError("relationship_id must be a non-empty string")
+
+        if not from_entity_id:
+            raise JupiterOneClientError("from_entity_id is required")
+        self._validate_entity_id(from_entity_id, "from_entity_id")
+
+        if not to_entity_id:
+            raise JupiterOneClientError("to_entity_id is required")
+        self._validate_entity_id(to_entity_id, "to_entity_id")
+
+        variables: Dict[str, Any] = {
+            "relationshipId": relationship_id,
+            "fromEntityId": from_entity_id,
+            "toEntityId": to_entity_id,
+        }
 
         response = self._execute_query(DELETE_RELATIONSHIP, variables=variables)
         return response["data"]["deleteRelationship"]

jupiterone/constants.py

@@ -117,8 +117,16 @@
 }
 """
 DELETE_RELATIONSHIP = """
-  mutation DeleteRelationship($relationshipId: String! $timestamp: Long) {
-    deleteRelationship (relationshipId: $relationshipId, timestamp: $timestamp) {
+  mutation DeleteRelationship(
+    $relationshipId: String!
+    $fromEntityId: String!
+    $toEntityId: String!
+  ) {
+    deleteRelationship(
+      relationshipId: $relationshipId
+      fromEntityId: $fromEntityId
+      toEntityId: $toEntityId
+    ) {
       relationship {
         _id
       }

setup.py

@@ -5,7 +5,7 @@
 
 setup(
     name="jupiterone",
-    version="2.1.0",
+    version="2.2.0",
     description="A Python client for the JupiterOne API",
     license="MIT License",
     author="JupiterOne",

tests/test_delete_relationship.py

@@ -3,49 +3,71 @@
 import responses
 
 from jupiterone.client import JupiterOneClient
-from jupiterone.constants import CREATE_ENTITY
+from jupiterone.errors import JupiterOneClientError
 
 
-@responses.activate
-def test_tree_query_v1():
+MOCK_RESPONSE = {
+    'data': {
+        'deleteRelationship': {
+            'relationship': {
+                '_id': '1'
+            },
+            'edge': {
+                'id': '1',
+                'toVertexId': '1',
+                'fromVertexId': '2',
+                'relationship': {
+                    '_id': '1'
+                },
+                'properties': {}
+            }
+        }
+    }
+}
+
 
+def _make_callback(captured_requests):
+    """Return a callback that records request bodies and returns a success response."""
     def request_callback(request):
-        headers = {
-            'Content-Type': 'application/json'
-        }
+        captured_requests.append(json.loads(request.body))
+        return (200, {'Content-Type': 'application/json'}, json.dumps(MOCK_RESPONSE))
+    return request_callback
 
-        response = {
-            'data': {
-                'deleteRelationship': {
-                    'relationship': {
-                        '_id': '1'
-                    }, 
-                    'edge': {
-                        'id': '1',
-                        'toVertexId': '1',
-                        'fromVertexId': '2',
-                        'relationship': {
-                            '_id': '1'
-                        },
-                        'properties': {}
-                    }
-                }
-            }
-        }
 
-        return (200, headers, json.dumps(response))
-    
+@responses.activate
+def test_delete_relationship_sends_required_variables():
+    captured = []
     responses.add_callback(
         responses.POST, 'https://graphql.us.jupiterone.io',
-        callback=request_callback,
+        callback=_make_callback(captured),
         content_type='application/json',
     )
 
     j1 = JupiterOneClient(account='testAccount', token='testToken1234567890')
-    response = j1.delete_relationship('1')
+    response = j1.delete_relationship(
+        relationship_id='1',
+        from_entity_id='2222222222',
+        to_entity_id='3333333333'
+    )
+
+    assert len(captured) == 1
+    variables = captured[0]['variables']
+    assert variables['relationshipId'] == '1'
+    assert variables['fromEntityId'] == '2222222222'
+    assert variables['toEntityId'] == '3333333333'
 
-    assert type(response) == dict
-    assert type(response['relationship']) == dict
     assert response['relationship']['_id'] == '1'
     assert response['edge']['toVertexId'] == '1'
     assert response['edge']['fromVertexId'] == '2'
+
+
+def test_delete_relationship_raises_without_from_entity_id():
+    j1 = JupiterOneClient(account='testAccount', token='testToken1234567890')
+    with pytest.raises(JupiterOneClientError, match="from_entity_id is required"):
+        j1.delete_relationship(relationship_id='1', to_entity_id='3333333333')
+
+
+def test_delete_relationship_raises_without_to_entity_id():
+    j1 = JupiterOneClient(account='testAccount', token='testToken1234567890')
+    with pytest.raises(JupiterOneClientError, match="to_entity_id is required"):
+        j1.delete_relationship(relationship_id='1', from_entity_id='2222222222')