4823: Working on porting more StreamReader tests to ArvadosFile.
[arvados.git] / sdk / python / arvados / collection.py
index 24572edf56ab62b6c3a08ed4cc15825370adf87b..ea18123f65474ff67bd6b4c137ecb15208694f34 100644 (file)
@@ -2,13 +2,15 @@ import functools
 import logging
 import os
 import re
+import errno
 
 from collections import deque
 from stat import *
 
-from .arvfile import ArvadosFileBase, split
+from .arvfile import ArvadosFileBase, split, ArvadosFile, ArvadosFileWriter, ArvadosFileReader, BlockManager
 from keep import *
-from .stream import StreamReader, normalize_stream
+from .stream import StreamReader, normalize_stream, locator_block_size
+from .ranges import Range, LocatorAndRange
 import config
 import errors
 import util
@@ -181,7 +183,7 @@ class CollectionReader(CollectionBase):
                 if filename not in streams[streamname]:
                     streams[streamname][filename] = []
                 for r in f.segments:
-                    streams[streamname][filename].extend(s.locators_and_ranges(r[0], r[1]))
+                    streams[streamname][filename].extend(s.locators_and_ranges(r.locator, r.range_size))
 
         self._streams = [normalize_stream(s, streams[s])
                          for s in sorted(streams)]
@@ -636,3 +638,413 @@ class ResumableCollectionWriter(CollectionWriter):
             raise errors.AssertionError(
                 "resumable writer can't accept unsourced data")
         return super(ResumableCollectionWriter, self).write(data)
