diff --git a/README.md b/README.md index ef67a8a..a7cf46a 100644 --- a/README.md +++ b/README.md @@ -246,6 +246,8 @@ j1.create_relationship( # Basic relationship update j1.update_relationship( relationship_id='', + from_entity_id='', + to_entity_id='', properties={ "": "", }, @@ -254,6 +256,8 @@ j1.update_relationship( # Update relationship with complex properties j1.update_relationship( relationship_id='', + from_entity_id='', + to_entity_id='', properties={ 'accessLevel': 'write', 'lastModified': int(time.time()) * 1000, @@ -265,12 +269,25 @@ j1.update_relationship( # Update relationship with tags j1.update_relationship( relationship_id='', + from_entity_id='', + to_entity_id='', properties={ 'tag.Status': 'active', 'tag.Priority': 'high', 'tag.ReviewRequired': 'true' } ) + +# Update relationship with custom timestamp +j1.update_relationship( + relationship_id='', + from_entity_id='', + to_entity_id='', + properties={ + 'lastUpdated': int(time.time()) * 1000 + }, + timestamp=int(time.time()) * 1000 # Custom timestamp +) ``` ##### Delete a relationship diff --git a/examples/03_relationship_management.py b/examples/03_relationship_management.py index ff3c058..73037cb 100644 --- a/examples/03_relationship_management.py +++ b/examples/03_relationship_management.py @@ -113,7 +113,7 @@ def create_relationship_examples(j1, from_entity_id, to_entity_id): return basic_relationship, relationship_with_props, complex_relationship -def update_relationship_examples(j1, relationship_id): +def update_relationship_examples(j1, relationship_id, from_entity_id, to_entity_id): """Demonstrate relationship update operations.""" print("=== Relationship Update Examples ===\n") @@ -122,6 +122,8 @@ def update_relationship_examples(j1, relationship_id): print("1. Updating basic relationship properties:") basic_update = j1.update_relationship( relationship_id=relationship_id, + from_entity_id=from_entity_id, + to_entity_id=to_entity_id, properties={ 'accessLevel': 'write', 'lastModified': int(time.time()) * 1000 @@ -133,6 +135,8 @@ def update_relationship_examples(j1, relationship_id): print("2. Updating with complex properties:") j1.update_relationship( relationship_id=relationship_id, + from_entity_id=from_entity_id, + to_entity_id=to_entity_id, properties={ 'accessLevel': 'admin', 'lastModified': int(time.time()) * 1000, @@ -151,6 +155,8 @@ def update_relationship_examples(j1, relationship_id): print("3. Updating relationship tags:") j1.update_relationship( relationship_id=relationship_id, + from_entity_id=from_entity_id, + to_entity_id=to_entity_id, properties={ 'tag.Status': 'active', 'tag.Priority': 'high', @@ -159,6 +165,19 @@ def update_relationship_examples(j1, relationship_id): } ) print(f"Updated relationship tags\n") + + # 4. Update with custom timestamp + print("4. Updating with custom timestamp:") + j1.update_relationship( + relationship_id=relationship_id, + from_entity_id=from_entity_id, + to_entity_id=to_entity_id, + properties={ + 'lastUpdated': int(time.time()) * 1000 + }, + timestamp=int(time.time()) * 1000 # Custom timestamp + ) + print(f"Updated with custom timestamp\n") def delete_relationship_examples(j1, relationship_id): """Demonstrate relationship deletion.""" @@ -366,7 +385,7 @@ def main(): basic_rel, props_rel, complex_rel = create_relationship_examples(j1, from_entity_id, to_entity_id) # Update examples (using the relationship with properties) - update_relationship_examples(j1, props_rel['relationship']['_id']) + update_relationship_examples(j1, props_rel['relationship']['_id'], from_entity_id, to_entity_id) # Complete lifecycle example relationship_lifecycle_example(j1, from_entity_id, to_entity_id) diff --git a/jupiterone/client.py b/jupiterone/client.py index 076fd04..ec4fa52 100644 --- a/jupiterone/client.py +++ b/jupiterone/client.py @@ -24,7 +24,7 @@ DELETE_ENTITY, UPDATE_ENTITY, CREATE_RELATIONSHIP, - UPDATE_RELATIONSHIPV2, + UPDATE_RELATIONSHIP, DELETE_RELATIONSHIP, CURSOR_QUERY_V1, DEFERRED_RESPONSE_QUERY, @@ -552,21 +552,21 @@ def update_relationship(self, **kwargs) -> Dict: args: relationship_id (str): Unique _id of the relationship + from_entity_id (str): Unique _id of the source entity + to_entity_id (str): Unique _id of the target entity properties (dict): Dictionary of key/value relationship properties + timestamp (int, optional): Timestamp for the update (defaults to current time) """ - now_dt = datetime.now() - variables = { - "relationship": {"_id": kwargs.pop("relationship_id")}, - "timestamp": int(datetime.now().timestamp() * 1000), + "relationshipId": kwargs.pop("relationship_id"), + "fromEntityId": kwargs.pop("from_entity_id"), + "toEntityId": kwargs.pop("to_entity_id"), + "timestamp": kwargs.pop("timestamp", int(datetime.now().timestamp() * 1000)), + "properties": kwargs.pop("properties", None) } - properties = kwargs.pop("properties", None) - if properties: - variables["relationship"].update(properties) - - response = self._execute_query(query=UPDATE_RELATIONSHIPV2, variables=variables) - return response["data"]["updateRelationshipV2"] + response = self._execute_query(query=UPDATE_RELATIONSHIP, variables=variables) + return response["data"]["updateRelationship"] def delete_relationship(self, relationship_id: str = None): """Deletes a relationship between two entities. diff --git a/jupiterone/constants.py b/jupiterone/constants.py index fb78d8c..4f741d9 100644 --- a/jupiterone/constants.py +++ b/jupiterone/constants.py @@ -83,16 +83,36 @@ } } """ -UPDATE_RELATIONSHIPV2 = """ -mutation UpdateRelationshipV2 ( - $relationship: JSON! +UPDATE_RELATIONSHIP = """ +mutation UpdateRelationship( + $relationshipId: String! + $fromEntityId: String! + $toEntityId: String! $timestamp: Long + $properties: JSON ) { - updateRelationshipV2 ( - relationship: $relationship, + updateRelationship( + relationshipId: $relationshipId, + fromEntityId: $fromEntityId, + toEntityId: $toEntityId, timestamp: $timestamp, + properties: $properties ) { - relationship + relationship { + _id + _key + _type + _class + _fromEntityId + _toEntityId + displayName + } + edge { + id + fromVertexId + toVertexId + properties + } } } """ diff --git a/setup.py b/setup.py index 368833e..57b1d2d 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name="jupiterone", - version="1.7.0", + version="2.0.0", description="A Python client for the JupiterOne API", license="MIT License", author="JupiterOne", diff --git a/tests/test_update_relationship.py b/tests/test_update_relationship.py index 6ccf24d..a4e0c55 100644 --- a/tests/test_update_relationship.py +++ b/tests/test_update_relationship.py @@ -6,7 +6,7 @@ from datetime import datetime from jupiterone.client import JupiterOneClient -from jupiterone.constants import UPDATE_RELATIONSHIPV2 +from jupiterone.constants import UPDATE_RELATIONSHIP from jupiterone.errors import JupiterOneApiError @@ -22,16 +22,20 @@ def test_update_relationship_basic(self, mock_execute_query): """Test basic relationship update""" mock_response = { "data": { - "updateRelationshipV2": { + "updateRelationship": { "relationship": { "_id": "rel-123", "_type": "test_relationship", - "_class": "TestRelationship" + "_class": "TestRelationship", + "_fromEntityId": "entity-1", + "_toEntityId": "entity-2", + "displayName": "test relationship" }, "edge": { "id": "edge-123", "toVertexId": "entity-2", - "fromVertexId": "entity-1" + "fromVertexId": "entity-1", + "properties": {"status": "active", "updated": True} } } } @@ -40,65 +44,88 @@ def test_update_relationship_basic(self, mock_execute_query): result = self.client.update_relationship( relationship_id="rel-123", + from_entity_id="entity-1", + to_entity_id="entity-2", properties={"status": "active", "updated": True} ) # Verify the method was called with correct parameters mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args - assert call_args[1]['query'] == UPDATE_RELATIONSHIPV2 + assert call_args[1]['query'] == UPDATE_RELATIONSHIP variables = call_args[1]['variables'] - assert variables["relationship"]["_id"] == "rel-123" - assert variables["relationship"]["status"] == "active" - assert variables["relationship"]["updated"] is True + assert variables["relationshipId"] == "rel-123" + assert variables["fromEntityId"] == "entity-1" + assert variables["toEntityId"] == "entity-2" + assert variables["properties"]["status"] == "active" + assert variables["properties"]["updated"] is True assert "timestamp" in variables # Verify the result - assert result == mock_response["data"]["updateRelationshipV2"] + assert result == mock_response["data"]["updateRelationship"] @patch.object(JupiterOneClient, '_execute_query') def test_update_relationship_without_properties(self, mock_execute_query): """Test relationship update without properties""" mock_response = { "data": { - "updateRelationshipV2": { + "updateRelationship": { "relationship": { - "_id": "rel-123" + "_id": "rel-123", + "_fromEntityId": "entity-1", + "_toEntityId": "entity-2" }, "edge": { - "id": "edge-123" + "id": "edge-123", + "fromVertexId": "entity-1", + "toVertexId": "entity-2" } } } } mock_execute_query.return_value = mock_response - result = self.client.update_relationship(relationship_id="rel-123") + result = self.client.update_relationship( + relationship_id="rel-123", + from_entity_id="entity-1", + to_entity_id="entity-2" + ) # Verify the method was called with correct parameters mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args variables = call_args[1]['variables'] - assert variables["relationship"]["_id"] == "rel-123" - assert len(variables["relationship"]) == 1 # Only _id should be present + assert variables["relationshipId"] == "rel-123" + assert variables["fromEntityId"] == "entity-1" + assert variables["toEntityId"] == "entity-2" + assert variables["properties"] is None assert "timestamp" in variables # Verify the result - assert result == mock_response["data"]["updateRelationshipV2"] + assert result == mock_response["data"]["updateRelationship"] @patch.object(JupiterOneClient, '_execute_query') def test_update_relationship_with_complex_properties(self, mock_execute_query): """Test relationship update with complex property types""" mock_response = { "data": { - "updateRelationshipV2": { + "updateRelationship": { "relationship": { "_id": "rel-123", - "nested": {"key": "value"}, - "list": [1, 2, 3], - "boolean": True, - "number": 42 + "_fromEntityId": "entity-1", + "_toEntityId": "entity-2" + }, + "edge": { + "id": "edge-123", + "fromVertexId": "entity-1", + "toVertexId": "entity-2", + "properties": { + "nested": {"key": "value"}, + "list": [1, 2, 3], + "boolean": True, + "number": 42 + } } } } @@ -114,6 +141,8 @@ def test_update_relationship_with_complex_properties(self, mock_execute_query): result = self.client.update_relationship( relationship_id="rel-123", + from_entity_id="entity-1", + to_entity_id="entity-2", properties=properties ) @@ -121,21 +150,23 @@ def test_update_relationship_with_complex_properties(self, mock_execute_query): mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args variables = call_args[1]['variables'] - assert variables["relationship"]["_id"] == "rel-123" - assert variables["relationship"]["nested"] == {"key": "value"} - assert variables["relationship"]["list"] == [1, 2, 3] - assert variables["relationship"]["boolean"] is True - assert variables["relationship"]["number"] == 42 + assert variables["relationshipId"] == "rel-123" + assert variables["fromEntityId"] == "entity-1" + assert variables["toEntityId"] == "entity-2" + assert variables["properties"]["nested"] == {"key": "value"} + assert variables["properties"]["list"] == [1, 2, 3] + assert variables["properties"]["boolean"] is True + assert variables["properties"]["number"] == 42 # Verify the result - assert result == mock_response["data"]["updateRelationshipV2"] + assert result == mock_response["data"]["updateRelationship"] @patch.object(JupiterOneClient, '_execute_query') def test_update_relationship_timestamp_generation(self, mock_execute_query): """Test that timestamp is properly generated""" mock_response = { "data": { - "updateRelationshipV2": { + "updateRelationship": { "relationship": {"_id": "rel-123"} } } @@ -148,6 +179,8 @@ def test_update_relationship_timestamp_generation(self, mock_execute_query): self.client.update_relationship( relationship_id="rel-123", + from_entity_id="entity-1", + to_entity_id="entity-2", properties={"test": "value"} ) @@ -168,6 +201,8 @@ def test_update_relationship_api_error(self, mock_execute_query): with pytest.raises(JupiterOneApiError, match="API Error"): self.client.update_relationship( relationship_id="rel-123", + from_entity_id="entity-1", + to_entity_id="entity-2", properties={"test": "value"} ) @@ -181,6 +216,8 @@ def test_update_relationship_missing_relationship_id(self): with pytest.raises(JupiterOneApiError): self.client.update_relationship( relationship_id=None, + from_entity_id="entity-1", + to_entity_id="entity-2", properties={"test": "value"} ) @@ -189,7 +226,7 @@ def test_update_relationship_empty_properties(self, mock_execute_query): """Test relationship update with empty properties dict""" mock_response = { "data": { - "updateRelationshipV2": { + "updateRelationship": { "relationship": {"_id": "rel-123"} } } @@ -198,6 +235,8 @@ def test_update_relationship_empty_properties(self, mock_execute_query): result = self.client.update_relationship( relationship_id="rel-123", + from_entity_id="entity-1", + to_entity_id="entity-2", properties={} ) @@ -205,18 +244,20 @@ def test_update_relationship_empty_properties(self, mock_execute_query): mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args variables = call_args[1]['variables'] - assert variables["relationship"]["_id"] == "rel-123" - assert len(variables["relationship"]) == 1 # Only _id should be present + assert variables["relationshipId"] == "rel-123" + assert variables["fromEntityId"] == "entity-1" + assert variables["toEntityId"] == "entity-2" + assert variables["properties"] == {} # Verify the result - assert result == mock_response["data"]["updateRelationshipV2"] + assert result == mock_response["data"]["updateRelationship"] @patch.object(JupiterOneClient, '_execute_query') def test_update_relationship_with_none_properties(self, mock_execute_query): """Test relationship update with None properties""" mock_response = { "data": { - "updateRelationshipV2": { + "updateRelationship": { "relationship": {"_id": "rel-123"} } } @@ -225,6 +266,8 @@ def test_update_relationship_with_none_properties(self, mock_execute_query): result = self.client.update_relationship( relationship_id="rel-123", + from_entity_id="entity-1", + to_entity_id="entity-2", properties=None ) @@ -232,8 +275,10 @@ def test_update_relationship_with_none_properties(self, mock_execute_query): mock_execute_query.assert_called_once() call_args = mock_execute_query.call_args variables = call_args[1]['variables'] - assert variables["relationship"]["_id"] == "rel-123" - assert len(variables["relationship"]) == 1 # Only _id should be present + assert variables["relationshipId"] == "rel-123" + assert variables["fromEntityId"] == "entity-1" + assert variables["toEntityId"] == "entity-2" + assert variables["properties"] is None # Verify the result - assert result == mock_response["data"]["updateRelationshipV2"] \ No newline at end of file + assert result == mock_response["data"]["updateRelationship"] \ No newline at end of file