local.py

#
# Copyright (C) 2021 Associated Universities, Inc. Washington DC, USA.
#
# This file is part of NRAO Workspaces.
#
# Workspaces is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Workspaces is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Workspaces.  If not, see <https://www.gnu.org/licenses/>.
import os
import pathlib
import shutil
import tempfile
from typing import BinaryIO, Dict

from .interfaces import DeliveryContextIF, Destination, DestinationTempFile


class LocalDestination(Destination):
    """
    LocalDestination is for delivering to a local directory on the filesystem.
    """

    def __init__(self, context: DeliveryContextIF, path: pathlib.Path):
        super().__init__(context)
        self.path = path

    def add_file(self, file: pathlib.Path, relative_path: str):
        """
        Copy contents of file to new file with path relative_path; creates directories if they don't exist

        :param file: Source file whose contents are getting copied
        :param relative_path: Relative path to new file in destination
        """
        # ensure that the directory exists. we must check, because we get an exception
        # if we try to create one that already exists
        if not (self.path / relative_path).parent.exists():
            (self.path / relative_path).parent.mkdir(parents=True)

        # now copy the file
        # if we don't care about file metadata we could use copy() instead
        shutil.copy2(file, self.path / relative_path)

    def create_file(self, relative_path: str) -> DestinationTempFile:
        """
        Creates a temporary file in the destination which will eventually be added at the supplied relative path.

        :param relative_path: the path to the eventual file
        :return: a DestinationTempFile you can write to
        """
        # make a temporary file in the local destination

        # the temporary file will have a prefix of just the filename
        prefix_name = relative_path.split("/")[-1]

        # make the named temporary file
        rawfile = tempfile.NamedTemporaryFile(prefix=prefix_name, dir=self.path, delete=False)

        # Any of these files created during delivery will need broad permissions:
        local_file_wrapper = LocalDestinationTempFile(self, relative_path, rawfile)
        local_file_wrapper.chmod(0o777)

        # hand off to the LocalDestinationTempFile, which handles the rest
        return local_file_wrapper

    def results(self) -> Dict:
        """
        Return a file:/// URL for this location
        """

        # we could also supply a file:/// URL, which would be constructed like so:
        # "url": pathlib.Path(self.path.absolute()).as_uri()
        return {"delivered_to": str(self.path)}

    def __str__(self):
        return f"local destination {self.path}"


class LocalDestinationTempFile(DestinationTempFile):
    """
    Implements the DestinationTempFile functionality for local destinations. Presumably if we wind up with other
    kinds of destination (globus or streaming or something like that) we will need to greatly reconsider how to
    implement this.
    """

    def __init__(
        self,
        destination: LocalDestination,
        relative_path: str,
        tempfile: tempfile.NamedTemporaryFile,
    ):
        self.destination = destination
        self.relative_path = relative_path
        self.tempfile = tempfile

    def close(self):
        # The key idea here is that after we close the temp file but before we delete it, we add it to the destination
        self.tempfile.close()

        # now that the file is finalized, we can add it to the destination using the path to the named temporary file
        # and the relative path the user originally requested
        self.destination.add_file(pathlib.Path(self.tempfile.name), self.relative_path)

        # now that the file has been delivered, we can remove it safely
        os.unlink(self.tempfile.name)

    def file(self) -> BinaryIO:
        return self.tempfile

    def filename(self) -> str:
        return self.tempfile.name

    def chmod(self, mode):
        temp_path = pathlib.Path(self.filename())
        temp_path.chmod(mode)