Properties types labels loading inspect

CHANGELOG.md

@@ -1,10 +1,10 @@
 # Changelog - [0.5.0] - yyyy-MM-dd
 
 ### New Features
-
+1. Support projecting specific relationship types for heterogeneous graphs.
 
 ### Bug Fixes
-
+1. Fix bug where get_relationship_properties, get_node_properties, and get_node_labels might not returning entire answer set
 
 ### Other Changes
 

mcp_server/src/mcp_server_neo4j_gds/gds.py

@@ -152,7 +152,10 @@ def projected_graph(gds, node_labels=None, relationship_types=None, undirected=F
 def get_node_labels(gds: GraphDataScience):
     query = """
                 MATCH (n)
-                RETURN DISTINCT labels(n) as labels
+                WITH DISTINCT labels(n) AS labels
+                UNWIND labels AS label
+                WITH DISTINCT label
+                RETURN COLLECT(label) AS labels
             """
     df = gds.run_cypher(query)
     if df.empty:
@@ -169,9 +172,14 @@ def get_node_properties_keys(gds: GraphDataScience, node_labels=None):
     if node_labels is None:
         node_labels = []
     nodelabels_query = create_node_cypher_match_query(node_labels)
-    query = (
-        nodelabels_query + """RETURN DISTINCT keys(properties(n)) AS properties_keys"""
-    )
+    property_extractor = """
+                            WITH keys(properties(n)) AS prop_keys_list
+                            UNWIND prop_keys_list AS prop_keys
+                            WITH DISTINCT prop_keys
+                            RETURN COLLECT(prop_keys) AS properties_keys
+                        """
+
+    query = nodelabels_query + property_extractor
     df = gds.run_cypher(query)
     if df.empty:
         return []
@@ -182,7 +190,13 @@ def get_relationship_properties_keys(gds: GraphDataScience, relationshipTypes=No
     if relationshipTypes is None:
         relationshipTypes = []
     rel_query = create_relationship_cypher_match_query([], relationshipTypes)
-    query = rel_query + " RETURN DISTINCT keys(properties(r)) AS properties_keys"
+    property_extractor = """
+                            WITH keys(properties(r)) AS prop_keys_list
+                            UNWIND prop_keys_list AS prop_keys
+                            WITH DISTINCT prop_keys
+                            RETURN COLLECT(prop_keys) AS properties_keys
+                        """
+    query = rel_query + property_extractor
 
     df = gds.run_cypher(query)
     if df.empty:
@@ -193,10 +207,12 @@ def get_relationship_properties_keys(gds: GraphDataScience, relationshipTypes=No
 def get_relationship_types(gds: GraphDataScience, node_labels=None):
     if node_labels is None:
         node_labels = []
-    query = (
-        create_relationship_cypher_match_query(node_labels, [])
-        + " RETURN DISTINCT TYPE(r) as relationship_types"
-    )
+    type_extractor = """
+                     WITH type(r) AS type
+                     WITH DISTINCT type
+                     RETURN COLLECT(type) AS relationship_types
+                     """
+    query = create_relationship_cypher_match_query(node_labels, []) + type_extractor
 
     df = gds.run_cypher(query)
     if df.empty:

mcp_server/tests/test_projection.py → mcp_server/tests/test_gds_functions.py

@@ -331,3 +331,56 @@ def test_projection_with_labels_and_types(neo4j_container):
 
     list_result = gds.graph.list()
     assert len(list_result) == 0
+
+
+@pytest.mark.asyncio
+def test_get_labels_and_types_and_properties(neo4j_container):
+    """Import test data into Neo4j."""
+    # Set environment variables for the import script
+    os.environ["NEO4J_URI"] = neo4j_container
+    os.environ["NEO4J_USERNAME"] = NEO4J_USER
+    os.environ["NEO4J_PASSWORD"] = NEO4J_PASSWORD
+
+    driver = GraphDatabase.driver(neo4j_container, auth=(NEO4J_USER, NEO4J_PASSWORD))
+    existing_count2 = -2
+    gds = GraphDataScience(driver)
+    with driver.session() as session:
+        session.run(
+            "CREATE (:Foo{prop1:1, prop3:5})-[:R1{relprop1:2}]->(:Bar:SpareBar), (:Bar)-[:R2{relprop2:2}]->(:Bar{prop2:2})"
+        )
+
+    from mcp_server.src.mcp_server_neo4j_gds.gds import (
+        get_node_labels,
+        get_relationship_types,
+        get_relationship_properties_keys,
+        get_node_properties_keys,
+    )
+
+    node_labels = get_node_labels(gds)
+    rel_types = get_relationship_types(gds)
+    rel_props = get_relationship_properties_keys(gds)
+    node_props = get_node_properties_keys(gds)
+
+    with driver.session() as session:
+        session.run("MATCH (n:Foo)  DETACH DELETE n")
+        session.run("MATCH (n:Bar)  DETACH DELETE n")
+
+        res = session.run(
+            "MATCH (n) WHERE 'Foo' IN labels(n) OR 'Bar' IN labels(n) OR 'SpareBar' IN labels(n) RETURN count(n) as count"
+        )
+        existing_count2 = res.single()["count"]
+
+    driver.close()
+    # assertions at the end to ensure failures do not affect other tests
+    assert existing_count2 == 0
+
+    assert "Foo" in node_labels
+    assert "Bar" in node_labels
+    assert "SpareBar" in node_labels
+    assert "R1" in rel_types
+    assert "R2" in rel_types
+    assert "relprop1" in rel_props
+    assert "relprop2" in rel_props
+    assert "prop1" in node_props
+    assert "prop2" in node_props
+    assert "prop3" in node_props