♻️ Refactor: migrate to asyncpg and improve database connection handling in studies dispatcher

Author: pcrespov

services/web/server/src/simcore_service_webserver/studies_dispatcher/_catalog.py

@@ -4,9 +4,6 @@
 from dataclasses import dataclass
 
 import sqlalchemy as sa
-from aiohttp import web
-from aiopg.sa.connection import SAConnection
-from aiopg.sa.engine import Engine
 from models_library.groups import EVERYONE_GROUP_ID
 from models_library.services import ServiceKey, ServiceVersion
 from models_library.services_constants import (
@@ -22,11 +19,12 @@
 from simcore_postgres_database.models.services_consume_filetypes import (
     services_consume_filetypes,
 )
+from simcore_postgres_database.utils_repos import pass_or_acquire_connection
 from simcore_postgres_database.utils_services import create_select_latest_services_query
+from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
 
-from ..db.plugin import get_database_engine_legacy
 from ._errors import ServiceNotFoundError
-from .settings import StudiesDispatcherSettings, get_plugin_settings
+from .settings import StudiesDispatcherSettings
 
 LARGEST_PAGE_SIZE = 1000
 
@@ -44,22 +42,28 @@ class ServiceMetaData:
     file_extensions: list[str]
 
 
-async def _get_service_filetypes(conn: SAConnection) -> dict[ServiceKey, list[str]]:
+async def _get_service_filetypes(
+    engine: AsyncEngine,
+    connection: AsyncConnection | None = None,
+) -> dict[ServiceKey, list[str]]:
     query = sa.select(
         services_consume_filetypes.c.service_key,
         sa.func.array_agg(
             sa.func.distinct(services_consume_filetypes.c.filetype)
         ).label("list_of_file_types"),
     ).group_by(services_consume_filetypes.c.service_key)
 
-    result = await conn.execute(query)
-    rows = await result.fetchall()
+    async with pass_or_acquire_connection(engine, connection) as conn:
+        result = await conn.execute(query)
+        rows = result.fetchall()
 
-    return {row.service_key: row.list_of_file_types for row in rows}
+        return {row.service_key: row.list_of_file_types for row in rows}
 
 
 async def iter_latest_product_services(
-    app: web.Application,
+    settings: StudiesDispatcherSettings,
+    engine: AsyncEngine,
+    connection: AsyncConnection | None = None,
     *,
     product_name: str,
     page_number: PositiveInt = 1,  # 1-based
@@ -68,9 +72,6 @@ async def iter_latest_product_services(
     assert page_number >= 1  # nosec
     assert ((page_number - 1) * page_size) >= 0  # nosec
 
-    engine: Engine = get_database_engine_legacy(app)
-    settings: StudiesDispatcherSettings = get_plugin_settings(app)
-
     # Select query for latest version of the service
     latest_services = create_select_latest_services_query().alias("latest_services")
 
@@ -109,10 +110,10 @@ async def iter_latest_product_services(
     # pagination
     query = query.limit(page_size).offset((page_number - 1) * page_size)
 
-    async with engine.acquire() as conn:
-        service_filetypes = await _get_service_filetypes(conn)
+    async with pass_or_acquire_connection(engine, connection) as conn:
+        service_filetypes = await _get_service_filetypes(engine, conn)
 
-        async for row in await conn.execute(query):
+        async for row in await conn.stream(query):
             yield ServiceMetaData(
                 key=row.key,
                 version=row.version,
@@ -135,14 +136,13 @@ class ValidService:
 
 @log_decorator(_logger, level=logging.DEBUG)
 async def validate_requested_service(
-    app: web.Application,
+    engine: AsyncEngine,
+    connection: AsyncConnection | None = None,
     *,
     service_key: ServiceKey,
     service_version: ServiceVersion,
 ) -> ValidService:
-    engine: Engine = get_database_engine_legacy(app)
-
-    async with engine.acquire() as conn:
+    async with pass_or_acquire_connection(engine, connection) as conn:
         query = sa.select(
             services_meta_data.c.name,
             services_meta_data.c.key,
@@ -153,7 +153,7 @@ async def validate_requested_service(
         )
 
         result = await conn.execute(query)
-        row = await result.fetchone()
+        row = result.one_or_none()
 
         if row is None:
             raise ServiceNotFoundError(

services/web/server/src/simcore_service_webserver/studies_dispatcher/_controller/rest/nih.py

@@ -9,10 +9,12 @@
 )
 
 from ...._meta import API_VTAG
+from ....db.plugin import get_asyncpg_engine
 from ....products import products_web
 from ....utils_aiohttp import envelope_json_response
 from ... import _service
 from ..._catalog import iter_latest_product_services
+from ...settings import get_plugin_settings
 from .nih_schemas import ServiceGet, Viewer
 
 _logger = logging.getLogger(__name__)
@@ -26,9 +28,12 @@ async def list_latest_services(request: Request):
     """Returns a list latest version of services"""
     product_name = products_web.get_product_name(request)
 
+    plugin_settings = get_plugin_settings(request.app)
+    engine = get_asyncpg_engine(request.app)
+
     services = []
     async for service_data in iter_latest_product_services(
-        request.app, product_name=product_name
+        plugin_settings, engine, product_name=product_name
     ):
         try:
             service = ServiceGet.create(service_data, request)

services/web/server/src/simcore_service_webserver/studies_dispatcher/_controller/rest/redirects.py

@@ -7,6 +7,7 @@
 from models_library.projects_nodes_io import NodeID
 from servicelib.aiohttp.requests_validation import parse_request_query_parameters_as
 
+from ....db.plugin import get_asyncpg_engine
 from ....dynamic_scheduler import api as dynamic_scheduler_service
 from ....products import products_web
 from ....utils_aiohttp import create_redirect_to_page_response, get_api_base_url
@@ -133,7 +134,7 @@ async def get_redirection_to_viewer(request: web.Request):
         service_params_ = query_params
 
         valid_service: ValidService = await validate_requested_service(
-            app=request.app,
+            get_asyncpg_engine(request.app),
             service_key=service_params_.viewer_key,
             service_version=service_params_.viewer_version,
         )

services/web/server/tests/unit/with_dbs/04/studies_dispatcher/test_studies_dispatcher_repository.py

@@ -7,19 +7,35 @@
 from collections.abc import AsyncIterator
 
 import pytest
+from models_library.groups import EVERYONE_GROUP_ID
+from models_library.services import ServiceKey, ServiceVersion
 from pytest_simcore.helpers.faker_factories import (
+    random_service_access_rights,
     random_service_consume_filetype,
     random_service_meta_data,
 )
 from pytest_simcore.helpers.postgres_tools import insert_and_get_row_lifespan
-from simcore_postgres_database.models.services import services_meta_data
+from simcore_postgres_database.models.services import (
+    services_access_rights,
+    services_meta_data,
+)
 from simcore_postgres_database.models.services_consume_filetypes import (
     services_consume_filetypes,
 )
+from simcore_service_webserver.studies_dispatcher._catalog import (
+    ServiceMetaData,
+    ValidService,
+    iter_latest_product_services,
+    validate_requested_service,
+)
+from simcore_service_webserver.studies_dispatcher._errors import ServiceNotFoundError
 from simcore_service_webserver.studies_dispatcher._models import ViewerInfo
 from simcore_service_webserver.studies_dispatcher._repository import (
     StudiesDispatcherRepository,
 )
+from simcore_service_webserver.studies_dispatcher.settings import (
+    StudiesDispatcherSettings,
+)
 from sqlalchemy.ext.asyncio import AsyncEngine
 
 
@@ -69,6 +85,37 @@ async def consume_filetypes_in_db(
         yield row
 
 
+@pytest.fixture
+async def service_access_rights_in_db(
+    asyncpg_engine: AsyncEngine, service_metadata_in_db: dict
+):
+    """Pre-populate services access rights table with test data."""
+    access_data = random_service_access_rights(
+        key=service_metadata_in_db["key"],
+        version=service_metadata_in_db["version"],
+        gid=EVERYONE_GROUP_ID,
+        execute_access=True,
+        product_name="osparc",
+    )
+
+    # pylint: disable=contextmanager-generator-missing-cleanup
+    async with insert_and_get_row_lifespan(
+        asyncpg_engine,
+        table=services_access_rights,
+        values=access_data,
+        pk_col=services_access_rights.c.key,
+        pk_value=access_data["key"],
+    ) as row:
+        yield row
+
+
+@pytest.fixture
+def studies_dispatcher_settings() -> StudiesDispatcherSettings:
+    return StudiesDispatcherSettings(
+        STUDIES_DEFAULT_SERVICE_THUMBNAIL="https://example.com/default-thumbnail.png"
+    )
+
+
 @pytest.fixture
 def studies_dispatcher_repository(
     asyncpg_engine: AsyncEngine,
@@ -201,3 +248,89 @@ async def test_find_compatible_viewer_not_found(
 
     # Assert
     assert viewer_wrong_filetype is None
+
+
+async def test_iter_latest_product_services(
+    asyncpg_engine: AsyncEngine,
+    studies_dispatcher_settings: StudiesDispatcherSettings,
+    service_metadata_in_db: dict,
+    service_access_rights_in_db: dict,
+    consume_filetypes_in_db: dict,
+):
+    """Test iterating through latest product services."""
+    # Act
+    services = []
+    async for service in iter_latest_product_services(
+        studies_dispatcher_settings, asyncpg_engine, product_name="osparc"
+    ):
+        services.append(service)
+
+    # Assert
+    assert len(services) == 1
+    service = services[0]
+    assert isinstance(service, ServiceMetaData)
+    assert service.key == service_metadata_in_db["key"]
+    assert service.version == service_metadata_in_db["version"]
+    assert service.title == service_metadata_in_db["name"]
+    assert service.description == service_metadata_in_db["description"]
+    assert service.file_extensions == [consume_filetypes_in_db["filetype"]]
+
+
+async def test_iter_latest_product_services_with_pagination(
+    asyncpg_engine: AsyncEngine,
+    studies_dispatcher_settings: StudiesDispatcherSettings,
+    service_metadata_in_db: dict,
+    service_access_rights_in_db: dict,
+):
+    """Test iterating through services with pagination."""
+    # Act
+    services = []
+    async for service in iter_latest_product_services(
+        studies_dispatcher_settings,
+        asyncpg_engine,
+        product_name="osparc",
+        page_number=1,
+        page_size=1,
+    ):
+        services.append(service)
+
+    # Assert
+    assert len(services) == 1
+
+
+async def test_validate_requested_service_success(
+    asyncpg_engine: AsyncEngine,
+    service_metadata_in_db: dict,
+    consume_filetypes_in_db: dict,
+):
+    """Test validating a service that exists and is valid."""
+    # Act
+    valid_service = await validate_requested_service(
+        engine=asyncpg_engine,
+        service_key=ServiceKey(service_metadata_in_db["key"]),
+        service_version=ServiceVersion(service_metadata_in_db["version"]),
+    )
+
+    # Assert
+    assert isinstance(valid_service, ValidService)
+    assert valid_service.key == service_metadata_in_db["key"]
+    assert valid_service.version == service_metadata_in_db["version"]
+    assert valid_service.title == service_metadata_in_db["name"]
+    assert valid_service.is_public == consume_filetypes_in_db["is_guest_allowed"]
+    assert valid_service.thumbnail is None  # No valid URL in test data
+
+
+async def test_validate_requested_service_not_found(
+    asyncpg_engine: AsyncEngine,
+):
+    """Test validating a service that doesn't exist."""
+    # Act & Assert
+    with pytest.raises(ServiceNotFoundError) as exc_info:
+        await validate_requested_service(
+            asyncpg_engine,
+            service_key=ServiceKey("simcore/services/dynamic/nonexistent"),
+            service_version=ServiceVersion("1.0.0"),
+        )
+
+    assert exc_info.value.service_key == "simcore/services/dynamic/nonexistent"
+    assert exc_info.value.service_version == "1.0.0"