From ff65b686c471ef779cdd3d81b0d1b2200685cd23 Mon Sep 17 00:00:00 2001
From: "David H. Irving" <david.irving@noirlab.edu>
Date: Wed, 31 Jan 2024 16:15:28 -0700
Subject: [PATCH] Allow RemoteButler to be used for cutout backend

Create Butlers using LabeledButlerFactory and provide the Gafaelfawr access token.  This allows us to use Butler client/server if the environment is configured to do so.
---
 Dockerfile.worker        |  2 +-
 src/vocutouts/actors.py  |  3 +++
 src/vocutouts/policy.py  |  1 +
 src/vocutouts/workers.py | 26 +++++++++++++++-----------
 tests/conftest.py        |  2 ++
 5 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/Dockerfile.worker b/Dockerfile.worker
index c464cd4..19ad863 100644
--- a/Dockerfile.worker
+++ b/Dockerfile.worker
@@ -2,7 +2,7 @@
 # are based on stack containers and install any required supporting code
 # for the image cutout backend, Dramatiq, and the backend worker definition.
 
-FROM lsstsqre/centos:7-stack-lsst_distrib-w_2022_40
+FROM lsstsqre/centos:7-stack-lsst_distrib-w_2024_04
 
 # Reset the user to root since we need to do system install tasks.
 USER root
diff --git a/src/vocutouts/actors.py b/src/vocutouts/actors.py
index de7cc0e..7313122 100644
--- a/src/vocutouts/actors.py
+++ b/src/vocutouts/actors.py
@@ -57,6 +57,7 @@ def cutout(
     job_id: str,
     dataset_ids: list[str],
     stencils: list[dict[str, Any]],
+    access_token: str,
 ) -> list[dict[str, str]]:
     """Stub for a circle cutout.
 
@@ -77,6 +78,8 @@ def cutout(
         JSON-serializable (a requirement for Dramatiq) representations of the
         `~vocutouts.models.stencils.Stencil` objects corresponding to the
         user's request.
+    access_token
+        Gafaelfawr access token used to authenticate to Butler server.
 
     Returns
     -------
diff --git a/src/vocutouts/policy.py b/src/vocutouts/policy.py
index 20918cd..51c7d42 100644
--- a/src/vocutouts/policy.py
+++ b/src/vocutouts/policy.py
@@ -66,6 +66,7 @@ def dispatch(self, job: Job, access_token: str) -> Message:
                 job.job_id,
                 cutout_params.ids,
                 [s.to_dict() for s in cutout_params.stencils],
+                access_token,
             ),
             time_limit=job.execution_duration * 1000,
             on_success=job_completed,
diff --git a/src/vocutouts/workers.py b/src/vocutouts/workers.py
index e5f23a3..8006491 100644
--- a/src/vocutouts/workers.py
+++ b/src/vocutouts/workers.py
@@ -28,7 +28,7 @@
 from dramatiq.middleware import CurrentMessage
 from dramatiq.results import Results
 from dramatiq.results.backends import RedisBackend
-from lsst.daf.butler import Butler
+from lsst.daf.butler import LabeledButlerFactory
 from lsst.image_cutout_backend import ImageCutoutBackend, projection_finders
 from lsst.image_cutout_backend.stencils import (
     SkyCircle,
@@ -37,8 +37,7 @@
 )
 from safir.logging import configure_logging
 
-BACKENDS: dict[str, ImageCutoutBackend] = {}
-"""Cache of image cutout backends by Butler repository label."""
+BUTLER_FACTORY = LabeledButlerFactory()
 
 configure_logging(
     name=os.getenv("SAFIR_LOGGER", "vocutouts"),
@@ -86,7 +85,7 @@ class TaskTransientError(Exception):
     """Some transient problem occurred."""
 
 
-def get_backend(butler_label: str) -> ImageCutoutBackend:
+def get_backend(butler_label: str, access_token: str) -> ImageCutoutBackend:
     """Given the Butler label, retrieve or build a backend.
 
     The dataset ID will be a URI of the form ``butler://<label>/<uuid>``.
@@ -103,15 +102,17 @@ def get_backend(butler_label: str) -> ImageCutoutBackend:
     lsst.image_cutout_backend.ImageCutoutBackend
         Backend to use.
     """
-    if butler_label in BACKENDS:
-        return BACKENDS[butler_label]
-    butler = Butler(butler_label)
+
+    butler = BUTLER_FACTORY.create_butler(
+        label=butler_label, access_token=access_token
+    )
+    # At present, projection finders and image cutout backend have no internal
+    # caching and are cheap to construct, so we just make a new one for each
+    # request.
     projection_finder = projection_finders.ProjectionFinder.make_default()
     output = os.environ["CUTOUT_STORAGE_URL"]
     tmpdir = os.environ.get("CUTOUT_TMPDIR", "/tmp")
-    backend = ImageCutoutBackend(butler, projection_finder, output, tmpdir)
-    BACKENDS[butler_label] = backend
-    return backend
+    return ImageCutoutBackend(butler, projection_finder, output, tmpdir)
 
 
 def parse_uri(uri: str) -> tuple[str, UUID]:
@@ -138,6 +139,7 @@ def cutout(
     job_id: str,
     dataset_ids: list[str],
     stencils: list[dict[str, Any]],
+    access_token: str,
 ) -> list[dict[str, str]]:
     """Perform a cutout.
 
@@ -160,6 +162,8 @@ def cutout(
         JSON-serializable (a requirement for Dramatiq) representations of the
         `~vocutouts.models.stencils.Stencil` objects corresponding to the
         user's request.
+    access_token
+        Gafaelfawr access token used to authenticate to Butler server.
 
     Returns
     -------
@@ -189,7 +193,7 @@ def cutout(
 
     # Parse the dataset ID and retrieve an appropriate backend.
     butler_label, uuid = parse_uri(dataset_ids[0])
-    backend = get_backend(butler_label)
+    backend = get_backend(butler_label, access_token)
 
     # Convert the stencils to SkyStencils.
     sky_stencils: list[SkyStencil] = []
diff --git a/tests/conftest.py b/tests/conftest.py
index 831d0cc..00c403a 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -32,11 +32,13 @@ def cutout_test(
     job_id: str,
     dataset_ids: list[str],
     stencils: list[dict[str, Any]],
+    access_token: str,
 ) -> list[dict[str, Any]]:
     message = CurrentMessage.get_current_message()
     now = isodatetime(datetime.now(tz=timezone.utc))
     job_started.send(job_id, message.message_id, now)
     assert len(dataset_ids) == 1
+    assert access_token == "sometoken"
     return [
         {
             "result_id": "cutout",