diff --git a/flowquery-py/src/parsing/parser.py b/flowquery-py/src/parsing/parser.py index 87e3290..b1c9a97 100644 --- a/flowquery-py/src/parsing/parser.py +++ b/flowquery-py/src/parsing/parser.py @@ -214,7 +214,7 @@ def _parse_with(self) -> Optional[With]: distinct = True self.set_next_token() self._expect_and_skip_whitespace_and_comments() - expressions = list(self._parse_expressions(AliasOption.REQUIRED)) + expressions = self._parse_expressions(AliasOption.REQUIRED) if len(expressions) == 0: raise ValueError("Expected expression") if distinct or any(expr.has_reducers() for expr in expressions): @@ -254,7 +254,7 @@ def _parse_return(self) -> Optional[Return]: distinct = True self.set_next_token() self._expect_and_skip_whitespace_and_comments() - expressions = list(self._parse_expressions(AliasOption.OPTIONAL)) + expressions = self._parse_expressions(AliasOption.OPTIONAL) if len(expressions) == 0: raise ValueError("Expected expression") if distinct or any(expr.has_reducers() for expr in expressions): @@ -353,7 +353,7 @@ def _parse_call(self) -> Optional[Call]: self._expect_previous_token_to_be_whitespace_or_comment() self.set_next_token() self._expect_and_skip_whitespace_and_comments() - expressions = list(self._parse_expressions(AliasOption.OPTIONAL)) + expressions = self._parse_expressions(AliasOption.OPTIONAL) if len(expressions) == 0: raise ValueError("Expected at least one expression") call.yielded = expressions # type: ignore[assignment] @@ -791,16 +791,30 @@ def _parse_order_by(self) -> Optional[OrderBy]: def _parse_expressions( self, alias_option: AliasOption = AliasOption.NOT_ALLOWED - ) -> Iterator[Expression]: + ) -> List[Expression]: + """Parse a comma-separated list of expressions with deferred variable + registration. Aliases set by earlier expressions in the same clause + won't shadow variables needed by later expressions + (e.g. ``RETURN a.x AS a, a.y AS b``).""" + parsed = list(self.__parse_expressions(alias_option)) + for expression, variable_name in parsed: + if variable_name is not None: + self._state.variables[variable_name] = expression + return [expression for expression, _ in parsed] + + def __parse_expressions( + self, alias_option: AliasOption + ) -> Iterator[Tuple[Expression, Optional[str]]]: while True: expression = self._parse_expression() if expression is not None: + variable_name: Optional[str] = None alias = self._parse_alias() if isinstance(expression.first_child(), Reference) and alias is None: reference = expression.first_child() assert isinstance(reference, Reference) # For type narrowing expression.set_alias(reference.identifier) - self._state.variables[reference.identifier] = expression + variable_name = reference.identifier elif (alias_option == AliasOption.REQUIRED and alias is None and not isinstance(expression.first_child(), Reference)): @@ -809,8 +823,8 @@ def _parse_expressions( raise ValueError("Alias not allowed") elif alias_option in (AliasOption.OPTIONAL, AliasOption.REQUIRED) and alias is not None: expression.set_alias(alias.get_alias()) - self._state.variables[alias.get_alias()] = expression - yield expression + variable_name = alias.get_alias() + yield expression, variable_name else: break self._skip_whitespace_and_comments() diff --git a/flowquery-py/tests/compute/test_runner.py b/flowquery-py/tests/compute/test_runner.py index fca339e..675f5e8 100644 --- a/flowquery-py/tests/compute/test_runner.py +++ b/flowquery-py/tests/compute/test_runner.py @@ -4339,4 +4339,59 @@ async def test_delete_virtual_node_leaves_other_nodes_intact(self): match = Runner("MATCH (n:PyKeepNode) RETURN n") await match.run() assert len(match.results) == 1 - assert match.results[0]["n"]["name"] == "Keep" \ No newline at end of file + assert match.results[0]["n"]["name"] == "Keep" + + @pytest.mark.asyncio + async def test_return_alias_shadowing_graph_variable(self): + """Test that RETURN alias doesn't shadow graph variable in same clause. + + When RETURN mentor.displayName AS mentor is followed by mentor.jobTitle, + the alias 'mentor' should not overwrite the graph node variable before + subsequent expressions are parsed. + """ + await Runner( + """ + CREATE VIRTUAL (:PyMentorUser) AS { + UNWIND [ + {id: 1, displayName: 'Alice Smith', jobTitle: 'Senior Engineer', department: 'Engineering'}, + {id: 2, displayName: 'Bob Jones', jobTitle: 'Staff Engineer', department: 'Engineering'}, + {id: 3, displayName: 'Chloe Dubois', jobTitle: 'Junior Engineer', department: 'Engineering'} + ] AS record + RETURN record.id AS id, record.displayName AS displayName, record.jobTitle AS jobTitle, record.department AS department + } + """ + ).run() + + await Runner( + """ + CREATE VIRTUAL (:PyMentorUser)-[:PY_MENTORS]-(:PyMentorUser) AS { + UNWIND [ + {left_id: 1, right_id: 3}, + {left_id: 2, right_id: 3} + ] AS record + RETURN record.left_id AS left_id, record.right_id AS right_id + } + """ + ).run() + + runner = Runner( + """ + MATCH (mentor:PyMentorUser)-[:PY_MENTORS]->(mentee:PyMentorUser) + WHERE mentee.displayName = "Chloe Dubois" + RETURN mentor.displayName AS mentor, mentor.jobTitle AS mentorJobTitle, mentor.department AS mentorDepartment + """ + ) + await runner.run() + results = runner.results + + assert len(results) == 2 + assert results[0] == { + "mentor": "Alice Smith", + "mentorJobTitle": "Senior Engineer", + "mentorDepartment": "Engineering", + } + assert results[1] == { + "mentor": "Bob Jones", + "mentorJobTitle": "Staff Engineer", + "mentorDepartment": "Engineering", + } \ No newline at end of file diff --git a/src/parsing/parser.ts b/src/parsing/parser.ts index 0266c26..b152d3d 100644 --- a/src/parsing/parser.ts +++ b/src/parsing/parser.ts @@ -224,7 +224,7 @@ class Parser extends BaseParser { this.setNextToken(); this.expectAndSkipWhitespaceAndComments(); } - const expressions = Array.from(this.parseExpressions(AliasOption.REQUIRED)); + const expressions = this.parseExpressions(AliasOption.REQUIRED); if (expressions.length === 0) { throw new Error("Expected expression"); } @@ -279,7 +279,7 @@ class Parser extends BaseParser { this.setNextToken(); this.expectAndSkipWhitespaceAndComments(); } - const expressions = Array.from(this.parseExpressions(AliasOption.OPTIONAL)); + const expressions = this.parseExpressions(AliasOption.OPTIONAL); if (expressions.length === 0) { throw new Error("Expected expression"); } @@ -393,7 +393,7 @@ class Parser extends BaseParser { this.expectPreviousTokenToBeWhitespaceOrComment(); this.setNextToken(); this.expectAndSkipWhitespaceAndComments(); - const expressions = Array.from(this.parseExpressions(AliasOption.OPTIONAL)); + const expressions = this.parseExpressions(AliasOption.OPTIONAL); if (expressions.length === 0) { throw new Error("Expected at least one expression"); } @@ -903,36 +903,52 @@ class Parser extends BaseParser { return new OrderBy(fields); } - private *parseExpressions( - alias_option: AliasOption = AliasOption.NOT_ALLOWED - ): IterableIterator { + /** + * Parses a comma-separated list of expressions with deferred variable + * registration. Aliases set by earlier expressions in the same clause + * won't shadow variables needed by later expressions + * (e.g. `RETURN a.x AS a, a.y AS b`). + */ + private parseExpressions(alias_option: AliasOption = AliasOption.NOT_ALLOWED): Expression[] { + const parsed = Array.from(this._parseExpressions(alias_option)); + for (const [expression, variableName] of parsed) { + if (variableName !== null) { + this._state.variables.set(variableName, expression); + } + } + return parsed.map(([expression]) => expression); + } + + private *_parseExpressions( + alias_option: AliasOption + ): IterableIterator<[Expression, string | null]> { while (true) { const expression: Expression | null = this.parseExpression(); - if (expression !== null) { - const alias = this.parseAlias(); - if (expression.firstChild() instanceof Reference && alias === null) { - const reference: Reference = expression.firstChild() as Reference; - expression.setAlias(reference.identifier); - this._state.variables.set(reference.identifier, expression); - } else if ( - alias_option === AliasOption.REQUIRED && - alias === null && - !(expression.firstChild() instanceof Reference) - ) { - throw new Error("Alias required"); - } else if (alias_option === AliasOption.NOT_ALLOWED && alias !== null) { - throw new Error("Alias not allowed"); - } else if ( - [AliasOption.OPTIONAL, AliasOption.REQUIRED].includes(alias_option) && - alias !== null - ) { - expression.setAlias(alias.getAlias()); - this._state.variables.set(alias.getAlias(), expression); - } - yield expression; - } else { + if (expression === null) { break; } + let variableName: string | null = null; + const alias = this.parseAlias(); + if (expression.firstChild() instanceof Reference && alias === null) { + const reference: Reference = expression.firstChild() as Reference; + expression.setAlias(reference.identifier); + variableName = reference.identifier; + } else if ( + alias_option === AliasOption.REQUIRED && + alias === null && + !(expression.firstChild() instanceof Reference) + ) { + throw new Error("Alias required"); + } else if (alias_option === AliasOption.NOT_ALLOWED && alias !== null) { + throw new Error("Alias not allowed"); + } else if ( + [AliasOption.OPTIONAL, AliasOption.REQUIRED].includes(alias_option) && + alias !== null + ) { + expression.setAlias(alias.getAlias()); + variableName = alias.getAlias(); + } + yield [expression, variableName]; this.skipWhitespaceAndComments(); if (!this.token.isComma()) { break; @@ -1360,7 +1376,7 @@ class Parser extends BaseParser { this.setNextToken(); this.expectAndSkipWhitespaceAndComments(); } - func.parameters = Array.from(this.parseExpressions(AliasOption.NOT_ALLOWED)); + func.parameters = this.parseExpressions(AliasOption.NOT_ALLOWED); this.skipWhitespaceAndComments(); if (!this.token.isRightParenthesis()) { throw new Error("Expected right parenthesis"); @@ -1394,7 +1410,7 @@ class Parser extends BaseParser { this.setNextToken(); // skip function name this.setNextToken(); // skip left parenthesis this.skipWhitespaceAndComments(); - asyncFunc.parameters = Array.from(this.parseExpressions(AliasOption.NOT_ALLOWED)); + asyncFunc.parameters = this.parseExpressions(AliasOption.NOT_ALLOWED); this.skipWhitespaceAndComments(); if (!this.token.isRightParenthesis()) { throw new Error("Expected right parenthesis"); diff --git a/tests/compute/runner.test.ts b/tests/compute/runner.test.ts index 5e94407..1b6ff46 100644 --- a/tests/compute/runner.test.ts +++ b/tests/compute/runner.test.ts @@ -4050,3 +4050,50 @@ test("Test delete virtual node leaves other nodes intact", async () => { expect(match.results.length).toBe(1); expect(match.results[0].n.name).toBe("Keep"); }); + +test("Test RETURN alias shadowing graph variable in same RETURN clause", async () => { + // Create User nodes with displayName, jobTitle, and department + await new Runner(` + CREATE VIRTUAL (:MentorUser) AS { + UNWIND [ + {id: 1, displayName: 'Alice Smith', jobTitle: 'Senior Engineer', department: 'Engineering'}, + {id: 2, displayName: 'Bob Jones', jobTitle: 'Staff Engineer', department: 'Engineering'}, + {id: 3, displayName: 'Chloe Dubois', jobTitle: 'Junior Engineer', department: 'Engineering'} + ] AS record + RETURN record.id AS id, record.displayName AS displayName, record.jobTitle AS jobTitle, record.department AS department + } + `).run(); + + // Create MENTORS relationships + await new Runner(` + CREATE VIRTUAL (:MentorUser)-[:MENTORS]-(:MentorUser) AS { + UNWIND [ + {left_id: 1, right_id: 3}, + {left_id: 2, right_id: 3} + ] AS record + RETURN record.left_id AS left_id, record.right_id AS right_id + } + `).run(); + + // This query aliases mentor.displayName AS mentor, which shadows the graph variable "mentor". + // Subsequent expressions mentor.jobTitle and mentor.department should still reference the graph node. + const runner = new Runner(` + MATCH (mentor:MentorUser)-[:MENTORS]->(mentee:MentorUser) + WHERE mentee.displayName = "Chloe Dubois" + RETURN mentor.displayName AS mentor, mentor.jobTitle AS mentorJobTitle, mentor.department AS mentorDepartment + `); + await runner.run(); + const results = runner.results; + + expect(results.length).toBe(2); + expect(results[0]).toEqual({ + mentor: "Alice Smith", + mentorJobTitle: "Senior Engineer", + mentorDepartment: "Engineering", + }); + expect(results[1]).toEqual({ + mentor: "Bob Jones", + mentorJobTitle: "Staff Engineer", + mentorDepartment: "Engineering", + }); +});