Add support for go workspaces

cachi2/core/package_managers/gomod.py

@@ -1,3 +1,4 @@
+import json
 import logging
 import os
 import re
@@ -50,6 +51,8 @@
 GOMOD_INPUT_DOC = f"{GOMOD_DOC}#specifying-modules-to-process"
 VENDORING_DOC = f"{GOMOD_DOC}#vendoring"
 
+ModuleDict = dict[str, Any]
+
 
 class _ParsedModel(pydantic.BaseModel):
     """Attributes automatically get PascalCase aliases to make parsing Golang JSON easier.
@@ -440,6 +443,7 @@ def _create_modules_from_parsed_data(
     parsed_modules: Iterable[ParsedModule],
     modules_in_go_sum: frozenset[ModuleID],
     version_resolver: "ModuleVersionResolver",
+    go_work_path: Optional[RootedPath] = None,
 ) -> list[Module]:
     def _create_module(module: ParsedModule) -> Module:
         mod_id = _get_module_id(module)
@@ -452,7 +456,11 @@ def _create_module(module: ParsedModule) -> Module:
             real_path = name
 
             if mod_id not in modules_in_go_sum:
-                missing_hash_in_file = main_module_dir.subpath_from_root / "go.sum"
+                if go_work_path:
+                    missing_hash_in_file = go_work_path.subpath_from_root / "go.work.sum"
+                else:
+                    missing_hash_in_file = main_module_dir.subpath_from_root / "go.sum"
+
                 log.warning("checksum not found in %s: %s@%s", missing_hash_in_file, name, version)
         else:
             # module/name v1.0.0 => ./local/path
@@ -571,9 +579,11 @@ def fetch_gomod_source(request: Request) -> RequestOutput:
             log.info(f'Fetching the gomod dependencies at the "{subpath}" directory')
 
             main_module_dir = request.source_dir.join_within_root(subpath)
+            go_work_path = _get_go_work_path(main_module_dir)
+
             try:
                 resolve_result = _resolve_gomod(
-                    main_module_dir, request, Path(tmp_dir), version_resolver
+                    main_module_dir, request, Path(tmp_dir), version_resolver, go_work_path
                 )
             except PackageManagerError:
                 log.error("Failed to fetch gomod dependencies")
@@ -591,6 +601,7 @@ def fetch_gomod_source(request: Request) -> RequestOutput:
                     resolve_result.parsed_modules,
                     resolve_result.modules_in_go_sum,
                     version_resolver,
+                    go_work_path,
                 )
             )
 
@@ -798,7 +809,11 @@ def _setup_go_toolchain(go_mod_file: RootedPath) -> Go:
 
 
 def _resolve_gomod(
-    app_dir: RootedPath, request: Request, tmp_dir: Path, version_resolver: "ModuleVersionResolver"
+    app_dir: RootedPath,
+    request: Request,
+    tmp_dir: Path,
+    version_resolver: "ModuleVersionResolver",
+    go_work_path: Optional[RootedPath] = None,
 ) -> ResolvedGoModule:
     """
     Resolve and fetch gomod dependencies for given app source archive.
