Merge pull request #23 from microsoft/fix_node_and_relationship_references

Fix node and relationship references

Author: niclasko

flowquery-py/src/parsing/parser.py

@@ -474,6 +474,12 @@ def _parse_node(self) -> Optional[Node]:
             self._variables[identifier] = node
         elif identifier is not None:
             reference = self._variables.get(identifier)
+            # Resolve through Expression -> Reference -> Node (e.g., after WITH)
+            ref_child = reference.first_child() if isinstance(reference, Expression) else None
+            if isinstance(ref_child, Reference):
+                inner = ref_child.referred
+                if isinstance(inner, Node):
+                    reference = inner
             if reference is None or not isinstance(reference, Node):
                 raise ValueError(f"Undefined node reference: {identifier}")
             node = NodeReference(node, reference)
@@ -524,6 +530,13 @@ def _parse_relationship(self) -> Optional[Relationship]:
             self._variables[variable] = relationship
         elif variable is not None:
             reference = self._variables.get(variable)
+            # Resolve through Expression -> Reference -> Relationship (e.g., after WITH)
+            if isinstance(reference, Expression) and isinstance(
+                reference.first_child(), Reference
+            ):
+                inner = reference.first_child().referred
+                if isinstance(inner, Relationship):
+                    reference = inner
             if reference is None or not isinstance(reference, Relationship):
                 raise ValueError(f"Undefined relationship reference: {variable}")
             relationship = RelationshipReference(relationship, reference)

flowquery-py/tests/compute/test_runner.py

@@ -1669,4 +1669,44 @@ async def test_reserved_keywords_as_relationship_types_and_labels(self):
         await runner.run()
         results = runner.results
         assert len(results) == 1
-        assert results[0] == {"name1": "Node 1", "name2": "Node 2"}
+        assert results[0] == {"name1": "Node 1", "name2": "Node 2"}
+
+    @pytest.mark.asyncio
+    async def test_match_with_node_reference_passed_through_with(self):
+        """Test that node variables passed through WITH can be re-referenced in subsequent MATCH."""
+        await Runner("""
+            CREATE VIRTUAL (:WithRefUser) AS {
+                UNWIND [
+                    {id: 1, name: 'Alice', mail: 'alice@test.com', jobTitle: 'CEO'},
+                    {id: 2, name: 'Bob', mail: 'bob@test.com', jobTitle: 'VP'},
+                    {id: 3, name: 'Carol', mail: 'carol@test.com', jobTitle: 'VP'},
+                    {id: 4, name: 'Dave', mail: 'dave@test.com', jobTitle: 'Engineer'}
+                ] AS record
+                RETURN record.id AS id, record.name AS name, record.mail AS mail, record.jobTitle AS jobTitle
+            }
+        """).run()
+        await Runner("""
+            CREATE VIRTUAL (:WithRefUser)-[:MANAGES]-(:WithRefUser) AS {
+                UNWIND [
+                    {left_id: 1, right_id: 2},
+                    {left_id: 1, right_id: 3},
+                    {left_id: 2, right_id: 4}
+                ] AS record
+                RETURN record.left_id AS left_id, record.right_id AS right_id
+            }
+        """).run()
+        runner = Runner("""
+            MATCH (ceo:WithRefUser)-[:MANAGES]->(dr1:WithRefUser)
+            WHERE ceo.jobTitle = 'CEO'
+            WITH ceo, dr1
+            MATCH (ceo)-[:MANAGES]->(dr2:WithRefUser)
+            WHERE dr1.mail <> dr2.mail
+            RETURN ceo.name AS ceo, dr1.name AS dr1, dr2.name AS dr2
+        """)
+        await runner.run()
+        results = runner.results
+        # CEO (Alice) manages Bob and Carol. All distinct pairs:
+        # (Alice, Bob, Carol) and (Alice, Carol, Bob)
+        assert len(results) == 2
+        assert results[0] == {"ceo": "Alice", "dr1": "Bob", "dr2": "Carol"}
+        assert results[1] == {"ceo": "Alice", "dr1": "Carol", "dr2": "Bob"}

src/parsing/expressions/reference.ts

@@ -3,9 +3,9 @@ import Identifier from "./identifier";
 
 /**
  * Represents a reference to a previously defined variable or expression.
- * 
+ *
  * References point to values defined earlier in the query (e.g., in WITH or LOAD statements).
- * 
+ *
  * @example
  * ```typescript
  * const ref = new Reference("myVar", previousNode);
@@ -14,17 +14,20 @@ import Identifier from "./identifier";
  */
 class Reference extends Identifier {
     private _referred: ASTNode | undefined = undefined;
-    
+
     /**
      * Creates a new Reference to a variable.
-     * 
+     *
      * @param value - The identifier name
      * @param referred - The node this reference points to (optional)
      */
     constructor(value: string, referred: ASTNode | undefined = undefined) {
         super(value);
         this._referred = referred;
     }
+    public get referred(): ASTNode | undefined {
+        return this._referred;
+    }
     public set referred(node: ASTNode) {
         this._referred = node;
     }
@@ -39,4 +42,4 @@ class Reference extends Identifier {
     }
 }
 
