tower-of-babel-public/tasks/operations/git.py

from pyinfra.operations import apk, files, server, git, systemd, python
from pyinfra import host

from pyinfra.facts.files import Directory
from pyinfra.facts.git import GitBranch

from pyinfra.api import deploy, operation, FactBase
from pyinfra.operations.util.files import chown, unix_path_join


class CoolerGitBranch(FactBase):
    requires_command = "git"

    @staticmethod
    def command(repo):
        # TODO should inject _sudo / _sudo_user if user is provided in repo()
        return "! test -d {0} || (cd {0} && git rev-parse --abbrev-ref HEAD)".format(
            repo
        )


class GitFetch(FactBase):
    def command(self, repo: str):
        return f"git -C {repo} fetch"

    def process(self, output):
        return output


class GitRevListComparison(FactBase):
    def command(self, repo: str, branch: str):
        return f"git -C {repo} rev-list HEAD..origin/{branch} | wc -l"

    def process(self, output):
        return output


class RawCommandOutput(FactBase):
    """
    Returns the raw output of a command.
    """

    def command(self, command):
        return command

    def process(self, output):
        return "\n".join(output)  # re-join and return the output lines


@operation(
    pipeline_facts={
        "git_branch": "target",
    }
)
def repo(
    src,
    dest,
    branch=None,
    pull=True,
    rebase=False,
    user=None,
    group=None,
    ssh_keyscan=False,
    update_submodules=False,
    recursive_submodules=False,
):
    """
    Clone/pull git repositories.

    + src: the git source URL
    + dest: directory to clone to
    + branch: branch to pull/checkout
    + pull: pull any changes for the branch
    + rebase: when pulling, use ``--rebase``
    + user: chown files to this user after
    + group: chown files to this group after
    + ssh_keyscan: keyscan the remote host if not in known_hosts before clone/pull
    + update_submodules: update any git submodules
    + recursive_submodules: update git submodules recursively

    Example:

    .. code:: python

        git.repo(
            name='Clone repo',
            src='https://github.com/Fizzadar/pyinfra.git',
            dest='/usr/local/src/pyinfra',
        )
    """

    # Ensure our target directory exists
    yield from files.directory(dest)

    if ssh_keyscan:
        raise NotImplementedError("TODO copypaste ssh_keyscan code")

    # Store git commands for directory prefix
    git_commands = []
    git_dir = unix_path_join(dest, ".git")
    is_repo = host.get_fact(Directory, path=git_dir)

    # Cloning new repo?
    if not is_repo:
        if branch:
            git_commands.append("clone {0} --branch {1} .".format(src, branch))
        else:
            git_commands.append("clone {0} .".format(src))

        host.create_fact(GitBranch, kwargs={"repo": dest}, data=branch)
        host.create_fact(CoolerGitBranch, kwargs={"repo": dest}, data=branch)
        host.create_fact(
            Directory,
            kwargs={"path": git_dir},
            data={"user": user, "group": group},
        )

    # Ensuring existing repo
    else:
        current_branch = host.get_fact(CoolerGitBranch, repo=dest)

        # always fetch upstream branches (that way we can compare if the latest
        # commit has changed, and then we don't need to execute anything!)
        host.get_fact(GitFetch, repo=dest)
        stdout = host.get_fact(GitRevListComparison, repo=dest, branch=branch)
        repository_has_updates = stdout[0] != "0"

        # since we immediately always fetch, we will always be modifying the
        # .git folder, and that folder MUST be owned by the correct user afterwards.
        if user or group:
            chown_command = chown(dest, user, group, recursive=True)
            host.get_fact(RawCommandOutput, command=chown_command.get_masked_value())

        if branch and current_branch != branch:
            git_commands.append("checkout {0}".format(branch))
            host.create_fact(GitBranch, kwargs={"repo": dest}, data=branch)
            host.create_fact(CoolerGitBranch, kwargs={"repo": dest}, data=branch)
            repository_has_updates = True

        if repository_has_updates and pull:
            if rebase:
                git_commands.append("pull --rebase")
            else:
                git_commands.append("pull")

    if update_submodules:
        if recursive_submodules:
            git_commands.append("submodule update --init --recursive")
        else:
            git_commands.append("submodule update --init")

    # Attach prefixes for directory
    command_prefix = "cd {0} && git".format(dest)
    git_commands = [
        "{0} {1}".format(command_prefix, command) for command in git_commands
    ]

    for cmd in git_commands:
        yield cmd

    # Apply any user or group if we did anything
    if git_commands and (user or group):
        yield chown(dest, user, group, recursive=True)