stacks: update build_stack_graph to use new endpoint (bug 1787516)

- use updated response from differential.revision.search endpoint
- translate `stackGraph` field to phids and edges
- update `tests.mocks` such that revision response includes new field

Author: zzzeid

landoapi/stacks.py

@@ -27,10 +27,6 @@ def build_stack_graph(
 ) -> Tuple[Set[str], Set[Tuple[str, str]]]:
     """Return a graph representation of a revision stack.
 
-    This function is expensive and can make up to approximately
-    n/2 calls to phabricator for a linear stack where n is the
-    number of revisions in the stack.
-
     Args:
         phab: A landoapi.phabricator.PhabricatorClient.
         revision_phid: String PHID of a revision in the stack.
@@ -42,35 +38,20 @@ def build_stack_graph(
         between two nodes. `child` and `parent` are also string
         PHIDs.
     """
-    phids = set()
-    new_phids = {revision_phid}
-    edges = []
-
-    # Repeatedly request all related edges, adding connected revisions
-    # each time until no new revisions are found.
-    while new_phids:
-        phids.update(new_phids)
-        edges = phab.call_conduit(
-            "edge.search",
-            types=["revision.parent", "revision.child"],
-            sourcePHIDs=[phid for phid in phids],
-            limit=10000,
-        )
-        edges = phab.expect(edges, "data")
-        new_phids = set()
-        for edge in edges:
-            new_phids.add(edge["sourcePHID"])
-            new_phids.add(edge["destinationPHID"])
-
-        new_phids = new_phids - phids
-
-    # Treat the stack like a commit DAG, we only care about edges going
-    # from child to parent. This is enough to represent the graph.
-    edges = {
-        (edge["sourcePHID"], edge["destinationPHID"])
-        for edge in edges
-        if edge["edgeType"] == "revision.parent"
-    }
+    result = phab.call_conduit(
+        "differential.revision.search",
+        constraints={"phids": [revision_phid]},
+        limit=1,
+    )
+
+    revision = PhabricatorClient.single(result, "data")
+    stack_graph = PhabricatorClient.expect(revision, "fields", "stackGraph")
+    phids = set(stack_graph.keys())
+    edges = set()
+
+    for source_node, destination_nodes in stack_graph.items():
+        for predecessor_node in destination_nodes:
+            edges.add((source_node, predecessor_node))
 
     return phids, edges
 

tests/mocks.py

@@ -4,6 +4,7 @@
 import hashlib
 import json
 from copy import deepcopy
+from collections import defaultdict
 
 from landoapi.phabricator import (
     PhabricatorAPIException,
@@ -96,6 +97,51 @@ def validate_change(change):
     return True
 
 
+def get_stack(_phid, phabdouble):
+    phids = set()
+    new_phids = {_phid}
+    edges = []
+
+    # Repeatedly request all related edges, adding connected revisions
+    # each time until no new revisions are found.
+    # NOTE: this was adapted from previous implementation of build_stack_graph.
+    while new_phids:
+        phids.update(new_phids)
+        edges = [
+            e
+            for e in phabdouble._edges
+            if e["sourcePHID"] in phids
+            and e["edgeType"] in ("revision.parent", "revision.child")
+        ]
+        new_phids = set()
+        for edge in edges:
+            new_phids.add(edge["sourcePHID"])
+            new_phids.add(edge["destinationPHID"])
+
+        new_phids = new_phids - phids
+
+    # Treat the stack like a commit DAG, we only care about edges going
+    # from child to parent. This is enough to represent the graph.
+    edges = {
+        (edge["sourcePHID"], edge["destinationPHID"])
+        for edge in edges
+        if edge["edgeType"] == "revision.parent"
+    }
+
+    stack_graph = defaultdict(list)
+    sources = [edge[0] for edge in edges]
+    for source, dest in edges:
+        # Check that destination phid has a corresponding source phid.
+        if dest not in sources:
+            # We are at a root node.
+            stack_graph[dest] = []
+        stack_graph[source].append(dest)
+    if not stack_graph:
+        # There is only one node, the root node.
+        stack_graph[_phid] = []
+    return dict(stack_graph)
+
+
 class PhabricatorDouble:
     """Phabricator test double.
 
@@ -898,6 +944,7 @@ def to_response(i):
                 "fields": {
                     "title": i["title"],
                     "authorPHID": i["authorPHID"],
+                    "stackGraph": i["stack_graph"],
                     "status": {
                         "value": i["status"].value,
                         "name": i["status"].output_name,
@@ -948,7 +995,10 @@ def to_response(i):
 
             return deepcopy(resp)
 
-        items = [r for r in self._revisions]
+        items = []
+        for r in self._revisions:
+            r["stack_graph"] = get_stack(r["phid"], self)
+            items.append(r)
 
         if constraints and "ids" in constraints:
             items = [i for i in items if i["id"] in constraints["ids"]]