+
+
+class Collection(CollectionBase):
+    def __init__(self, manifest_locator_or_text=None, parent=None, api_client=None,
+                 keep_client=None, num_retries=0, block_manager=None):
+
+        self.parent = parent
+        self._items = None
+        self._api_client = api_client
+        self._keep_client = keep_client
+        self._block_manager = block_manager
+
+        self.num_retries = num_retries
+        self._manifest_locator = None
+        self._manifest_text = None
+        self._api_response = None
+
+        if manifest_locator_or_text:
+            if re.match(util.keep_locator_pattern, manifest_locator_or_text):
+                self._manifest_locator = manifest_locator_or_text
+            elif re.match(util.collection_uuid_pattern, manifest_locator_or_text):
+                self._manifest_locator = manifest_locator_or_text
+            elif re.match(util.manifest_pattern, manifest_locator_or_text):
+                self._manifest_text = manifest_locator_or_text
+            else:
+                raise errors.ArgumentError(
+                    "Argument to CollectionReader must be a manifest or a collection UUID")
+
+    def _my_api(self):
+        if self._api_client is None:
+            if self.parent is not None:
+                return self.parent._my_api()
+            self._api_client = arvados.api('v1')
+            self._keep_client = None  # Make a new one with the new api.
+        return self._api_client
+
+    def _my_keep(self):
+        if self._keep_client is None:
+            if self.parent is not None:
+                return self.parent._my_keep()
+            self._keep_client = KeepClient(api_client=self._my_api(),
+                                           num_retries=self.num_retries)
+        return self._keep_client
+
+    def _my_block_manager(self):
+        if self._block_manager is None:
+            if self.parent is not None:
+                return self.parent._my_block_manager()
+            self._block_manager = BlockManager(self._my_keep())
+        return self._block_manager
+
+    def _populate_from_api_server(self):
+        # As in KeepClient itself, we must wait until the last
+        # possible moment to instantiate an API client, in order to
+        # avoid tripping up clients that don't have access to an API
+        # server.  If we do build one, make sure our Keep client uses
+        # it.  If instantiation fails, we'll fall back to the except
+        # clause, just like any other Collection lookup
+        # failure. Return an exception, or None if successful.
+        try:
+            self._api_response = self._my_api().collections().get(
+                uuid=self._manifest_locator).execute(
+                    num_retries=self.num_retries)
+            self._manifest_text = self._api_response['manifest_text']
+            return None
+        except Exception as e:
+            return e
+
+    def _populate_from_keep(self):
+        # Retrieve a manifest directly from Keep. This has a chance of
+        # working if [a] the locator includes a permission signature
+        # or [b] the Keep services are operating in world-readable
+        # mode. Return an exception, or None if successful.
+        try:
+            self._manifest_text = self._my_keep().get(
+                self._manifest_locator, num_retries=self.num_retries)
+        except Exception as e:
+            return e
+
+    def _populate(self):
+        self._items = {}
+        if self._manifest_locator is None and self._manifest_text is None:
+            return
+        error_via_api = None
+        error_via_keep = None
+        should_try_keep = ((self._manifest_text is None) and
+                           util.keep_locator_pattern.match(
+                self._manifest_locator))
+        if ((self._manifest_text is None) and
+            util.signed_locator_pattern.match(self._manifest_locator)):
+            error_via_keep = self._populate_from_keep()
+        if self._manifest_text is None:
+            error_via_api = self._populate_from_api_server()
+            if error_via_api is not None and not should_try_keep:
+                raise error_via_api
+        if ((self._manifest_text is None) and
+            not error_via_keep and
+            should_try_keep):
+            # Looks like a keep locator, and we didn't already try keep above
+            error_via_keep = self._populate_from_keep()
+        if self._manifest_text is None:
+            # Nothing worked!
+            raise arvados.errors.NotFoundError(
+                ("Failed to retrieve collection '{}' " +
+                 "from either API server ({}) or Keep ({})."
+                 ).format(
+                    self._manifest_locator,
+                    error_via_api,
+                    error_via_keep))
+        # populate
+        import_manifest(self._manifest_text, self)
+
+    def _populate_first(orig_func):
+        # Decorator for methods that read actual Collection data.
+        @functools.wraps(orig_func)
+        def wrapper(self, *args, **kwargs):
+            if self._items is None:
+                self._populate()
+            return orig_func(self, *args, **kwargs)
+        return wrapper
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.save(no_locator=True)
+        if self._block_manager is not None:
+            self._block_manager.stop_threads()
+
+    @_populate_first
+    def find(self, path, create=False, create_collection=False):
+        p = path.split("/")
+        if p[0] == '.':
+            del p[0]
+
+        if len(p) > 0:
+            item = self._items.get(p[0])
+            if len(p) == 1:
+                # item must be a file
+                if item is None and create:
+                    # create new file
+                    if create_collection:
+                        item = Collection(parent=self, num_retries=self.num_retries)
+                    else:
+                        item = ArvadosFile(self)
+                    self._items[p[0]] = item
+                return item
+            else:
+                if item is None and create:
+                    # create new collection
+                    item = Collection(parent=self, num_retries=self.num_retries)
+                    self._items[p[0]] = item
+                del p[0]
+                return item.find("/".join(p), create=create)
+        else:
+            return self
+
+    @_populate_first
+    def api_response(self):
+        """api_response() -> dict or None
+
+        Returns information about this Collection fetched from the API server.
+        If the Collection exists in Keep but not the API server, currently
+        returns None.  Future versions may provide a synthetic response.
+        """
+        return self._api_response
+
+    def open(self, path, mode):
+        mode = mode.replace("b", "")
+        if len(mode) == 0 or mode[0] not in ("r", "w", "a"):
+            raise ArgumentError("Bad mode '%s'" % mode)
+        create = (mode != "r")
+
+        f = self.find(path, create=create)
+        if f is None:
+            raise IOError((errno.ENOENT, "File not found"))
+        if not isinstance(f, ArvadosFile):
+            raise IOError((errno.EISDIR, "Path must refer to a file."))
+
+        if mode[0] == "w":
+            f.truncate(0)
+
+        if mode == "r":
+            return ArvadosFileReader(f, path, mode)
+        else:
+            return ArvadosFileWriter(f, path, mode)
+
+    @_populate_first
+    def modified(self):
+        for k,v in self._items.items():
+            if v.modified():
+                return True
+        return False
+
+    @_populate_first
+    def set_unmodified(self):
+        for k,v in self._items.items():
+            v.set_unmodified()
+
+    @_populate_first
+    def __iter__(self):
+        return self._items.iterkeys()
+
+    @_populate_first
+    def iterkeys(self):
+        return self._items.iterkeys()
+
+    @_populate_first
+    def __getitem__(self, k):
+        return self._items[k]
+
+    @_populate_first
+    def __contains__(self, k):
+        return k in self._items
+
+    @_populate_first
+    def __len__(self):
+       return len(self._items)
+
+    @_populate_first
+    def __delitem__(self, p):
+        del self._items[p]
+
+    @_populate_first
+    def keys(self):
+        return self._items.keys()
+
+    @_populate_first
+    def values(self):
+        return self._items.values()
+
+    @_populate_first
+    def items(self):
+        return self._items.items()
+
+    @_populate_first
+    def exists(self, path):
+        return self.find(path) != None
+
+    @_populate_first
+    def remove(self, path):
+        p = path.split("/")
+        if p[0] == '.':
+            del p[0]
+
+        if len(p) > 0:
+            item = self._items.get(p[0])
+            if item is None:
+                raise IOError((errno.ENOENT, "File not found"))
+            if len(p) == 1:
+                del self._items[p[0]]
+            else:
+                del p[0]
+                item.remove("/".join(p))
+        else:
+            raise IOError((errno.ENOENT, "File not found"))
+
+    @_populate_first
+    def manifest_text(self, strip=False, normalize=False):
+        if self.modified() or self._manifest_text is None or normalize:
+            return export_manifest(self, stream_name=".", portable_locators=strip)
+        else:
+            if strip:
+                return self.stripped_manifest()
+            else:
+                return self._manifest_text
+
+    def portable_data_hash(self):
+        stripped = self.manifest_text(strip=True)
+        return hashlib.md5(stripped).hexdigest() + '+' + str(len(stripped))
+
+    @_populate_first
+    def save(self, no_locator=False):
+        if self.modified():
+            self._my_block_manager().commit_all()
+            self._my_keep().put(self.manifest_text(strip=True))
+            if self._manifest_locator is not None and re.match(util.collection_uuid_pattern, self._manifest_locator):
+                self._api_response = self._my_api().collections().update(
+                    uuid=self._manifest_locator,
+                    body={'manifest_text': self.manifest_text(strip=False)}
+                    ).execute(
+                        num_retries=self.num_retries)
+            elif not no_locator:
+                raise AssertionError("Collection manifest_locator must be a collection uuid.  Use save_as() for new collections.")
+            self.set_unmodified()
+
+    @_populate_first
+    def save_as(self, name, owner_uuid=None, ensure_unique_name=False):
+        self._my_block_manager().commit_all()
+        self._my_keep().put(self.manifest_text(strip=True))
+        body = {"manifest_text": self.manifest_text(strip=False),
+                "name": name}
+        if owner_uuid:
+            body["owner_uuid"] = owner_uuid
+        self._api_response = self._my_api().collections().create(ensure_unique_name=ensure_unique_name, body=body).execute(num_retries=self.num_retries)
+        self._manifest_locator = self._api_response["uuid"]
+        self.set_unmodified()
+
+    @_populate_first
+    def rename(self, old, new):
+        old_path, old_fn = os.path.split(old)
+        old_col = self.find(path)
+        if old_col is None:
+            raise IOError((errno.ENOENT, "File not found"))
+        if not isinstance(old_p, Collection):
+            raise IOError((errno.ENOTDIR, "Parent in path is a file, not a directory"))
+        if old_fn in old_col:
+            new_path, new_fn = os.path.split(new)
+            new_col = self.find(new_path, create=True, create_collection=True)
+            if not isinstance(new_col, Collection):
+                raise IOError((errno.ENOTDIR, "Destination is a file, not a directory"))
+            ent = old_col[old_fn]
+            del old_col[old_fn]
+            ent.parent = new_col
+            new_col[new_fn] = ent
+        else:
+            raise IOError((errno.ENOENT, "File not found"))
+
+def import_manifest(manifest_text, into_collection=None, api_client=None, keep=None, num_retries=None):
+    if into_collection is not None:
+        if len(into_collection) > 0:
+            raise ArgumentError("Can only import manifest into an empty collection")
+        c = into_collection
+    else:
+        c = Collection(api_client=api_client, keep_client=keep, num_retries=num_retries)
+
+    STREAM_NAME = 0
+    BLOCKS = 1
+    SEGMENTS = 2
+
+    stream_name = None
+    state = STREAM_NAME
+
+    for n in re.finditer(r'(\S+)(\s+|$)', manifest_text):
+        tok = n.group(1)
+        sep = n.group(2)
+
+        if state == STREAM_NAME:
+            # starting a new stream
+            stream_name = tok.replace('\\040', ' ')
+            blocks = []
+            segments = []
+            streamoffset = 0L
+            state = BLOCKS
+            continue
+
+        if state == BLOCKS:
+            s = re.match(r'[0-9a-f]{32}\+(\d+)(\+\S+)*', tok)
+            if s:
+                blocksize = long(s.group(1))
+                blocks.append(Range(tok, streamoffset, blocksize))
+                streamoffset += blocksize
+            else:
+                state = SEGMENTS
+
+        if state == SEGMENTS:
+            s = re.search(r'^(\d+):(\d+):(\S+)', tok)
+            if s:
+                pos = long(s.group(1))
+                size = long(s.group(2))
+                name = s.group(3).replace('\\040', ' ')
+                f = c.find("%s/%s" % (stream_name, name), create=True)
+                f.add_segment(blocks, pos, size)
+            else:
+                # error!
+                raise errors.SyntaxError("Invalid manifest format")
+
+        if sep == "\n":
+            stream_name = None
+            state = STREAM_NAME
+
+    c.set_unmodified()
+    return c
+
+def export_manifest(item, stream_name=".", portable_locators=False):
+    buf = ""
+    if isinstance(item, Collection):
+        stream = {}
+        sorted_keys = sorted(item.keys())
+        for k in [s for s in sorted_keys if isinstance(item[s], ArvadosFile)]:
+            v = item[k]
+            st = []
+            for s in v.segments:
+                loc = s.locator
+                if loc.startswith("bufferblock"):
+                    loc = v.parent._my_block_manager()._bufferblocks[loc].locator()
+                if portable_locators:
+                    loc = KeepLocator(loc).stripped()
+                st.append(LocatorAndRange(loc, locator_block_size(loc),
+                                     s.segment_offset, s.range_size))
+            stream[k] = st
+        if stream:
+            buf += ' '.join(normalize_stream(stream_name, stream))
+            buf += "\n"
+        for k in [s for s in sorted_keys if isinstance(item[s], Collection)]:
+            buf += export_manifest(item[k], stream_name=os.path.join(stream_name, k), portable_locators=portable_locators)
+    elif isinstance(item, ArvadosFile):
+        st = []
+        for s in item.segments:
+            loc = s.locator
+            if loc.startswith("bufferblock"):
+                loc = item._bufferblocks[loc].calculate_locator()
+            if portable_locators:
+                loc = KeepLocator(loc).stripped()
+            st.append(LocatorAndRange(loc, locator_block_size(loc),
+                                 s.segment_offset, s.range_size))
+        stream[stream_name] = st
+        buf += ' '.join(normalize_stream(stream_name, stream))
+        buf += "\n"
+    return buf