@@ -813,7 +828,6 @@ def _resolve_gomod(
     :raises PackageManagerError: if fetching dependencies fails
     """
     _protect_against_symlinks(app_dir)
-    modules_in_go_sum = _parse_go_sum(app_dir)
 
     config = get_config()
 
@@ -838,6 +852,11 @@ def _resolve_gomod(
 
     run_params = {"env": env, "cwd": app_dir}
 
+    if go_work_path:
+        modules_in_go_sum = _parse_go_sum_from_workspaces(go_work_path, go, run_params)
+    else:
+        modules_in_go_sum = _parse_go_sum(app_dir.join_within_root("go.sum"))
+
     # Vendor dependencies if the gomod-vendor flag is set
     flags = request.flags
     should_vendor, can_make_changes = _should_vendor_deps(
@@ -860,11 +879,8 @@ def _resolve_gomod(
         # Make Go ignore the vendor dir even if there is one
         go_list.extend(["-mod", "readonly"])
 
-    main_module_name = go([*go_list, "-m"], run_params).rstrip()
-    main_module = ParsedModule(
-        path=main_module_name,
-        version=version_resolver.get_golang_version(main_module_name, app_dir),
-        main=True,
+    main_module, workspace_modules = _parse_local_modules(
+        go, go_list, run_params, app_dir, version_resolver
     )
 
     def go_list_deps(pattern: Literal["./...", "all"]) -> Iterator[ParsedPackage]:
@@ -878,10 +894,10 @@ def go_list_deps(pattern: Literal["./...", "all"]) -> Iterator[ParsedPackage]:
         cmd = [*go_list, "-deps", "-json=ImportPath,Module,Standard,Deps", pattern]
         return map(ParsedPackage.model_validate, load_json_stream(go(cmd, run_params)))
 
-    package_modules = (
+    package_modules = [
         module for pkg in go_list_deps("all") if (module := pkg.module) and not module.main
-    )
-
+    ]
+    package_modules.extend(workspace_modules)
     all_modules = _deduplicate_resolved_modules(package_modules, downloaded_modules)
 
     log.info("Retrieving the list of packages")
@@ -892,13 +908,146 @@ def go_list_deps(pattern: Literal["./...", "all"]) -> Iterator[ParsedPackage]:
     return ResolvedGoModule(main_module, all_modules, all_packages, modules_in_go_sum)
 
 
-def _parse_go_sum(module_dir: RootedPath) -> frozenset[ModuleID]:
-    """Return the set of modules present in the go.sum file in the specified directory.
+def _parse_local_modules(
+    go: Go,
+    go_list: list[str],
+    run_params: dict[str, Any],
+    app_dir: RootedPath,
+    version_resolver: "ModuleVersionResolver",
+) -> tuple[ParsedModule, list[ParsedModule]]:
+    """
+    Identify and parse the main module and all workspace modules, if they exist.
+
+    :return: A tuple containing the main module and a list of workspaces
+    """
+    modules_json_stream = go([*go_list, "-m", "-json"], run_params).rstrip()
+    main_module_dict, workspace_dict_list = _process_modules_json_stream(
+        app_dir, modules_json_stream
+    )
+
+    main_module_path = main_module_dict["Path"]
+    main_module_version = version_resolver.get_golang_version(main_module_path, app_dir)
+
+    main_module = ParsedModule(
+        path=main_module_path,
+        version=main_module_version,
+        main=True,
+    )
+
+    workspace_modules = [
+        _parse_workspace_module(app_dir, workspace_dict, main_module_version)
+        for workspace_dict in workspace_dict_list
+    ]
+
+    return main_module, workspace_modules
+
+
+def _process_modules_json_stream(
+    app_dir: RootedPath, modules_json_stream: str
+) -> tuple[ModuleDict, list[ModuleDict]]:
+    """Process the json stream returned by "go list -m -json".
+
+    The stream will contain the module currently being processed, or a list of all workspaces in
+    case a go.work file is present in the repository.
+
+    :param app_dir: the path to the module directory
+    :param modules_json_stream: the json stream returned by "go list -m -json"
+    :return: A tuple containing the main module and a list of workspaces
+    """
+    module_list = []
+    main_module = None
+
+    for module in load_json_stream(modules_json_stream):
+        if module["Dir"] == str(app_dir):
+            main_module = module
+        else:
+            module_list.append(module)
+
+    # should never happen, since the main module will always be a part of the json stream
+    if not main_module:
+        raise RuntimeError('Failed to find the main module info in the "go list -m -json" output.')
+
+    return main_module, module_list
+
+
+def _parse_workspace_module(
+    app_dir: RootedPath, module: ModuleDict, main_module_version: str
+) -> ParsedModule:
+    """Create a ParsedModule from a listed workspace.
+
+    The replacement info returned will always be relative to the module currently being processed.
+    """
+    # We can't use pathlib since corresponding method PurePath.relative_to('foo', walk_up=True)
+    # is available only in Python 3.12 and later.
+    relative_dir = os.path.relpath(module["Dir"], app_dir)
+
+    # We need to prepend "./" to a workspace that is direct child of app_dir so it can be treated
+    # the same way as a locally replaced module when converting it into a Module.
+    if not relative_dir.startswith("."):
+        relative_dir = f"./{relative_dir}"
+
+    replaced_module = ParsedModule(path=relative_dir)
+
+    return ParsedModule(
+        path=module["Path"],
+        replace=replaced_module,
+    )
+
+
+def _get_go_work_path(app_dir: RootedPath) -> Optional[RootedPath]:
+    """Get the directory that contains the go.work file, if it exists."""
+    go = Go()
+    go_work_file = go(["env", "GOWORK"], {"cwd": app_dir}).rstrip()
+
+    if not go_work_file:
+        return None
+
+    go_work_path = Path(go_work_file).parent
+
+    # make sure that the path to go.work is within the request's root
+    return app_dir.join_within_root(go_work_path)
+
+
+def _parse_go_sum_from_workspaces(
+    go_work_path: RootedPath,
+    go: Go,
+    run_params: dict[str, Any],
+) -> frozenset[ModuleID]:
+    """Return the set of modules present in all go.sum files across the existing workspaces."""
+    go_sum_files = _get_go_sum_files(go_work_path, go, run_params)
+
+    modules: frozenset[ModuleID] = frozenset()
+
+    for go_sum_file in go_sum_files:
+        modules = modules | _parse_go_sum(go_sum_file)
+
+    return modules
+
+
+def _get_go_sum_files(
+    go_work_path: RootedPath,
+    go: Go,
+    run_params: dict[str, Any],
+) -> list[RootedPath]:
+    """Find all go.sum files present in the related workspaces."""
+    go_work_json = go(["work", "edit", "-json"], run_params).rstrip()
+    go_work = json.loads(go_work_json)
+
+    go_sums = [
+        go_work_path.join_within_root(f"{module['DiskPath']}/go.sum") for module in go_work["Use"]
+    ]
+
+    go_sums.append(go_work_path.join_within_root("go.work.sum"))
+
+    return go_sums
+
+
+def _parse_go_sum(go_sum: RootedPath) -> frozenset[ModuleID]:
+    """Return the set of modules present in the specified go.sum file.
 
     A module is considered present if the checksum for its .zip file is present. The go.mod file
     checksums are not relevant for our purposes.
     """
-    go_sum = module_dir.join_within_root("go.sum")
     if not go_sum.path.exists():
         return frozenset()
 

hack/mock-unittest-data/gomod.sh

@@ -10,6 +10,7 @@ banner-end
 mocked_data_dir=${1:-tests/unit/data/gomod-mocks}
 mkdir -p "$mocked_data_dir/non-vendored"
 mkdir -p "$mocked_data_dir/vendored"
+mkdir -p "$mocked_data_dir/workspaces"
 mocked_data_dir_abspath=$(realpath "$mocked_data_dir")
 
 tmpdir=$(dirname "$(mktemp --dry-run)")
@@ -25,6 +26,38 @@ $(
     cd "$tmpdir/gomod-pandemonium"
     export GOMODCACHE="$tmpdir/cachi2-mock-gomodcache"
 
+    git switch workspaces
+
+    echo "generating $mocked_data_dir/workspaces/go.sum"
+    cp go.sum "$mocked_data_dir_abspath/workspaces/go.sum"
+
+    echo "generating $mocked_data_dir/workspaces/go_list_modules.json"
+    go work edit -json > \
+        "$mocked_data_dir_abspath/workspaces/go_work.json"
+
+    echo "generating $mocked_data_dir/workspaces/go_list_modules.json"
+    go list -m -json > \
+        "$mocked_data_dir_abspath/workspaces/go_list_modules.json"
+
+    echo "generating $mocked_data_dir/workspaces/go_mod_download.json"
+    go mod download -json > \
+        "$mocked_data_dir_abspath/workspaces/go_mod_download.json"
+
+    echo "generating $mocked_data_dir/workspaces/go_list_deps_all.json"
+    go list -deps -json=ImportPath,Module,Standard,Deps all > \
+        "$mocked_data_dir_abspath/workspaces/go_list_deps_all.json"
+
+    echo "generating $mocked_data_dir/workspaces/go_list_deps_threedot.json"
+    go list -deps -json=ImportPath,Module,Standard,Deps ./... > \
+        "$mocked_data_dir_abspath/workspaces/go_list_deps_threedot.json"
+
+    git restore .
+    git switch main
+
+    echo "generating $mocked_data_dir/non-vendored/go_list_modules.json"
+    go list -m -json > \
+        "$mocked_data_dir_abspath/non-vendored/go_list_modules.json"
+
     echo "generating $mocked_data_dir/non-vendored/go_mod_download.json"
     go mod download -json > \
         "$mocked_data_dir_abspath/non-vendored/go_mod_download.json"
@@ -59,7 +92,7 @@ $(
 --------------------------------------------------------------------------------
 banner-end
 
-find "$mocked_data_dir/non-vendored" "$mocked_data_dir/vendored" -type f |
+find "$mocked_data_dir/non-vendored" "$mocked_data_dir/vendored" "$mocked_data_dir/workspaces" -type f |
     while read -r f; do
         sed "s|$tmpdir.cachi2-mock-gomodcache|{gomodcache_dir}|" --in-place "$f"
         sed "s|$tmpdir.gomod-pandemonium|{repo_dir}|" --in-place "$f"

tests/integration/test_data/gomod_workspaces/.build-config.yaml

@@ -0,0 +1,10 @@
+environment_variables:
+- name: GOCACHE
+  value: ${output_dir}/deps/gomod
+- name: GOMODCACHE
+  value: ${output_dir}/deps/gomod/pkg/mod
+- name: GOPATH
+  value: ${output_dir}/deps/gomod
+- name: GOPROXY
+  value: file://${GOMODCACHE}/cache/download
+project_files: []