kt: Add checkout command

Usage: kt checkout <kernel>

This will set up the working directory for a kernel.
It will create git worktrees for dist-git-tree and kernel-src-tree repos.
If used if the working directory for a kernel is created,
the two git worktrees are updated.

An --extra option was added if people need 2 working environments for the
same kernel at the same time.

Example without --extra:

$ kt checkout lts9-2
Output:
$KERNELS_DIR/lts-9.2 will be created with 2 git worktrees:
kernel-src-tree -- for branch {USER}/ciqlts9_2
kernel-dist-git -- for branch {USER}/ciqlts9_2

Example with --extra:

$ kt checkout lts9-2 --extra CVE_X
Output:
$KERNELS_DIR/lts-9.2_CVE_X will be created with 2 git worktrees:
kernel-src-tree -- for branch {USER}/ciqlts9_2_CVE_X
kernel-dist-git -- for branch {USER}/ciqlts9_2_CVE_X

TODO: Autocomplete suggests only the kernel in data/kernels.yaml, not the
kernel workspace at the moment. This will be improved later, since
kt checkout may be used to just update en existing kernel workspace.

Signed-off-by: Roxana Nicolescu <rnicolescu@ciq.com>

Author: roxanan1996

bin/kt

@@ -4,6 +4,7 @@ import logging
 
 import click
 
+from kt.commands.checkout.command import checkout
 from kt.commands.list_kernels.command import list_kernels
 from kt.commands.setup.command import setup
 
@@ -24,6 +25,7 @@ def main():
 
     cli.add_command(list_kernels)
     cli.add_command(setup)
+    cli.add_command(checkout)
     cli()
 
 

kt/KT.md

@@ -172,3 +172,58 @@ If config.base_path = ~/ciq, these will be created:
 
 If there's a repo that needs to be cloned relevant for any future command,
 this is when it should be cloned.
+
+### kt checkout
+Prepares the working directory for a kernel.
+It uses the KernelInfo dataclass created based on the kernels.yaml file.
+
+The working directory location is based on configuration:
+`<config.kernels_dir>/<kernel>/`
+
+2 git worktrees are created for a kernel:
+
+- kernel-dist-git
+
+- kernel-src-tree
+
+They will point out to their root sources. Check kt setup for more info.
+They should be located in <config.base_path>.
+The worktrees reference the remote <branch> from kernels.yaml.
+The local branch is `{<user>}/<branch>`.
+
+If `--change-dir` or `-c` option is used, it will also go to the working
+directory of the kernel.
+
+If `--cleanup` option is used, it will delete the worktree and the local branch
+before creating it from scratch again.
+
+#### Example:
+```
+$ kt checkout lts9_4
+```
+
+For this configuration
+```
+{
+    "base_path": "~/ciq",
+    "kernels_dir": "~/ciq/kernels",
+    "images_source_dir": "~/ciq/default_test_images",
+    "images_dir": "~/ciq/tmp/virt-images",
+    "ssh_key": "~/.ssh/id_ed25519_generic.pub",
+}
+```
+
+This is the working directory for this kernel:
+`~/ciq/kernels/lts9_4`.
+
+2 git worktrees are created:
+
+1. kernel-dist-git
+
+    This representes branch `{<user>}/lts9_4:origin/lts9_4`.
+    The source repo is ~/ciq/dist-git-tree-lts
+
+2. kernel-src-tree
+
+    This representes branch `{<user>}/ciqlts9_4:origin/ciqlts9_4`
+    The source repo is ~/ciq/kernel-src-tree

kt/commands/checkout/command.py

