feat(butler): Implement butler pool factory

tcjennings · tcjennings · commit 7f7275b67dc0 · 2025-03-31T16:55:03.000-05:00
diff --git a/src/lsst/cmservice/common/butler.py b/src/lsst/cmservice/common/butler.py
@@ -1,15 +1,38 @@
-"""Utility functions for working with butler commands"""
+"""Module in support of working with Butlers.
+
+A CM Service that has a configured ``BUTLER_REPO_INDEX`` provides a
+factory function that manages a pool of Butler instances for each of these
+repos. This factory function is inspired by the
+`lsst.daf.butler.LabeledButlerFactory` and provides ``clone()`` instances of
+available Butlers when asked to provide one.
+
+Notes
+-----
+The butler "factory" follows a "global" pattern where it is assigned to
+a module-level variable at import-time. Other modules can import this factory
+and use it to produce on-demand butler clones. It is not necessary to use a
+butler factory as an injected dependency when using this pattern, but this
+module should be imported as early as possible in the application startup;
+it does not depend on a running event loop.
+"""
 
-from functools import partial
+from collections.abc import Callable
+from functools import cache, cached_property, partial
+from pathlib import Path
+from typing import TYPE_CHECKING
 
 import yaml
-from anyio import Path, to_thread
+from anyio import to_thread
+from botocore.exceptions import ClientError
 from sqlalchemy.engine import url
+from sqlalchemy.exc import OperationalError
 
 from lsst.daf.butler import Butler, ButlerConfig, ButlerRepoIndex
 from lsst.daf.butler._exceptions import MissingCollectionError
+from lsst.daf.butler.direct_butler import DirectButler
+from lsst.daf.butler.registry import CollectionArgType
 from lsst.resources import ResourcePathExpression
-from lsst.utils.db_auth import DbAuth
+from lsst.utils.db_auth import DbAuth, DbAuthError
 
 from ..config import config
 from . import errors
@@ -24,51 +47,170 @@
 """
 
 
-async def get_butler_config(repo: str, *, without_datastore: bool = False) -> ButlerConfig:
-    """Create a butler config object for a repo known to the service's
-    environment.
+class ButlerFactory:
+    """When created, the LabeledButlerFactory will create an instance of each
+    Butler known to the application configuration. This occurs synchronously
+    so it is best performed at application startup. After initializing, the
+    factory can hand out ``clone()`` copies of available Butlers.
     """
 
