fix: use protocol for SessionTasks (#1165)

Closes #1163.

This should improve type safety with the SessionTasks Python class.

Author: leafty

bases/renku_data_services/data_tasks/dependencies.py

@@ -59,12 +59,9 @@ def from_env(cls, cfg: Config | None = None) -> "DependencyManager":
             metrics=metrics,
             authz=authz,
         )
-        session_repo = SessionRepository(
+        session_environment_repo = SessionRepository.make_session_environment_repo(
             session_maker=cfg.db.async_session_maker,
             project_authz=authz,
-            resource_pools=None,  # type: ignore
-            shipwright_client=None,
-            builds_config=None,  # type: ignore
         )
         syncer = UsersSync(
             cfg.db.async_session_maker,
@@ -73,7 +70,7 @@ def from_env(cls, cfg: Config | None = None) -> "DependencyManager":
             metrics=metrics,
             authz=authz,
         )
-        session_tasks = SessionTasks(session_repo=session_repo)
+        session_tasks = SessionTasks(session_environment_repo=session_environment_repo)
         kc_api: IKeycloakAPI
         if cfg.dummy_stores:
             dummy_users = [

components/renku_data_services/session/db.py

@@ -5,7 +5,7 @@
 from collections.abc import Callable
 from contextlib import AbstractAsyncContextManager, nullcontext
 from datetime import UTC, datetime
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Protocol
 
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -28,7 +28,21 @@
     from renku_data_services.session.config import BuildsConfig
 
 
-class SessionRepository:
+class SessionEnvironmentRepositoryProtocol(Protocol):
+    """Protocol for operations on session environments."""
+
+    async def get_environments(self, include_archived: bool = False) -> list[models.Environment]:
+        """Get all global session environments from the database."""
+        ...
+
+    async def insert_environment(
+        self, user: base_models.APIUser, environment: models.UnsavedEnvironment
+    ) -> models.Environment:
+        """Insert a new session environment."""
+        ...
+
+
+class SessionRepository(SessionEnvironmentRepositoryProtocol):
     """Repository for sessions."""
 
     def __init__(
@@ -1160,3 +1174,22 @@ async def _get_environment_authorization(
         if launcher:
             authorized = await self.project_authz.has_permission(user, ResourceType.project, launcher.project_id, scope)
         return authorized
+
+    @classmethod
+    def make_session_environment_repo(
+        cls,
+        session_maker: Callable[..., AsyncSession],
+        project_authz: Authz,
+    ) -> SessionEnvironmentRepositoryProtocol:
+        """Create an instance of SessionEnvironmentRepositoryProtocol."""
+        # NOTE: resource_pools, shipwright_client and builds_config are set to None
+        # because the SessionEnvironmentRepositoryProtocol only exposes database
+        # operations for session environments.
+        instance = cls(
+            session_maker,
+            project_authz=project_authz,
+            resource_pools=None,  # type: ignore
+            shipwright_client=None,
+            builds_config=None,  # type: ignore
+        )
+        return instance

components/renku_data_services/session/tasks.py

@@ -6,29 +6,29 @@
 
 from renku_data_services.app_config import logging
 from renku_data_services.base_models.core import InternalServiceAdmin
-from renku_data_services.session.db import SessionRepository
+from renku_data_services.session.db import SessionEnvironmentRepositoryProtocol
 from renku_data_services.session.models import EnvironmentImageSource, EnvironmentKind, UnsavedEnvironment
 
 
 @dataclass(kw_only=True)
 class SessionTasks:
     """Task definitions for sessions."""
 
-    session_repo: SessionRepository
+    session_environment_repo: SessionEnvironmentRepositoryProtocol
 
     async def initialize_session_environments_task(self, requested_by: InternalServiceAdmin) -> None:
         """Initialize session environments."""
         logger = logging.getLogger(self.__class__.__name__)
 
         # Skip this task if global session environments already exist
-        existing_envs = await self.session_repo.get_environments()
+        existing_envs = await self.session_environment_repo.get_environments()
         if existing_envs:
             logger.debug("Global session environments are already initialized.")
             return None
 
         for env in self._get_default_session_environments():
             try:
-                await self.session_repo.insert_environment(user=requested_by, environment=env)
+                await self.session_environment_repo.insert_environment(user=requested_by, environment=env)
                 logger.info(f"Added global environment with image {env.container_image}")
             except Exception as err:
                 logger.error(f"Failed to create global environment with image {env.container_image} because {err}")