~launchpad-pqm/launchpad/devel

# Copyright 2009 Canonical Ltd.  This software is licensed under the

# GNU Affero General Public License version 3 (see the file LICENSE).

__all__ = ['DiskPoolEntry', 'DiskPool', 'poolify', 'unpoolify']

import os

import tempfile

from canonical.librarian.utils import (

    copy_and_close,

    sha1_from_path,

    )

from lp.archivepublisher import HARDCODED_COMPONENT_ORDER

from lp.services.propertycache import cachedproperty

from lp.soyuz.interfaces.publishing import (

    MissingSymlinkInPool,

    NotInPool,

    PoolFileOverwriteError,

    )

def poolify(source, component):

    """Poolify a given source and component name."""

    if source.startswith("lib"):

        return os.path.join(component, source[:4], source)

    else:

        return os.path.join(component, source[:1], source)

def unpoolify(self, path):

    """Take a path and unpoolify it.

    Return a tuple of component, source, filename

    """

    p = path.split("/")

    if len(p) < 3 or len(p) > 4:

        raise ValueError("Path %s is not in a valid pool form" % path)

    if len(p) == 4:

        return p[0], p[2], p[3]

    return p[0], p[2], None

def relative_symlink(src_path, dst_path):

    """os.symlink replacement that creates relative symbolic links."""

    path_sep = os.path.sep

    src_path = os.path.normpath(src_path)

    dst_path = os.path.normpath(dst_path)

    src_path_elems = src_path.split(path_sep)

    dst_path_elems = dst_path.split(path_sep)

    if os.path.isabs(src_path):

        if not os.path.isabs(dst_path):

            dst_path = os.path.abspath(dst_path)

        common_prefix = os.path.commonprefix([src_path_elems, dst_path_elems])

        backward_elems = ['..'] * (len(dst_path_elems)-len(common_prefix)-1)

        forward_elems = src_path_elems[len(common_prefix):]

        src_path = path_sep.join(backward_elems + forward_elems)

    os.symlink(src_path, dst_path)

class FileAddActionEnum:

    """Possible actions taken when adding a file.

    FILE_ADDED: we added the actual file to the disk

    SYMLINK_ADDED: we created a symlink to another copy of the same file

    NONE: no action was necessary or taken.

    """

    FILE_ADDED = "file_added"

    SYMLINK_ADDED = "symlink_added"

    NONE = "none"

class _diskpool_atomicfile:

    """Simple file-like object used by the pool to atomically move into place

    a file after downloading from the librarian.

    This class is designed to solve a very specific problem encountered in

    the publisher. Namely that should the publisher crash during the process

    of publishing a file to the pool, an empty or incomplete file would be

    present in the pool. Its mere presence would fool the publisher into

    believing it had already downloaded that file to the pool, resulting

    in failures in the apt-ftparchive stage.

    By performing a rename() when the file is guaranteed to have been

    fully written to disk (after the fd.close()) we can be sure that if

    the filename is present in the pool, it is definitely complete.

    """

    def __init__(self, targetfilename, mode, rootpath="/tmp"):

        # atomicfile implements the file object interface, but it is only

        # really used (or useful) for writing binary files, which is why we

        # keep the mode constructor argument but assert it's sane below.

        if mode == "w":

            mode = "wb"

        assert mode == "wb"

        assert not os.path.exists(targetfilename)

        self.targetfilename = targetfilename

        fd, name = tempfile.mkstemp(prefix="temp-download.", dir=rootpath)

        self.fd = os.fdopen(fd, mode)

        self.tempname = name

        self.write = self.fd.write

    def close(self):

        """Make the atomic move into place having closed the temp file."""

        self.fd.close()

        os.chmod(self.tempname, 0644)

        # Note that this will fail if the target and the temp dirs are on

        # different filesystems.

        os.rename(self.tempname, self.targetfilename)

