ongoing

Author: sanderegg

services/storage/src/simcore_service_storage/api/rest/health.py

@@ -9,16 +9,17 @@
 from aiohttp import web
 from aws_library.s3 import S3AccessError
 from common_library.json_serialization import json_dumps
+from fastapi import Request
 from models_library.api_schemas_storage import HealthCheck, S3BucketName
 from models_library.app_diagnostics import AppStatusCheck
 from pydantic import TypeAdapter
+from servicelib.db_asyncpg_utils import check_postgres_liveness
+from servicelib.fastapi.db_asyncpg_engine import get_engine
 from servicelib.rest_constants import RESPONSE_MODEL_POLICY
+from simcore_postgres_database.utils_aiosqlalchemy import get_pg_engine_stateinfo
 
 from ..._meta import API_VERSION, API_VTAG, PROJECT_NAME, VERSION
-from ...constants import APP_CONFIG_KEY
-from ...core.settings import ApplicationSettings
-from ...modules.db.db import get_engine_state
-from ...modules.db.db import is_service_responsive as is_pg_responsive
+from ...core.settings import get_application_settings
 from ...modules.s3 import get_s3_client
 
 _logger = logging.getLogger(__name__)
@@ -43,10 +44,10 @@ async def get_health(request: web.Request) -> web.Response:
 
 
 @routes.get(f"/{API_VTAG}/status", name="get_status")
-async def get_status(request: web.Request) -> web.Response:
+async def get_status(request: Request) -> web.Response:
     # NOTE: all calls here must NOT raise
     assert request.app  # nosec
-    app_settings: ApplicationSettings = request.app[APP_CONFIG_KEY]
+    app_settings = get_application_settings(request.app)
     s3_state = "disabled"
     if app_settings.STORAGE_S3:
         try:
@@ -63,9 +64,12 @@ async def get_status(request: web.Request) -> web.Response:
             s3_state = "failed"
 
     postgres_state = "disabled"
+
     if app_settings.STORAGE_POSTGRES:
         postgres_state = (
-            "connected" if await is_pg_responsive(request.app) else "failed"
+            "connected"
+            if await check_postgres_liveness(get_engine(request.app))
+            else "failed"
         )
 
     status = AppStatusCheck.model_validate(
@@ -75,7 +79,7 @@ async def get_status(request: web.Request) -> web.Response:
             "services": {
                 "postgres": {
                     "healthy": postgres_state,
-                    "pool": get_engine_state(request.app),
+                    "pool": get_pg_engine_stateinfo(get_engine(request.app)),
                 },
                 "s3": {"healthy": s3_state},
             },

services/storage/src/simcore_service_storage/core/application.py

@@ -24,7 +24,6 @@
     APP_NAME,
     APP_STARTED_BANNER_MSG,
 )
-from ..api.rest.utils import dsm_exception_handler
 from ..dsm import setup_dsm
 from ..dsm_cleaner import setup_dsm_cleaner
 from ..modules.db.db import setup_db
@@ -83,7 +82,7 @@ def create_app(settings: ApplicationSettings) -> FastAPI:
         setup_redis(app)
         setup_dsm_cleaner(app)
 
-    app.middlewares.append(dsm_exception_handler)
+    # app.middlewares.append(dsm_exception_handler)
 
     if settings.STORAGE_PROFILING:
         app.add_middleware(ProfilerMiddleware)

services/storage/src/simcore_service_storage/core/settings.py

@@ -1,5 +1,6 @@
 from typing import Annotated, Self
 
+from fastapi import FastAPI
 from pydantic import (
     AliasChoices,
     Field,
@@ -108,3 +109,8 @@ def _ensure_settings_consistency(self) -> Self:
             )
             raise ValueError(msg)
         return self
+
+
+def get_application_settings(app: FastAPI) -> ApplicationSettings:
+    assert isinstance(app.state.settings, ApplicationSettings)  # nosec
+    return app.state.settings

services/storage/src/simcore_service_storage/dsm.py

@@ -1,17 +1,18 @@
 import logging
+from typing import cast
 
 from fastapi import FastAPI
 
-from .constants import APP_DSM_KEY
 from .datcore_dsm import DatCoreDataManager, create_datcore_data_manager
 from .dsm_factory import DataManagerProvider
+from .exceptions.errors import ConfigurationError
 from .simcore_s3_dsm import SimcoreS3DataManager, create_simcore_s3_data_manager
 
 logger = logging.getLogger(__name__)
 
 
 def setup_dsm(app: FastAPI):
-    async def _cleanup_context(app: FastAPI):
+    async def _on_startup(app: FastAPI) -> None:
         dsm_provider = DataManagerProvider(app)
         dsm_provider.register_builder(
             SimcoreS3DataManager.get_location_id(),
@@ -23,17 +24,22 @@ async def _cleanup_context(app: FastAPI):
             create_datcore_data_manager,
             DatCoreDataManager,
         )
-        app[APP_DSM_KEY] = dsm_provider
+        app.state.dsm_provider = dsm_provider
 