-    try:
-        repo_uri: ResourcePathExpression = BUTLER_REPO_INDEX.get_repo_uri(label=repo)
-    except KeyError:
-        # No such repo known to the service
-        logger.warning("Butler repo %s not known to environment.", repo)
-        repo_uri = repo
-
-    try:
-        bc_f = partial(
-            ButlerConfig,
-            other=repo_uri,
-            without_datastore=without_datastore,
-        )
-        bc = await to_thread.run_sync(bc_f)
-    except FileNotFoundError:
-        # No such repo known to Butler
-        logger.error("Butler repo %s not known.", repo)
-        raise RuntimeError("Unknown Butler Repo %s", repo)
-
-    try:
-        db_auth_info = yaml.safe_load(await Path(config.butler.authentication_file).read_text())
-    except FileNotFoundError:
-        logger.error("No Butler Registry authentication secrets found.")
-        # delegate db auth info discovery to normal toolchain
+    def __init__(self) -> None:
+        """Initialize a butler factory by creating butler pool instances for
+        each known repository.
+        """
+        # forced property lookup; cache the registry auth file contents
+        _ = self.butler_auth_config
+
+        # create and cache any butler factories known to the service
+        for label in BUTLER_REPO_INDEX.get_known_repos():
+            _ = self.get_butler_factory(label)
+        return None
+
+    @cached_property
+    def butler_auth_config(self) -> list[dict[str, str]] | None:
+        """Read a butler authentication file for secrets.
+
+        Notes
+        -----
+        This is a `functools.cached_property` that should be first accessed
+        during instance initialization so any IO is performed synchronously
+        at application startup and cached for the lifetime of the object and/or
+        the application.
+        """
+        try:
+            # FIXME ought to validate the loaded value against expectations
+            return yaml.safe_load(Path(config.butler.authentication_file).read_text())
+        except FileNotFoundError:
+            logger.warning("No Butler Registry authentication secrets file found.")
+            # delegate db auth info discovery to normal toolchain
+            return None
+
+    def get_butler(self, label: str, collections: list[str] | None = None) -> Butler | None:
+        """Get a butler clone from the factory.
+
+        Notes
+        -----
+        This is the primary public interface to the factory object.
+        """
+        factory = self.get_butler_factory(label)
+        if factory is None:
+            return None
+        return factory(collections=collections)
+
+    @cache
+    def get_butler_factory(self, label: str) -> Callable[..., Butler] | None:
+        """Return a factory function that creates a butler clone.
+
+        Notes
+        -----
+        This method is backed by a `functools.cache`, a threadsafe cache.
+
+        If the return value is None, the service could not create a Butler for
+        the desired label or no such Butler is configured. In the former case,
+        the service log should include exception information related to the
+        failed Butler creation.
+
+        Returns
+        -------
+        `lsst.daf.butler` or `None`
+            A cloned instance of a ``Butler`` or None if the labelled Butler
+            could not be created from the configuration inputs.
+        """
+        try:
+            _butler_config = self.get_butler_config(label=label)
+            _butler = Butler.from_config(_butler_config)
+            if TYPE_CHECKING:
+                assert isinstance(_butler, DirectButler)
+            _butler._preload_cache(load_dimension_record_cache=False)
+        except (ClientError, OperationalError):
+            # Case that configured butler was unable to be created
+            logger.exception()
+            return None
+        except KeyError:
+            # Case that no such butler was configured
+            logger.warning(f"No such butler configured: {label}")
+            return None
+
+        def factory(collections: CollectionArgType) -> Butler:
+            return _butler.clone(collections=collections)
+
+        return factory
+
+    def update_butler_url(self, bc: ButlerConfig) -> ButlerConfig:
+        """Update a butler config with registry secrets.
+
+        Returns
+        -------
+        ``ButlerConfig``
+            The same configuration object passed as input is returned, with
+            secrets applied or not, depending on whether secrets were available
+            for the input config's registry URL.
+
+        Notes
+        -----
+        This method makes use of the cached ``butler_auth_config`` property
+        which contains the set of butler registry secrets known to the service.
+        Missing secrets may produce non-working Butlers, but not all butlers
+        have associated secrets, so failure to locate a secret is not an error.
+        """
+        if self.butler_auth_config is None:
+            return bc
+
+        db_url = url.make_url(bc["registry"]["db"])
+        try:
+            db_auth = DbAuth(authList=self.butler_auth_config).getAuth(
+                dialectname=db_url.drivername,
+                username=db_url.username,
+                host=db_url.host,
+                port=db_url.port,
+                database=db_url.database,
+            )
+            bc[".registry.username"] = db_auth[0]
+            bc[".registry.password"] = db_auth[1]
+        except DbAuthError as e:
+            logger.warning(f"Could not parse db auth from provided url[{db_url}]: {e}")
         return bc
 