class DiskPoolEntry:

    """Represents a single file in the pool, across all components.

    Creating a DiskPoolEntry performs disk reads, so don't create an

    instance of this class unless you need to know what's already on

    the disk for this file.

    'tempath' must be in the same filesystem as 'rootpath', it will be

    used to store the instalation candidate while it is being downloaded

    from the Librarian.

    Remaining files in the 'temppath' indicated installation failures and

    require manual removal after further investigation.

    """

    def __init__(self, rootpath, temppath, source, filename, logger):

        self.rootpath = rootpath

        self.temppath = temppath

        self.source = source

        self.filename = filename

        self.logger = logger

        self.file_component = None

        self.symlink_components = set()

        for component in HARDCODED_COMPONENT_ORDER:

            path = self.pathFor(component)

            if os.path.islink(path):

                self.symlink_components.add(component)

            elif os.path.isfile(path):

                assert not self.file_component

                self.file_component = component

        if self.symlink_components:

            assert self.file_component

    def debug(self, *args, **kwargs):

        self.logger.debug(*args, **kwargs)

    def pathFor(self, component):

        """Return the path for this file in the given component."""

        return os.path.join(self.rootpath,

                            poolify(self.source, component),

                            self.filename)

    def preferredComponent(self, add=None, remove=None):

        """Return the appropriate component for the real file.

        If add is passed, add it to the list before calculating.

        If remove is passed, remove it before calculating.

        Thus, we can calcuate which component should contain the main file

        after the addition or removal we are working on.

        """

        components = set()

        if self.file_component:

            components.add(self.file_component)

        components = components.union(self.symlink_components)

        if add is not None:

            components.add(add)

        if remove is not None and remove in components:

            components.remove(remove)

        for component in HARDCODED_COMPONENT_ORDER:

            if component in components:

                return component

    @cachedproperty

    def file_hash(self):

        """Return the SHA1 sum of this file."""

        targetpath = self.pathFor(self.file_component)

        return sha1_from_path(targetpath)

    def addFile(self, component, sha1, contents):

        """See DiskPool.addFile."""

        assert component in HARDCODED_COMPONENT_ORDER

        targetpath = self.pathFor(component)

        if not os.path.exists(os.path.dirname(targetpath)):

            os.makedirs(os.path.dirname(targetpath))

        if self.file_component:

            # There's something on disk. Check hash.

            if sha1 != self.file_hash:

                raise PoolFileOverwriteError('%s != %s for %s' %

                    (sha1, self.file_hash,

                     self.pathFor(self.file_component)))

            if (component == self.file_component

                or component in self.symlink_components):

                # The file is already here

                return FileAddActionEnum.NONE

            else:

                # The file is present in a different component,

                # make a symlink.

                relative_symlink(

                    self.pathFor(self.file_component), targetpath)

                self.symlink_components.add(component)

                # Then fix to ensure the right component is linked.

                self._sanitiseLinks()

                return FileAddActionEnum.SYMLINK_ADDED

        # If we get to here, we want to write the file.

        assert not os.path.exists(targetpath)

        self.debug("Making new file in %s for %s/%s" %

                   (component, self.source, self.filename))

        file_to_write = _diskpool_atomicfile(

            targetpath, "wb", rootpath=self.temppath)

        contents.open()

        copy_and_close(contents, file_to_write)

        self.file_component = component

        return FileAddActionEnum.FILE_ADDED

    def removeFile(self, component):

        """Remove a file from a given component; return bytes freed.

        This method handles three situations:

        1) Remove a symlink

        2) Remove the main file and there are no symlinks left.

        3) Remove the main file and there are symlinks left.

        """

        if not self.file_component:

            raise NotInPool(

                "File for removing %s %s/%s is not in pool, skipping." %

                (component, self.source, self.filename))

        # Okay, it's there, if it's a symlink then we need to remove

        # it simply.

        if component in self.symlink_components:

            self.debug("Removing %s %s/%s as it is a symlink"

                       % (component, self.source, self.filename))

            # ensure we are removing a symbolic link and

            # it is published in one or more components

            link_path = self.pathFor(component)

            assert os.path.islink(link_path)

            return self._reallyRemove(component)

        if component != self.file_component:

            raise MissingSymlinkInPool(

                "Symlink for %s/%s in %s is missing, skipping." %

                (self.source, self.filename, component))

        # It's not a symlink, this means we need to check whether we

        # have symlinks or not.

        if len(self.symlink_components) == 0:

            self.debug("Removing %s/%s from %s" %

                       (self.source, self.filename, component))

        else:

            # The target for removal is the real file, and there are symlinks

            # pointing to it. In order to avoid breakage, we need to first

            # shuffle the symlinks, so that the one we want to delete will

            # just be one of the links, and becomes safe.

            targetcomponent = self.preferredComponent(remove=component)

            self._shufflesymlinks(targetcomponent)

        return self._reallyRemove(component)

    def _reallyRemove(self, component):

        """Remove file and return file size.

        Remove the file from the filesystem and from our data

        structures.

        """

        fullpath = self.pathFor(component)

        assert os.path.exists(fullpath)

        if component == self.file_component:

            # Deleting the master file is only allowed if there

            # are no symlinks left.

            assert not self.symlink_components

            self.file_component = None

        elif component in self.symlink_components:

            self.symlink_components.remove(component)

        size = os.lstat(fullpath).st_size

        os.remove(fullpath)

        return size

    def _shufflesymlinks(self, targetcomponent):

        """Shuffle the symlinks for filename so that targetcomponent contains

        the real file and the rest are symlinks to the right place..."""

        if targetcomponent == self.file_component:

            # We're already in the right place.

            return

        if targetcomponent not in self.symlink_components:

            raise ValueError(

                "Target component '%s' is not a symlink for %s" %

                             (targetcomponent, self.filename))

        self.debug("Shuffling symlinks so primary for %s is in %s" %

                   (self.filename, targetcomponent))

        # Okay, so first up, we unlink the targetcomponent symlink.

        targetpath = self.pathFor(targetcomponent)

        os.remove(targetpath)

        # Now we rename the source file into the target component.

        sourcepath = self.pathFor(self.file_component)

        # XXX cprov 2006-05-26: if it fails the symlinks are severely broken

        # or maybe we are writing them wrong. It needs manual fix !

        # Nonetheless, we carry on checking other candidates.

        # Use 'find -L . -type l' on pool to find out broken symlinks

        # Normally they only can be fixed by remove the broken links and

        # run a careful (-C) publication.

        # ensure targetpath doesn't exists and  the sourcepath exists

        # before rename them.

        assert not os.path.exists(targetpath)

        assert os.path.exists(sourcepath)

        os.rename(sourcepath, targetpath)

        # XXX cprov 2006-06-12: it may cause problems to the database, since

        # ZTM isn't handled properly in scripts/publish-distro.py. Things are

        # commited mid-procedure & bare exception is caught.

        # Update the data structures.

        self.symlink_components.add(self.file_component)

        self.symlink_components.remove(targetcomponent)

        self.file_component = targetcomponent

        # Now we make the symlinks on the filesystem.

        for comp in self.symlink_components:

            newpath = self.pathFor(comp)

            try:

                os.remove(newpath)

            except OSError:

                # Do nothing because it's almost certainly a not found.

                pass

            relative_symlink(targetpath, newpath)

    def _sanitiseLinks(self):

        """Ensure the real file is in the most preferred component.

        If this file is in more than one component, ensure the real

        file is in the most preferred component and the other components

        use symlinks.

        It's important that the real file be in the most preferred

        component because partial mirrors may only take a subset of

        components, and these partial mirrors must not have broken

        symlinks where they should have working files.

        """

        component = self.preferredComponent()

        if not self.file_component == component:

            self._shufflesymlinks(component)

