♻️ Replace aiopg in catalog service (ITISFoundation#2869)

* catalog uses sqlalchemy in async mode
* use async sqlalchemy in tests too
* added caching of director-v0 services

Author: sanderegg

packages/postgres-database/src/simcore_postgres_database/utils_aiosqlalchemy.py

@@ -0,0 +1,32 @@
+from typing import Any, Dict
+
+import sqlalchemy as sa
+from sqlalchemy.ext.asyncio import AsyncEngine
+
+from .utils_migration import get_current_head
+
+
+async def get_pg_engine_stateinfo(engine: AsyncEngine) -> Dict[str, Any]:
+    return {
+        "current pool connections": f"{engine.pool.checkedin()=},{engine.pool.checkedout()=}",
+    }
+
+
+class DBMigrationError(RuntimeError):
+    pass
+
+
+async def raise_if_migration_not_ready(engine: AsyncEngine):
+    """Ensures db migration is complete
+
+    :raises DBMigrationError
+    """
+    async with engine.connect() as conn:
+        version_num = await conn.scalar(
+            sa.DDL('SELECT "version_num" FROM "alembic_version"')
+        )
+        head_version_num = get_current_head()
+        if version_num != head_version_num:
+            raise DBMigrationError(
+                f"Migration is incomplete, expected {head_version_num} but got {version_num}"
+            )

packages/pytest-simcore/src/pytest_simcore/postgres_service.py

@@ -4,7 +4,6 @@
 import logging
 from typing import AsyncIterator, Dict, Iterator, List
 
-import aiopg.sa
 import pytest
 import sqlalchemy as sa
 import tenacity
@@ -168,7 +167,7 @@ def postgres_db(
 @pytest.fixture(scope="function")
 async def aiopg_engine(
     postgres_db: sa.engine.Engine,
-) -> AsyncIterator[aiopg.sa.engine.Engine]:
+) -> AsyncIterator:
     """An aiopg engine connected to an initialized database"""
     from aiopg.sa import create_engine
 
@@ -181,6 +180,22 @@ async def aiopg_engine(
         await engine.wait_closed()
 
 
+@pytest.fixture(scope="function")
+async def sqlalchemy_async_engine(
+    postgres_db: sa.engine.Engine,
+) -> AsyncIterator:
+    # NOTE: prevent having to import this if latest sqlalchemy not installed
+    from sqlalchemy.ext.asyncio import create_async_engine
+
+    engine = create_async_engine(
+        f"{postgres_db.url}".replace("postgresql", "postgresql+asyncpg")
+    )
+    assert engine
+    yield engine
+
+    await engine.dispose()
+
+
 @pytest.fixture(scope="function")
 def postgres_host_config(postgres_dsn: Dict[str, str], monkeypatch) -> Dict[str, str]:
     monkeypatch.setenv("POSTGRES_USER", postgres_dsn["user"])
@@ -195,7 +210,7 @@ def postgres_host_config(postgres_dsn: Dict[str, str], monkeypatch) -> Dict[str,
 
 
 @pytest.fixture(scope="module")
-def postgres_session(postgres_db: sa.engine.Engine) -> sa.orm.session.Session:
+def postgres_session(postgres_db: sa.engine.Engine) -> Iterator[sa.orm.session.Session]:
     from sqlalchemy.orm.session import Session
 
     Session_cls = sessionmaker(postgres_db)

packages/settings-library/src/settings_library/postgres.py

@@ -33,7 +33,12 @@ class PostgresSettings(BaseCustomSettings):
     POSTGRES_CLIENT_NAME: Optional[str] = Field(
         None,
         description="Name of the application connecting the postgres database, will default to use the host hostname (hostname on linux)",
-        env=["HOST", "HOSTNAME", "POSTGRES_CLIENT_NAME"],
+        env=[
+            "POSTGRES_CLIENT_NAME",
+            # This is useful when running inside a docker container, then the hostname is set each client gets a different name
+            "HOST",
+            "HOSTNAME",
+        ],
     )
 
     @validator("POSTGRES_MAXSIZE")
@@ -57,6 +62,11 @@ def dsn(self) -> str:
         )
         return dsn
 
+    @cached_property
+    def dsn_with_async_sqlalchemy(self) -> str:
+        dsn = self.dsn.replace("postgresql", "postgresql+asyncpg")
+        return dsn
+
     @cached_property
     def dsn_with_query(self) -> str:
         """Some clients do not support queries in the dsn"""

services/catalog/requirements/_base.in

@@ -21,12 +21,14 @@ fastapi[all]
 pydantic[dotenv]
 
 # database
-aiopg[sa]
+asyncpg
+sqlalchemy[asyncio]
 
 # web client
 httpx
 
 # other
+aiocache[redis,msgpack]
 tenacity
 packaging
 pyyaml

services/catalog/requirements/_base.txt

@@ -4,6 +4,8 @@
 #
 #    pip-compile --output-file=requirements/_base.txt --strip-extras requirements/_base.in
 #
+aiocache==0.11.1
+    # via -r requirements/_base.in
 aiodebug==1.1.2
     # via
     #   -c requirements/../../../packages/service-library/requirements/./_base.in
@@ -12,8 +14,15 @@ aiofiles==0.5.0
     # via
     #   -c requirements/../../../packages/service-library/requirements/./_base.in
     #   -r requirements/../../../packages/service-library/requirements/_base.in
-aiopg==1.3.3
-    # via -r requirements/_base.in
+aioredis==1.3.1
+    # via
+    #   -c requirements/../../../packages/models-library/requirements/../../../requirements/constraints.txt
+    #   -c requirements/../../../packages/postgres-database/requirements/../../../requirements/constraints.txt
+    #   -c requirements/../../../packages/service-library/requirements/../../../requirements/constraints.txt
+    #   -c requirements/../../../packages/service-library/requirements/./../../../requirements/constraints.txt
+    #   -c requirements/../../../packages/settings-library/requirements/../../../requirements/constraints.txt
+    #   -c requirements/../../../requirements/constraints.txt
+    #   aiocache
 alembic==1.7.4
     # via -r requirements/../../../packages/postgres-database/requirements/_base.in
 anyio==3.5.0
@@ -23,7 +32,9 @@ anyio==3.5.0
 asgiref==3.4.1
     # via uvicorn
 async-timeout==4.0.2
-    # via aiopg
+    # via aioredis
+asyncpg==0.25.0
+    # via -r requirements/_base.in
 certifi==2020.12.5
     # via
     #   httpcore
@@ -56,6 +67,8 @@ h11==0.12.0
     # via
     #   httpcore
     #   uvicorn
+hiredis==2.0.0
+    # via aioredis
 httpcore==0.14.4
     # via httpx
 httptools==0.2.0
@@ -102,6 +115,8 @@ markupsafe==1.1.1
     # via
     #   jinja2
     #   mako
+msgpack==1.0.3
+    # via aiocache
 multidict==5.1.0
     # via yarl
 opentracing==2.4.0
@@ -113,9 +128,7 @@ orjson==3.4.8
 packaging==20.9
     # via -r requirements/_base.in
 psycopg2-binary==2.8.6
-    # via
-    #   aiopg
-    #   sqlalchemy
+    # via sqlalchemy
 pydantic==1.9.0
     # via
     #   -c requirements/../../../packages/models-library/requirements/../../../requirements/constraints.txt
@@ -180,7 +193,7 @@ sqlalchemy==1.4.31
     #   -c requirements/../../../packages/settings-library/requirements/../../../requirements/constraints.txt
     #   -c requirements/../../../requirements/constraints.txt
     #   -r requirements/../../../packages/postgres-database/requirements/_base.in
-    #   aiopg
+    #   -r requirements/_base.in
     #   alembic
 starlette==0.17.1
     # via fastapi

services/catalog/requirements/_test.in

@@ -14,7 +14,8 @@ jsonschema
 
 # testing
 pytest
-pytest-aiohttp  # incompatible with pytest-asyncio. See https://github.com/pytest-dev/pytest-asyncio/issues/76
+pytest-aiohttp
+pytest-benchmark
 pytest-cov
 pytest-mock
 pytest-runner

services/catalog/requirements/_test.txt

@@ -166,6 +166,8 @@ ptvsd==4.3.2
     # via -r requirements/_test.in
 py==1.11.0
     # via pytest
+py-cpuinfo==8.0.0
+    # via pytest-benchmark
 pycparser==2.21
     # via cffi
 pylint==2.12.2
@@ -183,13 +185,16 @@ pytest==6.2.5
     #   -r requirements/_test.in
     #   pytest-aiohttp
     #   pytest-asyncio
+    #   pytest-benchmark
     #   pytest-cov
     #   pytest-docker
     #   pytest-mock
 pytest-aiohttp==1.0.4
     # via -r requirements/_test.in
 pytest-asyncio==0.18.1
     # via pytest-aiohttp
+pytest-benchmark==3.4.1
+    # via -r requirements/_test.in
 pytest-cov==3.0.0
     # via -r requirements/_test.in
 pytest-docker==0.10.3

services/catalog/src/simcore_service_catalog/api/dependencies/database.py

@@ -1,38 +1,32 @@
 import logging
 from typing import AsyncGenerator, Callable, Type
 
-from aiopg.sa import Engine
 from fastapi import Depends
 from fastapi.requests import Request
+from sqlalchemy.ext.asyncio import AsyncEngine
 
 from ...db.repositories import BaseRepository
 
 logger = logging.getLogger(__name__)
 
 
-def _get_db_engine(request: Request) -> Engine:
+def _get_db_engine(request: Request) -> AsyncEngine:
     return request.app.state.engine
 
 
 def get_repository(repo_type: Type[BaseRepository]) -> Callable:
     async def _get_repo(
-        engine: Engine = Depends(_get_db_engine),
+        engine: AsyncEngine = Depends(_get_db_engine),
     ) -> AsyncGenerator[BaseRepository, None]:
         # NOTE: 2 different ideas were tried here with not so good
         # 1st one was acquiring a connection per repository which lead to the following issue https://github.com/ITISFoundation/osparc-simcore/pull/1966
         # 2nd one was acquiring a connection per request which works but blocks the director-v2 responsiveness once
         # the max amount of connections is reached
-        # now the current solution is to acquire connection when needed.
-        available_engines = engine.maxsize - (engine.size - engine.freesize)
-        if available_engines <= 1:
-            logger.warning(
-                "Low pg connections available in pool: pool size=%d, acquired=%d, free=%d, reserved=[%d, %d]",
-                engine.size,
-                engine.size - engine.freesize,
-                engine.freesize,
-                engine.minsize,
-                engine.maxsize,
-            )
+        # now the current solution is to connect connection when needed.
+        logger.info(
+            "%s",
+            f"current pool connections {engine.pool.checkedin()=},{engine.pool.checkedout()=}",
+        )
         yield repo_type(db_engine=engine)
 
     return _get_repo