@@ -0,0 +1,72 @@
+import click
+
+from kt.commands.checkout.impl import main
+from kt.ktlib.shell_completion import ShellCompletion
+
+epilog = """
+Prepares the working directory for a kernel.
+It uses the KernelInfo dataclass created based on the kernels.yaml file.
+
+The working directory location is based on configuration:
+<config.kernels_dir>/<kernel>/
+
+2 git worktrees are created for a kernel:
+
+- dist-git-tree
+
+- kernel-src-tree
+
+They will point out to their root sources. Check kt setup for more info.
+They should be located in <config.base_path>.
+
+The worktrees reference the remote <branch> from kernels.yaml.
+The local branch is {<user>}/<branch>.
+
+Examples:
+
+\b
+$ kt checkout lts9_4
+\b
+$ kt checkout lts9_4 --cleanup
+\b
+$ kt checkout lts9_4 --cleanup -c
+\b
+$ kt checkout lts9_4 --cleanup -change-dir
+\b
+$ kt checkout lts9_4 -e CVE-2022-49909
+Will create folder lts9_4_CVE-2022-49909 instead of lts9_4.
+"""
+
+
+@click.command(epilog=epilog)
+@click.option(
+    "-c",
+    "--change-dir",
+    is_flag=True,
+    help="Change directory to the kernel directory",
+)
+@click.option(
+    "--override",
+    is_flag=True,
+    help="Delete existing worktree for the kernel and start again",
+)
+@click.option(
+    "--cleanup",
+    is_flag=True,
+    help="Delete existing worktree for the kernel",
+)
+@click.option(
+    "-e",
+    "--extra",
+    type=str,
+    help="Feature you'll be working on",
+)
+@click.argument("kernel", required=True, type=str, shell_complete=ShellCompletion.show_kernels)
+def checkout(kernel, change_dir, override, cleanup, extra):
+    main(
+        name=kernel,
+        change_dir=change_dir,
+        override=override,
+        cleanup=cleanup,
+        extra=extra,
+    )

kt/commands/checkout/impl.py

@@ -0,0 +1,33 @@
+import logging
+import os
+import shutil
+
+from kt.ktlib.config import Config
+from kt.ktlib.kernel_workspace import KernelWorkspace
+from kt.ktlib.kernels import KernelsInfo
+
+
+def main(name: str, change_dir: bool, cleanup: bool, override: bool, extra: str):
+    config = Config.load()
+    kernels = KernelsInfo.from_yaml(config=config).kernels
+    if name not in kernels:
+        raise ValueError(f"Invalid param: {name} does not exist")
+
+    kernel_info = kernels[name]
+    kernel_workspace = KernelWorkspace.load(name=name, config=config, kernel_info=kernel_info, extra=extra)
+    if cleanup:
+        kernel_workspace.cleanup()
+        return
+
+    if override:
+        kernel_workspace.cleanup()
+
+    kernel_workspace.setup()
+    if change_dir:
+        shell_exec = os.environ["SHELL"]
+        if not shutil.which(shell_exec):
+            raise Exception(f"No executable {shell_exec}; cannot start a new one")
+
+        logging.info(f"Changing directory to {kernel_workspace.folder}")
+        os.chdir(kernel_workspace.folder)
+        os.execve(shell_exec, [shell_exec], os.environ)

kt/ktlib/kernel_workspace.py