class DiskPool:

    """Scan a pool on the filesystem and record information about it.

    Its constructor receives 'rootpath', which is the pool path where the

    files will be installed, and the 'temppath', which is a temporary

    directory used to store the installation candidate from librarian.

    'rootpath' and 'temppath' must be in the same filesystem, see

    DiskPoolEntry for further information.

    """

    results = FileAddActionEnum

    def __init__(self, rootpath, temppath, logger):

        self.rootpath = rootpath

        if not rootpath.endswith("/"):

            self.rootpath += "/"

        self.temppath = temppath

        if not temppath.endswith("/"):

            self.temppath += "/"

        self.entries = {}

        self.logger = logger

    def _getEntry(self, sourcename, file):

        """Return a new DiskPoolEntry for the given sourcename and file."""

        return DiskPoolEntry(

            self.rootpath, self.temppath, sourcename, file, self.logger)

    def pathFor(self, comp, source, file=None):

        """Return the path for the given pool folder or file.

        If file is none, the path to the folder containing all packages

        for the given component and source package name will be returned.

        If file is specified, the path to the specific package file will

        be returned.

        """

        path = os.path.join(

            self.rootpath, poolify(source, comp))

        if file:

            return os.path.join(path, file)

        return path

    def addFile(self, component, sourcename, filename, sha1, contents):

        """Add a file with the given contents to the pool.

        Component, sourcename and filename are used to calculate the

        on-disk location.

        sha1 is used to compare with the existing file's checksum, if

        a file already exists for any component.

        contents is a file-like object containing the contents we want

        to write.

        There are four possible outcomes:

        - If the file doesn't exist in the pool for any component, it will

        be written from the given contents and results.ADDED_FILE will be

        returned.

        - If the file already exists in the pool, in this or any other

        component, the checksum of the file on disk will be calculated and

        compared with the checksum provided. If they fail to match,

        PoolFileOverwriteError will be raised.

        - If the file already exists but not in this component, and the

        checksum test above passes, a symlink will be added, and

        results.SYMLINK_ADDED will be returned. Also, the symlinks will be

        checked and sanitised, to ensure the real copy of the file is in the

        most preferred component, according to HARDCODED_COMPONENT_ORDER.

        - If the file already exists and is already in this component,

        either as a file or a symlink, and the checksum check passes,

        results.NONE will be returned and nothing will be done.

        """

        entry = self._getEntry(sourcename, filename)

        return entry.addFile(component, sha1, contents)

    def removeFile(self, component, sourcename, filename):

        """Remove the specified file from the pool.

        There are three possible outcomes:

        - If the specified file does not exist, NotInPool will be raised.

        - If the specified file exists and is a symlink, or is the only

        copy of the file in the pool, it will simply be deleted, and its

        size will be returned.

        - If the specified file is a real file and there are symlinks

        referencing it, the symlink in the next most preferred component

        will be deleted, and the file will be moved to replace it. The

        size of the deleted symlink will be returned.

        """

        entry = self._getEntry(sourcename, filename)

        return entry.removeFile(component)