adding truss-transfer in truss (#1327)

* add truss transfer pr

* adding unit test

* fmt

* add exceptioon for test

* add lazy data resolver

pyproject.toml

@@ -167,6 +167,7 @@ opentelemetry-exporter-otlp = ">=1.25.0"
 opentelemetry-sdk = ">=1.25.0"
 uvicorn = ">=0.24.0"
 uvloop = ">=0.17.0"
+truss_transfer="0.0.1rc3"
 
 [build-system]
 build-backend = "poetry.core.masonry.api"

truss/templates/shared/lazy_data_resolver.py

@@ -15,6 +15,13 @@
 except ModuleNotFoundError:
     from truss.templates.shared.util import BLOB_DOWNLOAD_TIMEOUT_SECS
 
+try:
+    import truss_transfer
+
+    TRUSS_TRANSFER_AVAILABLE = True
+except ImportError:
+    TRUSS_TRANSFER_AVAILABLE = False
+
 LAZY_DATA_RESOLVER_PATH = Path("/bptr/bptr-manifest")
 NUM_WORKERS = 4
 CACHE_DIR = Path("/cache/org/artifacts")
@@ -44,6 +51,11 @@ class BasetenPointerManifest(pydantic.BaseModel):
 
 
 class LazyDataResolver:
+    """Deprecation warning: This class is deprecated and will be removed in a future release.
+
+    Please use LazyDataResolverV2 instead (using the `truss_transfer` package).
+    """
+
     def __init__(self, data_dir: Path):
         self._data_dir: Path = data_dir
         self._bptr_resolution: Dict[str, Tuple[str, str, int]] = _read_bptr_resolution()
@@ -131,6 +143,14 @@ def fetch(self):
         self._resolution_done = True
 
 
+class LazyDataResolverV2:
+    def __init__(self, data_dir: Path):
+        self._data_dir: Path = data_dir
+
+    def fetch(self):
+        truss_transfer.lazy_data_resolve(str(self._data_dir))
+
+
 def _read_bptr_resolution() -> Dict[str, Tuple[str, str, int]]:
     if not LAZY_DATA_RESOLVER_PATH.is_file():
         return {}

truss/tests/templates/core/server/test_lazy_data_resolver_v2.py

@@ -0,0 +1,72 @@
+import json
+import tempfile
+import time
+from pathlib import Path
+
+import pytest
+from truss.templates.shared.lazy_data_resolver import (
+    TRUSS_TRANSFER_AVAILABLE,
+    LazyDataResolverV2,
+)
+
+LAZY_DATA_RESOLVER_PATH = Path("/bptr/bptr-manifest")
+TARGET_FILE = Path("nested/config.json")
+
+
+def write_bptr_manifest_to_file(expiration_timestamp: int = 2683764059):
+    bptr_manifest = {
+        "pointers": [
+            {
+                "resolution": {
+                    "url": "https://raw.githubusercontent.com/basetenlabs/truss/00e01b679afbe353b0b2fe4de6b138d912bb7167/.circleci/config.yml",
+                    "expiration_timestamp": expiration_timestamp,
+                },
+                "uid": "8c6b2f215f0333437cdc3fe7c79be0c802847d2f2a0ccdc0bb251814e63cf375",
+                "file_name": TARGET_FILE.as_posix(),
+                "hashtype": "blake3",
+                "hash": "8c6b2f215f0333437cdc3fe7c79be0c802847d2f2a0ccdc0bb251814e63cf375",
+                "size": 1482,
+            }
+        ]
+    }
+    # write to LAZY_DATA_RESOLVER_PATH
+    with open(LAZY_DATA_RESOLVER_PATH, "w") as f:
+        json.dump(bptr_manifest, f)
+
+
+@pytest.mark.skipif(not TRUSS_TRANSFER_AVAILABLE, reason="Truss Transfer not available")
+def test_lazy_data_resolver_v2():
+    # truss_transfer reads from LAZY_DATA_RESOLVER_PATH
+    if LAZY_DATA_RESOLVER_PATH.exists():
+        LAZY_DATA_RESOLVER_PATH.unlink()
+    try:
+        LAZY_DATA_RESOLVER_PATH.mkdir(parents=True, exist_ok=True)
+    except Exception as e:
+        pytest.skip(f"Unable to create {LAZY_DATA_RESOLVER_PATH}: {e}")
+
+    # without LAZY_DATA_RESOLVER_PATH -> does not create folder / file
+    with tempfile.TemporaryDirectory() as tempdir:
+        data_dir = Path(tempdir)
+        resolver = LazyDataResolverV2(data_dir)
+        assert not (data_dir / TARGET_FILE).exists()
+
+    # with LAZY_DATA_RESOLVER_PATH -> fetches data
+    with tempfile.TemporaryDirectory() as tempdir:
+        data_dir = Path(tempdir)
+        write_bptr_manifest_to_file()
+        resolver = LazyDataResolverV2(data_dir)
+        resolver.fetch()
+        assert (data_dir / TARGET_FILE).exists()
+        assert (data_dir / TARGET_FILE).stat().st_size == 1482
+
+    # with expired LAZY_DATA_RESOLVER_PATH -> raises exception
+    with tempfile.TemporaryDirectory() as tempdir:
+        data_dir = Path(tempdir)
+        write_bptr_manifest_to_file(expiration_timestamp=int(time.time()) - 1)
+        resolver = LazyDataResolverV2(data_dir)
+        with pytest.raises(Exception):
+            resolver.fetch()
+
+
+if __name__ == "__main__":
+    test_lazy_data_resolver_v2()