-    db_url = url.make_url(bc["registry"]["db"])
-    db_auth = DbAuth(authList=db_auth_info).getAuth(
-        dialectname=db_url.drivername,
-        username=db_url.username,
-        host=db_url.host,
-        port=db_url.port,
-        database=db_url.database,
-    )
-
-    bc[".registry.username"] = db_auth[0]
-    bc[".registry.password"] = db_auth[1]
-    return bc
+    def get_butler_config(self, label: str, *, without_datastore: bool = True) -> ButlerConfig:
+        """Create a butler config object for a repo known to the service's
+        environment.
+
+        Returns
+        -------
+        lsst.daf.butler.ButlerConfig
+            Unless otherwise specified, the ButlerConfig returned by this
+            method is configured ``without_datastore``.
+        """
+
+        try:
+            repo_uri: ResourcePathExpression = BUTLER_REPO_INDEX.get_repo_uri(label=label)
+        except KeyError:
+            logger.warning("Butler repo %s not known to environment.", label)
+            repo_uri = label
+
+        try:
+            bc_f = partial(
+                ButlerConfig,
+                other=repo_uri,
+                without_datastore=without_datastore,
+            )
+            bc = bc_f()
+        except FileNotFoundError:
+            logger.error("Butler repo %s not known.", label)
+            raise RuntimeError("Unknown Butler Repo %s", label)
+
+        # Any generated Butler config is hydrated with application secrets via
+        # this tail call.
+        return self.update_butler_url(bc)
+
+
+BUTLER_FACTORY = ButlerFactory()
+"""A module level butler factory created at module import, available for use
+in other modules.
+"""
 
 
+# TODO: deprecate these functions that attempt to "remove" data from Butlers.
 async def remove_run_collections(
     butler_repo: str,
     collection_name: str,
diff --git a/src/lsst/cmservice/daemon.py b/src/lsst/cmservice/daemon.py
@@ -10,6 +10,7 @@
 from safir.logging import configure_uvicorn_logging
 
 from . import __version__
+from .common.butler import BUTLER_FACTORY  # noqa: F401
 from .common.daemon import daemon_iteration
 from .common.logging import LOGGER
 from .common.panda import get_panda_token
@@ -30,19 +31,20 @@ async def lifespan(app: FastAPI) -> AsyncGenerator:
     os.environ |= config.panda.model_dump(by_alias=True, exclude_none=True)
     os.environ |= config.htcondor.model_dump(by_alias=True, exclude_none=True)
     app.state.tasks = set()
-    daemon = create_task(main_loop(), name="daemon")
+    daemon = create_task(main_loop(app=app), name="daemon")
     app.state.tasks.add(daemon)
     yield
     # stop
 
 
-async def main_loop() -> None:
+async def main_loop(app: FastAPI) -> None:
     """Daemon execution loop.
 
     With a database session, perform a single daemon interation and then sleep
     until the next daemon appointment.
     """
     engine = create_database_engine(config.db.url, config.db.password)
+
     sleep_time = config.daemon.processing_interval
 
     async with engine.begin():
diff --git a/src/lsst/cmservice/handlers/elements.py b/src/lsst/cmservice/handlers/elements.py
@@ -8,11 +8,10 @@
 from anyio import to_thread
 from sqlalchemy.ext.asyncio import async_scoped_session
 
-from lsst.daf.butler import Butler
-
-from ..common.butler import get_butler_config
+from ..common.butler import BUTLER_FACTORY
 from ..common.enums import StatusEnum
 from ..common.errors import CMMissingScriptInputError, test_type_and_raise
+from ..common.logging import LOGGER
 from ..config import config
 from ..db.campaign import Campaign
 from ..db.element import ElementMixin
@@ -21,6 +20,8 @@
 from ..db.script import Script
 from .script_handler import FunctionHandler
 
+logger = LOGGER.bind(module=__name__)
+
 
 class RunElementScriptHandler(FunctionHandler):
     """Shared base class to handling running and
@@ -221,14 +222,10 @@ async def split(
         if mock_butler:
             sorted_field_values = np.arange(10)
         else:
-            butler_config = await get_butler_config(butler_repo, without_datastore=True)
-            butler_f = partial(
-                Butler.from_config,
-                butler_config,
-                collections=[input_coll, campaign_input_coll],
-                without_datastore=True,
-            )
-            butler = await to_thread.run_sync(butler_f)
+            butler = BUTLER_FACTORY.get_butler(butler_repo, collections=[input_coll, campaign_input_coll])
+            if butler is None:
+                logger.error(f"butler repo {butler_repo} is not known to the application.")
+                raise RuntimeError("No such butler")
             itr_q_f = partial(
                 butler.registry.queryDataIds,
                 [split_field],
diff --git a/tests/common/test_butler.py b/tests/common/test_butler.py
