Repository: All-Hands-AI/OpenHands
File: openhands/app_server/services/db_session_injector.py (Python)

Code before:

        return AsyncAdapt_asyncpg_connection(
            engine.dialect.dbapi,
            await self._create_async_gcp_db_connection(),
            prepared_statement_cache_size=100,
        )

    async def _create_async_gcp_engine(self):
        from sqlalchemy.dialects.postgresql.asyncpg import (
            AsyncAdapt_asyncpg_connection,
        )

        base_engine = self._create_gcp_engine()
        dbapi = base_engine.dialect.dbapi

        def adapted_creator():
            return AsyncAdapt_asyncpg_connection(
                dbapi,
                await_only(self._create_async_gcp_db_connection()),
                prepared_statement_cache_size=100,
            )

        return create_async_engine(
            'postgresql+asyncpg://',
            creator=adapted_creator,
            poolclass=NullPool,
        )

    async def get_async_db_engine(self) -> AsyncEngine:
        async_engine = self._async_engine
        if async_engine:
            return async_engine
        if self.gcp_db_instance:  # GCP environments
            async_engine = await self._create_async_gcp_engine()
        else:
            if self.host:
                try:
                    import asyncpg  # noqa: F401
                except Exception as e:
                    raise RuntimeError(
                        "PostgreSQL driver 'asyncpg' is required for async connections but is not installed."
                    ) from e
                password = self.password.get_secret_value() if self.password else None
                url = URL.create(
                    'postgresql+asyncpg',
                    username=self.user or '',
                    password=password,
                    host=self.host,
                    port=self.port,
                    database=self.name,
                )

Review comment (performance):
The NullPool is a cause for concern to me here - it disables connection pooling which slows things down a bit

Code after:

        return AsyncAdapt_asyncpg_connection(
            engine.dialect.dbapi,
            await self._create_async_gcp_db_connection(),
            prepared_statement_cache_size=100,
        )

    async def _create_async_gcp_engine(self):
        from sqlalchemy.dialects.postgresql.asyncpg import (
            AsyncAdapt_asyncpg_connection,
        )

        base_engine = self._create_gcp_engine()
        dbapi = base_engine.dialect.dbapi

        def adapted_creator():
            return AsyncAdapt_asyncpg_connection(
                dbapi,
                await_only(self._create_async_gcp_db_connection()),
                prepared_statement_cache_size=100,
            )

        return create_async_engine(
            'postgresql+asyncpg://',
            creator=adapted_creator,
            pool_size=self.pool_size,
            max_overflow=self.max_overflow,
            pool_pre_ping=True,
        )

    async def get_async_db_engine(self) -> AsyncEngine:
        async_engine = self._async_engine
        if async_engine:
            return async_engine
        if self.gcp_db_instance:  # GCP environments
            async_engine = await self._create_async_gcp_engine()
        else:
            if self.host:
                try:
                    import asyncpg  # noqa: F401
                except Exception as e:
                    raise RuntimeError(
                        "PostgreSQL driver 'asyncpg' is required for async connections but is not installed."
                    ) from e
                password = self.password.get_secret_value() if self.password else None
                url = URL.create(
                    'postgresql+asyncpg',
                    username=self.user or '',
                    password=password,
                    host=self.host,
                    port=self.port,