Merge pull request #188 from American-Institutes-for-Research/HEA-764/Add-an-AssetDownloadView-to-Django

Hea 764/add an asset download view to django

Author: rhunwicks

apps/common/views.py

@@ -1,6 +1,10 @@
 import os
 
+import fsspec
+from dagster import AssetKey, DagsterEventType, DagsterInstance, EventRecordsFilter
 from django.contrib.auth.mixins import LoginRequiredMixin, PermissionRequiredMixin
+from django.http import Http404, HttpResponseRedirect
+from django.views import View
 from revproxy.views import ProxyView
 
 
@@ -9,3 +13,45 @@ class DagsterProxyView(LoginRequiredMixin, PermissionRequiredMixin, ProxyView):
     upstream = f"{os.environ.get('DAGSTER_WEBSERVER_URL')}/{os.environ.get('DAGSTER_WEBSERVER_PREFIX')}/"
     permission_required = "common.access_dagster_ui"
     raise_exception = False
+
+
+class AssetDownloadView(LoginRequiredMixin, PermissionRequiredMixin, View):
+    """
+    A view that generates the s3 url for direct download of materialized assets
+    """
+
+    permission_required = "common.access_dagster_ui"
+    login_url = "/admin/login/"
+    raise_exception = False
+
+    def get(self, request, asset_name, partition_name=None):
+        try:
+            with DagsterInstance.get() as instance:
+                filter_kwargs = {
+                    "event_type": DagsterEventType.ASSET_MATERIALIZATION,
+                    "asset_key": AssetKey([asset_name]),
+                }
+                if partition_name:
+                    filter_kwargs["asset_partitions"] = [partition_name]
+
+                records = list(instance.get_event_records(EventRecordsFilter(**filter_kwargs), limit=1))
+                if not records:
+                    raise Http404(f"No materialization found for asset '{asset_name}'")
+
+                metadata = records[0].asset_materialization.metadata or {}
+                asset_uri = metadata.get("path")
+                if not asset_uri:
+                    raise Http404(f"Asset '{asset_name}' does not contain a 'path' in metadata")
+
+                # Generate signed URL for S3
+                protocol, path = asset_uri.value.split("://", 1)
+                fs = fsspec.filesystem(protocol)
+
+                if not hasattr(fs, "sign"):
+                    raise NotImplementedError(f"Filesystem '{protocol}' does not support pre-signed URLs")
+
+                signed_url = fs.sign(path)
+                return HttpResponseRedirect(signed_url)
+
+        except Exception as e:
+            raise Http404(f"Failed to locate or sign asset '{asset_name}': {str(e)}")

docker-compose.yml

@@ -111,6 +111,7 @@ services:
       SUPPORT_EMAIL_ADDRESS: ${SUPPORT_EMAIL_ADDRESS}
       DJANGO_MIGRATE: 1
       GOOGLE_APPLICATION_CREDENTIALS: ${GOOGLE_APPLICATION_CREDENTIALS}
+      BASELINE_EXPLORER_API_ROOT_URL: ${BASELINE_EXPLORER_API_ROOT_URL}
     command:
       - --timeout=3600
       - --workers=12

env.example

@@ -54,4 +54,7 @@ BSS_FILES_STORAGE_OPTIONS='{"token": "service_account", "access": "read_only", "
 # LAUNCHER can be used to configure a wrapper program around the Python process
 # For example, to add ddtrace or debugpy
 # Use the VSCode debugger as a launcher
-# LAUNCHER = "python3 -m debugpy --listen 0.0.0.0:5679"
+# LAUNCHER = "python3 -m debugpy --listen 0.0.0.0:5679"
+
+# Baseline Explorer API Root URL
+BASELINE_EXPLORER_API_ROOT_URL=http://localhost:8000

hea/urls.py

@@ -49,7 +49,7 @@
     WealthGroupViewSet,
     WildFoodGatheringViewSet,
 )
