From 794c90046fa8eaeb161f0ef878d35408931caffa Mon Sep 17 00:00:00 2001 From: "xunalei.lin" Date: Tue, 16 Sep 2025 08:54:39 +0800 Subject: [PATCH 1/3] feat: enhance PrimaryIdAsAttribute validation for schema retrieval from DB --- tigergraphx/core/managers/schema_manager.py | 22 +++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/tigergraphx/core/managers/schema_manager.py b/tigergraphx/core/managers/schema_manager.py index d17618a..ec8fc54 100644 --- a/tigergraphx/core/managers/schema_manager.py +++ b/tigergraphx/core/managers/schema_manager.py @@ -63,13 +63,13 @@ def create_schema(self, drop_existing_graph=False) -> bool: # Add vector attributes gsql_add_vector_attr = self._create_gsql_add_vector_attr() if gsql_add_vector_attr: - logger.info(f"Adding vector attribute(s) for graph: {self._graph_name}...") + logger.info( + f"Adding vector attribute(s) for graph: {self._graph_name}..." + ) result = self._tigergraph_api.gsql(gsql_add_vector_attr) logger.debug(f"GSQL response: {result}") if f"Using graph '{self._graph_name}'" not in result: - error_msg = ( - f"Failed to use graph '{self._graph_name}'. GSQL response: {result}" - ) + error_msg = f"Failed to use graph '{self._graph_name}'. GSQL response: {result}" logger.error(error_msg) raise RuntimeError(error_msg) if "Successfully created schema change jobs" not in result: @@ -96,7 +96,9 @@ def create_schema(self, drop_existing_graph=False) -> bool: return True - logger.debug(f"Graph '{self._graph_name}' already exists. Skipping graph creation.") + logger.debug( + f"Graph '{self._graph_name}' already exists. Skipping graph creation." + ) return False def drop_graph(self) -> None: @@ -366,10 +368,14 @@ def get_schema_from_db( # Construct nodes dictionary nodes = {} for vertex in raw_schema.get("VertexTypes", []): - primary_id = vertex["PrimaryId"] - if not primary_id["PrimaryIdAsAttribute"]: + primary_id = vertex.get("PrimaryId", {}) + primary_id_as_attr = primary_id.get("PrimaryIdAsAttribute") + + if primary_id_as_attr is not True: raise ValueError( - f"PrimaryIdAsAttribute must be set to True for node type {vertex['Name']}." + f"The node type '{vertex.get('Name', '')}' has PrimaryIdAsAttribute unset. " + f"TigerGraphX requires the primary ID to be used as an attribute for certain methods. " + f"Please set PrimaryIdAsAttribute to True for this node type." ) # Extract regular attributes From 07a8a58378db80e7bae1852a8c0f7662f358dce6 Mon Sep 17 00:00:00 2001 From: "xunalei.lin" Date: Tue, 16 Sep 2025 11:03:11 +0800 Subject: [PATCH 2/3] feat: add support for defining and handling reverse edges --- tests/integration/core/graph_test.py | 21 ++++++------------- .../config/graph_db/schema/edge_schema.py | 5 ++++- .../config/graph_db/schema/graph_schema.py | 13 ++++++++++++ tigergraphx/core/graph.py | 4 ++-- tigergraphx/core/managers/schema_manager.py | 14 ++++++++++++- 5 files changed, 38 insertions(+), 19 deletions(-) diff --git a/tests/integration/core/graph_test.py b/tests/integration/core/graph_test.py index 8abe351..4dd8175 100644 --- a/tests/integration/core/graph_test.py +++ b/tests/integration/core/graph_test.py @@ -256,12 +256,8 @@ def test_add_edge_without_edge_type(self): def test_has_edges(self): # Test edge existence - assert not self.G.has_edge( - "User_A", 2, "User", "purchased", "Product" - ) - assert not self.G.has_edge( - "User_A", 3, "User", "purchased", "Product" - ) + assert not self.G.has_edge("User_A", 2, "User", "purchased", "Product") + assert not self.G.has_edge("User_A", 3, "User", "purchased", "Product") assert self.G.has_edge("User_C", 1, "User", "purchased", "Product") assert self.G.has_edge("User_C", 2, "User", "purchased", "Product") assert self.G.has_edge("User_C", 3, "User", "purchased", "Product") @@ -269,9 +265,7 @@ def test_has_edges(self): def test_get_edge_data(self): # Test fetching edge data edge_data = self.time_execution( - lambda: self.G.get_edge_data( - "User_C", 2, "User", "purchased", "Product" - ), + lambda: self.G.get_edge_data("User_C", 2, "User", "purchased", "Product"), "get_edge_data", ) assert edge_data["purchase_date"] == "2024-01-12 00:00:00" @@ -292,16 +286,12 @@ def test_degree(self): assert degree == 2 degree = self.time_execution( - lambda: self.G.degree( - 1, "Product", ["reverse_purchased", "similar_to"] - ), + lambda: self.G.degree(1, "Product", ["reverse_purchased", "similar_to"]), "degree", ) assert degree == 4 - degree = self.time_execution( - lambda: self.G.degree(1, "Product", []), "degree" - ) + degree = self.time_execution(lambda: self.G.degree(1, "Product", []), "degree") assert degree == 4 degree = self.time_execution( @@ -581,6 +571,7 @@ def test_get_schema(self): "edges": { "relationship": { "is_directed_edge": True, + "reverse_edge_name": "reverse_relationship", "from_node_type": "Entity", "to_node_type": "Entity", "attributes": { diff --git a/tigergraphx/config/graph_db/schema/edge_schema.py b/tigergraphx/config/graph_db/schema/edge_schema.py index f690d36..6dab7e2 100644 --- a/tigergraphx/config/graph_db/schema/edge_schema.py +++ b/tigergraphx/config/graph_db/schema/edge_schema.py @@ -5,7 +5,7 @@ # Permission is granted to use, copy, modify, and distribute this software # under the License. The software is provided "AS IS", without warranty. -from typing import Any, Dict, Set +from typing import Any, Dict, Optional, Set from pydantic import Field, model_validator from .attribute_schema import AttributeSchema, AttributesType, create_attribute_schema @@ -22,6 +22,9 @@ class EdgeSchema(BaseConfig): is_directed_edge: bool = Field( default=False, description="Whether the edge is directed." ) + reverse_edge_name: Optional[str] = Field( + default=None, description="The name of the reverse edge." + ) from_node_type: str = Field(description="The type of the source node.") to_node_type: str = Field(description="The type of the target node.") discriminator: Set[str] | str = Field( diff --git a/tigergraphx/config/graph_db/schema/graph_schema.py b/tigergraphx/config/graph_db/schema/graph_schema.py index e0e0fad..774a0cb 100644 --- a/tigergraphx/config/graph_db/schema/graph_schema.py +++ b/tigergraphx/config/graph_db/schema/graph_schema.py @@ -83,3 +83,16 @@ def validate_reserved_keywords(self) -> "GraphSchema": ) return self + + @model_validator(mode="after") + def set_reverse_edge_names(self) -> "GraphSchema": + """ + For directed edges without a reverse_edge_name, set it to 'reverse_[edge_name]'. + + Returns: + The updated graph schema. + """ + for edge_name, edge in self.edges.items(): + if edge.is_directed_edge and not edge.reverse_edge_name: + edge.reverse_edge_name = f"reverse_{edge_name}" + return self diff --git a/tigergraphx/core/graph.py b/tigergraphx/core/graph.py index 5239827..9ef3e03 100644 --- a/tigergraphx/core/graph.py +++ b/tigergraphx/core/graph.py @@ -79,8 +79,8 @@ def __init__( self.edge_types: Set[str] = set() for edge_name, edge in self._context.graph_schema.edges.items(): self.edge_types.add(edge_name) - if edge.is_directed_edge: - self.edge_types.add(f"reverse_{edge_name}") + if edge.is_directed_edge and edge.reverse_edge_name: + self.edge_types.add(edge.reverse_edge_name) logger.debug(f"self.name: {self.name}") logger.debug(f"self.node_types: {self.node_types}") logger.debug(f"self.edge_types: {self.edge_types}") diff --git a/tigergraphx/core/managers/schema_manager.py b/tigergraphx/core/managers/schema_manager.py index ec8fc54..ccc29d5 100644 --- a/tigergraphx/core/managers/schema_manager.py +++ b/tigergraphx/core/managers/schema_manager.py @@ -190,7 +190,7 @@ def _create_gsql_graph_schema(self) -> str: # Construct the edge definition, with conditional attribute string and direction edge_type_str = "DIRECTED" if edge_schema.is_directed_edge else "UNDIRECTED" reverse_edge_clause = ( - f' WITH REVERSE_EDGE="reverse_{edge_name}"' + f' WITH REVERSE_EDGE="{edge_schema.reverse_edge_name}"' if edge_schema.is_directed_edge else "" ) @@ -430,8 +430,20 @@ def get_schema_from_db( for attr in edge.get("Attributes", []) if attr.get("IsDiscriminator") } + + # Determine reverse edge name + reverse_edge_name = edge.get("Config", {}).get("REVERSE_EDGE") + + # Precheck: reverse_edge_name must not be None for directed edges + if edge.get("IsDirected") and not reverse_edge_name: + raise ValueError( + f"Directed edge '{edge['Name']}' must have a reverse edge name defined " + f"in 'Config.REVERSE_EDGE'." + ) + edges[edge["Name"]] = { "is_directed_edge": edge["IsDirected"], + "reverse_edge_name": reverse_edge_name, "from_node_type": edge["FromVertexTypeName"], "to_node_type": edge["ToVertexTypeName"], "discriminator": discriminator, From 1bee1a4f9d963cda98b4fd3e0cc3c4d7e86ca1c6 Mon Sep 17 00:00:00 2001 From: "xunalei.lin" Date: Mon, 29 Sep 2025 08:30:08 +0800 Subject: [PATCH 3/3] feat: add schema change support for nodes, edges, and attributes --- tests/integration/core/graph/__init__.py | 0 .../core/{ => graph}/base_graph_fixture.py | 2 +- .../core/graph/graph_schema_change_test.py | 426 ++++++++++++++++++ .../core/{ => graph}/graph_test.py | 2 +- .../core/{ => graph}/tigervector_test.py | 0 .../core/tigergraph_api/database_api_test.py | 171 +++++++ ...gergraph_api_test.py => graph_api_test.py} | 166 ------- .../tigergraph_api/graph_schema_api_test.py | 245 ++++++++++ .../schema/schema_change_builder_test.py | 98 ++++ tigergraphx/config/endpoint_definitions.yaml | 25 + .../config/graph_db/schema/edge_schema.py | 11 + .../config/graph_db/schema/graph_schema.py | 3 +- tigergraphx/core/__init__.py | 2 + tigergraphx/core/graph.py | 225 +++++++++ tigergraphx/core/managers/__init__.py | 2 + .../core/managers/schema_change/__init__.py | 12 + .../schema_change/schema_change_builder.py | 203 +++++++++ tigergraphx/core/managers/schema_manager.py | 275 ++++++++++- .../core/tigergraph_api/api/schema_api.py | 69 ++- .../core/tigergraph_api/tigergraph_api.py | 72 ++- 20 files changed, 1822 insertions(+), 187 deletions(-) create mode 100644 tests/integration/core/graph/__init__.py rename tests/integration/core/{ => graph}/base_graph_fixture.py (89%) create mode 100644 tests/integration/core/graph/graph_schema_change_test.py rename tests/integration/core/{ => graph}/graph_test.py (99%) rename tests/integration/core/{ => graph}/tigervector_test.py (100%) create mode 100644 tests/integration/core/tigergraph_api/database_api_test.py rename tests/integration/core/tigergraph_api/{tigergraph_api_test.py => graph_api_test.py} (50%) create mode 100644 tests/integration/core/tigergraph_api/graph_schema_api_test.py create mode 100644 tests/unit/core/managers/schema/schema_change_builder_test.py create mode 100644 tigergraphx/core/managers/schema_change/__init__.py create mode 100644 tigergraphx/core/managers/schema_change/schema_change_builder.py diff --git a/tests/integration/core/graph/__init__.py b/tests/integration/core/graph/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/integration/core/base_graph_fixture.py b/tests/integration/core/graph/base_graph_fixture.py similarity index 89% rename from tests/integration/core/base_graph_fixture.py rename to tests/integration/core/graph/base_graph_fixture.py index a7966e9..a2153dd 100644 --- a/tests/integration/core/base_graph_fixture.py +++ b/tests/integration/core/graph/base_graph_fixture.py @@ -11,7 +11,7 @@ class BaseGraphFixture: @pytest.fixture(scope="class", autouse=True) def load_connection_config(self, request): """Load connection config from YAML and attach to class.""" - config_path = Path(__file__).parent / "config" / "tigergraph_connection.yaml" + config_path = Path(__file__).parent.parent / "config" / "tigergraph_connection.yaml" with open(config_path, "r") as f: config = yaml.safe_load(f) request.cls.tigergraph_connection_config = config diff --git a/tests/integration/core/graph/graph_schema_change_test.py b/tests/integration/core/graph/graph_schema_change_test.py new file mode 100644 index 0000000..0937f46 --- /dev/null +++ b/tests/integration/core/graph/graph_schema_change_test.py @@ -0,0 +1,426 @@ +import pytest +from typing import Any, Dict, List, Optional + +from .base_graph_fixture import BaseGraphFixture + +from tigergraphx.core import Graph + + +class TestGraphSchemaChange(BaseGraphFixture): + def setup_graph(self): + """Set up the graph.""" + graph_schema = { + "graph_name": "SchemaChangeGraph", + "nodes": { + "Employee": { + "primary_key": "id", + "attributes": { + "id": "STRING", + }, + } + }, + "edges": {}, + } + self.G = Graph( + graph_schema=graph_schema, + tigergraph_connection_config=self.tigergraph_connection_config, + ) + + @pytest.fixture(autouse=True) + def setup_graph_and_run_test_case(self): + """Add nodes and edges before each test case.""" + self.setup_graph() + yield + + @pytest.fixture(scope="class", autouse=True) + def drop_graph(self): + """Drop the graph after all tests are done in the session.""" + yield + self.setup_graph() + self.G.drop_graph() + + def test_apply_schema_changes_no_change(self): + """Test applying schema changes when no changes are provided.""" + # Should execute successfully and return True even if nothing to do + result = self.G.apply_schema_changes() + assert result is False + + # Schema should remain unchanged + self._assert_schema_contains(nodes=["Employee"], should_exist=True) + + def test_apply_schema_changes_add_and_drop_nodes_and_edges(self): + """Test apply_schema_changes with adding and dropping nodes and edges, + including edge cases like existing node/edge names and empty changes. + The schema before and after the test remains the same. + """ + nodes_to_add: Dict[str, Any] = { + "Department": { + "primary_key": "id", + "attributes": {"id": "STRING", "name": "STRING"}, + }, + } + edges_to_add: Dict[str, Any] = { + "mentors": { + "is_directed_edge": True, + "from_node_type": "Employee", + "to_node_type": "Employee", + }, + } + + # Use try/finally to ensure cleanup + try: + # Add nodes and edges + result = self.G.apply_schema_changes( + add_nodes=nodes_to_add, add_edges=edges_to_add + ) + assert result is True + self._assert_schema_contains( + nodes=list(nodes_to_add.keys()), + edges=list(edges_to_add.keys()), + should_exist=True, + ) + + # Attempt to add nodes/edges that already exist (should raise ValueError) + with pytest.raises(ValueError): + self.G.apply_schema_changes(add_nodes=nodes_to_add) + with pytest.raises(ValueError): + self.G.apply_schema_changes(add_edges=edges_to_add) + finally: + # Drop the nodes and edges to restore schema + result = self.G.apply_schema_changes( + drop_nodes=list(nodes_to_add), drop_edges=list(edges_to_add) + ) + assert result is True + self._assert_schema_contains( + nodes=list(nodes_to_add.keys()), + edges=list(edges_to_add.keys()), + should_exist=False, + ) + + def test_apply_schema_changes_drop_nonexistent_node(self): + """Test dropping a node that does not exist raises ValueError.""" + with pytest.raises( + ValueError, match="Node type 'NonexistentNode' does not exist" + ): + self.G.apply_schema_changes(drop_nodes=["NonexistentNode"]) + + def test_apply_schema_changes_drop_nonexistent_edge(self): + """Test dropping an edge that does not exist raises ValueError.""" + with pytest.raises( + ValueError, match="Edge type 'NonexistentEdge' does not exist" + ): + self.G.apply_schema_changes(drop_edges=["NonexistentEdge"]) + + def test_apply_schema_changes_add_and_drop_node_attributes(self): + """Test adding and dropping node attributes.""" + node_attrs_to_add: Dict[str, Any] = { + "Employee": { + "name": {"data_type": "STRING"}, + "age": {"data_type": "INT"}, + } + } + + # Use try/finally to ensure cleanup + try: + # Add attributes to Employee + result = self.G.apply_schema_changes(add_node_attributes=node_attrs_to_add) + assert result is True + self._assert_schema_contains( + node_attributes={"Employee": ["name", "age"]}, should_exist=True + ) + + # Attempt to add attribute that already exists (should raise ValueError) + with pytest.raises(ValueError): + self.G.apply_schema_changes(add_node_attributes=node_attrs_to_add) + finally: + # Drop attributes to restore schema + node_attrs_to_drop = {"Employee": ["name", "age"]} + result = self.G.apply_schema_changes( + drop_node_attributes=node_attrs_to_drop + ) + assert result is True + self._assert_schema_contains( + node_attributes={"Employee": ["name", "age"]}, should_exist=False + ) + + def test_apply_schema_changes_add_and_drop_edge_attributes(self): + """Test adding and dropping edge attributes.""" + edge_to_add: Dict[str, Any] = { + "mentors": { + "is_directed_edge": True, + "from_node_type": "Employee", + "to_node_type": "Employee", + } + } + + # Use try/finally to ensure cleanup + try: + # Ensure the edge exists first + self.G.apply_schema_changes(add_edges=edge_to_add) + + # Add attributes to edge + edge_attrs_to_add: Dict[str, Any] = { + "mentors": { + "since": {"data_type": "DATETIME"}, + "level": {"data_type": "STRING"}, + } + } + result = self.G.apply_schema_changes(add_edge_attributes=edge_attrs_to_add) + assert result is True + self._assert_schema_contains( + edge_attributes={"mentors": ["since", "level"]}, should_exist=True + ) + + # Attempt to add attribute that already exists (should raise ValueError) + with pytest.raises(ValueError): + self.G.apply_schema_changes(add_edge_attributes=edge_attrs_to_add) + finally: + # Drop attributes and edge to restore schema + edge_attrs_to_drop = {"mentors": ["since", "level"]} + self.G.apply_schema_changes(drop_edge_attributes=edge_attrs_to_drop) + self._assert_schema_contains( + edge_attributes={"mentors": ["since", "level"]}, should_exist=False + ) + self.G.apply_schema_changes(drop_edges=list(edge_to_add.keys())) + + def test_add_and_drop_node_type(self): + """Test adding and dropping a single node type.""" + node_schema = { + "primary_key": "id", + "attributes": {"id": "STRING", "name": "STRING"}, + } + node_name = "Department" + + # Use try/finally to ensure cleanup + try: + result = self.G.add_node_type(node_name, node_schema) + assert result is True + self._assert_schema_contains( + nodes=[node_name], + should_exist=True, + ) + finally: + result = self.G.drop_node_type(node_name) + assert result is True + self._assert_schema_contains( + nodes=[node_name], + should_exist=False, + ) + + def test_add_and_drop_node_types(self): + """Test adding and dropping multiple node types.""" + nodes: Dict[str, Any] = { + "Company": { + "primary_key": "id", + "attributes": {"id": "STRING", "name": "STRING"}, + }, + "Project": { + "primary_key": "id", + "attributes": {"id": "STRING", "name": "STRING"}, + }, + } + + # Use try/finally to ensure cleanup + try: + result = self.G.add_node_types(nodes) + assert result is True + self._assert_schema_contains( + nodes=list(nodes.keys()), + should_exist=True, + ) + finally: + result = self.G.drop_node_types(list(nodes.keys())) + assert result is True + self._assert_schema_contains( + nodes=list(nodes.keys()), + should_exist=False, + ) + + def test_add_and_drop_edge_type(self): + """Test adding and dropping a single edge type using only Employee nodes.""" + edge_schema = { + "is_directed_edge": True, + "from_node_type": "Employee", + "to_node_type": "Employee", + "attributes": {"since": {"data_type": "DATETIME"}}, + } + edge_name = "mentors" + + # Use try/finally to ensure cleanup + try: + result = self.G.add_edge_type(edge_name, edge_schema) + assert result is True + self._assert_schema_contains( + edges=[edge_name], + should_exist=True, + ) + finally: + result = self.G.drop_edge_type(edge_name) + assert result is True + self._assert_schema_contains( + edges=[edge_name], + should_exist=False, + ) + + def test_add_and_drop_edge_types(self): + """Test adding and dropping multiple edge types using only Employee nodes.""" + edges: Dict[str, Any] = { + "collaborates_with": { + "is_directed_edge": True, + "from_node_type": "Employee", + "to_node_type": "Employee", + }, + "reports_to": { + "is_directed_edge": True, + "from_node_type": "Employee", + "to_node_type": "Employee", + }, + } + + # Use try/finally to ensure cleanup + try: + result = self.G.add_edge_types(edges) + assert result is True + self._assert_schema_contains( + edges=list(edges.keys()), + should_exist=True, + ) + finally: + result = self.G.drop_edge_types(list(edges.keys())) + assert result is True + self._assert_schema_contains( + edges=list(edges.keys()), + should_exist=False, + ) + + def test_add_and_drop_node_attributes(self): + """Test adding and dropping multiple attributes on a node type.""" + node_attrs: Dict[str, Dict[str, Any]] = { + "Employee": { + "name": {"data_type": "STRING"}, + "age": {"data_type": "INT"}, + } + } + + # Use try/finally to ensure cleanup + try: + result = self.G.add_node_attributes(node_attrs) + assert result is True + self._assert_schema_contains( + node_attributes={"Employee": ["name", "age"]}, + should_exist=True, + ) + finally: + result = self.G.drop_node_attributes({"Employee": ["name", "age"]}) + assert result is True + self._assert_schema_contains( + node_attributes={"Employee": ["name", "age"]}, + should_exist=False, + ) + + def test_add_and_drop_edge_attributes(self): + """Test adding and dropping multiple attributes on an edge type.""" + # Ensure the edge exists first + self.G.add_edge_type( + "mentors", + { + "is_directed_edge": True, + "from_node_type": "Employee", + "to_node_type": "Employee", + }, + ) + + edge_attrs: Dict[str, Dict[str, Any]] = { + "mentors": { + "since": {"data_type": "DATETIME"}, + "level": {"data_type": "STRING"}, + } + } + + # Use try/finally to ensure cleanup + try: + result = self.G.add_edge_attributes(edge_attrs) + assert result is True + self._assert_schema_contains( + edge_attributes={"mentors": ["since", "level"]}, + should_exist=True, + ) + finally: + result = self.G.drop_edge_attributes({"mentors": ["since", "level"]}) + assert result is True + self._assert_schema_contains( + edge_attributes={"mentors": ["since", "level"]}, + should_exist=False, + ) + # Clean up edge itself + self.G.drop_edge_type("mentors") + + def _assert_schema_contains( + self, + nodes: Optional[List[str]] = None, + edges: Optional[List[str]] = None, + node_attributes: Optional[Dict[str, List[str]]] = None, + edge_attributes: Optional[Dict[str, List[str]]] = None, + should_exist: bool = True, + ): + """Private helper to assert that nodes/edges and their attributes exist + (or not) in both current and fresh graph schema, and schemas are identical. + """ + nodes = nodes or [] + edges = edges or [] + node_attributes = node_attributes or {} + edge_attributes = edge_attributes or {} + + schema = self.G.get_schema(format="dict") + assert isinstance(schema, dict) + G_fresh = Graph.from_db( + graph_name=self.G.name, + tigergraph_connection_config=self.tigergraph_connection_config, + ) + fresh_schema = G_fresh.get_schema(format="dict") + assert isinstance(fresh_schema, dict) + + # Assert that schema and fresh_schema are identical + assert schema == fresh_schema + + # Check node existence + for node in nodes: + if should_exist: + assert node in schema.get("nodes", {}) + assert node in fresh_schema.get("nodes", {}) + else: + assert node not in schema.get("nodes", {}) + assert node not in fresh_schema.get("nodes", {}) + + # Check edge existence + for edge in edges: + if should_exist: + assert edge in schema.get("edges", {}) + assert edge in fresh_schema.get("edges", {}) + else: + assert edge not in schema.get("edges", {}) + assert edge not in fresh_schema.get("edges", {}) + + # Check node attributes + for node, attrs in node_attributes.items(): + node_dict = schema.get("nodes", {}).get(node, {}) + fresh_node_dict = fresh_schema.get("nodes", {}).get(node, {}) + for attr in attrs: + if should_exist: + assert attr in node_dict.get("attributes", {}) + assert attr in fresh_node_dict.get("attributes", {}) + else: + assert attr not in node_dict.get("attributes", {}) + assert attr not in fresh_node_dict.get("attributes", {}) + + # Check edge attributes + for edge, attrs in edge_attributes.items(): + edge_dict = schema.get("edges", {}).get(edge, {}) + fresh_edge_dict = fresh_schema.get("edges", {}).get(edge, {}) + for attr in attrs: + if should_exist: + assert attr in edge_dict.get("attributes", {}) + assert attr in fresh_edge_dict.get("attributes", {}) + else: + assert attr not in edge_dict.get("attributes", {}) + assert attr not in fresh_edge_dict.get("attributes", {}) diff --git a/tests/integration/core/graph_test.py b/tests/integration/core/graph/graph_test.py similarity index 99% rename from tests/integration/core/graph_test.py rename to tests/integration/core/graph/graph_test.py index 4dd8175..32e3657 100644 --- a/tests/integration/core/graph_test.py +++ b/tests/integration/core/graph/graph_test.py @@ -462,7 +462,7 @@ def test_bfs(self): class TestGraph2(BaseGraphFixture): def setup_graph(self): - """Set up the graph and add nodes and edges.""" + """Set up the graph.""" graph_schema = { "graph_name": "ERGraph", "nodes": { diff --git a/tests/integration/core/tigervector_test.py b/tests/integration/core/graph/tigervector_test.py similarity index 100% rename from tests/integration/core/tigervector_test.py rename to tests/integration/core/graph/tigervector_test.py diff --git a/tests/integration/core/tigergraph_api/database_api_test.py b/tests/integration/core/tigergraph_api/database_api_test.py new file mode 100644 index 0000000..6b30e2f --- /dev/null +++ b/tests/integration/core/tigergraph_api/database_api_test.py @@ -0,0 +1,171 @@ +import pytest +from pathlib import Path +import yaml + +from tigergraphx.core.tigergraph_api import TigerGraphAPI + + +class TestDatabaseAPIs: + @pytest.fixture(autouse=True) + def init(self): + config_path = ( + Path(__file__).parent.parent / "config" / "tigergraph_connection.yaml" + ) + with open(config_path, "r") as f: + self.tigergraph_connection_config = yaml.safe_load(f) + + # Initialize the TigerGraphAPI + self.api = TigerGraphAPI(self.tigergraph_connection_config) + + # ------------------------------ Admin ------------------------------ + def test_ping(self): + """ + Integration test for the TigerGraph ping endpoint. + """ + result = self.api.ping() + + assert isinstance(result, str), "Response should be a str." + assert result == "pong", "Response should be 'pong'." + + def test_get_version(self): + """ + Integration test for the TigerGraph get_version endpoint. + """ + result = self.api.get_version() + + assert isinstance(result, str), "Response should be a str." + assert result.startswith("3.") or result.startswith("4."), ( + f"Unexpected version format: {result}" + ) + + # ------------------------------ GSQL ------------------------------ + def test_gsql(self): + """ + Integration test for the TigerGraph gsql endpoint. + """ + result = self.api.gsql("ls") + + assert isinstance(result, str), "Response should be a string." + assert "Global vertices, edges, and all graphs" in result + + # ------------------------------ Data Source ------------------------------ + def test_data_source_CRUD(self): + data_source_name = "data_source_test" + data_source_type = "s3" + + # Create initial data source + result = self.api.create_data_source( + name=data_source_name, + data_source_type=data_source_type, + ) + assert isinstance(result, str), f"Expected str, got {type(result)}" + assert f"Data source {data_source_name} is created" in result, ( + f"Unexpected response: {result}" + ) + + try: + # Get data source + result = self.api.get_data_source(data_source_name) + assert isinstance(result, dict), f"Expected dict, got {type(result)}" + assert result.get("name") == data_source_name, ( + f"Expected name '{data_source_name}', got {result.get('name')}" + ) + assert result.get("type") == data_source_type.upper(), ( + f"Expected type '{data_source_type}', got {result.get('type')}" + ) + + # Get all data sources + all_sources = self.api.get_all_data_sources() + assert isinstance(all_sources, list), ( + f"Expected list, got {type(all_sources)}" + ) + assert any(ds["name"] == data_source_name for ds in all_sources), ( + "Created data source not found in list" + ) + + # Update the data source + result = self.api.update_data_source( + name=data_source_name, + data_source_type=data_source_type, + access_key="aaaaaaaaaaaaaaaaaaaa", + secret_key="", + ) + assert isinstance(result, str) + assert ( + f"Data source {data_source_name} is created" in result + or "updated" in result.lower() + ) + + finally: + # Drop data source + drop_result = self.api.drop_all_data_sources() + assert isinstance(drop_result, str), ( + f"Expected str, got {type(drop_result)}" + ) + assert "All data sources is dropped successfully." in drop_result, ( + f"Unexpected drop response: {drop_result}" + ) + + @pytest.mark.skip( + reason=""" + Skipped by default. To enable this test, manually configure the data source by setting + values for data_source_type, access_key, secret_key, extra_config, and sample_path + (the path to the file to sample). + """ + ) + def test_preview_sample_data(self): + data_source_name = "data_source_1" + data_source_type = "s3" + access_key = "" + secret_key = "" + extra_config = { + "file.reader.settings.fs.s3a.aws.credentials.provider": "org.apache.hadoop.fs.s3a.AnonymousAWSCredentialsProvider" + } + sample_path = "s3a://" + + # Create data source + result = self.api.create_data_source( + name=data_source_name, + data_source_type=data_source_type, + access_key=access_key, + secret_key=secret_key, + extra_config=extra_config, + ) + assert isinstance(result, str), f"Expected str, got {type(result)}" + assert f"Data source {data_source_name} is created" in result, ( + f"Unexpected response: {result}" + ) + + try: + # Preview sample data + preview_result = self.api.preview_sample_data( + path=sample_path, + data_source_type=data_source_type, + data_source=data_source_name, + data_format="csv", + size=5, + has_header=True, + separator=",", + eol="\\n", + quote='"', + ) + assert isinstance(preview_result, dict), ( + f"Expected dict, got {type(preview_result)}" + ) + assert "data" in preview_result, "Key 'data' not found in preview result" + assert isinstance(preview_result["data"], list), ( + f"Expected 'data' to be list, got {type(preview_result['data'])}" + ) + assert len(preview_result["data"]) <= 5, ( + f"Preview data has more rows than expected: {len(preview_result['data'])}" + ) + + finally: + # Drop data source + drop_result = self.api.drop_data_source(name=data_source_name) + assert isinstance(drop_result, str), ( + f"Expected str, got {type(drop_result)}" + ) + assert f"Data source {data_source_name} is dropped" in drop_result, ( + f"Unexpected drop response: {drop_result}" + ) diff --git a/tests/integration/core/tigergraph_api/tigergraph_api_test.py b/tests/integration/core/tigergraph_api/graph_api_test.py similarity index 50% rename from tests/integration/core/tigergraph_api/tigergraph_api_test.py rename to tests/integration/core/tigergraph_api/graph_api_test.py index af0cf99..cad1e88 100644 --- a/tests/integration/core/tigergraph_api/tigergraph_api_test.py +++ b/tests/integration/core/tigergraph_api/graph_api_test.py @@ -201,169 +201,3 @@ def test_run_interpreted_query_syntax_error(self): match="line 5:2 no viable alternative at input ", ): self.api.run_interpreted_query(gsql_query) - - -class TestDatabaseAPIs: - @pytest.fixture(autouse=True) - def init(self): - config_path = ( - Path(__file__).parent.parent / "config" / "tigergraph_connection.yaml" - ) - with open(config_path, "r") as f: - self.tigergraph_connection_config = yaml.safe_load(f) - - # Initialize the TigerGraphAPI - self.api = TigerGraphAPI(self.tigergraph_connection_config) - - # ------------------------------ Admin ------------------------------ - def test_ping(self): - """ - Integration test for the TigerGraph ping endpoint. - """ - result = self.api.ping() - - assert isinstance(result, str), "Response should be a str." - assert result == "pong", "Response should be 'pong'." - - def test_get_version(self): - """ - Integration test for the TigerGraph get_version endpoint. - """ - result = self.api.get_version() - - assert isinstance(result, str), "Response should be a str." - assert result.startswith("3.") or result.startswith("4."), ( - f"Unexpected version format: {result}" - ) - - # ------------------------------ GSQL ------------------------------ - def test_gsql(self): - """ - Integration test for the TigerGraph gsql endpoint. - """ - result = self.api.gsql("ls") - - assert isinstance(result, str), "Response should be a string." - assert "Global vertices, edges, and all graphs" in result - - # ------------------------------ Data Source ------------------------------ - def test_data_source_CRUD(self): - data_source_name = "data_source_test" - data_source_type = "s3" - - # Create initial data source - result = self.api.create_data_source( - name=data_source_name, - data_source_type=data_source_type, - ) - assert isinstance(result, str), f"Expected str, got {type(result)}" - assert f"Data source {data_source_name} is created" in result, ( - f"Unexpected response: {result}" - ) - - try: - # Get data source - result = self.api.get_data_source(data_source_name) - assert isinstance(result, dict), f"Expected dict, got {type(result)}" - assert result.get("name") == data_source_name, ( - f"Expected name '{data_source_name}', got {result.get('name')}" - ) - assert result.get("type") == data_source_type.upper(), ( - f"Expected type '{data_source_type}', got {result.get('type')}" - ) - - # Get all data sources - all_sources = self.api.get_all_data_sources() - assert isinstance(all_sources, list), ( - f"Expected list, got {type(all_sources)}" - ) - assert any(ds["name"] == data_source_name for ds in all_sources), ( - "Created data source not found in list" - ) - - # Update the data source - result = self.api.update_data_source( - name=data_source_name, - data_source_type=data_source_type, - access_key="aaaaaaaaaaaaaaaaaaaa", - secret_key="", - ) - assert isinstance(result, str) - assert ( - f"Data source {data_source_name} is created" in result - or "updated" in result.lower() - ) - - finally: - # Drop data source - drop_result = self.api.drop_all_data_sources() - assert isinstance(drop_result, str), ( - f"Expected str, got {type(drop_result)}" - ) - assert "All data sources is dropped successfully." in drop_result, ( - f"Unexpected drop response: {drop_result}" - ) - - @pytest.mark.skip( - reason=""" - Skipped by default. To enable this test, manually configure the data source by setting - values for data_source_type, access_key, secret_key, extra_config, and sample_path - (the path to the file to sample). - """ - ) - def test_preview_sample_data(self): - data_source_name = "data_source_1" - data_source_type = "s3" - access_key = "" - secret_key = "" - extra_config = { - "file.reader.settings.fs.s3a.aws.credentials.provider": "org.apache.hadoop.fs.s3a.AnonymousAWSCredentialsProvider" - } - sample_path = "s3a://" - - # Create data source - result = self.api.create_data_source( - name=data_source_name, - data_source_type=data_source_type, - access_key=access_key, - secret_key=secret_key, - extra_config=extra_config, - ) - assert isinstance(result, str), f"Expected str, got {type(result)}" - assert f"Data source {data_source_name} is created" in result, ( - f"Unexpected response: {result}" - ) - - try: - # Preview sample data - preview_result = self.api.preview_sample_data( - path=sample_path, - data_source_type=data_source_type, - data_source=data_source_name, - data_format="csv", - size=5, - has_header=True, - separator=",", - eol="\\n", - quote='"', - ) - assert isinstance(preview_result, dict), ( - f"Expected dict, got {type(preview_result)}" - ) - assert "data" in preview_result, "Key 'data' not found in preview result" - assert isinstance(preview_result["data"], list), ( - f"Expected 'data' to be list, got {type(preview_result['data'])}" - ) - assert len(preview_result["data"]) <= 5, ( - f"Preview data has more rows than expected: {len(preview_result['data'])}" - ) - - finally: - # Drop data source - drop_result = self.api.drop_data_source(name=data_source_name) - assert isinstance(drop_result, str), ( - f"Expected str, got {type(drop_result)}" - ) - assert f"Data source {data_source_name} is dropped" in drop_result, ( - f"Unexpected drop response: {drop_result}" - ) diff --git a/tests/integration/core/tigergraph_api/graph_schema_api_test.py b/tests/integration/core/tigergraph_api/graph_schema_api_test.py new file mode 100644 index 0000000..fa7d6c4 --- /dev/null +++ b/tests/integration/core/tigergraph_api/graph_schema_api_test.py @@ -0,0 +1,245 @@ +import pytest +from pathlib import Path +import yaml + +from tigergraphx.core.tigergraph_api import TigerGraphAPI + + +class TestGraphSchemaAPIs: + def setup_graph(self): + """Set up the graph and add nodes and edges.""" + # Load config from YAML + config_path = ( + Path(__file__).parent.parent / "config" / "tigergraph_connection.yaml" + ) + with open(config_path, "r") as f: + self.tigergraph_connection_config = yaml.safe_load(f) + + @pytest.fixture(autouse=True) + def init(self): + """Add nodes and edges before each test case.""" + # Initialize the graph + self.setup_graph() + + # Initialize the TigerGraphAPI + self.api = TigerGraphAPI(self.tigergraph_connection_config) + + yield # The test case runs here + + # ------------------------------ Schema ------------------------------ + def test_create_and_drop_graph(self): + graph_name = "TestSchemaGraph" + + # Precheck: drop graph if it already exists + try: + schema = self.api.get_schema(graph_name) + if schema and schema.get("GraphName") == graph_name: + drop_result = self.api.drop_graph(graph_name=graph_name) + assert isinstance(drop_result, str), "Drop response should be a string." + assert "Successfully dropped graph" in drop_result + except Exception: + # If get_schema fails (e.g., graph doesn't exist), ignore + pass + + # Create new graph + result = self.api.create_empty_graph(graph_name=graph_name) + assert isinstance(result, str), "Create response should be a string." + assert "Successfully created graph" in result + + # Drop the graph + result = self.api.drop_graph(graph_name=graph_name) + assert isinstance(result, str), "Drop response should be a string." + assert "Successfully dropped graph" in result + + def test_local_schema_change_job(self): + graph_name = "TestSchemaChangeGraph" + job_name_prefix = "test_job" + + # Precheck: drop graph if it already exists + try: + schema = self.api.get_schema(graph_name) + if schema and schema.get("GraphName") == graph_name: + drop_result = self.api.drop_graph(graph_name=graph_name) + assert isinstance(drop_result, str), "Drop response should be a string." + assert "Successfully dropped graph" in drop_result + except Exception: + # Graph does not exist + pass + + # Step 1: Create an empty graph + result = self.api.create_empty_graph(graph_name=graph_name) + assert isinstance(result, str), "Create response should be a string." + assert "Successfully created graph" in result + + try: + # Step 2: Add schema (nodes and edges) + add_job_name = f"{job_name_prefix}_add" + add_payload = { + "addVertexTypes": [ + { + "Name": "Person", + "PrimaryId": { + "AttributeName": "id", + "AttributeType": {"Name": "STRING"}, + }, + "Attributes": [ + { + "AttributeName": "name", + "AttributeType": {"Name": "STRING"}, + } + ], + "Config": { + "STATS": "OUTDEGREE_BY_EDGETYPE", + "PRIMARY_ID_AS_ATTRIBUTE": "true", + }, + }, + { + "Name": "Company", + "PrimaryId": { + "AttributeName": "id", + "AttributeType": {"Name": "STRING"}, + }, + "Attributes": [], + "Config": { + "STATS": "OUTDEGREE_BY_EDGETYPE", + "PRIMARY_ID_AS_ATTRIBUTE": "true", + }, + }, + ], + "addEdgeTypes": [ + { + "Name": "works_at", + "IsDirected": True, + "FromVertexTypeName": "Person", + "ToVertexTypeName": "Company", + "Attributes": [ + { + "AttributeName": "since", + "AttributeType": {"Name": "DATETIME"}, + "IsDiscriminator": True + }, + { + "AttributeName": "country", + "AttributeType": {"Name": "STRING"}, + } + ], + "Config": { + "REVERSE_EDGE": "reverse_works_at", + }, + } + ], + } + result = self.api.create_local_schema_change_job( + graph_name=graph_name, job_name=add_job_name, payload=add_payload + ) + assert isinstance(result, str), ( + "Create schema change job response should be a string." + ) + assert "Successfully created schema change job" in result + try: + result = self.api.run_local_schema_change_job( + graph_name=graph_name, job_name=add_job_name + ) + assert isinstance(result, str), ( + "Run schema change job response should be a string." + ) + assert "Schema change job run successfully!" in result + finally: + result = self.api.drop_local_schema_change_job( + graph_name=graph_name, job_name=add_job_name + ) + assert isinstance(result, str), ( + "Drop schema change job response should be a string." + ) + assert "Successfully dropped schema change jobs" in result + + # Step 3: Modify schema (add attributes) + modify_job_name = f"{job_name_prefix}_modify" + modify_payload = { + "alterEdgeTypes": [ + { + "name": "works_at", + "addAttributes": [ + { + "AttributeName": "city", + "AttributeType": {"Name": "STRING"}, + } + ], + "dropAttributes": ["country"], + } + ], + "alterVertexTypes": [ + { + "name": "Person", + "addAttributes": [ + { + "AttributeName": "age", + "AttributeType": {"Name": "UINT"}, + "DefaultValue": "0", + } + ], + "dropAttributes": ["name"], + } + ], + } + result = self.api.create_local_schema_change_job( + graph_name=graph_name, job_name=modify_job_name, payload=modify_payload + ) + assert isinstance(result, str), ( + "Create schema change job response should be a string." + ) + assert "Successfully created schema change job" in result + try: + result = self.api.run_local_schema_change_job( + graph_name=graph_name, job_name=modify_job_name + ) + assert isinstance(result, str), ( + "Run schema change job response should be a string." + ) + assert "Schema change job run successfully!" in result + finally: + result = self.api.drop_local_schema_change_job( + graph_name=graph_name, job_name=modify_job_name + ) + assert isinstance(result, str), ( + "Drop schema change job response should be a string." + ) + assert "Successfully dropped schema change jobs" in result + + # Step 4: Delete schema elements + delete_job_name = f"{job_name_prefix}_delete" + delete_payload = { + "dropEdgeTypes": ["works_at"], + "dropVertexTypes": ["Company"], + } + result = self.api.create_local_schema_change_job( + graph_name=graph_name, + job_name=delete_job_name, + payload=delete_payload, + ) + assert isinstance(result, str), ( + "Create schema change job response should be a string." + ) + assert "Successfully created schema change job" in result + try: + result = self.api.run_local_schema_change_job( + graph_name=graph_name, job_name=delete_job_name + ) + assert isinstance(result, str), ( + "Run schema change job response should be a string." + ) + assert "Schema change job run successfully!" in result + finally: + result = self.api.drop_local_schema_change_job( + graph_name=graph_name, job_name=delete_job_name + ) + assert isinstance(result, str), ( + "Drop schema change job response should be a string." + ) + assert "Successfully dropped schema change jobs" in result + + finally: + # Step 5: Drop the graph + result = self.api.drop_graph(graph_name=graph_name) + assert isinstance(result, str), "Drop response should be a string." + assert "Successfully dropped graph" in result diff --git a/tests/unit/core/managers/schema/schema_change_builder_test.py b/tests/unit/core/managers/schema/schema_change_builder_test.py new file mode 100644 index 0000000..f7f04ec --- /dev/null +++ b/tests/unit/core/managers/schema/schema_change_builder_test.py @@ -0,0 +1,98 @@ +import pytest + +from tigergraphx.config import ( + NodeSchema, + EdgeSchema, + AttributeSchema, + DataType, +) +from tigergraphx.core.managers.schema.schema_change_builder import SchemaChangeBuilder + + +class TestSchemaChangeBuilder: + @pytest.fixture(autouse=True) + def setup(self): + self.builder = SchemaChangeBuilder() + + def test_add_node_type(self): + node_schema = NodeSchema( + primary_key="id", + attributes={"id": AttributeSchema(data_type=DataType.STRING)}, + ) + self.builder.add_node_type("Person", node_schema) + payload = self.builder.build_payload() + assert "addVertexTypes" in payload + assert payload["addVertexTypes"][0]["Name"] == "Person" + assert payload["addVertexTypes"][0]["PrimaryId"]["AttributeName"] == "id" + + def test_drop_node_type(self): + self.builder.drop_node_type("Person") + payload = self.builder.build_payload() + assert "dropVertexTypes" in payload + assert "Person" in payload["dropVertexTypes"] + + def test_add_node_attribute(self): + attr_schema = AttributeSchema(data_type=DataType.STRING) + self.builder.add_node_attribute("Person", "name", attr_schema) + payload = self.builder.build_payload() + alter_node = payload["alterVertexTypes"][0] + assert alter_node["name"] == "Person" + assert alter_node["addAttributes"][0]["AttributeName"] == "name" + assert alter_node["addAttributes"][0]["AttributeType"]["Name"] == "STRING" + + def test_drop_node_attribute(self): + self.builder.drop_node_attribute("Person", "name") + payload = self.builder.build_payload() + alter_node = payload["alterVertexTypes"][0] + assert alter_node["name"] == "Person" + assert "name" in alter_node["dropAttributes"] + + def test_add_edge_type(self): + edge_schema = EdgeSchema( + from_node_type="Person", to_node_type="Company", is_directed_edge=True + ) + self.builder.add_edge_type("works_at", edge_schema) + payload = self.builder.build_payload() + edge = payload["addEdgeTypes"][0] + assert edge["Name"] == "works_at" + assert edge["IsDirected"] is True + assert edge["FromVertexTypeName"] == "Person" + assert edge["ToVertexTypeName"] == "Company" + + def test_drop_edge_type(self): + self.builder.drop_edge_type("works_at") + payload = self.builder.build_payload() + assert "works_at" in payload["dropEdgeTypes"] + + def test_add_edge_attribute(self): + attr_schema = AttributeSchema(data_type=DataType.INT) + self.builder.add_edge_attribute("works_at", "since", attr_schema) + payload = self.builder.build_payload() + alter_edge = payload["alterEdgeTypes"][0] + assert alter_edge["name"] == "works_at" + assert alter_edge["addAttributes"][0]["AttributeName"] == "since" + assert alter_edge["addAttributes"][0]["AttributeType"]["Name"] == "INT" + + def test_drop_edge_attribute(self): + self.builder.drop_edge_attribute("works_at", "since") + payload = self.builder.build_payload() + alter_edge = payload["alterEdgeTypes"][0] + assert "since" in alter_edge["dropAttributes"] + + def test_complex_node_and_edge_changes(self): + # add node + vector attr + edge + edge attr + drops + node_schema = NodeSchema( + primary_key="id", + attributes={"id": AttributeSchema(data_type=DataType.STRING)}, + ) + self.builder.add_node_type("Person", node_schema) + edge_schema = EdgeSchema(from_node_type="Person", to_node_type="Company") + self.builder.add_edge_type("works_at", edge_schema) + + payload = self.builder.build_payload() + # Node + node = payload["addVertexTypes"][0] + assert node["Name"] == "Person" + # Edge + edge = payload["addEdgeTypes"][0] + assert edge["Name"] == "works_at" diff --git a/tigergraphx/config/endpoint_definitions.yaml b/tigergraphx/config/endpoint_definitions.yaml index 7baeb20..17d6a34 100644 --- a/tigergraphx/config/endpoint_definitions.yaml +++ b/tigergraphx/config/endpoint_definitions.yaml @@ -36,6 +36,31 @@ endpoints: # 3.x: "/gsqlserver/gsql/schema" 4.x: "/gsql/v1/schema/graphs/{graph_name}" + create_empty_graph: + path: + 4.x: "/gsql/v1/schema/graphs?graphName={graph_name}&gsql=false" + method: "POST" + + drop_graph: + path: + 4.x: "/gsql/v1/schema/graphs/{graph_name}?cascade=true" + method: "DELETE" + + create_local_schema_change_job: + path: + 4.x: "/gsql/v1/schema/jobs/{job_name}?graph={graph_name}" + method: "POST" + + run_local_schema_change_job: + path: + 4.x: "/gsql/v1/schema/jobs/{job_name}?graph={graph_name}" + method: "PUT" + + drop_local_schema_change_job: + path: + 4.x: "/gsql/v1/schema/jobs?graph={graph_name}&jobName={job_name}" + method: "DELETE" + # ------------------------------ Data Source ------------------------------ create_data_source: path: diff --git a/tigergraphx/config/graph_db/schema/edge_schema.py b/tigergraphx/config/graph_db/schema/edge_schema.py index 6dab7e2..6ca5578 100644 --- a/tigergraphx/config/graph_db/schema/edge_schema.py +++ b/tigergraphx/config/graph_db/schema/edge_schema.py @@ -90,6 +90,17 @@ def validate_reserved_keywords(self) -> "EdgeSchema": return self + def set_default_reverse_edge(self, edge_name: str) -> None: + """ + Set the default reverse edge name if the edge is directed and + no reverse edge name is provided. + + Args: + edge_name: The name of the edge. + """ + if self.is_directed_edge and not self.reverse_edge_name: + self.reverse_edge_name = f"reverse_{edge_name}" + def create_edge_schema( is_directed_edge: bool, diff --git a/tigergraphx/config/graph_db/schema/graph_schema.py b/tigergraphx/config/graph_db/schema/graph_schema.py index 774a0cb..35ce28c 100644 --- a/tigergraphx/config/graph_db/schema/graph_schema.py +++ b/tigergraphx/config/graph_db/schema/graph_schema.py @@ -93,6 +93,5 @@ def set_reverse_edge_names(self) -> "GraphSchema": The updated graph schema. """ for edge_name, edge in self.edges.items(): - if edge.is_directed_edge and not edge.reverse_edge_name: - edge.reverse_edge_name = f"reverse_{edge_name}" + edge.set_default_reverse_edge(edge_name) return self diff --git a/tigergraphx/core/__init__.py b/tigergraphx/core/__init__.py index 16391ff..ee38cc9 100644 --- a/tigergraphx/core/__init__.py +++ b/tigergraphx/core/__init__.py @@ -8,6 +8,7 @@ from .graph import Graph from .tigergraph_api import TigerGraphAPI, TigerGraphAPIError from .tigergraph_database import TigerGraphDatabase +from .managers import SchemaChangeBuilder __all__ = [ @@ -15,4 +16,5 @@ "TigerGraphAPI", "TigerGraphAPIError", "TigerGraphDatabase", + "SchemaChangeBuilder", ] diff --git a/tigergraphx/core/graph.py b/tigergraphx/core/graph.py index 9ef3e03..6136c94 100644 --- a/tigergraphx/core/graph.py +++ b/tigergraphx/core/graph.py @@ -14,6 +14,9 @@ TigerGraphConnectionConfig, GraphSchema, LoadingJobConfig, + NodeSchema, + EdgeSchema, + AttributeSchema, ) from tigergraphx.core.graph_context import GraphContext @@ -173,6 +176,228 @@ def drop_graph(self) -> None: """ return self._schema_manager.drop_graph() + def apply_schema_changes( + self, + add_nodes: Optional[Dict[str, NodeSchema | Dict | str | Path]] = None, + drop_nodes: Optional[List[str]] = None, + add_node_attributes: Optional[ + Dict[str, Dict[str, AttributeSchema | Dict | str | Path]] + ] = None, + drop_node_attributes: Optional[Dict[str, List[str]]] = None, + add_edges: Optional[Dict[str, EdgeSchema | Dict | str | Path]] = None, + drop_edges: Optional[List[str]] = None, + add_edge_attributes: Optional[ + Dict[str, Dict[str, AttributeSchema | Dict | str | Path]] + ] = None, + drop_edge_attributes: Optional[Dict[str, List[str]]] = None, + ) -> bool: + """ + Apply multiple schema changes in a single schema change job. + + This can add or drop nodes, edges, or their attributes. + The method executes all changes as one schema change job. + + Args: + add_nodes: Nodes to add with their schemas. + drop_nodes: Names of nodes to drop. + add_node_attributes: Attributes to add to existing nodes. + drop_node_attributes: Attributes to drop from existing nodes. + add_edges: Edges to add with their schemas. + drop_edges: Names of edges to drop. + add_edge_attributes: Attributes to add to existing edges. + drop_edge_attributes: Attributes to drop from existing edges. + + Returns: + True if the schema change job was executed successfully, + False if no changes were applied. + """ + return self._schema_manager.apply_schema_changes( + add_nodes=add_nodes, + drop_nodes=drop_nodes, + add_node_attributes=add_node_attributes, + drop_node_attributes=drop_node_attributes, + add_edges=add_edges, + drop_edges=drop_edges, + add_edge_attributes=add_edge_attributes, + drop_edge_attributes=drop_edge_attributes, + ) + + def add_node_type( + self, + name: str, + schema: NodeSchema | Dict | str | Path, + ) -> bool: + """Add a new node type to the graph schema. + + Args: + name: The name of the node type. + schema: The schema definition for the node type. + + Returns: + True if the schema change succeeds. + """ + return self._schema_manager.add_node_type(name, schema) + + def add_node_types( + self, + nodes: Dict[str, NodeSchema | Dict | str | Path], + ) -> bool: + """Add multiple node types to the graph schema. + + Args: + nodes: A dictionary of node type names to their schemas. + + Returns: + True if the schema change succeeds. + """ + return self._schema_manager.add_node_types(nodes) + + def drop_node_type( + self, + name: str, + ) -> bool: + """Drop a node type from the graph schema. + + Args: + name: The name of the node type to drop. + + Returns: + True if the schema change succeeds. + """ + return self._schema_manager.drop_node_types([name]) + + def drop_node_types( + self, + node_names: List[str], + ) -> bool: + """Drop multiple node types from the graph schema. + + Args: + node_names: A list of node type names to drop. + + Returns: + True if the schema change succeeds. + """ + return self._schema_manager.drop_node_types(node_names) + + def add_edge_type( + self, + name: str, + schema: EdgeSchema | Dict | str | Path, + ) -> bool: + """Add a new edge type to the graph schema. + + Args: + name: The name of the edge type. + schema: The schema definition for the edge type. + + Returns: + True if the schema change succeeds. + """ + return self._schema_manager.add_edge_type(name, schema) + + def add_edge_types( + self, + edges: Dict[str, EdgeSchema | Dict | str | Path], + ) -> bool: + """Add multiple edge types to the graph schema. + + Args: + edges: A dictionary of edge type names to their schemas. + + Returns: + True if the schema change succeeds. + """ + return self._schema_manager.add_edge_types(edges) + + def drop_edge_type( + self, + name: str, + ) -> bool: + """Drop a single edge type from the graph schema. + + Args: + name: The name of the edge type to drop. + + Returns: + True if the schema change succeeds. + """ + return self._schema_manager.drop_edge_types([name]) + + def drop_edge_types( + self, + edge_names: List[str], + ) -> bool: + """Drop multiple edge types from the graph schema. + + Args: + edge_names: A list of edge type names to drop. + + Returns: + True if the schema change succeeds. + """ + return self._schema_manager.drop_edge_types(edge_names) + + def add_node_attributes( + self, + node_attributes: Dict[str, Dict[str, AttributeSchema | Dict | str | Path]], + ) -> bool: + """Add attributes to nodes in the graph schema. + + Args: + node_attributes: A dictionary where keys are node type names and values + are dictionaries of attribute names mapped to their schema. + + Returns: + True if the schema change succeeds. + """ + return self._schema_manager.add_node_attributes(node_attributes) + + def drop_node_attributes( + self, + node_attributes: Dict[str, List[str]], + ) -> bool: + """Drop attributes from nodes in the graph schema. + + Args: + node_attributes: A dictionary where keys are node type names and values + are lists of attribute names to drop. + + Returns: + True if the schema change succeeds. + """ + return self._schema_manager.drop_node_attributes(node_attributes) + + def add_edge_attributes( + self, + edge_attributes: Dict[str, Dict[str, AttributeSchema | Dict | str | Path]], + ) -> bool: + """Add attributes to edges in the graph schema. + + Args: + edge_attributes: A dictionary where keys are edge type names and values + are dictionaries of attribute names mapped to their schema. + + Returns: + True if the schema change succeeds. + """ + return self._schema_manager.add_edge_attributes(edge_attributes) + + def drop_edge_attributes( + self, + edge_attributes: Dict[str, List[str]], + ) -> bool: + """Drop attributes from edges in the graph schema. + + Args: + edge_attributes: A dictionary where keys are edge type names and values + are lists of attribute names to drop. + + Returns: + True if the schema change succeeds. + """ + return self._schema_manager.drop_edge_attributes(edge_attributes) + # ------------------------------ Data Loading Operations ------------------------------ def load_data( self, loading_job_config: LoadingJobConfig | Dict | str | Path diff --git a/tigergraphx/core/managers/__init__.py b/tigergraphx/core/managers/__init__.py index f2166f9..ff46ce5 100644 --- a/tigergraphx/core/managers/__init__.py +++ b/tigergraphx/core/managers/__init__.py @@ -12,6 +12,7 @@ from .statistics_manager import StatisticsManager from .query_manager import QueryManager from .vector_manager import VectorManager +from .schema_change import SchemaChangeBuilder __all__ = [ "SchemaManager", @@ -21,4 +22,5 @@ "StatisticsManager", "QueryManager", "VectorManager", + "SchemaChangeBuilder", ] diff --git a/tigergraphx/core/managers/schema_change/__init__.py b/tigergraphx/core/managers/schema_change/__init__.py new file mode 100644 index 0000000..09a4f2a --- /dev/null +++ b/tigergraphx/core/managers/schema_change/__init__.py @@ -0,0 +1,12 @@ +# Copyright 2025 TigerGraph Inc. +# Licensed under the Apache License, Version 2.0. +# See the LICENSE file or https://www.apache.org/licenses/LICENSE-2.0 +# +# Permission is granted to use, copy, modify, and distribute this software +# under the License. The software is provided "AS IS", without warranty. + +from .schema_change_builder import SchemaChangeBuilder + +__all__ = [ + "SchemaChangeBuilder", +] diff --git a/tigergraphx/core/managers/schema_change/schema_change_builder.py b/tigergraphx/core/managers/schema_change/schema_change_builder.py new file mode 100644 index 0000000..abc634d --- /dev/null +++ b/tigergraphx/core/managers/schema_change/schema_change_builder.py @@ -0,0 +1,203 @@ +# Copyright 2025 TigerGraph Inc. +# Licensed under the Apache License, Version 2.0. +# See the LICENSE file or https://www.apache.org/licenses/LICENSE-2.0 +# +# Permission is granted to use, copy, modify, and distribute this software +# under the License. The software is provided "AS IS", without warranty. + +from typing import Any, Dict, List +from tigergraphx.config import ( + NodeSchema, + EdgeSchema, + AttributeSchema, +) + + +class SchemaChangeBuilder: + def __init__(self): + # Node operations + self.nodes_to_add: Dict[str, NodeSchema] = {} + self.nodes_to_drop: List[str] = [] + + # Alter node operations + self.nodes_to_alter_add_attrs: Dict[str, Dict[str, AttributeSchema]] = {} + self.nodes_to_alter_drop_attrs: Dict[str, List[str]] = {} + + # Edge operations + self.edges_to_add: Dict[str, EdgeSchema] = {} + self.edges_to_drop: List[str] = [] + + # Alter edge operations + self.edges_to_alter_add_attrs: Dict[str, Dict[str, AttributeSchema]] = {} + self.edges_to_alter_drop_attrs: Dict[str, List[str]] = {} + + # ---------------- Node/Edge add/drop methods ---------------- + def add_node_type(self, name: str, schema: NodeSchema): + self.nodes_to_add[name] = schema + + def drop_node_type(self, name: str): + self.nodes_to_drop.append(name) + + def add_node_attribute( + self, + node_name: str, + attr_name: str, + schema: AttributeSchema, + ): + self.nodes_to_alter_add_attrs.setdefault(node_name, {})[attr_name] = schema + + def drop_node_attribute(self, node_name: str, attr_name: str): + self.nodes_to_alter_drop_attrs.setdefault(node_name, []).append(attr_name) + + def add_edge_type(self, name: str, schema: EdgeSchema): + self.edges_to_add[name] = schema + + def drop_edge_type(self, name: str): + self.edges_to_drop.append(name) + + def add_edge_attribute( + self, + edge_name: str, + attr_name: str, + schema: AttributeSchema, + ): + self.edges_to_alter_add_attrs.setdefault(edge_name, {})[attr_name] = schema + + def drop_edge_attribute(self, edge_name: str, attr_name: str): + self.edges_to_alter_drop_attrs.setdefault(edge_name, []).append(attr_name) + + # ---------------- Payload conversion ---------------- + def build_payload(self) -> Dict: + payload = {} + + # Add nodes + if self.nodes_to_add: + payload["addVertexTypes"] = [ + self._node_to_payload(name, schema) + for name, schema in self.nodes_to_add.items() + ] + + # Add edges + if self.edges_to_add: + payload["addEdgeTypes"] = [ + self._edge_to_payload(name, schema) + for name, schema in self.edges_to_add.items() + ] + + # Alter nodes + if any( + [ + self.nodes_to_alter_add_attrs, + self.nodes_to_alter_drop_attrs, + ] + ): + payload["alterVertexTypes"] = self._build_alter_vertex_payload() + + # Alter edges + if any([self.edges_to_alter_add_attrs, self.edges_to_alter_drop_attrs]): + payload["alterEdgeTypes"] = self._build_alter_edge_payload() + + # Drop nodes/edges + if self.nodes_to_drop: + payload["dropVertexTypes"] = self.nodes_to_drop + if self.edges_to_drop: + payload["dropEdgeTypes"] = self.edges_to_drop + + return payload + + # ---------------- Helpers ---------------- + def _node_to_payload(self, name: str, schema: NodeSchema) -> Dict: + """Convert a NodeSchema into the payload format for schema change job.""" + primary_key_name = schema.primary_key + primary_key_type = schema.attributes[primary_key_name].data_type.value + + attributes = [ + {"AttributeName": k, "AttributeType": {"Name": v.data_type.value}} + for k, v in schema.attributes.items() + if k != primary_key_name + ] + + return { + "Name": name, + "PrimaryId": { + "AttributeName": primary_key_name, + "AttributeType": {"Name": primary_key_type}, + }, + "Attributes": attributes, + "Config": { + "STATS": "OUTDEGREE_BY_EDGETYPE", + "PRIMARY_ID_AS_ATTRIBUTE": "true", + }, + } + + def _edge_to_payload(self, name: str, schema: EdgeSchema) -> Dict: + """Convert an EdgeSchema into the payload format for schema change job.""" + attributes_payload = [] + + for attr_name, attr_schema in schema.attributes.items(): + attr_dict = { + "AttributeName": attr_name, + "AttributeType": {"Name": attr_schema.data_type.value}, + } + if ( + isinstance(schema.discriminator, set) + and attr_name in schema.discriminator + ): + attr_dict["IsDiscriminator"] = True + elif ( + isinstance(schema.discriminator, str) + and attr_name == schema.discriminator + ): + attr_dict["IsDiscriminator"] = True + attributes_payload.append(attr_dict) + + config = ( + {"REVERSE_EDGE": schema.reverse_edge_name} + if schema.reverse_edge_name + else {} + ) + + return { + "Name": name, + "IsDirected": schema.is_directed_edge, + "FromVertexTypeName": schema.from_node_type, + "ToVertexTypeName": schema.to_node_type, + "Attributes": attributes_payload, + "Config": config, + } + + def _build_alter_vertex_payload(self) -> List[Dict]: + alter_payload = [] + node_names = set( + list(self.nodes_to_alter_add_attrs.keys()) + + list(self.nodes_to_alter_drop_attrs.keys()) + ) + for node in node_names: + entry: Dict[str, Any] = {"name": node} + if node in self.nodes_to_alter_add_attrs: + entry["addAttributes"] = [ + {"AttributeName": k, "AttributeType": {"Name": v.data_type.value}} + for k, v in self.nodes_to_alter_add_attrs[node].items() + ] + if node in self.nodes_to_alter_drop_attrs: + entry["dropAttributes"] = self.nodes_to_alter_drop_attrs[node] + alter_payload.append(entry) + return alter_payload + + def _build_alter_edge_payload(self) -> List[Dict]: + alter_payload = [] + edge_names = set( + list(self.edges_to_alter_add_attrs.keys()) + + list(self.edges_to_alter_drop_attrs.keys()) + ) + for edge in edge_names: + entry: Dict[str, Any] = {"name": edge} + if edge in self.edges_to_alter_add_attrs: + entry["addAttributes"] = [ + {"AttributeName": k, "AttributeType": {"Name": v.data_type.value}} + for k, v in self.edges_to_alter_add_attrs[edge].items() + ] + if edge in self.edges_to_alter_drop_attrs: + entry["dropAttributes"] = self.edges_to_alter_drop_attrs[edge] + alter_payload.append(entry) + return alter_payload diff --git a/tigergraphx/core/managers/schema_manager.py b/tigergraphx/core/managers/schema_manager.py index ccc29d5..ea05cd6 100644 --- a/tigergraphx/core/managers/schema_manager.py +++ b/tigergraphx/core/managers/schema_manager.py @@ -6,13 +6,21 @@ # under the License. The software is provided "AS IS", without warranty. import logging -from typing import Dict, Literal, Optional +from typing import Dict, List, Literal, Optional from pathlib import Path from .base_manager import BaseManager from tigergraphx.core.graph_context import GraphContext -from tigergraphx.config import GraphSchema, TigerGraphConnectionConfig +from tigergraphx.config import ( + GraphSchema, + TigerGraphConnectionConfig, + NodeSchema, + EdgeSchema, + AttributeSchema, +) + +from .schema_change import SchemaChangeBuilder logger = logging.getLogger(__name__) @@ -103,15 +111,216 @@ def create_schema(self, drop_existing_graph=False) -> bool: def drop_graph(self) -> None: logger.info(f"Dropping graph: {self._graph_name}...") - gsql_script = self._create_gsql_drop_graph() - result = self._tigergraph_api.gsql(gsql_script) + result = self._tigergraph_api.drop_graph(self._graph_name) logger.debug(result) - if f"The graph {self._graph_name} is dropped" not in result: + if "Successfully dropped graph" not in result: error_msg = f"Failed to drop the graph. GSQL response: {result}" logger.error(error_msg) raise RuntimeError(error_msg) logger.info("Graph dropped successfully.") + # -------- Schema Change Operations -------- + def apply_schema_changes( + self, + add_nodes: Optional[Dict[str, NodeSchema | Dict | str | Path]] = None, + drop_nodes: Optional[List[str]] = None, + add_node_attributes: Optional[ + Dict[str, Dict[str, AttributeSchema | Dict | str | Path]] + ] = None, + drop_node_attributes: Optional[Dict[str, List[str]]] = None, + add_edges: Optional[Dict[str, EdgeSchema | Dict | str | Path]] = None, + drop_edges: Optional[List[str]] = None, + add_edge_attributes: Optional[ + Dict[str, Dict[str, AttributeSchema | Dict | str | Path]] + ] = None, + drop_edge_attributes: Optional[Dict[str, List[str]]] = None, + ) -> bool: + """ + Apply multiple schema changes in a single schema change job. + + Returns: + True if the schema change job was executed successfully, False if no changes were applied. + """ + builder = SchemaChangeBuilder() + + add_nodes = add_nodes or {} + drop_nodes = drop_nodes or [] + add_node_attributes = add_node_attributes or {} + drop_node_attributes = drop_node_attributes or {} + add_edges = add_edges or {} + drop_edges = drop_edges or [] + add_edge_attributes = add_edge_attributes or {} + drop_edge_attributes = drop_edge_attributes or {} + + if not ( + add_nodes + or drop_nodes + or add_node_attributes + or drop_node_attributes + or add_edges + or drop_edges + or add_edge_attributes + or drop_edge_attributes + ): + return False + + # Add nodes + for node_name, schema in add_nodes.items(): + if node_name in self._graph_schema.nodes: + raise ValueError( + f"Node type '{node_name}' already exists in the graph schema." + ) + schema = NodeSchema.ensure_config(schema) + builder.add_node_type(node_name, schema) + self._graph_schema.nodes[node_name] = schema + + # Drop nodes + for node_name in drop_nodes: + if node_name not in self._graph_schema.nodes: + raise ValueError( + f"Node type '{node_name}' does not exist in the graph schema." + ) + builder.drop_node_type(node_name) + del self._graph_schema.nodes[node_name] + + # Add node attributes + for node_name, attrs in add_node_attributes.items(): + if node_name not in self._graph_schema.nodes: + raise ValueError( + f"Node type '{node_name}' does not exist to add attributes." + ) + schema = self._graph_schema.nodes[node_name] + for attr_name, attr_type in attrs.items(): + if attr_name in schema.attributes: + raise ValueError( + f"Attribute '{attr_name}' already exists in node '{node_name}'." + ) + attr_type = AttributeSchema.ensure_config(attr_type) + schema.attributes[attr_name] = attr_type + builder.add_node_attribute(node_name, attr_name, attr_type) + + # Drop node attributes + for node_name, attrs in drop_node_attributes.items(): + if node_name not in self._graph_schema.nodes: + raise ValueError( + f"Node type '{node_name}' does not exist to drop attributes." + ) + schema = self._graph_schema.nodes[node_name] + for attr_name in attrs: + if attr_name not in schema.attributes: + raise ValueError( + f"Attribute '{attr_name}' does not exist in node '{node_name}'." + ) + del schema.attributes[attr_name] + builder.drop_node_attribute(node_name, attr_name) + + # Add edges + for edge_name, schema in add_edges.items(): + if edge_name in self._graph_schema.edges: + raise ValueError( + f"Edge type '{edge_name}' already exists in the graph schema." + ) + schema = EdgeSchema.ensure_config(schema) + schema.set_default_reverse_edge(edge_name) + builder.add_edge_type(edge_name, schema) + self._graph_schema.edges[edge_name] = schema + + # Drop edges + for edge_name in drop_edges: + if edge_name not in self._graph_schema.edges: + raise ValueError( + f"Edge type '{edge_name}' does not exist in the graph schema." + ) + builder.drop_edge_type(edge_name) + del self._graph_schema.edges[edge_name] + + # Add edge attributes + for edge_name, attrs in add_edge_attributes.items(): + if edge_name not in self._graph_schema.edges: + raise ValueError( + f"Edge type '{edge_name}' does not exist to add attributes." + ) + schema = self._graph_schema.edges[edge_name] + for attr_name, attr_type in attrs.items(): + if attr_name in schema.attributes: + raise ValueError( + f"Attribute '{attr_name}' already exists in edge '{edge_name}'." + ) + attr_type = AttributeSchema.ensure_config(attr_type) + schema.attributes[attr_name] = attr_type + builder.add_edge_attribute(edge_name, attr_name, attr_type) + + # Drop edge attributes + for edge_name, attrs in drop_edge_attributes.items(): + if edge_name not in self._graph_schema.edges: + raise ValueError( + f"Edge type '{edge_name}' does not exist to drop attributes." + ) + schema = self._graph_schema.edges[edge_name] + for attr_name in attrs: + if attr_name not in schema.attributes: + raise ValueError( + f"Attribute '{attr_name}' does not exist in edge '{edge_name}'." + ) + del schema.attributes[attr_name] + builder.drop_edge_attribute(edge_name, attr_name) + + # Execute the schema change job + job_name = f"schema_change_{self._graph_name}" + return self._execute_schema_change_job(builder, job_name) + + def add_node_type(self, name: str, schema: NodeSchema | Dict | str | Path) -> bool: + """Add a single node type to the graph schema.""" + return self.apply_schema_changes(add_nodes={name: schema}) + + def add_node_types(self, nodes: Dict[str, NodeSchema | Dict | str | Path]) -> bool: + """Add multiple node types to the graph schema.""" + return self.apply_schema_changes(add_nodes=nodes) + + def drop_node_type(self, name: str) -> bool: + """Drop a single node type from the graph schema.""" + return self.apply_schema_changes(drop_nodes=[name]) + + def drop_node_types(self, node_names: List[str]) -> bool: + """Drop multiple node types from the graph schema.""" + return self.apply_schema_changes(drop_nodes=node_names) + + def add_edge_type(self, name: str, schema: EdgeSchema | Dict | str | Path) -> bool: + """Add a single edge type to the graph schema.""" + return self.apply_schema_changes(add_edges={name: schema}) + + def add_edge_types(self, edges: Dict[str, EdgeSchema | Dict | str | Path]) -> bool: + """Add multiple edge types to the graph schema.""" + return self.apply_schema_changes(add_edges=edges) + + def drop_edge_type(self, name: str) -> bool: + """Drop a single edge type from the graph schema.""" + return self.apply_schema_changes(drop_edges=[name]) + + def drop_edge_types(self, edge_names: List[str]) -> bool: + """Drop multiple edge types from the graph schema.""" + return self.apply_schema_changes(drop_edges=edge_names) + + def add_node_attributes( + self, node_attributes: Dict[str, Dict[str, AttributeSchema | Dict | str | Path]] + ) -> bool: + """Add attributes to nodes in the graph schema.""" + return self.apply_schema_changes(add_node_attributes=node_attributes) + + def drop_node_attributes(self, node_attributes: Dict[str, List[str]]) -> bool: + """Drop attributes from nodes in the graph schema.""" + return self.apply_schema_changes(drop_node_attributes=node_attributes) + + def add_edge_attributes( + self, edge_attributes: Dict[str, Dict[str, AttributeSchema | Dict | str | Path]] + ) -> bool: + """Add attributes to edges in the graph schema.""" + return self.apply_schema_changes(add_edge_attributes=edge_attributes) + + def drop_edge_attributes(self, edge_attributes: Dict[str, List[str]]) -> bool: + """Drop attributes from edges in the graph schema.""" + return self.apply_schema_changes(drop_edge_attributes=edge_attributes) + def _check_graph_exists(self) -> bool: """Check if the specified graph name exists in the gsql_script.""" result = self._tigergraph_api.gsql(f"USE Graph {self._graph_name}") @@ -122,16 +331,6 @@ def _check_graph_exists(self) -> bool: ) return "Using graph" in result - def _create_gsql_drop_graph(self) -> str: - # Generating the gsql script to drop graph - gsql_script = f""" -USE GRAPH {self._graph_name} -DROP QUERY * -DROP JOB * -DROP GRAPH {self._graph_name} -""" - return gsql_script.strip() - def _create_gsql_graph_schema(self) -> str: # Extracting node attributes graph_schema = self._graph_schema @@ -347,6 +546,52 @@ def _create_gsql_add_vector_attr(self) -> str: logger.debug("GSQL script for adding vector attributes: %s", gsql_script) return gsql_script.rstrip() + def _execute_schema_change_job( + self, builder: SchemaChangeBuilder, job_name: str + ) -> bool: + """Create, run, and drop a local schema change job, logging all steps.""" + logger.info( + f"Running schema change job '{job_name}' on graph: {self._graph_name}..." + ) + try: + # Create job + create_result = self._tigergraph_api.create_local_schema_change_job( + self._graph_name, job_name, builder.build_payload() + ) + logger.debug(f"Create job result: {create_result}") + if "Successfully created schema change job" not in create_result: + raise RuntimeError( + f"Failed to create schema change job: {create_result}" + ) + + # Run job + run_result = self._tigergraph_api.run_local_schema_change_job( + self._graph_name, job_name + ) + logger.debug(f"Run job result: {run_result}") + if "Schema change job run successfully!" not in run_result: + raise RuntimeError(f"Failed to run schema change job: {run_result}") + + logger.info(f"Schema change job '{job_name}' executed successfully.") + return True + + except Exception as e: + logger.error(f"Schema change job '{job_name}' failed: {e}") + raise + + finally: + # Always attempt to drop the job + try: + drop_result = self._tigergraph_api.drop_local_schema_change_job( + self._graph_name, job_name + ) + logger.debug(f"Drop job result: {drop_result}") + except Exception as cleanup_error: + logger.warning( + f"Failed to drop schema change job '{job_name}' during cleanup: " + f"{cleanup_error}" + ) + @staticmethod def get_schema_from_db( graph_name: str, diff --git a/tigergraphx/core/tigergraph_api/api/schema_api.py b/tigergraphx/core/tigergraph_api/api/schema_api.py index d79c6e3..57ea2c7 100644 --- a/tigergraphx/core/tigergraph_api/api/schema_api.py +++ b/tigergraphx/core/tigergraph_api/api/schema_api.py @@ -5,7 +5,7 @@ # Permission is granted to use, copy, modify, and distribute this software # under the License. The software is provided "AS IS", without warranty. -from typing import Dict +from typing import Any, Dict from .base_api import BaseAPI @@ -18,3 +18,70 @@ def get_schema(self, graph_name: str) -> Dict: if not isinstance(result, dict): raise TypeError(f"Expected dict, but got {type(result).__name__}: {result}") return result + + def create_empty_graph(self, graph_name: str) -> str: + """ + Create an empty graph. + """ + data = '{"VertexTypes":[], "EdgeTypes":[]}' + + result = self._request( + endpoint_name="create_empty_graph", graph_name=graph_name, data=data + ) + if not isinstance(result, str): + raise TypeError(f"Expected str, but got {type(result).__name__}: {result}") + return result + + def drop_graph(self, graph_name: str) -> str: + """ + Drop a graph. + """ + result = self._request( + endpoint_name="drop_graph", + graph_name=graph_name, + ) + if not isinstance(result, str): + raise TypeError(f"Expected str, but got {type(result).__name__}: {result}") + return result + + def create_local_schema_change_job( + self, graph_name: str, job_name: str, payload: Dict[str, Any] + ) -> str: + """ + Create a local schema change job + """ + result = self._request( + endpoint_name="create_local_schema_change_job", + graph_name=graph_name, + job_name=job_name, + json=payload, + ) + if not isinstance(result, str): + raise TypeError(f"Expected str, but got {type(result).__name__}: {result}") + return result + + def run_local_schema_change_job(self, graph_name: str, job_name: str) -> str: + """ + Run a local schema change job + """ + result = self._request( + endpoint_name="run_local_schema_change_job", + graph_name=graph_name, + job_name=job_name, + ) + if not isinstance(result, str): + raise TypeError(f"Expected str, but got {type(result).__name__}: {result}") + return result + + def drop_local_schema_change_job(self, graph_name: str, job_name: str) -> str: + """ + Drop a local schema change job + """ + result = self._request( + endpoint_name="drop_local_schema_change_job", + graph_name=graph_name, + job_name=job_name, + ) + if not isinstance(result, str): + raise TypeError(f"Expected str, but got {type(result).__name__}: {result}") + return result diff --git a/tigergraphx/core/tigergraph_api/tigergraph_api.py b/tigergraphx/core/tigergraph_api/tigergraph_api.py index e8a04c4..0b5fede 100644 --- a/tigergraphx/core/tigergraph_api/tigergraph_api.py +++ b/tigergraphx/core/tigergraph_api/tigergraph_api.py @@ -163,7 +163,9 @@ def create_token( Returns: The generated authentication token as a string. """ - return self._security_api.create_token(secret_alias, graph_name, lifetime_seconds) + return self._security_api.create_token( + secret_alias, graph_name, lifetime_seconds + ) def drop_token( self, @@ -349,6 +351,74 @@ def get_schema(self, graph_name: str) -> Dict: """ return self._schema_api.get_schema(graph_name) + def create_empty_graph(self, graph_name: str) -> str: + """ + Create a new empty graph with the given name. + + Args: + graph_name: The name of the graph. + + Returns: + The response message. + """ + return self._schema_api.create_empty_graph(graph_name) + + def drop_graph(self, graph_name: str) -> str: + """ + Drop a graph with the given name. + + Args: + graph_name: The name of the graph. + + Returns: + The response message. + """ + return self._schema_api.drop_graph(graph_name) + + def create_local_schema_change_job( + self, graph_name: str, job_name: str, payload: Dict[str, Any] + ) -> str: + """ + Create a local schema change job for the given graph. + + Args: + graph_name: The name of the graph. + job_name: The name of the job. + payload: The schema change details. + + Returns: + The response message. + """ + return self._schema_api.create_local_schema_change_job( + graph_name, job_name, payload + ) + + def run_local_schema_change_job(self, graph_name: str, job_name: str) -> str: + """ + Run a local schema change job for the given graph. + + Args: + graph_name: The name of the graph. + job_name: The name of the job. + + Returns: + The response message. + """ + return self._schema_api.run_local_schema_change_job(graph_name, job_name) + + def drop_local_schema_change_job(self, graph_name: str, job_name: str) -> str: + """ + Drop a local schema change job for the given graph. + + Args: + graph_name: The name of the graph. + job_name: The name of the job. + + Returns: + The response message. + """ + return self._schema_api.drop_local_schema_change_job(graph_name, job_name) + # ------------------------------ Node ------------------------------ def retrieve_a_node(self, graph_name: str, node_type: str, node_id: str) -> List: """