-        yield
-
-        logger.info("Shuting down %s", f"{dsm_provider=}")
+    async def _on_shutdown() -> None:
+        if app.state.dsm_provider:
+            # nothing to do
+            ...
 
     # ------
 
-    app.cleanup_ctx.append(_cleanup_context)
+    app.add_event_handler("startup", _on_startup)
+    app.add_event_handler("shutdown", _on_shutdown)
 
 
 def get_dsm_provider(app: FastAPI) -> DataManagerProvider:
-    dsm_provider: DataManagerProvider = app[APP_DSM_KEY]
-    return dsm_provider
+    if not app.state.dsm_provider:
+        raise ConfigurationError(
+            msg="DSM provider not available. Please check the configuration."
+        )
+    return cast(DataManagerProvider, app.state.dsm_provider)

services/storage/src/simcore_service_storage/exceptions/errors.py

@@ -5,6 +5,10 @@ class StorageRuntimeError(OsparcErrorMixin, RuntimeError):
     ...
 
 
+class ConfigurationError(StorageRuntimeError):
+    msg_template: str = "Application misconfiguration: {msg}"
+
+
 class DatabaseAccessError(StorageRuntimeError):
     msg_template: str = "Unexpected error while accessing database backend"
 

services/storage/src/simcore_service_storage/modules/db/db.py

@@ -1,76 +1,23 @@
 import logging
-from typing import Any
 
-from aiopg.sa.engine import Engine
 from fastapi import FastAPI
-from servicelib.aiohttp.aiopg_utils import is_pg_responsive
-from servicelib.common_aiopg_utils import DataSourceName, create_pg_engine
+from servicelib.db_async_engine import close_db_connection
+from servicelib.fastapi.db_asyncpg_engine import connect_to_db
 from servicelib.retry_policies import PostgresRetryPolicyUponInitialization
-from settings_library.postgres import PostgresSettings
-from simcore_postgres_database.utils_aiopg import (
-    get_pg_engine_stateinfo,
-    raise_if_migration_not_ready,
-)
 from tenacity import retry
 
-from ...constants import APP_AIOPG_ENGINE_KEY
+from ...core.settings import get_application_settings
 
 _logger = logging.getLogger(__name__)
 
 
-@retry(**PostgresRetryPolicyUponInitialization(_logger).kwargs)
-async def _ensure_pg_ready(dsn: DataSourceName, min_size: int, max_size: int) -> None:
-    _logger.info("Checking pg is ready %s", dsn)
+def setup_db(app: FastAPI) -> None:
+    @retry(**PostgresRetryPolicyUponInitialization(_logger).kwargs)
+    async def _on_startup(app: FastAPI) -> None:
+        await connect_to_db(app, get_application_settings(app).STORAGE_POSTGRES)
 
-    async with create_pg_engine(dsn, minsize=min_size, maxsize=max_size) as engine:
-        await raise_if_migration_not_ready(engine)
+    async def _on_shutdown() -> None:
+        await close_db_connection(app)
 
-
-async def postgres_cleanup_ctx(app: FastAPI):
-    pg_cfg: PostgresSettings = app[APP_CONFIG_KEY].STORAGE_POSTGRES
-    dsn = DataSourceName(
-        application_name=f"{__name__}_{id(app)}",
-        database=pg_cfg.POSTGRES_DB,
-        user=pg_cfg.POSTGRES_USER,
-        password=pg_cfg.POSTGRES_PASSWORD.get_secret_value(),
-        host=pg_cfg.POSTGRES_HOST,
-        port=pg_cfg.POSTGRES_PORT,
-    )
-
-    await _ensure_pg_ready(
-        dsn, min_size=pg_cfg.POSTGRES_MINSIZE, max_size=pg_cfg.POSTGRES_MAXSIZE
-    )
-    _logger.info("Creating pg engine for %s", dsn)
-    async with create_pg_engine(
-        dsn, minsize=pg_cfg.POSTGRES_MINSIZE, maxsize=pg_cfg.POSTGRES_MAXSIZE
-    ) as engine:
-        assert engine  # nosec
-        app[APP_AIOPG_ENGINE_KEY] = engine
-
-        _logger.info("Created pg engine for %s", dsn)
-        yield  # ----------
-        _logger.info("Deleting pg engine for %s", dsn)
-    _logger.info("Deleted pg engine for %s", dsn)
-
-
-async def is_service_responsive(app: FastAPI) -> bool:
-    """Returns true if the app can connect to db service"""
-    return await is_pg_responsive(engine=app[APP_AIOPG_ENGINE_KEY])
-
-
-def get_engine_state(app: FastAPI) -> dict[str, Any]:
-    engine: Engine | None = app.get(APP_AIOPG_ENGINE_KEY)
-    if engine:
-        engine_info: dict[str, Any] = get_pg_engine_stateinfo(engine)
-        return engine_info
-    return {}
-
-
-def setup_db(app: FastAPI):
-    app[APP_AIOPG_ENGINE_KEY] = None
-
-    # app is created at this point but not yet started
-    _logger.debug("Setting up %s [service: %s] ...", __name__, "postgres")
-
-    # async connection to db
-    app.cleanup_ctx.append(postgres_cleanup_ctx)
+    app.add_event_handler("startup", _on_startup)
+    app.add_event_handler("shutdown", _on_shutdown)

