diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index 4f7987443..5ea7f0556 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -109,7 +109,7 @@ jobs: run: >- python -m tests -n=4 --cov-report=xml --cov-report=term - ${{ env.extra_test_args }} + --enable-ssh ${{ env.extra_test_args }} - name: upload coverage report uses: codecov/codecov-action@v2.1.0 with: diff --git a/dvc/scm/asyn.py b/dvc/scm/asyn.py new file mode 100644 index 000000000..f95c4dad3 --- /dev/null +++ b/dvc/scm/asyn.py @@ -0,0 +1,48 @@ +"""DVC re-implementation of fsspec's dedicated async event loop.""" +import asyncio +import os +import threading +from typing import List, Optional + +from fsspec.asyn import ( # noqa: F401, pylint:disable=unused-import + _selector_policy, + sync, + sync_wrapper, +) + +# dedicated async IO thread +iothread: List[Optional[threading.Thread]] = [None] +# global DVC event loop +default_loop: List[Optional[asyncio.AbstractEventLoop]] = [None] +lock = threading.Lock() + + +def get_loop() -> asyncio.AbstractEventLoop: + """Create or return the global DVC event loop.""" + if default_loop[0] is None: + with lock: + if default_loop[0] is None: + with _selector_policy(): + default_loop[0] = asyncio.new_event_loop() + loop = default_loop[0] + th = threading.Thread( + target=loop.run_forever, # type: ignore[attr-defined] + name="dvcIO", + ) + th.daemon = True + th.start() + iothread[0] = th + assert default_loop[0] is not None + return default_loop[0] + + +class BaseAsyncObject: + def __init__(self, loop: Optional[asyncio.AbstractEventLoop] = None): + self._loop: asyncio.AbstractEventLoop = loop or get_loop() + self._pid = os.getpid() + + @property + def loop(self): + # AsyncMixin is not fork-safe + assert self._pid == os.getpid() + return self._loop diff --git a/dvc/scm/git/backend/base.py b/dvc/scm/git/backend/base.py index ad950574a..cbc571255 100644 --- a/dvc/scm/git/backend/base.py +++ b/dvc/scm/git/backend/base.py @@ -178,7 +178,7 @@ class BaseGitBackend(ABC): """ @abstractmethod - def iter_remote_refs(self, url: str, base: Optional[str] = None): + def iter_remote_refs(self, url: str, base: Optional[str] = None, **kwargs): """Iterate over all refs in the specified remote Git repo. If base is specified, only refs which begin with base will be yielded. @@ -197,6 +197,7 @@ class BaseGitBackend(ABC): dest: str, force: bool = False, on_diverged: Optional[Callable[[str, str], bool]] = None, + **kwargs, ): """Push refspec to a remote Git repo. @@ -222,6 +223,7 @@ class BaseGitBackend(ABC): refspecs: Iterable[str], force: Optional[bool] = False, on_diverged: Optional[Callable[[str, str], bool]] = None, + **kwargs, ): """Fetch refspecs from a remote Git repo. @@ -349,5 +351,5 @@ class BaseGitBackend(ABC): """ @abstractmethod - def validate_git_remote(self, url: str): + def validate_git_remote(self, url: str, **kwargs): """Verify that url is a valid git URL or remote name.""" diff --git a/dvc/scm/git/backend/dulwich.py b/dvc/scm/git/backend/dulwich/__init__.py similarity index 96% rename from dvc/scm/git/backend/dulwich.py rename to dvc/scm/git/backend/dulwich/__init__.py index 977f842f5..d4344d9e8 100644 --- a/dvc/scm/git/backend/dulwich.py +++ b/dvc/scm/git/backend/dulwich/__init__.py @@ -23,13 +23,13 @@ from dvc.progress import Tqdm from dvc.scm.base import GitAuthError, InvalidRemoteSCMRepo, SCMError from dvc.utils import relpath -from ..objects import GitObject -from .base import BaseGitBackend +from ...objects import GitObject +from ..base import BaseGitBackend if TYPE_CHECKING: from dvc.types import StrPath - from ..objects import GitCommit + from ...objects import GitCommit logger = logging.getLogger(__name__) @@ -80,6 +80,13 @@ class DulwichObject(GitObject): class DulwichBackend(BaseGitBackend): # pylint:disable=abstract-method """Dulwich Git backend.""" + from dulwich import client + + from .asyncssh_vendor import AsyncSSHVendor + + # monkeypatch dulwich client's default SSH vendor to use asyncssh + client.get_ssh_vendor = AsyncSSHVendor + # Dulwich progress will return messages equivalent to git CLI, # our pbars should just display the messages as formatted by dulwich BAR_FMT_NOTOTAL = "{desc}{bar:b}|{postfix[info]} [{elapsed}]" @@ -354,14 +361,14 @@ class DulwichBackend(BaseGitBackend): # pylint:disable=abstract-method else: yield os.fsdecode(key) - def iter_remote_refs(self, url: str, base: Optional[str] = None): + def iter_remote_refs(self, url: str, base: Optional[str] = None, **kwargs): from dulwich.client import HTTPUnauthorized, get_transport_and_path from dulwich.errors import NotGitRepository from dulwich.porcelain import get_remote_repo try: _remote, location = get_remote_repo(self.repo, url) - client, path = get_transport_and_path(location) + client, path = get_transport_and_path(location, **kwargs) except Exception as exc: raise InvalidRemoteSCMRepo(url) from exc @@ -389,6 +396,7 @@ class DulwichBackend(BaseGitBackend): # pylint:disable=abstract-method dest: str, force: bool = False, on_diverged: Optional[Callable[[str, str], bool]] = None, + **kwargs, ): from dulwich.client import HTTPUnauthorized, get_transport_and_path from dulwich.errors import NotGitRepository, SendPackError @@ -402,7 +410,7 @@ class DulwichBackend(BaseGitBackend): # pylint:disable=abstract-method try: _remote, location = get_remote_repo(self.repo, url) - client, path = get_transport_and_path(location) + client, path = get_transport_and_path(location, **kwargs) except Exception as exc: raise SCMError( f"'{url}' is not a valid Git remote or URL" @@ -476,6 +484,7 @@ class DulwichBackend(BaseGitBackend): # pylint:disable=abstract-method refspecs: Iterable[str], force: Optional[bool] = False, on_diverged: Optional[Callable[[str, str], bool]] = None, + **kwargs, ): from dulwich.client import get_transport_and_path from dulwich.objectspec import parse_reftuples @@ -504,7 +513,7 @@ class DulwichBackend(BaseGitBackend): # pylint:disable=abstract-method try: _remote, location = get_remote_repo(self.repo, url) - client, path = get_transport_and_path(location) + client, path = get_transport_and_path(location, **kwargs) except Exception as exc: raise SCMError( f"'{url}' is not a valid Git remote or URL" @@ -659,13 +668,13 @@ class DulwichBackend(BaseGitBackend): # pylint:disable=abstract-method ) -> Optional[str]: raise NotImplementedError - def validate_git_remote(self, url: str): + def validate_git_remote(self, url: str, **kwargs): from dulwich.client import LocalGitClient, get_transport_and_path from dulwich.porcelain import get_remote_repo try: _, location = get_remote_repo(self.repo, url) - client, path = get_transport_and_path(location) + client, path = get_transport_and_path(location, **kwargs) except Exception as exc: raise InvalidRemoteSCMRepo(url) from exc if isinstance(client, LocalGitClient) and not os.path.exists( diff --git a/dvc/scm/git/backend/dulwich/asyncssh_vendor.py b/dvc/scm/git/backend/dulwich/asyncssh_vendor.py new file mode 100644 index 000000000..3264c5c39 --- /dev/null +++ b/dvc/scm/git/backend/dulwich/asyncssh_vendor.py @@ -0,0 +1,92 @@ +"""asyncssh SSH vendor for Dulwich.""" +from typing import List, Optional + +from dulwich.client import SSHVendor + +from dvc.scm.asyn import BaseAsyncObject, sync_wrapper + + +class _StderrWrapper: + def __init__(self, stderr): + self.stderr = stderr + + async def _readlines(self): + lines = [] + while True: + line = await self.stderr.readline() + if not line: + break + lines.append(line) + return lines + + readlines = sync_wrapper(_readlines) + + +class AsyncSSHWrapper(BaseAsyncObject): + def __init__(self, conn, proc, **kwargs): + super().__init__(**kwargs) + self.conn = conn + self.proc = proc + self.stderr = _StderrWrapper(proc.stderr) + + def can_read(self) -> bool: + # pylint:disable=protected-access + return self.proc.stdout._session._recv_buf_len > 0 + + async def _read(self, n: Optional[int] = None) -> bytes: + if self.proc.stdout.at_eof(): + return b"" + + return await self.proc.stdout.read(n=n if n is not None else -1) + + read = sync_wrapper(_read) + + def write(self, data: bytes): + self.proc.stdin.write(data) + + def close(self): + self.conn.close() + + +class AsyncSSHVendor(BaseAsyncObject, SSHVendor): + def __init__(self, **kwargs): + super().__init__(**kwargs) + + async def _run_command( + self, + host: str, + command: List[str], + username: Optional[str] = None, + port: Optional[int] = None, + password: Optional[str] = None, + key_filename: Optional[str] = None, + **kwargs, + ): + """Connect to an SSH server. + + Run a command remotely and return a file-like object for interaction + with the remote command. + + Args: + host: Host name + command: Command to run (as argv array) + username: Optional ame of user to log in as + port: Optional SSH port to use + password: Optional ssh password for login or private key + key_filename: Optional path to private keyfile + """ + import asyncssh + + conn = await asyncssh.connect( + host, + port=port, + username=username, + password=password, + client_keys=[key_filename] if key_filename else [], + known_hosts=None, + encoding=None, + ) + proc = await conn.create_process(command, encoding=None) + return AsyncSSHWrapper(conn, proc) + + run_command = sync_wrapper(_run_command) diff --git a/dvc/scm/git/backend/gitpython.py b/dvc/scm/git/backend/gitpython.py index c982d91d9..2c790a295 100644 --- a/dvc/scm/git/backend/gitpython.py +++ b/dvc/scm/git/backend/gitpython.py @@ -438,7 +438,7 @@ class GitPythonBackend(BaseGitBackend): # pylint:disable=abstract-method for ref in Reference.iter_items(self.repo, common_path=base): yield ref.path - def iter_remote_refs(self, url: str, base: Optional[str] = None): + def iter_remote_refs(self, url: str, base: Optional[str] = None, **kwargs): raise NotImplementedError def get_refs_containing(self, rev: str, pattern: Optional[str] = None): @@ -465,6 +465,7 @@ class GitPythonBackend(BaseGitBackend): # pylint:disable=abstract-method dest: str, force: bool = False, on_diverged: Optional[Callable[[str, str], bool]] = None, + **kwargs, ): raise NotImplementedError @@ -474,6 +475,7 @@ class GitPythonBackend(BaseGitBackend): # pylint:disable=abstract-method refspecs: Iterable[str], force: Optional[bool] = False, on_diverged: Optional[Callable[[str, str], bool]] = None, + **kwargs, ): raise NotImplementedError @@ -627,5 +629,5 @@ class GitPythonBackend(BaseGitBackend): # pylint:disable=abstract-method raise SCMError("Merge failed") from exc return None - def validate_git_remote(self, url: str): + def validate_git_remote(self, url: str, **kwargs): raise NotImplementedError diff --git a/dvc/scm/git/backend/pygit2.py b/dvc/scm/git/backend/pygit2.py index 7971e4d4b..891f21d32 100644 --- a/dvc/scm/git/backend/pygit2.py +++ b/dvc/scm/git/backend/pygit2.py @@ -374,6 +374,7 @@ class Pygit2Backend(BaseGitBackend): # pylint:disable=abstract-method dest: str, force: bool = False, on_diverged: Optional[Callable[[str, str], bool]] = None, + **kwargs, ): raise NotImplementedError @@ -383,6 +384,7 @@ class Pygit2Backend(BaseGitBackend): # pylint:disable=abstract-method refspecs: Iterable[str], force: Optional[bool] = False, on_diverged: Optional[Callable[[str, str], bool]] = None, + **kwargs, ): raise NotImplementedError @@ -537,7 +539,7 @@ class Pygit2Backend(BaseGitBackend): # pylint:disable=abstract-method index.add(entry.path) index.write() - def iter_remote_refs(self, url: str, base: Optional[str] = None): + def iter_remote_refs(self, url: str, base: Optional[str] = None, **kwargs): raise NotImplementedError def status( @@ -584,5 +586,5 @@ class Pygit2Backend(BaseGitBackend): # pylint:disable=abstract-method self.repo.index.write() return None - def validate_git_remote(self, url: str): + def validate_git_remote(self, url: str, **kwargs): raise NotImplementedError diff --git a/tests/docker-compose.yml b/tests/docker-compose.yml index d5fb6ad80..2aa6b3a11 100644 --- a/tests/docker-compose.yml +++ b/tests/docker-compose.yml @@ -37,3 +37,14 @@ services: - 2222 volumes: - ./remotes/user.key.pub:/tmp/key + + git-server: + image: ghcr.io/linuxserver/openssh-server + environment: + - USER_NAME=user + - PUBLIC_KEY_FILE=/tmp/key + ports: + - 2222 + volumes: + - ./remotes/user.key.pub:/tmp/key + - ./remotes/git-init:/config/custom-cont-init.d diff --git a/tests/func/test_scm.py b/tests/func/test_scm.py index 85d65fe09..1c3829825 100644 --- a/tests/func/test_scm.py +++ b/tests/func/test_scm.py @@ -331,3 +331,52 @@ def test_git_stash_clear(tmp_dir, scm, ref): not os.path.exists(log_path) or not open(log_path, encoding="utf-8").read() ) + + +@pytest.mark.needs_internet +@pytest.mark.parametrize("server", [pytest.lazy_fixture("git_ssh")]) +def test_git_ssh(tmp_dir, scm, server): + from dulwich.repo import Repo as DulwichRepo + from sshfs import SSHFileSystem + + from dvc.utils.fs import remove + from tests.remotes.ssh import TEST_SSH_KEY_PATH, TEST_SSH_USER + + fs = SSHFileSystem( + host=server.host, + port=server.port, + username=TEST_SSH_USER, + client_keys=[TEST_SSH_KEY_PATH], + ) + server._ssh.execute("git init --bare test-repo.git") + url = f"ssh://{TEST_SSH_USER}@{server.host}:{server.port}/~/test-repo.git" + + tmp_dir.scm_gen("foo", "foo", commit="init") + rev = scm.get_rev() + + scm.push_refspec( + url, + "refs/heads/master", + "refs/heads/master", + force=True, + key_filename=TEST_SSH_KEY_PATH, + ) + + assert ( + rev.encode("ascii") + == fs.open("test-repo.git/refs/heads/master").read().strip() + ) + + remove(tmp_dir / ".git") + remove(tmp_dir / "foo") + DulwichRepo.init(str(tmp_dir)) + + scm.fetch_refspecs( + url, + ["refs/heads/master"], + force=True, + key_filename=TEST_SSH_KEY_PATH, + ) + assert rev == scm.get_ref("refs/heads/master") + scm.checkout("master") + assert "foo" == (tmp_dir / "foo").read_text() diff --git a/tests/remotes/__init__.py b/tests/remotes/__init__.py index 4b86530df..4e9b426a0 100644 --- a/tests/remotes/__init__.py +++ b/tests/remotes/__init__.py @@ -8,6 +8,7 @@ from .gdrive import ( # noqa: F401; noqa: F401 GDrive, gdrive, ) +from .git_server import git_server, git_ssh # noqa: F401 from .gs import ( # noqa: F401; noqa: F401 GCP, TEST_GCP_CREDS_FILE, diff --git a/tests/remotes/git-init/git.sh b/tests/remotes/git-init/git.sh new file mode 100644 index 000000000..f6e2f273a --- /dev/null +++ b/tests/remotes/git-init/git.sh @@ -0,0 +1,2 @@ +#!/bin/bash +apk add --no-cache git diff --git a/tests/remotes/git_server.py b/tests/remotes/git_server.py new file mode 100644 index 000000000..13f470266 --- /dev/null +++ b/tests/remotes/git_server.py @@ -0,0 +1,60 @@ +import pytest + +from .ssh import SSH, TEST_SSH_KEY_PATH, TEST_SSH_USER + + +class GitSSH(SSH): + @staticmethod + def get_url(host, port): # pylint: disable=arguments-differ + return f"ssh://{host}:{port}/tmp/data/git" + + +@pytest.fixture +def git_server(test_config, docker_compose, docker_services): + import asyncssh + from sshfs import SSHFileSystem + + test_config.requires("ssh") + conn_info = { + "host": "127.0.0.1", + "port": docker_services.port_for("git-server", 2222), + } + + def get_fs(): + return SSHFileSystem( + **conn_info, + username=TEST_SSH_USER, + client_keys=[TEST_SSH_KEY_PATH], + ) + + def _check(): + try: + fs = get_fs() + fs.exists("/") + fs.execute("git --version") + except asyncssh.Error: + return False + else: + return True + + docker_services.wait_until_responsive(timeout=30.0, pause=1, check=_check) + return conn_info + + +@pytest.fixture +def git_ssh_connection(git_server): + from sshfs import SSHFileSystem + + yield SSHFileSystem( + host=git_server["host"], + port=git_server["port"], + username=TEST_SSH_USER, + client_keys=[TEST_SSH_KEY_PATH], + ) + + +@pytest.fixture +def git_ssh(git_server, monkeypatch): + url = GitSSH(GitSSH.get_url(**git_server)) + url.mkdir(exist_ok=True, parents=True) + return url