Add storage upload and download methods

CHANGELOG.md

@@ -10,6 +10,7 @@
 - Saves inverse relations and relationship meta-data in local directory
 - Adds index.import_directory method
 - Allows search by short ID instead of UUID
+- Adds storage upload/download methods
 
 # v4.8.4
 

src/machinable/storage.py

@@ -1,11 +1,13 @@
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, List, Union
 
 import os
 
 from machinable import schema
-from machinable.element import Element, get_lineage
+from machinable.element import get_lineage
+from machinable.index import Index
 from machinable.interface import Interface
 from machinable.types import VersionType
+from machinable.utils import load_file
 
 
 class Storage(Interface):
@@ -25,11 +27,94 @@ def __init__(
             lineage=get_lineage(self),
         )
 
+    def upload(
+        self, interface: "Interface", related: Union[bool, int] = True
+    ) -> None:
+        """
+        Upload the interface to the storage
+
+        interface: Interface
+        related: bool | int
+            - 0/False: Do not save related interfaces
+            - 1/True: Save immediately related interfaces
+            - 2: Save related interfaces and their related interfaces
+        """
+        if self.contains(interface.uuid):
+            self.update(interface)
+        else:
+            self.commit(interface)
+
+        if related:
+            for r in interface.related(deep=int(related) == 2).all():
+                self.upload(r, related=False)
+
+    def download(
+        self,
+        uuid: str,
+        destination: Union[str, None] = None,
+        related: Union[bool, int] = True,
+    ) -> List[str]:
+        """
+        Download to destination
+
+        uuid: str
+            Primary interface UUID
+        destination: str | None
+            If None, download will be imported into active index. If directory filepath, download will be placed
+            in directory without import to index
+        related: bool | int
+            - 0/False: Do not upload related interfaces
+            - 1/True: Upload immediately related interfaces
+            - 2: Upload related interfaces and their related interfaces
+
+        Returns:
+            List of downloaded directories
+        """
+        index = None
+        download_directory = destination
+        if destination is None:
+            index = Index.get()
+            download_directory = index.config.directory
+
+        retrieved = []
+
+        # retrieve primary
+        local_directory = os.path.join(download_directory, uuid)
+        if self.retrieve(uuid, local_directory):
+            retrieved.append(local_directory)
+
+        # retrieve related
+        if related:
+            related_uuids = set()
+            for r in load_file(
+                [local_directory, "related", "metadata.jsonl"], []
+            ):
+                related_uuids.add(r["uuid"])
+                related_uuids.add(r["related_uuid"])
+
+            for r in related_uuids:
+                if r == uuid:
+                    continue
+                retrieved.extend(
+                    self.download(
+                        r,
+                        destination,
+                        related=2 if int(related) == 2 else False,
+                    )
+                )
+
+        # import to index
+        if index is not None:
+            for directory in retrieved:
+                index.import_directory(directory, relations=bool(related))
+
+        return retrieved
+
     def commit(self, interface: "Interface") -> None:
         ...
 
     def update(self, interface: "Interface") -> None:
-        ...
+        return self.commit(interface)
 
     def contains(self, uuid: str) -> bool:
         ...

tests/test_index.py

@@ -2,6 +2,7 @@
 import shutil
 import sqlite3
 
+import pytest
 from machinable import Component, index, schema
 
 
@@ -168,6 +169,18 @@ def test_index_find_related(tmp_path):
     assert _matches(q, [v1, v2])
 
 
+def test_index_find(tmp_path):
+    i = index.Index({"database": str(tmp_path / "index.sqlite")})
+    v = schema.Interface(module="machinable", predicate={"a": 0, "b": 0})
+    i.commit(v)
+    assert i.find_by_hash(v.hash) == i.find(v, by="hash")
+    assert i.find_by_id(v.uuid) == i.find(v, by="uuid")[0]
+    assert i.find_by_id(v.id) == i.find(v, by="id")[0]
+
+    with pytest.raises(ValueError):
+        i.find(v.hash, by="invalid")
+
+
 def test_index_import_directory(tmp_path):
     local_index = index.Index(str(tmp_path / "local"))
     remote_index = index.Index(str(tmp_path / "remote"))

tests/test_storage.py

@@ -1,34 +1,37 @@
 import os
 import shutil
 
-from machinable import Index, Storage, get
+from machinable import Storage
 
 
-def test_storage(tmp_path):
-    class CopyStorage(Storage):
-        class Config:
-            directory: str = ""
+class CopyStorage(Storage):
+    class Config:
+        directory: str = ""
+
+    def commit(self, interface) -> bool:
+        directory = os.path.join(self.config.directory, interface.uuid)
+        if not os.path.exists(directory):
+            os.makedirs(directory)
+            interface.to_directory(directory)
 
-        def commit(self, interface) -> bool:
-            directory = os.path.join(self.config.directory, interface.uuid)
-            if not os.path.exists(directory):
-                os.makedirs(directory)
-                interface.to_directory(directory)
+    def contains(self, uuid):
+        return os.path.exists(os.path.join(self.config.directory, uuid))
 
-        def contains(self, uuid):
-            return os.path.exists(os.path.join(self.config.directory, uuid))
+    def retrieve(self, uuid, local_directory) -> bool:
+        if not self.contains(uuid):
+            return False
 
-        def retrieve(self, uuid, local_directory) -> bool:
-            if not self.contains(uuid):
-                return False
+        shutil.copytree(
+            os.path.join(self.config.directory, uuid),
+            local_directory,
+            dirs_exist_ok=True,
+        )
 
-            shutil.copytree(
-                os.path.join(self.config.directory, uuid),
-                local_directory,
-                dirs_exist_ok=True,
-            )
+        return True
 
-            return True
+
+def test_storage(tmp_path):
+    from machinable import Index, get
 
     primary = str(tmp_path / "primary")
     secondary = str(tmp_path / "secondary")
@@ -64,3 +67,42 @@ def retrieve(self, uuid, local_directory) -> bool:
     st1.__exit__()
     st2.__exit__()
     i.__exit__()
+
+
+def test_storage_upload_and_download(tmp_path):
+    from machinable import Index, get
+
+    primary = str(tmp_path / "primary")
+    secondary = str(tmp_path / "secondary")
+
+    i = Index(primary).__enter__()
+    local = Index(str(tmp_path / "download"))
+
+    storage = CopyStorage({"directory": secondary})
+
+    project = get("machinable.project", "tests/samples/project").__enter__()
+
+    interface1 = get("dummy").launch()
+    interface2 = get("dummy", {"a": 5}, uses=interface1).launch()
+
+    assert not os.path.exists(tmp_path / "secondary" / interface2.uuid)
+    storage.upload(interface2)
+    assert os.path.exists(tmp_path / "secondary" / interface1.uuid)
+    assert os.path.exists(tmp_path / "secondary" / interface2.uuid)
+
+    assert not local.find(interface2)
+    with local:
+        downloads = storage.download(interface2.uuid, related=False)
+        assert len(downloads) == 1
+        assert os.path.exists(local.local_directory(interface2.uuid))
+        assert not os.path.exists(local.local_directory(interface1.uuid))
+        assert local.find(interface2)
+        assert not local.find(interface1)
+
+        downloads = storage.download(interface2.uuid, related=True)
+        assert len(downloads) == 2
+        assert os.path.exists(local.local_directory(interface1.uuid))
+        assert local.find(interface1)
+
+    project.__exit__()
+    i.__exit__()