services/storage/src/simcore_service_storage/modules/s3.py

@@ -1,29 +1,27 @@
 """Module to access s3 service"""
 
 import logging
-from collections.abc import AsyncGenerator
 from typing import cast
 
 from aws_library.s3 import SimcoreS3API
 from common_library.json_serialization import json_dumps
+from fastapi import FastAPI
 from servicelib.logging_utils import log_context
 from tenacity.asyncio import AsyncRetrying
 from tenacity.before_sleep import before_sleep_log
 from tenacity.wait import wait_fixed
 
-from ..constants import APP_CONFIG_KEY, APP_S3_KEY, RETRY_WAIT_SECS
-from ..core.settings import ApplicationSettings
+from ..constants import RETRY_WAIT_SECS
+from ..core.settings import get_application_settings
+from ..exceptions.errors import ConfigurationError
 
 _logger = logging.getLogger(__name__)
 
 
-async def setup_s3_client(app) -> AsyncGenerator[None, None]:
-    client = None
-
-    with log_context(_logger, logging.DEBUG, msg="setup.s3_client.cleanup_ctx"):
-        storage_settings: ApplicationSettings = app[APP_CONFIG_KEY]
-        storage_s3_settings = storage_settings.STORAGE_S3
-        assert storage_s3_settings  # nosec
+def setup_s3(app: FastAPI):
+    async def _on_startup(app: FastAPI) -> None:
+        app.state.s3_client = None
+        settings = get_application_settings(app)
 
         async for attempt in AsyncRetrying(
             wait=wait_fixed(RETRY_WAIT_SECS),
@@ -32,43 +30,34 @@ async def setup_s3_client(app) -> AsyncGenerator[None, None]:
         ):
             with attempt:
                 client = await SimcoreS3API.create(
-                    storage_s3_settings,
-                    storage_settings.STORAGE_S3_CLIENT_MAX_TRANSFER_CONCURRENCY,
+                    settings.STORAGE_S3,
+                    settings.STORAGE_S3_CLIENT_MAX_TRANSFER_CONCURRENCY,
                 )
                 _logger.info(
                     "S3 client %s successfully created [%s]",
                     f"{client=}",
                     json_dumps(attempt.retry_state.retry_object.statistics),
                 )
             assert client  # nosec
-            app[APP_S3_KEY] = client
-
-    yield
-
-    with log_context(_logger, logging.DEBUG, msg="teardown.s3_client.cleanup_ctx"):
-        if client:
-            await client.close()
+        app.state.s3_client = client
 
+        with log_context(_logger, logging.DEBUG, msg="setup.s3_bucket.cleanup_ctx"):
+            await client.create_bucket(
+                bucket=settings.STORAGE_S3.S3_BUCKET_NAME,
+                region=settings.STORAGE_S3.S3_REGION,
+            )
 
-async def setup_s3_bucket(app: FastAPI):
-    with log_context(_logger, logging.DEBUG, msg="setup.s3_bucket.cleanup_ctx"):
-        storage_s3_settings = app[APP_CONFIG_KEY].STORAGE_S3
-        client = get_s3_client(app)
-        await client.create_bucket(
-            bucket=storage_s3_settings.S3_BUCKET_NAME,
-            region=storage_s3_settings.S3_REGION,
-        )
-    yield
-
+    async def _on_shutdown() -> None:
+        if app.state.s3_client:
+            await cast(SimcoreS3API, app.state.s3_client).close()
 
-def setup_s3(app: FastAPI):
-    if setup_s3_client not in app.cleanup_ctx:
-        app.cleanup_ctx.append(setup_s3_client)
-    if setup_s3_bucket not in app.cleanup_ctx:
-        app.cleanup_ctx.append(setup_s3_bucket)
+    app.add_event_handler("startup", _on_startup)
+    app.add_event_handler("shutdown", _on_shutdown)
 
 
 def get_s3_client(app: FastAPI) -> SimcoreS3API:
-    assert app[APP_S3_KEY]  # nosec
-    assert isinstance(app[APP_S3_KEY], SimcoreS3API)  # nosec
-    return cast(SimcoreS3API, app[APP_S3_KEY])
+    if not app.state.s3_client:
+        raise ConfigurationError(
+            msg="S3 client is not available. Please check the configuration."
+        )
+    return cast(SimcoreS3API, app.state.s3_client)

services/storage/src/simcore_service_storage/routes.py

@@ -1,7 +1,7 @@
 import logging
 from pathlib import Path
 
-from aiohttp_swagger import setup_swagger  # type: ignore[import-untyped]
+from fastapi import FastAPI
 from servicelib.aiohttp.rest_middlewares import append_rest_middlewares
 from servicelib.aiohttp.rest_utils import (
     get_named_routes_as_message,