-export default Reference;
+export default Reference;

src/parsing/parser.ts

@@ -438,8 +438,15 @@ class Parser extends BaseParser {
             node.identifier = identifier;
             this.variables.set(identifier, node);
         } else if (identifier !== null) {
-            const reference = this.variables.get(identifier);
-            if (reference === undefined || reference.constructor !== Node) {
+            let reference = this.variables.get(identifier);
+            // Resolve through Expression -> Reference -> Node (e.g., after WITH)
+            if (reference instanceof Expression && reference.firstChild() instanceof Reference) {
+                const inner = (reference.firstChild() as Reference).referred;
+                if (inner instanceof Node) {
+                    reference = inner;
+                }
+            }
+            if (reference === undefined || !(reference instanceof Node)) {
                 throw new Error(`Undefined node reference: ${identifier}`);
             }
             node = new NodeReference(node, reference);
@@ -629,8 +636,15 @@ class Parser extends BaseParser {
             relationship.identifier = variable;
             this.variables.set(variable, relationship);
         } else if (variable !== null) {
-            const reference = this.variables.get(variable);
-            if (reference === undefined || reference.constructor !== Relationship) {
+            let reference = this.variables.get(variable);
+            // Resolve through Expression -> Reference -> Relationship (e.g., after WITH)
+            if (reference instanceof Expression && reference.firstChild() instanceof Reference) {
+                const inner = (reference.firstChild() as Reference).referred;
+                if (inner instanceof Relationship) {
+                    reference = inner;
+                }
+            }
+            if (reference === undefined || !(reference instanceof Relationship)) {
                 throw new Error(`Undefined relationship reference: ${variable}`);
             }
             relationship = new RelationshipReference(relationship, reference);

tests/compute/runner.test.ts

@@ -1556,3 +1556,49 @@ test("Test reserved keywords as relationship types and labels", async () => {
     expect(results.length).toBe(1);
     expect(results[0]).toEqual({ name1: "Node 1", name2: "Node 2" });
 });
+
+test("Test match with node reference passed through WITH", async () => {
+    await new Runner(`
+        CREATE VIRTUAL (:User) AS {
+            UNWIND [
+                {id: 1, name: 'Alice', mail: 'alice@test.com', jobTitle: 'CEO'},
+                {id: 2, name: 'Bob', mail: 'bob@test.com', jobTitle: 'VP'},
+                {id: 3, name: 'Carol', mail: 'carol@test.com', jobTitle: 'VP'},
+                {id: 4, name: 'Dave', mail: 'dave@test.com', jobTitle: 'Engineer'}
+            ] AS record
+            RETURN record.id AS id, record.name AS name, record.mail AS mail, record.jobTitle AS jobTitle
+        }
+    `).run();
+    await new Runner(`
+        CREATE VIRTUAL (:User)-[:MANAGES]-(:User) AS {
+            UNWIND [
+                {left_id: 1, right_id: 2},
+                {left_id: 1, right_id: 3},
+                {left_id: 2, right_id: 4}
+            ] AS record
+            RETURN record.left_id AS left_id, record.right_id AS right_id
+        }
+    `).run();
+    // Equivalent to:
+    //   MATCH (ceo:User)-[:MANAGES]->(dr1:User)
+    //   WHERE ceo.jobTitle = 'CEO'
+    //   WITH ceo, dr1
+    //   MATCH (ceo)-[:MANAGES]->(dr2:User)
+    //   WHERE dr1.mail <> dr2.mail
+    //   RETURN ceo, dr1, dr2
+    const match = new Runner(`
+        MATCH (ceo:User)-[:MANAGES]->(dr1:User)
+        WHERE ceo.jobTitle = 'CEO'
+        WITH ceo, dr1
+        MATCH (ceo)-[:MANAGES]->(dr2:User)
+        WHERE dr1.mail <> dr2.mail
+        RETURN ceo.name AS ceo, dr1.name AS dr1, dr2.name AS dr2
+    `);
+    await match.run();
+    const results = match.results;
+    // CEO (Alice) manages Bob and Carol. All distinct pairs:
+    // (Alice, Bob, Carol) and (Alice, Carol, Bob)
+    expect(results.length).toBe(2);
+    expect(results[0]).toEqual({ ceo: "Alice", dr1: "Bob", dr2: "Carol" });
+    expect(results[1]).toEqual({ ceo: "Alice", dr1: "Carol", dr2: "Bob" });
+});