diff --git a/revup/github/endpoint.py b/revup/github/endpoint.py index 2cbb9e5..6a14b8f 100644 --- a/revup/github/endpoint.py +++ b/revup/github/endpoint.py @@ -10,7 +10,6 @@ from revup.types import RevupForgeException, RevupRequestException TRANSIENT_STATUSES = frozenset({500, 502, 503, 504}) -RETRYABLE_GRAPHQL_ERRORS = frozenset({"RESOURCE_LIMITS_EXCEEDED"}) class GitHubEndpoint: @@ -127,7 +126,12 @@ async def _graphql_once(self, query: str, **kwargs: Any) -> Any: return r async def graphql( - self, query: str, *, max_retries: int = 3, base_delay: float = 1.0, **kwargs: Any + self, + query: str, + *, + max_retries: int = 3, + base_delay: float = 1.0, + **kwargs: Any, ) -> Any: for attempt in range(max_retries): try: @@ -138,10 +142,3 @@ async def graphql( msg = "GitHub returned {}".format(e.status) if not await self._should_retry(attempt, max_retries, base_delay, msg): raise - except RevupForgeException as e: - retryable = set(e.types) & RETRYABLE_GRAPHQL_ERRORS - if not retryable: - raise - msg = "GitHub GraphQL error ({})".format(", ".join(retryable)) - if not await self._should_retry(attempt, max_retries, base_delay, msg): - raise diff --git a/revup/github/github.py b/revup/github/github.py index 1ce147d..c7468ae 100644 --- a/revup/github/github.py +++ b/revup/github/github.py @@ -1,5 +1,5 @@ import logging -from typing import Any, Dict, Iterable, List, Optional, Set, Tuple +from typing import Any, Dict, List, Optional, Set, Tuple from revup.forge import ( MAX_COMMENTS_TO_QUERY, @@ -10,151 +10,10 @@ PrUpdate, ) from revup.github.endpoint import GitHubEndpoint -from revup.types import RevupForgeException +from revup.github.graphql import GraphqlQuery, SingleQuery +from revup.types import RevupForgeException, RevupRequestException - -def _get_args_dict(args: List[Any], prefix: str) -> Dict[str, Any]: - return {f"{prefix}{n}": arg for n, arg in enumerate(args)} - - -def _get_args_declaration(args: Dict[str, Any], typ: str) -> List[str]: - return [f"${var}: {typ}" for var in args] - - -def _get_result_args(num: int, prefix: str) -> List[str]: - return [f"{prefix}{n}" for n in range(num)] - - -def _zip_and_flatten(l1: Iterable[str], l2: Iterable[str]) -> List[str]: - ret: List[str] = [] - iter1 = iter(l1) - iter2 = iter(l2) - while True: - try: - ret.append(next(iter1)) - ret.append(next(iter2)) - except StopIteration: - break - return ret - - -class Github(Forge): - def __init__( - self, - endpoint: GitHubEndpoint, - repo_info: ForgeRepoInfo, - fork_info: ForgeRepoInfo, - ): - self.endpoint = endpoint - self.repo_info = repo_info - self.fork_info = fork_info - - @property - def repo_owner(self) -> str: - return self.fork_info.owner - - @property - def repo_name(self) -> str: - return self.repo_info.name - - @property - def is_fork(self) -> bool: - return self.fork_info.owner != self.repo_info.owner - - async def close(self) -> None: - await self.endpoint.close() - - async def query_everything( - self, - head_refs: List[str], - user_ids: List[str], - labels: List[str], - teams: List[Tuple[str, str]], - ) -> Tuple[ - str, - List[Optional[PrInfo]], - Dict[str, str], - Dict[str, str], - Dict[str, str], - Dict[str, str], - Dict[str, Optional[Set[str]]], - ]: - head_refs_args = _get_args_dict(head_refs, "pr") - user_id_args = _get_args_dict(user_ids, "user") - label_args = _get_args_dict(labels, "label") - team_org_args = _get_args_dict([t[0] for t in teams], "team_org") - team_slug_args = _get_args_dict([t[1] for t in teams], "team_slug") - - prs_out = _get_result_args(len(head_refs), "pr_out") - user_id_out = _get_result_args(len(user_ids), "user_out") - label_out = _get_result_args(len(labels), "label_out") - team_out = _get_result_args(len(teams), "team_out") - - arg_str = ", ".join( - _get_args_declaration(head_refs_args, "String!") - + _get_args_declaration(user_id_args, "String!") - + _get_args_declaration(label_args, "String!") - + _get_args_declaration(team_org_args, "String!") - + _get_args_declaration(team_slug_args, "String!") - ) - - # NOTE: There are possible limitations here because we depend on PRs being - # returned in order of OPEN prs, followed by MERGED prs in the order that - # they merged. github doesn't offer these options and it is excessively - # expensive to always fetch multiple prs and order them on this side. For now - # we hope that the most relevant PR will have the most recent update time. - request_str = "".join( - len(head_refs) - * [ - "{}: pullRequests (headRefName: ${}, states: [OPEN, MERGED], first: 1, " - "orderBy: {{direction: DESC, field:UPDATED_AT}}) {{" - "...PrResult" - "}}," - ] - ) - request_str = request_str.format(*_zip_and_flatten(prs_out, head_refs_args.keys())) - - user_str = "".join( - len(user_ids) * ["{}: assignableUsers (query: ${}, first: 25) {{...UserResult}},"] - ) - user_str = user_str.format(*_zip_and_flatten(user_id_out, user_id_args.keys())) - - label_str = "".join(len(labels) * ["{}: label (name: ${}) {{...LabelResult}},"]) - label_str = label_str.format(*_zip_and_flatten(label_out, label_args.keys())) - - team_str = "" - for i in range(len(teams)): - team_str += ( - f"{team_out[i]}: organization(login: ${list(team_org_args.keys())[i]}) " - f"{{team(slug: ${list(team_slug_args.keys())[i]}) " - f"{{id, members(first: 100) {{nodes {{login}}, totalCount}}}}}}," - ) - - multi_query_str = f""" - query GetPrResults($owner: String!, $name: String!, {arg_str}) {{ - repository(name: $name, owner: $owner) {{ - id - {request_str}{user_str}{label_str} - }} - {team_str} - }}""" - if user_str: - multi_query_str += """ - fragment UserResult on UserConnection { - nodes { - login - id - } - totalCount - }""" - if label_str: - multi_query_str += """ - fragment LabelResult on Label { - id - name - }""" - if request_str: - multi_query_str += f""" +PR_FRAGMENT = f""" fragment PrResult on PullRequestConnection {{ nodes {{ id @@ -257,175 +116,390 @@ async def query_everything( totalCount }}""" - pr_result = await self.endpoint.graphql( - multi_query_str, - owner=self.repo_info.owner, - name=self.repo_info.name, - **head_refs_args, - **user_id_args, - **label_args, - **team_org_args, - **team_slug_args, +USER_FRAGMENT = """ + fragment UserResult on UserConnection { + nodes { + login + id + } + totalCount + }""" + +LABEL_FRAGMENT = """ + fragment LabelResult on Label { + id + name + }""" + + +def _is_resource_limit_error(e: Exception) -> bool: + if isinstance(e, RevupForgeException): + return "RESOURCE_LIMITS_EXCEEDED" in e.types + if isinstance(e, RevupRequestException): + return e.status in {502, 503} + return False + + +def _add_pr_queries(q: GraphqlQuery, head_refs: List[str]) -> List[SingleQuery]: + return [ + q.add( + prefix="pr", + scope="repo", + field_template=( + "{}: pullRequests (headRefName: {}, states: [OPEN, MERGED], first: 1, " + "orderBy: {{direction: DESC, field:UPDATED_AT}}) {{...PrResult}}," + ), + var_types=["String!"], + values=[ref], + fragment=PR_FRAGMENT, + ) + for ref in head_refs + ] + + +def _add_user_queries(q: GraphqlQuery, user_ids: List[str]) -> List[SingleQuery]: + return [ + q.add( + prefix="user", + scope="repo", + field_template="{}: assignableUsers (query: {}, first: 25) {{...UserResult}},", + var_types=["String!"], + values=[uid], + fragment=USER_FRAGMENT, + ) + for uid in user_ids + ] + + +def _add_label_queries(q: GraphqlQuery, labels: List[str]) -> List[SingleQuery]: + return [ + q.add( + prefix="label", + scope="repo", + field_template="{}: label (name: {}) {{...LabelResult}},", + var_types=["String!"], + values=[label], + fragment=LABEL_FRAGMENT, ) + for label in labels + ] + + +def _add_team_queries(q: GraphqlQuery, teams: List[Tuple[str, str]]) -> List[SingleQuery]: + return [ + q.add( + prefix="team", + scope="top", + field_template=( + "{}: organization(login: {}) " + "{{team(slug: {}) " + "{{id, members(first: 100) {{nodes {{login}}, totalCount}}}}}}," + ), + var_types=["String!", "String!"], + values=[org, slug], + ) + for org, slug in teams + ] + + +def _parse_prs( + queries: List[SingleQuery], result: Any, head_refs: List[str] +) -> List[Optional[PrInfo]]: + raw = [sq.extract(result) for sq in queries] + prs: List[Optional[PrInfo]] = [] + for i, branch_name in enumerate(head_refs): + this_node = raw[i] + if len(this_node["nodes"]) == 1: + this_node = this_node["nodes"][0] + pr_labels: Set[str] = set() + pr_label_ids: Set[str] = set() + reviewers: Set[str] = set() + reviewer_ids: Set[str] = set() + reviewer_teams: Set[str] = set() + reviewer_team_ids: Set[str] = set() + assignees: Set[str] = set() + assignee_ids: Set[str] = set() + for label in this_node["labels"]["nodes"]: + pr_labels.add(label["name"]) + pr_label_ids.add(label["id"]) + for revs in this_node["reviewRequests"]["nodes"]: + requested = revs["requestedReviewer"] + if not requested: + continue + elif "slug" in requested: + reviewer_teams.add(f"{requested['organization']['login']}/{requested['slug']}") + reviewer_team_ids.add(requested["id"]) + elif "login" in requested: + reviewers.add(requested["login"]) + reviewer_ids.add(requested["id"]) + for revs in this_node["latestReviews"]["nodes"]: + if not revs["viewerDidAuthor"] and "login" in revs["author"]: + reviewers.add(revs["author"]["login"]) + reviewer_ids.add(revs["author"]["id"]) + for user in this_node["assignees"]["nodes"]: + assignees.add(user["login"]) + assignee_ids.add(user["id"]) + + headRefOid = ( + this_node["headCommit"]["nodes"][0]["commit"]["oid"] + if this_node["headCommit"]["nodes"] + else None + ) + baseRefOid = ( + this_node["baseCommit"]["nodes"][0]["commit"]["parents"]["nodes"][0]["oid"] + if this_node["baseCommit"]["nodes"] + else None + ) - prs: List[Optional[PrInfo]] = [] - for i, branch_name in enumerate(head_refs): - this_node = pr_result["data"]["repository"][prs_out[i]] - if len(this_node["nodes"]) == 1: - this_node = this_node["nodes"][0] - pr_labels: Set[str] = set() - pr_label_ids: Set[str] = set() - reviewers: Set[str] = set() - reviewer_ids: Set[str] = set() - reviewer_teams: Set[str] = set() - reviewer_team_ids: Set[str] = set() - assignees: Set[str] = set() - assignee_ids: Set[str] = set() - for label in this_node["labels"]["nodes"]: - pr_labels.add(label["name"]) - pr_label_ids.add(label["id"]) - for revs in this_node["reviewRequests"]["nodes"]: - requested = revs["requestedReviewer"] - if not requested: - continue - elif "slug" in requested: - reviewer_teams.add( - f"{requested['organization']['login']}/{requested['slug']}" - ) - reviewer_team_ids.add(requested["id"]) - elif "login" in requested: - reviewers.add(requested["login"]) - reviewer_ids.add(requested["id"]) - for revs in this_node["latestReviews"]["nodes"]: - # Ignore self reviews and bot reviews (without a login) - if not revs["viewerDidAuthor"] and "login" in revs["author"]: - reviewers.add(revs["author"]["login"]) - reviewer_ids.add(revs["author"]["id"]) - for user in this_node["assignees"]["nodes"]: - assignees.add(user["login"]) - assignee_ids.add(user["id"]) - - # The plain headRef and baseRef fields return the latest commit id associated with - # that branch name which may be newer than the PR itself if it was merged. We want - # the ids of the commits actually last associated with the PR, which we query from - # the commit list. This can also mean they are None if the PR has 0 commits. - headRefOid = ( - this_node["headCommit"]["nodes"][0]["commit"]["oid"] - if this_node["headCommit"]["nodes"] - else None + comments = [] + for c in this_node["comments"]["nodes"]: + comments.append(PrComment(c["body"], c["id"])) + + removed_reviewers: Set[str] = set() + removed_reviewer_ids: Set[str] = set() + removed_assignees: Set[str] = set() + removed_assignee_ids: Set[str] = set() + for event in this_node["timelineItems"]["nodes"]: + rr = event.get("requestedReviewer") + if rr and "login" in rr and rr["login"] not in reviewers: + removed_reviewers.add(rr["login"]) + removed_reviewer_ids.add(rr["id"]) + assignee = event.get("assignee") + if assignee and "login" in assignee and assignee["login"] not in assignees: + removed_assignees.add(assignee["login"]) + removed_assignee_ids.add(assignee["id"]) + + prs.append( + PrInfo( + id=this_node["id"], + url=this_node["url"], + baseRef=this_node["baseRefName"], + headRef=branch_name, + baseRefOid=baseRefOid, + headRefOid=headRefOid, + body=this_node["body"], + title=this_node["title"], + reviewers=reviewers, + reviewer_ids=reviewer_ids, + reviewer_teams=reviewer_teams, + reviewer_team_ids=reviewer_team_ids, + assignees=assignees, + assignee_ids=assignee_ids, + labels=pr_labels, + label_ids=pr_label_ids, + removed_reviewers=removed_reviewers, + removed_reviewer_ids=removed_reviewer_ids, + removed_assignees=removed_assignees, + removed_assignee_ids=removed_assignee_ids, + is_draft=this_node["isDraft"], + state=this_node["state"], + comments=comments, ) - baseRefOid = ( - this_node["baseCommit"]["nodes"][0]["commit"]["parents"]["nodes"][0]["oid"] - if this_node["baseCommit"]["nodes"] - else None + ) + else: + prs.append(None) + return prs + + +def _parse_users( + queries: List[SingleQuery], result: Any, user_ids: List[str] +) -> Tuple[Dict[str, str], Dict[str, str]]: + raw = [sq.extract(result) for sq in queries] + names_to_ids: Dict[str, str] = {} + names_to_logins: Dict[str, str] = {} + for i, user_id in enumerate(user_ids): + this_node = raw[i] + if len(this_node["nodes"]) == 0: + logging.warning("No matching user found for {}".format(user_id)) + else: + if this_node["totalCount"] > len(this_node["nodes"]): + logging.warning( + "Too many matching users found for {}, try being more specific".format(user_id) ) - - comments = [] - for c in this_node["comments"]["nodes"]: - comments.append(PrComment(c["body"], c["id"])) - - removed_reviewers: Set[str] = set() - removed_reviewer_ids: Set[str] = set() - removed_assignees: Set[str] = set() - removed_assignee_ids: Set[str] = set() - for event in this_node["timelineItems"]["nodes"]: - rr = event.get("requestedReviewer") - if rr and "login" in rr and rr["login"] not in reviewers: - removed_reviewers.add(rr["login"]) - removed_reviewer_ids.add(rr["id"]) - assignee = event.get("assignee") - if assignee and "login" in assignee and assignee["login"] not in assignees: - removed_assignees.add(assignee["login"]) - removed_assignee_ids.add(assignee["id"]) - - prs.append( - PrInfo( - id=this_node["id"], - url=this_node["url"], - baseRef=this_node["baseRefName"], - headRef=branch_name, - baseRefOid=baseRefOid, - headRefOid=headRefOid, - body=this_node["body"], - title=this_node["title"], - reviewers=reviewers, - reviewer_ids=reviewer_ids, - reviewer_teams=reviewer_teams, - reviewer_team_ids=reviewer_team_ids, - assignees=assignees, - assignee_ids=assignee_ids, - labels=pr_labels, - label_ids=pr_label_ids, - removed_reviewers=removed_reviewers, - removed_reviewer_ids=removed_reviewer_ids, - removed_assignees=removed_assignees, - removed_assignee_ids=removed_assignee_ids, - is_draft=this_node["isDraft"], - state=this_node["state"], - comments=comments, + shortest_name = this_node["nodes"][0]["login"] + names_to_ids[user_id] = this_node["nodes"][0]["id"] + found_match = False + for user in this_node["nodes"]: + if len(user["login"]) <= len(shortest_name) and user["login"].startswith(user_id): + shortest_name = user["login"] + names_to_ids[user_id] = user["id"] + names_to_logins[user_id] = user["login"] + found_match = True + if not found_match: + logging.warning( + "Couldn't find a prefixed match for {}, going with {} instead".format( + user_id, shortest_name ) ) + return names_to_ids, names_to_logins + + +def _parse_labels(queries: List[SingleQuery], result: Any, labels: List[str]) -> Dict[str, str]: + raw = [sq.extract(result) for sq in queries] + labels_to_ids: Dict[str, str] = {} + for i, label in enumerate(labels): + this_node = raw[i] + if this_node is not None: + labels_to_ids[label] = this_node["id"] + else: + logging.warning("Couldn't find an existing label named {}".format(label)) + return labels_to_ids + + +def _parse_teams( + queries: List[SingleQuery], result: Any, teams: List[Tuple[str, str]] +) -> Tuple[Dict[str, str], Dict[str, Optional[Set[str]]]]: + raw = [sq.extract(result) for sq in queries] + teams_to_ids: Dict[str, str] = {} + teams_to_members: Dict[str, Optional[Set[str]]] = {} + for i, (org, slug) in enumerate(teams): + team_node = raw[i] + if team_node is not None and team_node["team"] is not None: + team_ref = f"{org}/{slug}" + teams_to_ids[team_ref] = team_node["team"]["id"] + members_node = team_node["team"]["members"] + member_logins = {m["login"] for m in members_node["nodes"]} + if members_node["totalCount"] > len(members_node["nodes"]): + teams_to_members[team_ref] = None else: - prs.append(None) - - names_to_ids: Dict[str, str] = {} - names_to_logins: Dict[str, str] = {} - for i, user_id in enumerate(user_ids): - this_node = pr_result["data"]["repository"][user_id_out[i]] - if len(this_node["nodes"]) == 0: - logging.warning("No matching user found for {}".format(user_id)) - else: - if this_node["totalCount"] > len(this_node["nodes"]): - logging.warning( - "Too many matching users found for {}, try being more specific".format( - user_id - ) - ) - shortest_name = this_node["nodes"][0]["login"] - names_to_ids[user_id] = this_node["nodes"][0]["id"] - found_match = False - for user in this_node["nodes"]: - if len(user["login"]) <= len(shortest_name) and user["login"].startswith( - user_id - ): - shortest_name = user["login"] - names_to_ids[user_id] = user["id"] - names_to_logins[user_id] = user["login"] - found_match = True - if not found_match: - logging.warning( - "Couldn't find a prefixed match for {}, going with {} instead".format( - user_id, shortest_name - ) - ) + teams_to_members[team_ref] = member_logins + else: + logging.warning("Couldn't find a team matching {}/{}".format(org, slug)) + return teams_to_ids, teams_to_members + + +async def _refresh_new_comment_ids(endpoint: GitHubEndpoint, prs: List[PrUpdate]) -> None: + """Re-query comments for PRs with new (id=None) comments to avoid duplicates on retry.""" + prs_with_new = [pr for pr in prs if any(c.id is None for c in pr.comments)] + if not prs_with_new: + return + + query = GraphqlQuery() + node_queries = [ + query.add( + prefix="node", + scope="top", + field_template=( + "{}: node(id: {}) {{ ... on PullRequest {{ comments(first: " + + str(MAX_COMMENTS_TO_QUERY) + + ") {{ nodes {{ body id }} }} }} }}," + ), + var_types=["ID!"], + values=[pr.id], + ) + for pr in prs_with_new + ] + query_str, variables = query.build() - labels_to_ids: Dict[str, str] = {} - for i, label in enumerate(labels): - this_node = pr_result["data"]["repository"][label_out[i]] - if this_node is not None: - labels_to_ids[label] = this_node["id"] - else: - logging.warning("Couldn't find an existing label named {}".format(label)) - - teams_to_ids: Dict[str, str] = {} - teams_to_members: Dict[str, Optional[Set[str]]] = {} - for i, (org, slug) in enumerate(teams): - team_node = pr_result["data"][team_out[i]] - if team_node is not None and team_node["team"] is not None: - team_ref = f"{org}/{slug}" - teams_to_ids[team_ref] = team_node["team"]["id"] - members_node = team_node["team"]["members"] - member_logins = {m["login"] for m in members_node["nodes"]} - if members_node["totalCount"] > len(members_node["nodes"]): - # Team has more members than we fetched; we can't check membership reliably. - teams_to_members[team_ref] = None - else: - teams_to_members[team_ref] = member_logins - else: - logging.warning("Couldn't find a team matching {}/{}".format(org, slug)) + result = await endpoint.graphql(query_str, max_retries=1, **variables) + + raw = [sq.extract(result) for sq in node_queries] + for pr, pr_data in zip(prs_with_new, raw): + existing = pr_data.get("comments", {}).get("nodes", []) if pr_data else [] + existing_by_body = {c["body"]: c["id"] for c in existing} + for comment in pr.comments: + if comment.id is None and comment.text in existing_by_body: + comment.id = existing_by_body[comment.text] + logging.info("Comment already posted on PR, converting to edit") + + +class Github(Forge): + def __init__( + self, + endpoint: GitHubEndpoint, + repo_info: ForgeRepoInfo, + fork_info: ForgeRepoInfo, + ): + self.endpoint = endpoint + self.repo_info = repo_info + self.fork_info = fork_info + + @property + def repo_owner(self) -> str: + return self.fork_info.owner + + @property + def repo_name(self) -> str: + return self.repo_info.name + + @property + def is_fork(self) -> bool: + return self.fork_info.owner != self.repo_info.owner + + async def close(self) -> None: + await self.endpoint.close() + + def _make_query_everything( + self, + head_refs: List[str], + user_ids: List[str], + labels: List[str], + teams: List[Tuple[str, str]], + ) -> Tuple[ + GraphqlQuery, + List[SingleQuery], + List[SingleQuery], + List[SingleQuery], + List[SingleQuery], + ]: + q = GraphqlQuery(name="GetEverything") + q.add_fixed_var("owner", "String!", self.repo_info.owner) + q.add_fixed_var("name", "String!", self.repo_info.name) + q.fixed_repo_fields = "id\n" + + pr_queries = _add_pr_queries(q, head_refs) + user_queries = _add_user_queries(q, user_ids) + label_queries = _add_label_queries(q, labels) + team_queries = _add_team_queries(q, teams) + + return q, pr_queries, user_queries, label_queries, team_queries + + async def _execute_with_backoff(self, q: GraphqlQuery, max_retries: int = 3) -> Any: + query_str, variables = q.build() + try: + return await self.endpoint.graphql(query_str, max_retries=max_retries, **variables) + except (RevupForgeException, RevupRequestException) as e: + if not _is_resource_limit_error(e) or q.total_items() <= 1: + raise + left, right = q.split() + logging.warning( + "Request too complex ({} items), splitting into {} + {}".format( + q.total_items(), left.total_items(), right.total_items() + ) + ) + left_result = await self._execute_with_backoff(left, max_retries=max_retries) + right_result = await self._execute_with_backoff(right, max_retries=max_retries) + return _merge_results(left_result, right_result) + + async def query_everything( + self, + head_refs: List[str], + user_ids: List[str], + labels: List[str], + teams: List[Tuple[str, str]], + ) -> Tuple[ + str, + List[Optional[PrInfo]], + Dict[str, str], + Dict[str, str], + Dict[str, str], + Dict[str, str], + Dict[str, Optional[Set[str]]], + ]: + q, pr_queries, user_queries, label_queries, team_queries = self._make_query_everything( + head_refs, user_ids, labels, teams + ) + + result = await self._execute_with_backoff(q) + + repo_id = result["data"]["repository"]["id"] + prs = _parse_prs(pr_queries, result, head_refs) + names_to_ids, names_to_logins = _parse_users(user_queries, result, user_ids) + labels_to_ids = _parse_labels(label_queries, result, labels) + teams_to_ids, teams_to_members = _parse_teams(team_queries, result, teams) return ( - pr_result["data"]["repository"]["id"], + repo_id, prs, names_to_ids, names_to_logins, @@ -453,39 +527,58 @@ async def create_pull_requests(self, repo_id: str, prs: List[PrInfo]) -> None: "draft": pr.is_draft, } ) - inputs_args = _get_args_dict(inputs, "pr") - prs_out = _get_result_args(len(inputs), "pr_out") - - arg_str = ", ".join(_get_args_declaration(inputs_args, "CreatePullRequestInput!")) - request_str = "".join( - len(inputs) - * [ - """ - {}: createPullRequest(input: ${}) {{ + q = GraphqlQuery(operation="mutation") + pr_queries = [ + q.add( + prefix="pr", + scope="mutation", + field_template=""" + {}: createPullRequest(input: {}) {{ pullRequest {{ id url }} - }},""" - ] - ) - request_str = request_str.format(*_zip_and_flatten(prs_out, inputs_args.keys())) - - mutation_str = f""" - mutation ({arg_str}) {{ - {request_str} - }}""" + }},""", + var_types=["CreatePullRequestInput!"], + values=[inp], + ) + for inp in inputs + ] - # Creating a pull request can fail if the branch is already merged. - pr_results = await self.endpoint.graphql(mutation_str, require_success=False, **inputs_args) + pr_results = await self._execute_with_backoff(q) + raw = [sq.extract(pr_results) for sq in pr_queries] for i, pr in enumerate(prs): - result = pr_results["data"][prs_out[i]]["pullRequest"] - if result is not None: - pr.id = result["id"] - pr.url = result["url"] + result_node = raw[i]["pullRequest"] + if result_node is not None: + pr.id = result_node["id"] + pr.url = result_node["url"] async def update_pull_requests(self, prs: List[PrUpdate]) -> None: + q = self._build_update_mutation(prs) + query_str, variables = q.build() + try: + await self.endpoint.graphql(query_str, max_retries=1, **variables) + except (RevupForgeException, RevupRequestException) as e: + if isinstance(e, RevupForgeException) and "timeout" in e.message: + logging.warning( + "Github update request timed out! Most likely this is a false alarm and changes" + " actually succeeded. You may want to rerun this command to verify." + ) + return + if not _is_resource_limit_error(e) or len(prs) <= 1: + raise + mid = len(prs) // 2 + logging.warning( + "Update request too complex ({} PRs), splitting into {} + {}".format( + len(prs), mid, len(prs) - mid + ) + ) + await _refresh_new_comment_ids(self.endpoint, prs) + await self.update_pull_requests(prs[:mid]) + await self.update_pull_requests(prs[mid:]) + + def _build_update_mutation(self, prs: List[PrUpdate]) -> GraphqlQuery: inputs = [] labels = [] reviewers = [] @@ -515,7 +608,6 @@ async def update_pull_requests(self, prs: List[PrUpdate]) -> None: "labelableId": pr.id, } ) - if pr.reviewer_ids or pr.reviewer_team_ids: reviewers.append( { @@ -534,7 +626,6 @@ async def update_pull_requests(self, prs: List[PrUpdate]) -> None: "assignableId": pr.id, } ) - if pr.is_draft is not None: if pr.is_draft: convert_to_draft.append( @@ -550,7 +641,6 @@ async def update_pull_requests(self, prs: List[PrUpdate]) -> None: "pullRequestId": pr.id, } ) - for c in pr.comments: if c.id: edit_comments.append( @@ -569,165 +659,39 @@ async def update_pull_requests(self, prs: List[PrUpdate]) -> None: } ) - inputs_args = _get_args_dict(inputs, "pr") - prs_out = _get_result_args(len(inputs), "pr_out") - - labels_args = _get_args_dict(labels, "label") - labels_out = _get_result_args(len(labels), "label_out") - - reviewers_args = _get_args_dict(reviewers, "rev") - reviewers_out = _get_result_args(len(reviewers), "rev_out") - - assignees_args = _get_args_dict(assignees, "asn") - assignees_out = _get_result_args(len(assignees), "asn_out") - - to_draft_args = _get_args_dict(convert_to_draft, "to_d") - to_draft_out = _get_result_args(len(convert_to_draft), "to_d_out") - - from_draft_args = _get_args_dict(convert_from_draft, "from_d") - from_draft_out = _get_result_args(len(convert_from_draft), "from_d_out") - - comments_args = _get_args_dict(comments, "com") - comments_out = _get_result_args(len(comments), "com_out") - - edit_comments_args = _get_args_dict(edit_comments, "edit_com") - edit_comments_out = _get_result_args(len(edit_comments), "edit_com_out") - - arg_str = ", ".join( - _get_args_declaration(inputs_args, "UpdatePullRequestInput!") - + _get_args_declaration(labels_args, "AddLabelsToLabelableInput!") - + _get_args_declaration(reviewers_args, "RequestReviewsInput!") - + _get_args_declaration(assignees_args, "AddAssigneesToAssignableInput!") - + _get_args_declaration(to_draft_args, "ConvertPullRequestToDraftInput!") - + _get_args_declaration(from_draft_args, "MarkPullRequestReadyForReviewInput!") - + _get_args_declaration(comments_args, "AddCommentInput!") - + _get_args_declaration(edit_comments_args, "UpdateIssueCommentInput!") - ) - - update_str = "".join( - len(inputs) - * [ - """ - {}: updatePullRequest(input: ${}) {{ - clientMutationId - }},""" - ] - ) - update_str = update_str.format(*_zip_and_flatten(prs_out, inputs_args.keys())) - - request_reviewers_str = "".join( - len(reviewers_args) - * [ - """ - {}: requestReviews(input: ${}) {{ - clientMutationId - }},""" - ] - ) - request_reviewers_str = request_reviewers_str.format( - *_zip_and_flatten(reviewers_out, reviewers_args.keys()) - ) - assignees_str = "".join( - len(assignees_args) - * [ - """ - {}: addAssigneesToAssignable(input: ${}) {{ + q = GraphqlQuery(operation="mutation") + + def add_all(prefix: str, mutation: str, var_type: str, items: List[Any]) -> None: + for inp in items: + q.add( + prefix=prefix, + scope="mutation", + field_template=""" + {}: """ + + mutation + + """(input: {}) {{ clientMutationId - }},""" - ] - ) - assignees_str = assignees_str.format( - *_zip_and_flatten(assignees_out, assignees_args.keys()) - ) - - add_labels_str = "".join( - len(labels_args) - * [ - """ - {}: addLabelsToLabelable(input: ${}) {{ - clientMutationId - }},""" - ] - ) - add_labels_str = add_labels_str.format(*_zip_and_flatten(labels_out, labels_args.keys())) - - to_draft_str = "".join( - len(convert_to_draft) - * [ - """ - {}: convertPullRequestToDraft(input: ${}) {{ - clientMutationId - }},""" - ] - ) - to_draft_str = to_draft_str.format(*_zip_and_flatten(to_draft_out, to_draft_args.keys())) - - from_draft_str = "".join( - len(convert_from_draft) - * [ - """ - {}: markPullRequestReadyForReview(input: ${}) {{ - clientMutationId - }},""" - ] - ) - from_draft_str = from_draft_str.format( - *_zip_and_flatten(from_draft_out, from_draft_args.keys()) - ) - - add_comments_str = "".join( - len(comments_args) - * [ - """ - {}: addComment(input: ${}) {{ - clientMutationId - }},""" - ] - ) - add_comments_str = add_comments_str.format( - *_zip_and_flatten(comments_out, comments_args.keys()) - ) + }},""", + var_types=[var_type], + values=[inp], + ) - edit_comments_str = "".join( - len(edit_comments_args) - * [ - """ - {}: updateIssueComment(input: ${}) {{ - clientMutationId - }},""" - ] + add_all("com", "addComment", "AddCommentInput!", comments) + add_all("pr", "updatePullRequest", "UpdatePullRequestInput!", inputs) + add_all("rev", "requestReviews", "RequestReviewsInput!", reviewers) + add_all("asn", "addAssigneesToAssignable", "AddAssigneesToAssignableInput!", assignees) + add_all("label", "addLabelsToLabelable", "AddLabelsToLabelableInput!", labels) + add_all( + "to_d", "convertPullRequestToDraft", "ConvertPullRequestToDraftInput!", convert_to_draft ) - edit_comments_str = edit_comments_str.format( - *_zip_and_flatten(edit_comments_out, edit_comments_args.keys()) + add_all( + "from_d", + "markPullRequestReadyForReview", + "MarkPullRequestReadyForReviewInput!", + convert_from_draft, ) - - # Add comment mutations first to ensure comments are at the top of the PR - mutation_str = f""" - mutation ({arg_str}) {{ - {add_comments_str}{update_str}{request_reviewers_str}{assignees_str}{add_labels_str}\ -{to_draft_str}{from_draft_str}{edit_comments_str} - }}""" - - try: - await self.endpoint.graphql( - mutation_str, - **comments_args, - **inputs_args, - **reviewers_args, - **assignees_args, - **labels_args, - **to_draft_args, - **from_draft_args, - **edit_comments_args, - ) - except RevupForgeException as e: - if "timeout" in e.message: - logging.warning( - "Github update request timed out! Most likely this is a false alarm and changes" - " actually succeeded. You may want to rerun this command to verify." - ) - else: - raise e + add_all("edit_com", "updateIssueComment", "UpdateIssueCommentInput!", edit_comments) + return q async def query_pr_by_number(self, owner: str, name: str, number: int) -> Tuple[str, str]: result = await self.endpoint.graphql( @@ -746,3 +710,17 @@ async def query_pr_by_number(self, owner: str, name: str, number: int) -> Tuple[ ) pr = result["data"]["repository"]["pullRequest"] return pr["headRefName"], pr["baseRefName"] + + +def _merge_results(left: Any, right: Any) -> Any: + """Merge two GraphQL result dicts by combining their data keys.""" + merged: Dict[str, Any] = {"data": {}} + if "data" in left: + merged["data"].update(left["data"]) + if "data" in right: + for key, val in right["data"].items(): + if key == "repository" and "repository" in merged["data"]: + merged["data"]["repository"].update(val) + else: + merged["data"][key] = val + return merged diff --git a/revup/github/graphql.py b/revup/github/graphql.py new file mode 100644 index 0000000..b1b280e --- /dev/null +++ b/revup/github/graphql.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Dict, List, Tuple + + +@dataclass +class SingleQuery: + """A single aliased field in a GraphQL operation. + + Each query owns its field_template, variable values and a per-prefix index + that fixes its alias. Because the alias is stable, a GraphqlQuery can freely + move a SingleQuery between split halves without result aliases colliding. + """ + + prefix: str + scope: str # "repo" | "top" | "mutation" + field_template: str + var_types: List[str] + values: List[Any] + fragment: str = "" + index: int = 0 + + @property + def alias(self) -> str: + return f"{self.prefix}_out{self.index}" + + def var_name(self, var_idx: int) -> str: + if len(self.var_types) == 1: + return f"{self.prefix}{self.index}" + return f"{self.prefix}{self.index}_{var_idx}" + + def render_field(self) -> str: + var_names = [f"${self.var_name(j)}" for j in range(len(self.var_types))] + return self.field_template.format(self.alias, *var_names) + + def render_declarations(self) -> List[str]: + return [f"${self.var_name(j)}: {vtype}" for j, vtype in enumerate(self.var_types)] + + def render_variables(self) -> Dict[str, Any]: + return {self.var_name(j): val for j, val in enumerate(self.values)} + + def extract(self, result: Any) -> Any: + if self.scope == "repo": + return result["data"]["repository"][self.alias] + return result["data"][self.alias] + + +class GraphqlQuery: + """Builds a GraphQL query/mutation from a flat list of SingleQuery objects. + + Queries of different types (prefixes) coexist in one list, so split() can + divide the list in half regardless of type, unlike a per-type batch which + could only be split within a single type. + """ + + def __init__(self, operation: str = "query", name: str = ""): + self.operation = operation + self.name = name + self.fixed_vars: List[Tuple[str, str, Any]] = [] + self.fixed_repo_fields: str = "" + self.queries: List[SingleQuery] = [] + # Per-prefix counter so each type's aliases are 0, 1, 2, ... and unique. + self._prefix_counts: Dict[str, int] = {} + + def add_fixed_var(self, name: str, gql_type: str, value: Any) -> None: + self.fixed_vars.append((name, gql_type, value)) + + def add( + self, + *, + prefix: str, + scope: str, + field_template: str, + var_types: List[str], + values: List[Any], + fragment: str = "", + ) -> SingleQuery: + assert len(values) == len(var_types) + idx = self._prefix_counts.get(prefix, 0) + self._prefix_counts[prefix] = idx + 1 + sq = SingleQuery( + prefix=prefix, + scope=scope, + field_template=field_template, + var_types=list(var_types), + values=list(values), + fragment=fragment, + index=idx, + ) + self.queries.append(sq) + return sq + + def total_items(self) -> int: + return len(self.queries) + + def build(self) -> Tuple[str, Dict[str, Any]]: + all_decls: List[str] = [] + variables: Dict[str, Any] = {} + + for name, gql_type, value in self.fixed_vars: + all_decls.append(f"${name}: {gql_type}") + variables[name] = value + + for sq in self.queries: + all_decls.extend(sq.render_declarations()) + variables.update(sq.render_variables()) + + decl_str = ", ".join(all_decls) + name_str = f" {self.name}" if self.name else "" + + repo_fields = self.fixed_repo_fields + top_fields = "" + mutation_fields = "" + for sq in self.queries: + rendered = sq.render_field() + if sq.scope == "repo": + repo_fields += rendered + elif sq.scope == "top": + top_fields += rendered + else: + mutation_fields += rendered + + if self.operation == "query": + body = "" + if repo_fields: + body += f""" + repository(name: $name, owner: $owner) {{ + {repo_fields} + }}""" + body += top_fields + query_str = f""" + {self.operation}{name_str} ({decl_str}) {{{body} + }}""" + else: + query_str = f""" + {self.operation}{name_str} ({decl_str}) {{ + {mutation_fields} + }}""" + + fragments = "" + seen: set = set() + for sq in self.queries: + if sq.fragment and sq.fragment not in seen: + fragments += sq.fragment + seen.add(sq.fragment) + query_str += fragments + + return query_str, variables + + def _empty_clone(self) -> GraphqlQuery: + clone = GraphqlQuery(operation=self.operation, name=self.name) + clone.fixed_vars = list(self.fixed_vars) + clone.fixed_repo_fields = self.fixed_repo_fields + return clone + + def split(self) -> Tuple[GraphqlQuery, GraphqlQuery]: + """Split the flat query list in half. + + Each half may be heterogeneous in query type. Aliases are baked into each + SingleQuery, so results from the two halves never collide when merged. + """ + # Round up so an odd count keeps the extra item on the left (and a lone + # item lands left with an empty right). + mid = (len(self.queries) + 1) // 2 + left = self._empty_clone() + right = self._empty_clone() + left.queries = self.queries[:mid] + right.queries = self.queries[mid:] + return left, right