Feat 50: Dynamically generate job template (#73)

* generate, save, and download job template xlsx

* update UI link, add error handling

* unit tests for download job template

* update /api/job_upload_template to use StreamingResponse

* create job_template_configs.py

* create JobUploadTemplate class to handle template creation

* remove old upload_template_link

* add test suite for JobUploadTemplate and update unit tests for download endpoint

* code cleanup

* cleanup openpyxl

* address PR comments

src/aind_data_transfer_service/configs/job_upload_template.py

@@ -0,0 +1,116 @@
+"""Module to configure and create xlsx job upload template"""
+import datetime
+from io import BytesIO
+
+from aind_data_schema.core.data_description import Modality, Platform
+from openpyxl import Workbook
+from openpyxl.styles import Font
+from openpyxl.utils import get_column_letter
+from openpyxl.worksheet.datavalidation import DataValidation
+
+
+# TODO: convert to pydantic model
+class JobUploadTemplate:
+    """Class to configure and create xlsx job upload template"""
+
+    FILE_NAME = "job_upload_template.xlsx"
+    HEADERS = [
+        "platform",
+        "acq_datetime",
+        "subject_id",
+        "s3_bucket",
+        "modality0",
+        "modality0.source",
+        "modality1",
+        "modality1.source",
+    ]
+    SAMPLE_JOBS = [
+        [
+            Platform.BEHAVIOR.abbreviation,
+            datetime.datetime(2023, 10, 4, 4, 0, 0),
+            "123456",
+            "aind-behavior-data",
+            Modality.BEHAVIOR_VIDEOS.abbreviation,
+            "/allen/aind/stage/fake/dir",
+            Modality.BEHAVIOR.abbreviation,
+            "/allen/aind/stage/fake/dir",
+        ],
+        [
+            Platform.SMARTSPIM.abbreviation,
+            datetime.datetime(2023, 3, 4, 16, 30, 0),
+            "654321",
+            "aind-open-data",
+            Modality.SPIM.abbreviation,
+            "/allen/aind/stage/fake/dir",
+        ],
+        [
+            Platform.ECEPHYS.abbreviation,
+            datetime.datetime(2023, 1, 30, 19, 1, 0),
+            "654321",
+            "aind-ephys-data",
+            Modality.ECEPHYS.abbreviation,
+            "/allen/aind/stage/fake/dir",
+            Modality.BEHAVIOR_VIDEOS.abbreviation,
+            "/allen/aind/stage/fake/dir",
+        ],
+    ]
+    VALIDATORS = [
+        {
+            "name": "platform",
+            "options": [p().abbreviation for p in Platform._ALL],
+            "ranges": ["A2:A20"],
+        },
+        {
+            "name": "modality",
+            "options": [m().abbreviation for m in Modality._ALL],
+            "ranges": ["E2:E20", "G2:G20"],
+        },
+        {
+            "name": "s3_bucket",
+            "options": [
+                "aind-ephys-data",
+                "aind-ophys-data",
+                "aind-behavior-data",
+                "aind-private-data",
+            ],
+            "ranges": ["D2:D20"],
+        },
+    ]
+
+    @staticmethod
+    def create_job_template():
+        """Create job template as xlsx filestream"""
+        # job template
+        xl_io = BytesIO()
+        workbook = Workbook()
+        worksheet = workbook.active
+        worksheet.append(JobUploadTemplate.HEADERS)
+        for job in JobUploadTemplate.SAMPLE_JOBS:
+            worksheet.append(job)
+        # data validators
+        for validator in JobUploadTemplate.VALIDATORS:
+            dv = DataValidation(
+                type="list",
+                formula1=f'"{(",").join(validator["options"])}"',
+                allow_blank=True,
+                showErrorMessage=True,
+                showInputMessage=True,
+            )
+            dv.promptTitle = validator["name"]
+            dv.prompt = f'Select a {validator["name"]} from the dropdown'
+            dv.error = f'Invalid {validator["name"]}.'
+            for r in validator["ranges"]:
+                dv.add(r)
+            worksheet.add_data_validation(dv)
+        # formatting
+        bold = Font(bold=True)
+        for header in worksheet["A1:H1"]:
+            for cell in header:
+                cell.font = bold
+                worksheet.column_dimensions[
+                    get_column_letter(cell.column)
+                ].auto_size = True
+        # save file
+        workbook.save(xl_io)
+        workbook.close()
+        return xl_io

src/aind_data_transfer_service/server.py

@@ -7,10 +7,10 @@
 from asyncio import sleep
 from pathlib import PurePosixPath
 
-import openpyxl
 from fastapi import Request
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 from fastapi.templating import Jinja2Templates
+from openpyxl import load_workbook
 from pydantic import SecretStr
 from starlette.applications import Starlette
 from starlette.routing import Route
@@ -19,6 +19,9 @@
     BasicUploadJobConfigs,
     HpcJobConfigs,
 )
