diff --git a/lamindb/.agents b/lamindb/.agents index d109acd4d..22038d616 160000 --- a/lamindb/.agents +++ b/lamindb/.agents @@ -1 +1 @@ -Subproject commit d109acd4d939cbcf362aee56e599f9897706cc67 +Subproject commit 22038d616515b00fa868203e0d8bf2fae8845150 diff --git a/lamindb/core/storage/paths.py b/lamindb/core/storage/paths.py index 67923f2a2..c027ff8a9 100644 --- a/lamindb/core/storage/paths.py +++ b/lamindb/core/storage/paths.py @@ -2,6 +2,7 @@ import shutil from typing import TYPE_CHECKING +from urllib.parse import urlencode import fsspec from lamindb_setup.core import StorageSettings @@ -11,6 +12,7 @@ ) from lamindb.core._settings import settings +import lamindb_setup as ln_setup if TYPE_CHECKING: from lamindb_setup.types import AnyPath, AnyPathStr @@ -45,6 +47,58 @@ def auto_storage_key_from_artifact_uid( return storage_key +def join_uri_path(root: str, path: str) -> str: + normalized_root = root.rstrip("/") + normalized_path = path.lstrip("/") + return ( + f"{normalized_root}/{normalized_path}" if normalized_path else normalized_root + ) + + +def _serialize_name_param(name: str | None) -> str: + if name is None: + return "" + return f"?{urlencode({'name': name})}" + + +def create_download_link( + artifact: Artifact, + *, + name: str | None = None, + public_uid: str | None = None, + using_key: str | None = None, + origin: str | None = None, +) -> str: + storage_key = auto_storage_key_from_artifact(artifact) + filepath, storage_settings = filepath_from_artifact(artifact, using_key=using_key) + storage_root = storage_settings.root_as_str if storage_settings is not None else "" + + if storage_root.startswith("https://"): + filepath = join_uri_path(storage_root, storage_key) + + elif storage_root.startswith("gs://"): + gcs_root = ( + f"https://storage.googleapis.com/{storage_root.removeprefix('gs://')}" + ) + filepath = join_uri_path(gcs_root, storage_key) + + elif storage_root.startswith("s3://") and ln_setup.settings.instance.is_on_hub: + origin = origin or ln_setup.settings.instance.ui_url + search_params = _serialize_name_param(name) + common = ( + storage_root.removeprefix("s3://").rstrip("/") + + f"%2F/{storage_key}" + + search_params + ) + if public_uid is not None: + filepath = f"{origin or ''}/storage/public/{public_uid}/{common}" + if origin is not None: + filepath = f"{origin}/storage/s3/{common}" + else: + filepath = artifact.path + return str(filepath) + + def check_path_is_child_of_root(path: AnyPathStr, root: AnyPathStr) -> bool: if fsspec.utils.get_protocol(str(path)) != fsspec.utils.get_protocol(str(root)): return False diff --git a/lamindb/models/artifact.py b/lamindb/models/artifact.py index 286307df1..c1fcb14da 100644 --- a/lamindb/models/artifact.py +++ b/lamindb/models/artifact.py @@ -94,6 +94,7 @@ def _lazy_load_storage_module(): auto_storage_key_from_artifact, auto_storage_key_from_artifact_uid, check_path_is_child_of_root, + create_download_link, filepath_cache_key_from_artifact, filepath_from_artifact, ) @@ -106,6 +107,7 @@ def _lazy_load_storage_module(): auto_storage_key_from_artifact=auto_storage_key_from_artifact, auto_storage_key_from_artifact_uid=auto_storage_key_from_artifact_uid, check_path_is_child_of_root=check_path_is_child_of_root, + create_download_link=create_download_link, filepath_cache_key_from_artifact=filepath_cache_key_from_artifact, filepath_from_artifact=filepath_from_artifact, ) @@ -2113,6 +2115,28 @@ def path(self) -> UPath: filepath, _ = _s().filepath_from_artifact(self, using_key=settings._using_key) return filepath + def create_download_link( + self, + *, + name: str | None = None, + public_uid: str | None = None, + origin: str | None = None, + ) -> str: + """Create a download link with optional routing parameters. + + Args: + name: Optional display name query parameter in the download URL. + public_uid: Optional public UID for `/storage/public/...` routes. + origin: Optional origin for `/storage/s3/...` routes. + """ + return _s().create_download_link( + self, + name=name, + public_uid=public_uid, + origin=origin, + using_key=settings._using_key, + ) + @property def _cache_path(self) -> UPath: filepath, cache_key = _s().filepath_cache_key_from_artifact( diff --git a/tests/storage/test_artifact_storage.py b/tests/storage/test_artifact_storage.py index c2f9160bc..070c79217 100644 --- a/tests/storage/test_artifact_storage.py +++ b/tests/storage/test_artifact_storage.py @@ -206,3 +206,62 @@ def test_single_file_directory_preserved(tmp_path): assert [file.name for file in artifact.path.iterdir()] == ["only.txt"] artifact.delete(permanent=True) + + +def test_create_download_link_s3_routes(): + artifact = ln.Artifact("s3://lamindb-ci/test-data/test.parquet").save() + root = str(artifact.storage.root).removeprefix("s3://").rstrip("/") + expected_common = f"{root}%2F/{artifact.key}?name=my+table.parquet" + + public_link = artifact.create_download_link( + name="my table.parquet", + public_uid="pub123", + ) + proxy_link = artifact.create_download_link(name="my table.parquet") + native_link = artifact.path.as_posix() + + if ln.setup.settings.instance.is_on_hub: + origin = ln.setup.settings.instance.ui_url + if origin is None: + assert public_link == f"/storage/public/pub123/{expected_common}" + assert proxy_link == native_link + else: + assert public_link == f"{origin}/storage/s3/{expected_common}" + assert proxy_link == f"{origin}/storage/s3/{expected_common}" + else: + assert public_link == native_link + assert proxy_link == native_link + + artifact.delete(permanent=True, storage=False) + + +def test_create_download_link_gcs_root(): + artifact = ln.Artifact( + "gs://rxrx1-europe-west4/images/test/HEPG2-08/Plate1/B02_s1_w1.png", + description="Test GCP file for download link", + ).save() + root = str(artifact.storage.root).removeprefix("gs://").rstrip("/") + + link = artifact.create_download_link() + + assert ( + link + == f"https://storage.googleapis.com/{root}/{artifact.key.lstrip('/')}" + ) + + artifact.delete(permanent=True, storage=False) + + +def test_create_download_link_https_root_real_artifact(): + storage = ln.Storage("https://example.com").save() + artifact = ln.Artifact( + "https://example.com/files/document.txt", + storage=storage, + skip_check_exists=True, + ).save() + + assert artifact.storage.root == "https://example.com" + assert artifact.key == "files/document.txt" + assert artifact.create_download_link() == "https://example.com/files/document.txt" + + artifact.delete(permanent=True, storage=False)