Merge branch 'main' into willtai/use-exceptions-for-error-handling

Author: willtai

.github/workflows/pr-e2e-tests.yaml

@@ -71,4 +71,11 @@ jobs:
         run: |
           count=0; until curl -s --fail localhost:8080/v1/.well-known/ready; do ((count++)); [ $count -ge 10 ] && echo "Reached maximum retry limit" && exit 1; sleep 15; done
       - name: Run tests
-        run: poetry run pytest ./tests/e2e
+        shell: bash
+        run: |
+          if [[ "${{ matrix.neo4j-edition }}" == "community" ]]; then
+              poetry run pytest -m 'not enterprise_only' ./tests/e2e
+          else
+              poetry run pytest ./tests/e2e
+          fi
+

src/neo4j_genai/schema.py

@@ -13,6 +13,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from typing import Any, Optional
+from neo4j.exceptions import ClientError
 
 import neo4j
 
@@ -49,6 +50,12 @@
 RETURN {start: label, type: property, end: toString(other_node)} AS output
 """
 
+INDEX_QUERY = """
+CALL apoc.schema.nodes() YIELD label, properties, type, size, valuesSelectivity 
+WHERE type = "RANGE" RETURN *, 
+size * valuesSelectivity as distinctValues
+"""
+
 
 def query_database(
     driver: neo4j.Driver, query: str, params: Optional[dict] = None
@@ -82,6 +89,51 @@ def get_schema(
     Returns:
         str: the graph schema information in a serialized format.
     """