@@ -0,0 +1,181 @@
+import logging
+import os
+from dataclasses import dataclass
+
+from git import GitCommandError, Repo
+from pathlib3x import Path
+
+from kt.ktlib.config import Config
+from kt.ktlib.kernels import KernelInfo
+from kt.ktlib.util import Constants
+
+
+@dataclass
+class RepoWorktree:
+    source_root: Repo
+    folder: Path
+    remote: str
+    remote_branch: str
+    local_branch: str
+
+    @classmethod
+    def load_from_filepath(cls, folder: Path):
+        if not folder.exists():
+            raise RuntimeError(f"{folder} does not exist")
+
+        # check if folder is a git repo
+        repo = Repo(folder)
+
+        source_root = Path(repo.git.rev_parse("--path-format=absolute", "--git-common-dir")).parent
+        remote = repo.remotes.origin
+
+        remote_branch = repo.active_branch.tracking_branch()
+
+        local_branch = repo.active_branch.name
+
+        return cls(
+            source_root=source_root,
+            folder=folder,
+            remote=remote,
+            remote_branch=remote_branch,
+            local_branch=local_branch,
+        )
+
+    def setup(self):
+        """
+        First run: It will create the worktree
+        Second run: It will update the worktree
+        """
+
+        try:
+            remote_ref = f"{self.remote}/{self.remote_branch}"
+            self.source_root.git.worktree(
+                "add",
+                "--track",
+                "-b",
+                self.local_branch,
+                self.folder,
+                remote_ref,
+            )
+
+        except GitCommandError as e:
+            if "already exists" in e.stderr:
+                self.update()
+            else:
+                # Make sure the worktree is properly cleaned up
+                self.cleanup()
+                raise e
+
+    def update(self):
+        """
+        It will make sure the worktree is up-to-date with remote.
+        It assumes the worktree is already created and initialized.
+        """
+        logging.info("update")
+        repo = Repo(self.folder)
+
+        repo.remotes.origin.pull(rebase=True)
+
+    def cleanup(self):
+        # remove worktree, only if it exists
+        try:
+            self.source_root.git.worktree("remove", self.folder, "-f")
+        except GitCommandError as e:
+            if f"'{self.folder}' is not a working tree" not in e.stderr:
+                raise e
+
+        # remove local branch, only if it exists
+        try:
+            self.source_root.delete_head(self.local_branch, force=True)
+        except GitCommandError as e:
+            if f"branch '{self.local_branch}' not found" not in e.stderr:
+                raise e
+
+    def push(self, force: bool = False):
+        repo = Repo(self.folder)
+        origin = repo.remote(name=self.remote)
+        args = ["--set-upstream", origin.name, self.local_branch]
+        if force:
+            args.append("--force-with-lease")
+
+        repo.git.push(*args)
+
+
+@dataclass
+class KernelWorkspace:
+    folder: Path
+    dist_worktree: RepoWorktree
+    src_worktree: RepoWorktree
+
+    @classmethod
+    def load_from_filepath(cls, folder: Path):
+        if not folder.exists():
+            raise RuntimeError(f"{folder} does not exists")
+
+        ## Get the dist-git-tree path
+        dist_worktree_path = folder / Constants.DIST_TREE
+        dist_worktree = RepoWorktree.load_from_filepath(folder=dist_worktree_path)
+
+        src_worktree_path = folder / Constants.SRC_TREE
+        src_worktree = RepoWorktree.load_from_filepath(folder=src_worktree_path)
+
+        return cls(
+            folder=folder,
+            dist_worktree=dist_worktree,
+            src_worktree=src_worktree,
+        )
+
+    @classmethod
+    def load(cls, name: str, config: Config, kernel_info: KernelInfo, extra: str):
+        if extra:
+            name = name + "_" + extra
+
+        folder = config.kernels_dir / Path(name)
+        user = os.environ["USER"]
+        default_remote = "origin"
+
+        dist_folder = folder / Path(Constants.DIST_TREE)
+        dist_local_branch = f"{{{user}}}_{kernel_info.dist_git_branch}"
+        if extra:
+            dist_local_branch += f"_{extra}"
+
+        dist_worktree = RepoWorktree(
+            source_root=Repo(kernel_info.dist_git_root.folder),
+            folder=dist_folder,
+            remote=default_remote,
+            remote_branch=kernel_info.dist_git_branch,
+            local_branch=dist_local_branch,
+        )
+
+        src_folder = folder / Path(Constants.SRC_TREE)
+        src_local_branch = f"{{{user}}}_{kernel_info.src_tree_branch}"
+        if extra:
+            src_local_branch += f"_{extra}"
+
+        src_worktree = RepoWorktree(
+            source_root=Repo(kernel_info.src_tree_root.folder),
+            folder=src_folder,
+            remote=default_remote,
+            remote_branch=kernel_info.src_tree_branch,
+            local_branch=src_local_branch,
+        )
+
+        return cls(
+            folder=folder,
+            dist_worktree=dist_worktree,
+            src_worktree=src_worktree,
+        )
+
+    def setup(self):
+        # Make sure the folder is created
+        self.folder.mkdir(parents=True, exist_ok=True)
+
+        self.dist_worktree.setup()
+        self.src_worktree.setup()
+
+    def cleanup(self):
+        self.dist_worktree.cleanup()
+        self.src_worktree.cleanup()
+
+        # Remove working directory that includes the above git worktrees
+        self.folder.rmtree(ignore_errors=True)

kt/ktlib/shell_completion.py

@@ -0,0 +1,11 @@
+from kt.ktlib.config import Config
+from kt.ktlib.kernels import KernelsInfo
+
+
+class ShellCompletion:
+    @classmethod
+    def show_kernels(cls, ctx, param, incomplete):
+        config = Config.load()
+        kernels = KernelsInfo.from_yaml(config=config).kernels
+
+        return [kernel for kernel in kernels if kernel.startswith(incomplete)]

kt/ktlib/util.py

@@ -1,5 +1,7 @@
 class Constants:
     PRIVATE_REPOS_CONFIG_FILE = ".private_repos.yaml"
+    SRC_TREE = "kernel-src-tree"
+    DIST_TREE = "kernel-dist-git"
     DIST_GIT_ROOT = "dist_git_root"
     SRC_TREE_ROOT = "src_tree_root"
     COMMON_REPOS = "common_repos"