-from common.views import DagsterProxyView
+from common.views import AssetDownloadView, DagsterProxyView
 from common.viewsets import (
     ClassifiedProductViewSet,
     CountryViewSet,
@@ -141,8 +141,15 @@
         name="livelihood-zone-baseline-faceted-search",
     ),
     re_path(os.environ.get("DAGSTER_WEBSERVER_PREFIX", "pipelines") + r"/(?P<path>.*)", DagsterProxyView.as_view()),
+    path("assetdownload/<str:asset_name>/", AssetDownloadView.as_view(), name="asset_download"),
+    path(
+        "assetdownload/<str:asset_name>/<str:partition_name>/",
+        AssetDownloadView.as_view(),
+        name="asset_download_partitioned",
+    ),
 ]
 
+
 # Django's solution for translating JavaScript apps
 # Provides gettext translation functionality for Javascript clients (and ngettext, pgettext, iterpolate, etc.)
 # See: https://docs.djangoproject.com/en/4.2/topics/i18n/translation/#using-the-javascript-translation-catalog

pipelines/resources.py

@@ -1,7 +1,9 @@
 import json
+import os
 import pickle
 from abc import abstractmethod
 from typing import Any, Optional, Type
+from urllib.parse import urljoin
 
 import dagster._check as check
 import pandas as pd
@@ -10,15 +12,49 @@
     DagsterInvariantViolationError,
     InitResourceContext,
     InputContext,
+    MetadataValue,
     OutputContext,
 )
 from dagster._core.storage.upath_io_manager import UPathIOManager
 from dagster._utils import PICKLE_PROTOCOL
+from django.urls import reverse
 from pydantic.fields import Field
 from upath import UPath
 
 
-class PickleFilesystemIOManager(UPathIOManager):
+class DownloadAssetMixin:
+    """
+    Mixin that adds download URL of the asset to the metadata.
+    """
+
+    def get_download_url(self, context: OutputContext) -> str:
+        if not context or not getattr(context, "has_asset_key", False) or not context.has_asset_key:
+            return ""
+
+        asset_name = context.asset_key.path[-1]
+        root_url = os.getenv("BASELINE_EXPLORER_API_ROOT_URL", "http://localhost:8000")
+
+        kwargs = {"asset_name": asset_name}
+        if context.has_partition_key:
+            kwargs["partition_name"] = context.partition_key
+            relative_url = reverse("asset_download_partitioned", kwargs=kwargs)
+        else:
+            relative_url = reverse("asset_download", kwargs=kwargs)
+
+        return urljoin(root_url, relative_url)
+
+    def handle_output(self, context: OutputContext, obj):
+        download_url = self.get_download_url(context)
+        if download_url:
+            context.add_output_metadata(
+                {
+                    "download": MetadataValue.url(download_url),
+                }
+            )
+        super().handle_output(context, obj)
+
+
+class PickleFilesystemIOManager(DownloadAssetMixin, UPathIOManager):
     """
     Dagster I/O Manager that serializes Python objects to files using Pickle.
 
@@ -57,7 +93,7 @@ def load_from_path(self, context: InputContext, path: "UPath") -> Any:
             return pickle.load(file)
 
 
-class JSONFilesystemIOManager(UPathIOManager):
+class JSONFilesystemIOManager(DownloadAssetMixin, UPathIOManager):
     """
     Dagster I/O Manager that serializes Python objects to JSON files.
     """
@@ -77,7 +113,7 @@ def load_from_path(self, context: InputContext, path: "UPath") -> Any:
             return json.loads(file.read())
 
 
-class DataFrameCSVFilesystemIOManager(UPathIOManager):
+class DataFrameCSVFilesystemIOManager(DownloadAssetMixin, UPathIOManager):
     """
     Dagster I/O Manager that serializes DataFrames to CSV files.
     """
@@ -95,7 +131,7 @@ def load_from_path(self, context: InputContext, path: UPath) -> pd.DataFrame:
         return pd.read_csv(path)
 
 
-class DataFrameExcelFilesystemIOManager(UPathIOManager):
+class DataFrameExcelFilesystemIOManager(DownloadAssetMixin, UPathIOManager):
     """
     Dagster I/O Manager that serializes DataFrames to Excel .xlsx using the OpenPyXL engine.