+    structured_schema = get_structured_schema(driver)
+
+    def _format_props(props):
+        return ", ".join([f"{prop['property']}: {prop['type']}" for prop in props])
+
+    # Format node properties
+    formatted_node_props = [
+        f"{label} {{{_format_props(props)}}}"
+        for label, props in structured_schema["node_props"].items()
+    ]
+
+    # Format relationship properties
+    formatted_rel_props = [
+        f"{rel_type} {{{_format_props(props)}}}"
+        for rel_type, props in structured_schema["rel_props"].items()
+    ]
+
+    # Format relationships
+    formatted_rels = [
+        f"(:{element['start']})-[:{element['type']}]->(:{element['end']})"
+        for element in structured_schema["relationships"]
+    ]
+
+    return "\n".join(
+        [
+            "Node properties:",
+            "\n".join(formatted_node_props),
+            "Relationship properties:",
+            "\n".join(formatted_rel_props),
+            "The relationships:",
+            "\n".join(formatted_rels),
+        ]
+    )
+
+
+def get_structured_schema(driver: neo4j.Driver) -> dict[str, Any]:
+    """
+    Returns the structured schema of the graph.
+
+    Args:
+        driver (neo4j.Driver): Neo4j Python driver instance.
+
+    Returns:
+        dict[str, Any]: the graph schema information in a structured format.
+    """
     node_properties = [
         data["output"]
         for data in query_database(
@@ -97,6 +149,7 @@ def get_schema(
             driver, REL_PROPERTIES_QUERY, params={"EXCLUDED_LABELS": EXCLUDED_RELS}
         )
     ]
+
     relationships = [
         data["output"]
         for data in query_database(
@@ -106,35 +159,17 @@ def get_schema(
         )
     ]
 
-    # Format node properties
-    formatted_node_props = []
-    for element in node_properties:
-        props_str = ", ".join(
-            [f"{prop['property']}: {prop['type']}" for prop in element["properties"]]
-        )
-        formatted_node_props.append(f"{element['labels']} {{{props_str}}}")
-
-    # Format relationship properties
-    formatted_rel_props = []
-    for element in rel_properties:
-        props_str = ", ".join(
-            [f"{prop['property']}: {prop['type']}" for prop in element["properties"]]
-        )
-        formatted_rel_props.append(f"{element['type']} {{{props_str}}}")
-
-    # Format relationships
-    formatted_rels = [
-        f"(:{element['start']})-[:{element['type']}]->(:{element['end']})"
-        for element in relationships
-    ]
-
-    return "\n".join(
-        [
-            "Node properties:",
-            "\n".join(formatted_node_props),
-            "Relationship properties:",
-            "\n".join(formatted_rel_props),
-            "The relationships:",
-            "\n".join(formatted_rels),
-        ]
-    )
+    # Get constraints and indexes
+    try:
+        constraint = query_database(driver, "SHOW CONSTRAINTS")
+        index = query_database(driver, INDEX_QUERY)
+    except ClientError:
+        constraint = []
+        index = []
+
+    return {
+        "node_props": {el["labels"]: el["properties"] for el in node_properties},
+        "rel_props": {el["type"]: el["properties"] for el in rel_properties},
+        "relationships": relationships,
+        "metadata": {"constraint": constraint, "index": index},
+    }

tests/e2e/pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+markers = enterprise_only: marks tests as neo4j enterprise edition only

tests/e2e/test_schema_e2e.py

@@ -15,13 +15,12 @@
 
 import pytest
 from neo4j_genai.schema import (
-    BASE_ENTITY_LABEL,
-    EXCLUDED_LABELS,
-    EXCLUDED_RELS,
+    query_database,
+    get_structured_schema,
     NODE_PROPERTIES_QUERY,
+    BASE_ENTITY_LABEL,
     REL_PROPERTIES_QUERY,
     REL_QUERY,
-    query_database,
 )
 
 
@@ -78,41 +77,56 @@ def test_cypher_returns_correct_relationships(driver):
     )
 
 
-@pytest.mark.usefixtures("setup_neo4j_for_schema_query_with_excluded_labels")
-def test_filtering_labels_node_properties(driver):
-    node_properties = [
-        data["output"]
-        for data in query_database(
-            driver,
-            NODE_PROPERTIES_QUERY,
-            params={"EXCLUDED_LABELS": EXCLUDED_LABELS},
-        )
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query")
+def test_get_structured_schema_returns_correct_node_properties(driver):
+    result = get_structured_schema(driver)
+    assert result["node_props"]["LabelA"] == [
+        {"property": "property_a", "type": "STRING"}
     ]
 
-    assert node_properties == []
+
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query")
+def test_get_structured_schema_returns_correct_relationship_properties(driver):
+    result = get_structured_schema(driver)
+    assert result["rel_props"]["REL_TYPE"] == [
+        {"property": "rel_prop", "type": "STRING"}
+    ]
 
 
-@pytest.mark.usefixtures("setup_neo4j_for_schema_query_with_excluded_labels")
-def test_filtering_labels_relationship_properties(driver):
-    relationship_properties = [
-        data["output"]
-        for data in query_database(
-            driver, REL_PROPERTIES_QUERY, params={"EXCLUDED_LABELS": EXCLUDED_RELS}
-        )
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query")
+def test_get_structured_schema_returns_correct_relationships(driver):
+    result = get_structured_schema(driver)
+    assert sorted(result["relationships"], key=lambda x: x["end"]) == [
+        {"end": "LabelB", "start": "LabelA", "type": "REL_TYPE"},
+        {"end": "LabelC", "start": "LabelA", "type": "REL_TYPE"},
     ]
 
-    assert relationship_properties == []
 
+@pytest.mark.enterprise_only
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query")
+def test_get_structured_schema_returns_correct_constraints(driver):
+    query_database(driver, "DROP CONSTRAINT test_constraint IF EXISTS")
+    query_database(
+        driver,
+        "CREATE CONSTRAINT test_constraint IF NOT EXISTS FOR (n:LabelA) REQUIRE n.property_a IS NOT NULL",
+    )
 
-@pytest.mark.usefixtures("setup_neo4j_for_schema_query_with_excluded_labels")
-def test_filtering_labels_relationships(driver):
-    relationships = [
-        data["output"]
-        for data in query_database(
-            driver,
-            REL_QUERY,
-            params={"EXCLUDED_LABELS": EXCLUDED_LABELS + [BASE_ENTITY_LABEL]},
-        )
-    ]
+    result = get_structured_schema(driver)
+    assert result["metadata"]["constraint"][0].get("name") == "test_constraint"
+    assert result["metadata"]["constraint"][0].get("type") == "NODE_PROPERTY_EXISTENCE"
+    assert result["metadata"]["constraint"][0].get("entityType") == "NODE"
+    assert result["metadata"]["constraint"][0].get("labelsOrTypes") == ["LabelA"]
+    assert result["metadata"]["constraint"][0].get("properties") == ["property_a"]
+
+
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query")
+def test_get_structured_schema_returns_correct_indexes(driver):
+    query_database(driver, "DROP INDEX node_range_index IF EXISTS")
+    query_database(
+        driver, "CREATE INDEX node_range_index FOR (n:LabelA) ON (n.property_a)"
+    )
 
-    assert relationships == []
+    result = get_structured_schema(driver)
+    assert result["metadata"]["index"][0].get("label") == "LabelA"
+    assert result["metadata"]["index"][0].get("properties") == ["property_a"]
+    assert result["metadata"]["index"][0].get("type") == "RANGE"

tests/e2e/test_schema_filters_e2e.py

@@ -0,0 +1,65 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import pytest
+from neo4j_genai.schema import (
+    query_database,
+    NODE_PROPERTIES_QUERY,
+    BASE_ENTITY_LABEL,
+    EXCLUDED_LABELS,
+    EXCLUDED_RELS,
+    REL_PROPERTIES_QUERY,
+    REL_QUERY,
+)
+
+
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query_with_excluded_labels")
+def test_filtering_labels_node_properties(driver):
+    node_properties = [
+        data["output"]
+        for data in query_database(
+            driver,
+            NODE_PROPERTIES_QUERY,
+            params={"EXCLUDED_LABELS": EXCLUDED_LABELS},
+        )
+    ]
+
+    assert node_properties == []
+
+
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query_with_excluded_labels")
+def test_filtering_labels_relationship_properties(driver):
+    relationship_properties = [
+        data["output"]
+        for data in query_database(
+            driver, REL_PROPERTIES_QUERY, params={"EXCLUDED_LABELS": EXCLUDED_RELS}
+        )
+    ]
+
+    assert relationship_properties == []
+
+
+@pytest.mark.usefixtures("setup_neo4j_for_schema_query_with_excluded_labels")
+def test_filtering_labels_relationships(driver):
+    relationships = [
+        data["output"]
+        for data in query_database(
+            driver,
+            REL_QUERY,
+            params={"EXCLUDED_LABELS": EXCLUDED_LABELS + [BASE_ENTITY_LABEL]},
+        )
+    ]
+
+    assert relationships == []