+from aind_data_transfer_service.configs.job_upload_template import (
+    JobUploadTemplate,
+)
 from aind_data_transfer_service.hpc.client import HpcClient, HpcClientConfigs
 from aind_data_transfer_service.hpc.models import (
     HpcJobStatusResponse,
@@ -32,7 +35,6 @@
 templates = Jinja2Templates(directory=template_directory)
 
 # TODO: Add server configs model
-# UPLOAD_TEMPLATE_LINK
 # HPC_SIF_LOCATION
 # HPC_USERNAME
 # HPC_LOGGING_DIRECTORY
@@ -63,11 +65,13 @@ async def validate_csv(request: Request):
                 # byte chars. Adding "utf-8-sig" should remove them.
                 data = content.decode("utf-8-sig")
             else:
-                xlsx_sheet = openpyxl.load_workbook(io.BytesIO(content)).active
+                xlsx_book = load_workbook(io.BytesIO(content), read_only=True)
+                xlsx_sheet = xlsx_book.active
                 csv_io = io.StringIO()
                 csv_writer = csv.writer(csv_io)
                 for r in xlsx_sheet.rows:
                     csv_writer.writerow([cell.value for cell in r])
+                xlsx_book.close()
                 data = csv_io.getvalue()
             csv_reader = csv.DictReader(io.StringIO(data))
             for row in csv_reader:
@@ -274,7 +278,6 @@ async def index(request: Request):
         context=(
             {
                 "request": request,
-                "upload_template_link": os.getenv("UPLOAD_TEMPLATE_LINK"),
             }
         ),
     )
@@ -312,12 +315,37 @@ async def jobs(request: Request):
                 "request": request,
                 "job_status_list": job_status_list,
                 "num_of_jobs": len(job_status_list),
-                "upload_template_link": os.getenv("UPLOAD_TEMPLATE_LINK"),
             }
         ),
     )
 
 
+def download_job_template(request: Request):
+    """Get job template as xlsx filestream for download"""
+    try:
+        xl_io = JobUploadTemplate.create_job_template()
+    except Exception as e:
+        return JSONResponse(
+            content={
+                "message": "Error creating job template",
+                "data": {"error": f"{e.__class__.__name__}{e.args}"},
+            },
+            status_code=500,
+        )
+    return StreamingResponse(
+        io.BytesIO(xl_io.getvalue()),
+        media_type=(
+            "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
+        ),
+        headers={
+            "Content-Disposition": (
+                f"attachment; filename={JobUploadTemplate.FILE_NAME}"
+            )
+        },
+        status_code=200,
+    )
+
+
 routes = [
     Route("/", endpoint=index, methods=["GET", "POST"]),
     Route("/api/validate_csv", endpoint=validate_csv, methods=["POST"]),
@@ -326,6 +354,11 @@ async def jobs(request: Request):
     ),
     Route("/api/submit_hpc_jobs", endpoint=submit_hpc_jobs, methods=["POST"]),
     Route("/jobs", endpoint=jobs, methods=["GET"]),
+    Route(
+        "/api/job_upload_template",
+        endpoint=download_job_template,
+        methods=["GET"],
+    ),
 ]
 
 app = Starlette(routes=routes)

src/aind_data_transfer_service/templates/index.html

@@ -52,7 +52,7 @@
     <nav>
         <a href="/">Submit Jobs</a> |
         <a href="/jobs">Job Status</a> |
-        <a href= "{{ upload_template_link }}" >Job Submit Template</a>
+        <a title="Download job template as .xslx" href= "/api/job_upload_template" download>Job Submit Template</a>
     </nav>
     <h2>Submit Jobs</h2>
         <div>

src/aind_data_transfer_service/templates/job_status.html

