SERVER-94173 Fix Bazel Headers (#26849)

GitOrigin-RevId: ad87c3b1f1469d3d6a029089d7b679dd364617b4

2 files changed, 240 insertions, 0 deletions

diff --git a/buildscripts/fix_headers.py b/buildscripts/fix_headers.py
new file mode 100644
index 00000000000..e3be5101700
--- /dev/null
+++ b/buildscripts/fix_headers.py
@@ -0,0 +1,222 @@
+"""Cleanup Bazel target headers
+1. Evaluate expression into a list of cc_library targets.
+2. Identify headers defined outside of the package directory.
+3. Lookup target that should claim a given header.
+4. If said target exists, check for cycles by modifying BUILD.bazel and building.
+5. Print report with targets and buildozer commands to fix each one.
+"""
+
+# TODO(SERVER-94780) Add buildozer dep to poetry
+import subprocess
+import json
+from typing import List, Optional, Dict, Annotated, Tuple
+import typer
+import sys
+import os
+import pprint
+
+# Get relative imports to work when the package is not installed on the PYTHONPATH.
+if __name__ == "__main__" and __package__ is None:
+    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import buildscripts.util.buildozer_utils as bd_utils
+
+CC_LIB_SUFFIX = "_with_debug"
+
+
+def move_header(
+    fix_target: str, header: str, new_dep: Optional[str] = None, add_header: bool = False
+) -> None:
+    bd_utils.bd_remove([fix_target], "hdrs", [header])
+    if new_dep:
+        bd_utils.bd_add([fix_target], "deps", [new_dep])
+        if add_header:
+            bd_utils.bd_add([new_dep], "hdrs", [header])
+
+
+def undo_header_move(
+    fix_target: str, header: str, new_dep: Optional[str] = None, remove_header: bool = False
+) -> None:
+    if new_dep:
+        bd_utils.bd_remove([fix_target], "deps", [new_dep])
+        if remove_header:
+            bd_utils.bd_remove([new_dep], "hdrs", [header])
+    bd_utils.bd_add([fix_target], "hdrs", [header])
+
+
+def useful_print(fixes: Dict) -> None:
+    for target, target_fixes in fixes.items():
+        print("-", target)
+        print("  Fixes:\n")
+        for header, commands in target_fixes["fixes"].items():
+            print(f"    -{header}:")
+            for cmd in commands:
+                print("      ", cmd)
+
+
+class HeaderFixer:
+    def __init__(self):
+        # TODO(SERVER-94781) Remove SCons dep
+        subprocess.run(
+            [
+                sys.executable,
+                "buildscripts/scons.py",
+                "--build-profile=opt",
+                "--bazel-includes-info=dummy",  # TODO Allow no library to be passed.
+                "--libdeps-linting=off",
+                "--ninja=disabled",
+                "$BUILD_ROOT/scons/$VARIANT_DIR/sconf_temp",
+            ]
+        )
+        with open(".bazel_include_info.json") as f:
+            bazel_include_info = json.load(f)
+        self.bazel_exec = bazel_include_info["bazel_exec"]
+        self.bazel_config = bazel_include_info["config"]
+
+    def _query(
+        self, query: str, config: bool = False, args: List[str] = []
+    ) -> subprocess.CompletedProcess:
+        query_cmd = "cquery" if config else "query"
+        p = subprocess.run(
+            [self.bazel_exec, query_cmd] + self.bazel_config + args + [query],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        return p
+
+    def _build(self, target: str) -> subprocess.CompletedProcess:
+        p = subprocess.run(
+            [self.bazel_exec, "build"] + self.bazel_config + [target],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        return p
+
+    def _fix_package(self, package: str):
+        pass
+
+    def _create_header_target(self):
+        pass
+
+    def _find_misplaced_headers(self, target: str) -> List[str]:
+        p = self._query(f"labels(hdrs,{target}{CC_LIB_SUFFIX})")
+        misplaced_headers = []
+        target_package = target.split(":")[0] + ":"
+        for line in p.stdout.splitlines():
+            if not line.startswith("//"):
+                continue
+            if "." not in line:
+                continue
+            # skip if local
+            if line.startswith(target_package):
+                continue
+            misplaced_headers.append(line.split(" ")[0])
+
+        return misplaced_headers
+
+    def _find_header_target(self, header: str) -> Tuple[Optional[str], bool]:
+        potential_target = header.split(".")[0]
+        p = self._query(f"attr(srcs,{potential_target}.cpp,//...)")
+        target = None
+        for line in p.stdout.splitlines():
+            line = line.split()[0]
+            if not line.startswith("//"):
+                continue
+            if line.endswith(CC_LIB_SUFFIX):
+                target = line[: -len(CC_LIB_SUFFIX)]
+
+        if not target:
+            return None, False
+
+        p = self._query(f"filter('{header}',labels(hdrs,{target}{CC_LIB_SUFFIX}))")
+        filter_res = [line for line in p.stdout.splitlines() if line.startswith("//")]
+        if filter_res == []:
+            return target, False
+        return target, True
+
+    def _get_build_file(self, target: str) -> Optional[str]:
+        p = self._query(f"buildfiles({target})")
+        for line in p.stdout.splitlines():
+            if line.startswith("//src"):
+                return line.strip()
+
+        return None
+
+    def _check_dep_exists(self, fix_target: str, dep: str) -> bool:
+        p = self._query(f"filter({dep}$,deps({fix_target}))")
+        for line in p.stdout.splitlines():
+            if line.startswith("//") and line.split()[0] == dep:
+                return True
+
+        return False
+
+    def _fix_target(self, target: str) -> Dict:
+        target_fixes = {"fixes": {}, "cycles": {}}
+        orphaned_headers = []
+        for hdr in self._find_misplaced_headers(target):
+            new_dep, has_header = self._find_header_target(hdr)
+            if not new_dep:
+                orphaned_headers.append(hdr)
+                continue
+
+            buildozer_cmds = [f"buildozer 'remove hdrs {hdr}' {target}"]
+            if not has_header:
+                buildozer_cmds += [f"buildozer 'add hdrs {hdr}' {new_dep}"]
+            if self._check_dep_exists(target, new_dep):
+                new_dep = None
+            else:
+                buildozer_cmds += [f"buildozer 'add deps {new_dep}' {target}"]
+            move_header(target, hdr, new_dep, has_header)
+            p = self._build(target)
+            if p.returncode == 0:
+                target_fixes["fixes"][hdr] = buildozer_cmds
+            elif p.returncode == 1 and "cycle in dependency graph" in p.stdout:
+                target_fixes["cycles"][hdr] = buildozer_cmds
+            else:
+                print("Unexpected bazel failure.")
+            undo_header_move(target, hdr, new_dep, has_header)
+
+        print(f"Orphaned headers for {target}")
+        print("\n".join(orphaned_headers))
+        return target_fixes
+
+    def _evaluate_target_expression(self, target_exp: str) -> List[str]:
+        p = self._query(
+            f"filter('.*{CC_LIB_SUFFIX}$',kind(cc_library,deps({target_exp}, 1)))",
+            ["--noimplicit_deps"],
+        )
+        return [
+            line.split()[0][: -len(CC_LIB_SUFFIX)]
+            for line in p.stdout.splitlines()
+            if line.startswith("//")
+        ]
+
+    def fix_targets(self, target_exp: str) -> Dict:
+        fixes = {}
+        for target in self._evaluate_target_expression(target_exp):
+            fixes[target] = self._fix_target(target)
+
+        return fixes
+
+
+def main(
+    target_exp: Annotated[str, typer.Argument()],
+    output_file: Annotated[str, typer.Option()] = "",
+    copy_format: Annotated[bool, typer.Option()] = False,
+):
+    hf = HeaderFixer()
+    fixes = hf.fix_targets(target_exp)
+    json_output = pprint.pformat(json.dumps(fixes), compact=False).replace("'", '"')
+    if output_file:
+        with open(output_file, "w") as f:
+            print(json_output, filename, file=f)
+    elif copy_format:
+        useful_print(fixes)
+    else:
+        print(json_output)
+
+
+if __name__ == "__main__":
+    typer.run(main)
diff --git a/buildscripts/util/buildozer_utils.py b/buildscripts/util/buildozer_utils.py
new file mode 100644
index 00000000000..fd11175bf87
--- /dev/null
+++ b/buildscripts/util/buildozer_utils.py
@@ -0,0 +1,18 @@
+from typing import List
+import subprocess
+
+
+def _bd_command(cmd: str, labels: List[str]):
+    return subprocess.run(
+        ["buildozer"] + [cmd],
+        capture_output=True,
+        text=True,
+    )
+
+
+def bd_add(labels: List[str], attr: str, values: List[str]) -> None:
+    _bd_command(f'add {attr} {"".join(values)}', labels)
+
+
+def bd_remove(labels: List[str], attr: str, values: List[str]) -> None:
+    _bd_command(f'remove {attr} {"".join(values)}', labels)