Return a presigned URL

Author: blueraft

src/nomad_ml_workflows/actions/export_remote_entries/activities.py

@@ -105,13 +105,33 @@ def _unique_upload_filename(filename: str, upload_files: StagingUploadFiles) ->
         count += 1
 
 
-def _parse_s3_uri(remote_uri: str) -> tuple[str, str]:
-    """Parse an S3 URI into its bucket and object key/prefix."""
+DEFAULT_PRESIGNED_URL_EXPIRATION_SECONDS = 2 * 24 * 3600  # 2 days (172800s)
+
+
+def _parse_s3_uri(
+    remote_uri: str, default_bucket: str | None = None
+) -> tuple[str, str]:
+    """Parse an S3 URI or presigned URL into its bucket and object key/prefix."""
     parsed_uri = urlparse(remote_uri)
-    key = parsed_uri.path.lstrip('/')
-    if parsed_uri.scheme != 's3' or not parsed_uri.netloc or not key:
-        raise ValueError(f'Invalid S3 URI: {remote_uri}')
-    return parsed_uri.netloc, key
+    if parsed_uri.scheme == 's3':
+        key = parsed_uri.path.lstrip('/')
+        if not parsed_uri.netloc or not key:
+            raise ValueError(f'Invalid S3 URI: {remote_uri}')
+        return parsed_uri.netloc, key
+    elif parsed_uri.scheme in ('http', 'https'):
+        path = parsed_uri.path.lstrip('/')
+        if default_bucket and path.startswith(f'{default_bucket}/'):
+            key = path[len(default_bucket) + 1 :]
+            return default_bucket, key
+        elif default_bucket:
+            return default_bucket, path
+        parts = path.split('/', 1)
+        if len(parts) > 1:
+            return parts[0], parts[1]
+        elif parts and parts[0]:
+            return parsed_uri.netloc, parts[0]
+        raise ValueError(f'Could not determine S3 bucket/key from URL: {remote_uri}')
+    raise ValueError(f'Invalid S3 URI: {remote_uri}')
 
 
 def _upload_dataset_to_s3(
@@ -120,7 +140,7 @@ def _upload_dataset_to_s3(
     exportable_filepaths: list[Path],
     artifacts_subdirectory: Path,
 ) -> str:
-    """Upload exported dataset files to S3-compatible remote storage."""
+    """Upload exported dataset files to S3-compatible remote storage and return a presigned download URL."""
     import boto3
 
     client_kwargs = _build_boto3_client_kwargs(storage_settings)
@@ -136,12 +156,32 @@ def _upload_dataset_to_s3(
 
         object_key = _build_s3_key(prefix, f'{data.exportable_dir_name}.zip')
         s3_client.upload_file(zippath.as_posix(), bucket, object_key)
-        return f's3://{bucket}/{object_key}'
+        try:
+            return s3_client.generate_presigned_url(
+                'get_object',
+                Params={'Bucket': bucket, 'Key': object_key},
+                ExpiresIn=DEFAULT_PRESIGNED_URL_EXPIRATION_SECONDS,
+            )
+        except Exception:
+            return f's3://{bucket}/{object_key}'
 
     base_key_prefix = _build_s3_key(prefix, data.exportable_dir_name)
+    primary_object_key = None
     for filepath in exportable_filepaths:
         object_key = f'{base_key_prefix}/{filepath.name}'
         s3_client.upload_file(filepath.as_posix(), bucket, object_key)
+        if filepath.stem == DATA_FILE_NAME or primary_object_key is None:
+            primary_object_key = object_key
+
+    if primary_object_key:
+        try:
+            return s3_client.generate_presigned_url(
+                'get_object',
+                Params={'Bucket': bucket, 'Key': primary_object_key},
+                ExpiresIn=DEFAULT_PRESIGNED_URL_EXPIRATION_SECONDS,
+            )
+        except Exception:
+            pass
 
     return f's3://{bucket}/{base_key_prefix}/'
 
@@ -153,7 +193,7 @@ def _copy_dataset_from_s3_to_upload(
     """Download an S3 dataset and add it to a NOMAD staging upload."""
     import boto3
 
-    bucket, key = _parse_s3_uri(data.remote_uri)
+    bucket, key = _parse_s3_uri(data.remote_uri, default_bucket=storage_settings.bucket)
     client_kwargs = _build_boto3_client_kwargs(storage_settings)
     s3_client = boto3.client('s3', **client_kwargs)
     upload_files = _get_staging_upload_files(data.user_id, data.upload_id)

tests/actions/export_remote_entries/test_activities.py

@@ -48,6 +48,9 @@ async def test_upload_dataset_to_remote_storage_s3_zipped(
     mock_artifacts_dir.return_value = artifacts_dir.as_posix()
 
     mock_s3 = MagicMock()
+    mock_s3.generate_presigned_url.return_value = (
+        'https://my-bucket.s3.amazonaws.com/my-folder/export_entries_2026.zip?signed'
+    )
     mock_boto_client.return_value = mock_s3
 
     storage_settings = S3StorageSettings(
@@ -84,7 +87,15 @@ async def test_upload_dataset_to_remote_storage_s3_zipped(
     call_args = mock_s3.upload_file.call_args[0]
     assert call_args[1] == 'my-bucket'
     assert call_args[2] == 'my-folder/export_entries_2026.zip'
-    assert remote_uri == 's3://my-bucket/my-folder/export_entries_2026.zip'
+    mock_s3.generate_presigned_url.assert_called_once_with(
+        'get_object',
+        Params={'Bucket': 'my-bucket', 'Key': 'my-folder/export_entries_2026.zip'},
+        ExpiresIn=172800,
+    )
+    assert (
+        remote_uri
+        == 'https://my-bucket.s3.amazonaws.com/my-folder/export_entries_2026.zip?signed'
+    )
 
 
 @pytest.mark.asyncio
@@ -103,6 +114,9 @@ async def test_upload_dataset_to_remote_storage_s3_unzipped(
     mock_artifacts_dir.return_value = artifacts_dir.as_posix()
 
     mock_s3 = MagicMock()
+    mock_s3.generate_presigned_url.return_value = (
+        'https://my-bucket.s3.amazonaws.com/export_entries_2026/data.csv?signed'
+    )
     mock_boto_client.return_value = mock_s3
 
     storage_settings = S3StorageSettings(
@@ -125,7 +139,15 @@ async def test_upload_dataset_to_remote_storage_s3_unzipped(
     assert 'export_entries_2026/metadata.json' in uploaded_keys
     assert 'export_entries_2026/selected_entries.json' in uploaded_keys
     assert 'export_entries_2026/data.csv' in uploaded_keys
-    assert remote_uri == 's3://my-bucket/export_entries_2026/'
+    mock_s3.generate_presigned_url.assert_called_once_with(
+        'get_object',
+        Params={'Bucket': 'my-bucket', 'Key': 'export_entries_2026/data.csv'},
+        ExpiresIn=172800,
+    )
+    assert (
+        remote_uri
+        == 'https://my-bucket.s3.amazonaws.com/export_entries_2026/data.csv?signed'
+    )
 
 
 @pytest.mark.asyncio
@@ -144,7 +166,7 @@ async def test_copy_remote_dataset_to_upload_s3_zipped(
     data = CopyRemoteDatasetToUploadInput(
         user_id='user-123',
         upload_id='upload-123',
-        remote_uri='s3://my-bucket/exports/dataset.zip',
+        remote_uri='https://s3.example.com/my-bucket/exports/dataset.zip?X-Amz-Signature=123',
         storage_settings=S3StorageSettings(
             bucket='my-bucket',
             endpoint_url='https://s3.example.com',