@@ -26,7 +26,7 @@
     <nav>
         <a href="/">Submit Jobs</a> |
         <a href="/jobs">Job Status</a> |
-        <a href= "{{ upload_template_link }}" >Job Submit Template</a>
+        <a title="Download job template as .xslx" href= "/api/job_upload_template" download>Job Submit Template</a>
     </nav>
     <h2>Jobs Submitted: {{num_of_jobs}}</h2>
     <table>

tests/test_job_upload_template.py

@@ -0,0 +1,58 @@
+"""Module to test job upload template configs and generation"""
+
+import os
+import unittest
+from pathlib import Path
+
+from openpyxl import load_workbook
+
+from aind_data_transfer_service.configs.job_upload_template import (
+    JobUploadTemplate,
+)
+
+TEST_DIRECTORY = Path(os.path.dirname(os.path.realpath(__file__)))
+SAMPLE_JOB_TEMPLATE = TEST_DIRECTORY / "resources" / "job_upload_template.xlsx"
+
+
+class TestJobUploadTemplate(unittest.TestCase):
+    """Tests job upload template class"""
+
+    def read_xl_helper(self, source, return_validators=False):
+        """Helper function to read xlsx contents and validators"""
+        lines = []
+        workbook = load_workbook(source, read_only=(not return_validators))
+        worksheet = workbook.active
+        for row in worksheet.rows:
+            row_contents = [cell.value for cell in row]
+            lines.append(row_contents)
+        if return_validators:
+            validators = []
+            for dv in worksheet.data_validations.dataValidation:
+                validators.append(
+                    {
+                        "name": dv.promptTitle,
+                        "options": dv.formula1.strip('"').split(","),
+                        "ranges": str(dv.cells).split(" "),
+                    }
+                )
+            result = (lines, validators)
+        else:
+            result = lines
+        workbook.close()
+        return result
+
+    def test_create_job_template(self):
+        """Tests that xlsx job template is created with
+        correct contents and validators"""
+        expected_lines = self.read_xl_helper(SAMPLE_JOB_TEMPLATE)
+        (template_lines, template_validators) = self.read_xl_helper(
+            JobUploadTemplate.create_job_template(), True
+        )
+        self.assertEqual(expected_lines, template_lines)
+        self.assertCountEqual(
+            JobUploadTemplate.VALIDATORS, template_validators
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_server.py

@@ -4,6 +4,7 @@
 import os
 import unittest
 from copy import deepcopy
+from io import BytesIO
 from pathlib import Path, PurePosixPath
 from unittest.mock import MagicMock, patch
 
@@ -507,6 +508,45 @@ def test_jobs_failure(self, mock_get: MagicMock):
         self.assertEqual(response.status_code, 200)
         self.assertIn("Submit Jobs", response.text)
 
+    @patch(
+        "aind_data_transfer_service.configs.job_upload_template"
+        ".JobUploadTemplate.create_job_template"
+    )
+    def test_download_job_template(self, mock_create_template: MagicMock):
+        """Tests that job template downloads as xlsx file."""
+        mock_create_template.return_value = BytesIO(b"mock_template_stream")
+        with TestClient(app) as client:
+            response = client.get("/api/job_upload_template")
+        expected_file_name_header = (
+            "attachment; filename=job_upload_template.xlsx"
+        )
+        self.assertEqual(1, mock_create_template.call_count)
+        self.assertEqual(200, response.status_code)
+        self.assertEqual(
+            expected_file_name_header, response.headers["Content-Disposition"]
+        )
+
+    @patch(
+        "aind_data_transfer_service.configs.job_upload_template"
+        ".JobUploadTemplate.create_job_template"
+    )
+    def test_download_invalid_job_template(
+        self, mock_create_template: MagicMock
+    ):
+        """Tests that download invalid job template returns errors."""
+        mock_create_template.side_effect = Exception(
+            "mock invalid job template"
+        )
+        with TestClient(app) as client:
+            response = client.get("/api/job_upload_template")
+        expected_response = {
+            "message": "Error creating job template",
+            "data": {"error": "Exception('mock invalid job template',)"},
+        }
+        self.assertEqual(1, mock_create_template.call_count)
+        self.assertEqual(500, response.status_code)
+        self.assertEqual(expected_response, response.json())
+
 
 if __name__ == "__main__":
     unittest.main()