From 110da0c366b133a9981e89f221fe6dcd1ca166f2 Mon Sep 17 00:00:00 2001
From: "Mr. Outis" <mroutis@protonmail.com>
Date: Thu, 22 Aug 2019 01:00:16 -0500
Subject: [PATCH] ssh: adapt link logic and support symlinks

---
 dvc/remote/ssh/__init__.py   | 87 ++++++++++++++++++++++++++++++++++++
 dvc/remote/ssh/connection.py | 12 +++++
 2 files changed, 99 insertions(+)

diff --git a/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py
index 98ee172bc..f19e1df80 100644
--- a/dvc/remote/ssh/__init__.py
+++ b/dvc/remote/ssh/__init__.py
@@ -4,11 +4,13 @@ import errno
 import itertools
 import io
 import os
+import posixpath
 import getpass
 import logging
 import threading
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager, closing
+from copy import copy
 
 try:
     import paramiko
@@ -23,6 +25,7 @@ from dvc.remote.base import RemoteBASE
 from dvc.scheme import Schemes
 from dvc.remote.pool import get_connection
 from dvc.progress import Tqdm
+from dvc.exceptions import DvcException
 
 from .connection import SSHConnection
 
@@ -47,9 +50,21 @@ class RemoteSSH(RemoteBASE):
     # We use conservative setting of 4 instead to not exhaust max sessions.
     CHECKSUM_JOBS = 4
 
+    # XXX: Friendly reminder to add "reflink" to the default types
+    # after is implemented
+    DEFAULT_CACHE_TYPES = ["symlink", "copy"]
+
     def __init__(self, repo, config):
         super(RemoteSSH, self).__init__(repo, config)
 
+        types = config.get(Config.SECTION_CACHE_TYPE, None)
+        if types:
+            if isinstance(types, str):
+                types = [t.strip() for t in types.split(",")]
+            self.cache_types = types
+        else:
+            self.cache_types = copy(self.DEFAULT_CACHE_TYPES)
+
         url = config.get(Config.SECTION_REMOTE_URL)
         if url:
             parsed = urlparse(url)
@@ -278,3 +293,75 @@ class RemoteSSH(RemoteBASE):
                 ret = list(itertools.compress(checksums, in_remote))
                 pbar.update_desc("", 0)  # clear path name description
                 return ret
+
+    def link(self, from_info, to_info):
+        with self.ssh(from_info) as ssh:
+            self._link(from_info, to_info, self.cache_types, ssh)
+
+    def _link(self, from_info, to_info, link_types, ssh):
+        from_path = from_info.path
+        to_path = to_info.path
+
+        assert ssh.isfile(from_path)
+
+        dname = posixpath.dirname(to_path)
+
+        # XXX: We can remove the `exists` call, since `makedirs`
+        # will ignore it
+        if not ssh.exists(dname):
+            ssh.makedirs(dname)
+
+        # NOTE: just create an empty file for an empty cache
+        if ssh.getsize(from_path) == 0:
+            ssh.open(to_path, "w+").close()
+
+            msg = "Created empty file: {} -> {}".format(from_path, to_path)
+            logger.debug(msg)
+            return
+
+        self._try_links(from_info, to_info, link_types, ssh)
+
+    @classmethod
+    def _get_link_method(cls, link_type, ssh):
+        CACHE_TYPE_MAP = {"copy": ssh.cp, "symlink": ssh.symlink}
+
+        try:
+            return CACHE_TYPE_MAP[link_type]
+        except KeyError:
+            raise DvcException(
+                "Cache type: '{}' not supported!".format(link_type)
+            )
+
+    def _do_link(self, from_info, to_info, link_method, ssh):
+        # XXX: We are testing if file exists rather than if file is a link
+        if ssh.exists(to_info.path):
+            raise DvcException("Link '{}' already exists!".format(to_info))
+        else:
+            link_method(from_info.path, to_info.path)
+
+        if self.protected:
+            self.protect(to_info)
+
+        msg = "Created {}'{}': {} -> {}".format(
+            "protected " if self.protected else "",
+            self.cache_types[0],
+            from_info,
+            to_info,
+        )
+        logger.debug(msg)
+
+    def _try_links(self, from_info, to_info, link_types, ssh):
+        i = len(link_types)
+        while i > 0:
+            link_method = self._get_link_method(link_types[0], ssh)
+            try:
+                self._do_link(from_info, to_info, link_method, ssh)
+                return
+
+            except DvcException as exc:
+                msg = "Cache type '{}' is not supported: {}"
+                logger.debug(msg.format(link_types[0], str(exc)))
+                del link_types[0]
+                i -= 1
+
+        raise DvcException("no possible cache types left to try out.")
diff --git a/dvc/remote/ssh/connection.py b/dvc/remote/ssh/connection.py
index fbb0c2b7b..423ce6604 100644
--- a/dvc/remote/ssh/connection.py
+++ b/dvc/remote/ssh/connection.py
@@ -64,6 +64,12 @@ class SSHConnection:
 
         return 0
 
+    def getsize(self, path):
+        with ignore_file_not_found():
+            return self.sftp.stat(path).st_size
+
+        return 0
+
     def exists(self, path, sftp=None):
         return bool(self.st_mode(path))
 
@@ -282,3 +288,9 @@ class SSHConnection:
                         raise
                     break
         return self._sftp_channels
+
+    def open(self, *args, **kwargs):
+        return self.sftp.open(*args, **kwargs)
+
+    def symlink(self, *args, **kwargs):
+        return self.sftp.symlink(*args, **kwargs)