Merge pull request #1160 from lsst-sqre/tickets/DM-47716a

DM-47716: Avoid creating Firestore clients on every request

changelog.d/20241120_164500_rra_DM_47716_queue.md

@@ -0,0 +1,3 @@
+### Bug fixes
+
+- Avoid creating a Google Firestore client for every request, since it does authentication setup on creation. Instead, create a single client that will be used for all requests.

docs/documenteer.toml

@@ -35,6 +35,7 @@ nitpick_ignore = [
     # by having Sphinx complain about a new symbol.
     ["py:class", "dataclasses_avroschema.pydantic.main.AvroBaseModel"],
     ["py:class", "dataclasses_avroschema.main.AvroModel"],
+    ["py:class", "google.cloud.firestore_v1.async_client.AsyncClient"],
     ["py:class", "fastapi.applications.FastAPI"],
     ["py:class", "fastapi.datastructures.DefaultPlaceholder"],
     ["py:class", "fastapi.exceptions.HTTPException"],

src/gafaelfawr/factory.py

@@ -10,6 +10,7 @@
 import structlog
 from bonsai import LDAPClient
 from bonsai.asyncio import AIOConnectionPool
+from google.cloud import firestore
 from httpx import AsyncClient
 from kubernetes_asyncio.client import ApiClient
 from redis.asyncio import BlockingConnectionPool, Redis
@@ -82,6 +83,9 @@ class ProcessContext:
     config: Config
     """Gafaelfawr's configuration."""
 
+    firestore: firestore.AsyncClient | None
+    """Client to talk to Firestore, if configured."""
+
     http_client: AsyncClient
     """Shared HTTP client."""
 
@@ -126,6 +130,11 @@ async def from_config(cls, config: Config) -> Self:
         ProcessContext
             Shared context for a Gafaelfawr process.
         """
+        firestore_client = None
+        if config.firestore:
+            firestore_project = config.firestore.project
+            firestore_client = firestore.AsyncClient(project=firestore_project)
+
         ldap_pool = None
         if config.ldap:
             client = LDAPClient(str(config.ldap.url))
@@ -161,6 +170,7 @@ async def from_config(cls, config: Config) -> Self:
 
         return cls(
             config=config,
+            firestore=firestore_client,
             http_client=await http_client_dependency(),
             ldap_pool=ldap_pool,
             redis=redis_client,
@@ -352,9 +362,9 @@ def create_firestore_storage(self) -> FirestoreStorage:
         FirestoreStorage
             Newly-created Firestore storage.
         """
-        if not self._context.config.firestore:
+        if not self._context.firestore:
             raise NotConfiguredError("Firestore is not configured")
-        return FirestoreStorage(self._context.config.firestore, self._logger)
+        return FirestoreStorage(self._context.firestore, self._logger)
 
     def create_health_check_service(self) -> HealthCheckService:
         """Create a service for performing health checks.

src/gafaelfawr/storage/firestore.py

@@ -5,7 +5,6 @@
 from google.cloud import firestore
 from structlog.stdlib import BoundLogger
 
-from ..config import FirestoreConfig
 from ..constants import (
     GID_MAX,
     GID_MIN,
@@ -33,22 +32,26 @@ class FirestoreStorage:
     """Google Firestore storage layer.
 
     Gafaelfawr supports assigning UIDs and GIDs from Google Firestore rather
-    than getting them from LDAP or upstream authentication tokens.  This
-    module provides the read/write layer and transaction management for that
-    storage.  It's used from inside a per-process cache.
+    than getting them from LDAP or upstream authentication tokens. This module
+    provides the read/write layer and transaction management for that
+    storage. It's used from inside a per-process cache.
+
+    This class authenticates to Google on creation, so it should not be
+    created fresh for every request.
 
     Parameters
     ----------
-    config
-        Configuration for Google Firestore.
+    client
+        Firestore client to use.
     logger
         Logger for debug messages and errors.
     """
 
-    def __init__(self, config: FirestoreConfig, logger: BoundLogger) -> None:
-        self._config = config
+    def __init__(
+        self, client: firestore.AsyncClient, logger: BoundLogger
+    ) -> None:
+        self._client = client
         self._logger = logger
-        self._db = firestore.AsyncClient(project=config.project)
 
     async def get_gid(self, group: str) -> int:
         """Get the GID for a group.
@@ -73,9 +76,9 @@ async def get_gid(self, group: str) -> int:
         NoAvailableGidError
             Raised if no more GIDs are available in that range.
         """
-        transaction = self._db.transaction()
-        group_ref = self._db.collection("groups").document(group)
-        counter_ref = self._db.collection("counters").document("gid")
+        transaction = self._client.transaction()
+        group_ref = self._client.collection("groups").document(group)
+        counter_ref = self._client.collection("counters").document("gid")
         return await _get_or_assign_gid(
             transaction,
             group_name=group,
@@ -111,10 +114,10 @@ async def get_uid(self, username: str, *, bot: bool = False) -> int:
         NoAvailableUidError
             Raised if no more UIDs are available in that range.
         """
-        transaction = self._db.transaction()
-        user_ref = self._db.collection("users").document(username)
-        counter_name = "bot-uid" if bot else "uid"
-        counter_ref = self._db.collection("counters").document(counter_name)
+        transaction = self._client.transaction()
+        user_ref = self._client.collection("users").document(username)
+        counter = "bot-uid" if bot else "uid"
+        counter_ref = self._client.collection("counters").document(counter)
         return await _get_or_assign_uid(
             transaction,
             username=username,
@@ -131,10 +134,10 @@ async def initialize(self) -> None:
         silently do nothing.
         """
         counter_refs = {
-            n: self._db.collection("counters").document(n)
+            n: self._client.collection("counters").document(n)
             for n in _INITIAL_COUNTERS
         }
-        transaction = self._db.transaction()
+        transaction = self._client.transaction()
         await _initialize_in_transaction(
             transaction, counter_refs, self._logger
         )