Merge pull request #37 from microsoft/fixing_bugs_and_making_improvements

Fixing bugs and making improvements

Author: niclasko

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()

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"
+        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",
+        }

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<Expression> {
+    /**
+     * 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");

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",
+    });
+});