Fix connection creation leak

Author: maxachis

local_database/docker-compose.yml

@@ -12,6 +12,5 @@ services:
       - POSTGRES_PASSWORD=HanviliciousHamiltonHilltops
       - POSTGRES_USER=test_source_collector_user
       - POSTGRES_DB=source_collector_test_db
-    command: ['postgres', '-c', 'max_connections=160']
 volumes:
   dbscripts:

src/api/main.py

@@ -5,6 +5,7 @@
 from discord_poster import DiscordPoster
 from fastapi import FastAPI
 from pdap_access_manager import AccessManager
+from sqlalchemy.ext.asyncio import create_async_engine
 from starlette.responses import RedirectResponse
 
 from src.api.endpoints.agencies.routes import agencies_router
@@ -52,12 +53,9 @@ async def lifespan(app: FastAPI):
     env.read_env()
 
     # Initialize shared dependencies
-    db_client = DatabaseClient(
-        db_url=env_var_manager.get_postgres_connection_string()
-    )
-    adb_client = AsyncDatabaseClient(
-        db_url=env_var_manager.get_postgres_connection_string(is_async=True)
-    )
+
+    db_client = DatabaseClient()
+    adb_client = AsyncDatabaseClient()
     await setup_database(db_client)
     core_logger = AsyncCoreLogger(adb_client=adb_client)
 

src/db/client/async_.py

@@ -2,8 +2,8 @@
 from functools import wraps
 from typing import Optional, Type, Any, List, Sequence
 
-from sqlalchemy import select, func, Select, and_, update, Row, text
-from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
+from sqlalchemy import select, func, Select, and_, update, Row, text, Engine
+from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker, AsyncEngine
 from sqlalchemy.orm import selectinload
 
 from src.api.endpoints.annotate.all.get.models.response import GetNextURLForAllAnnotationResponse
@@ -103,15 +103,15 @@
 
 
 class AsyncDatabaseClient:
-    def __init__(self, db_url: str | None = None):
-        if db_url is None:
+    def __init__(self, engine: AsyncEngine | None = None):
+        if engine is None:
             db_url = EnvVarManager.get().get_postgres_connection_string(is_async=True)
-        self.db_url = db_url
-        echo = ConfigManager.get_sqlalchemy_echo()
-        self.engine = create_async_engine(
-            url=db_url,
-            echo=echo,
-        )
+            echo = ConfigManager.get_sqlalchemy_echo()
+            engine = create_async_engine(
+                url=db_url,
+                echo=echo,
+            )
+        self.engine = engine
         self.session_maker = async_sessionmaker(bind=self.engine, expire_on_commit=False)
         self.statement_composer = StatementComposer()
 

src/db/client/sync.py

@@ -1,7 +1,7 @@
 from functools import wraps
 from typing import List
 
-from sqlalchemy import create_engine, Select
+from sqlalchemy import create_engine, Select, Engine
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import sessionmaker, scoped_session, Session
 
@@ -28,15 +28,19 @@
 
 # Database Client
 class DatabaseClient:
-    def __init__(self, db_url: str | None = None):
+    def __init__(
+        self,
+        engine: Engine | None = None
+    ):
         """Initialize the DatabaseClient."""
-        if db_url is None:
+        if engine is None:
             db_url = EnvVarManager.get().get_postgres_connection_string(is_async=True)
+            engine = create_engine(
+                url=db_url,
+                echo=ConfigManager.get_sqlalchemy_echo(),
+            )
 
-        self.engine = create_engine(
-            url=db_url,
-            echo=ConfigManager.get_sqlalchemy_echo(),
-        )
+        self.engine = engine
         self.session_maker = scoped_session(sessionmaker(bind=self.engine))
         self.session = None
 

tests/automated/integration/readonly/conftest.py

@@ -3,6 +3,7 @@
 
 import pytest
 import pytest_asyncio
+from sqlalchemy import Engine
 from starlette.testclient import TestClient
 
 from src.db.helpers.connect import get_postgres_connection_string
@@ -33,8 +34,10 @@ async def california_readonly(
 async def readonly_helper(
     event_loop,
     client: TestClient,
+    engine: Engine
+
 ) -> AsyncGenerator[ReadOnlyTestHelper, Any]:
-    wipe_database(get_postgres_connection_string())
+    wipe_database(engine)
     db_data_creator = DBDataCreator()
     api_test_helper = APITestHelper(
         request_validator=RequestValidator(client=client),

tests/conftest.py

@@ -7,18 +7,19 @@
 import pytest_asyncio
 from aiohttp import ClientSession
 from alembic.config import Config
-from sqlalchemy import create_engine, inspect, MetaData
+from sqlalchemy import create_engine, inspect, MetaData, Engine
+from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 from sqlalchemy.orm import scoped_session, sessionmaker
 
 from src.core.env_var_manager import EnvVarManager
+from src.db.client.async_ import AsyncDatabaseClient
+from src.db.client.sync import DatabaseClient
+from src.db.helpers.connect import get_postgres_connection_string
+from src.db.models.impl.log.sqlalchemy import Log  # noqa: F401
 # Below are to prevent import errors
 from src.db.models.impl.missing import Missing  # noqa: F401
-from src.db.models.impl.log.sqlalchemy import Log  # noqa: F401
 from src.db.models.impl.task.error import TaskError  # noqa: F401
 from src.db.models.impl.url.checked_for_duplicate import URLCheckedForDuplicate  # noqa: F401
-from src.db.client.async_ import AsyncDatabaseClient
-from src.db.client.sync import DatabaseClient
-from src.db.helpers.connect import get_postgres_connection_string
 from src.util.helper_functions import load_from_environment
 from tests.helpers.alembic_runner import AlembicRunner
 from tests.helpers.data_creator.core import DBDataCreator
@@ -99,33 +100,55 @@ def setup_and_teardown():
         live_connection.close()
         engine.dispose()
 
+@pytest.fixture(scope="session")
+def engine():
+    conn = get_postgres_connection_string()
+    engine = create_engine(conn)
+    yield engine
+    engine.dispose()
+
+@pytest.fixture(scope="session")
+def async_engine():
+    conn = get_postgres_connection_string(is_async=True)
+    engine = create_async_engine(conn)
+    yield engine
+    engine.dispose()
+
 @pytest.fixture
-def wiped_database():
+def wiped_database(
+    engine: Engine
+):
     """Wipe all data from database."""
-    wipe_database(get_postgres_connection_string())
+    wipe_database(engine)
 
 
 
 @pytest.fixture
-def db_client_test(wiped_database) -> Generator[DatabaseClient, Any, None]:
+def db_client_test(
+    wiped_database,
+    engine
+) -> Generator[DatabaseClient, Any, None]:
     # Drop pre-existing table
-    conn = get_postgres_connection_string()
-    db_client = DatabaseClient(db_url=conn)
+    db_client = DatabaseClient(engine)
     yield db_client
     db_client.engine.dispose()
 
 @pytest_asyncio.fixture
-async def populated_database(wiped_database) -> None:
-    conn = get_postgres_connection_string(is_async=True)
-    adb_client = AsyncDatabaseClient(db_url=conn)
+async def populated_database(
+    wiped_database,
+    async_engine: AsyncEngine
+) -> None:
+    adb_client = AsyncDatabaseClient(async_engine)
     await populate_database(adb_client)
 
 @pytest_asyncio.fixture
-async def adb_client_test(wiped_database) -> AsyncGenerator[AsyncDatabaseClient, Any]:
-    conn = get_postgres_connection_string(is_async=True)
-    adb_client = AsyncDatabaseClient(db_url=conn)
+async def adb_client_test(
+    wiped_database,
+    async_engine: AsyncEngine
+) -> AsyncGenerator[AsyncDatabaseClient, Any]:
+    adb_client = AsyncDatabaseClient(async_engine)
     yield adb_client
-    adb_client.engine.dispose()
+    await adb_client.engine.dispose()
 
 @pytest.fixture
 def db_data_creator(

tests/helpers/setup/wipe.py

@@ -1,11 +1,10 @@
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, Engine
 
 from src.db.models.templates_.base import Base
 
 
-def wipe_database(connection_string: str) -> None:
+def wipe_database(engine: Engine) -> None:
     """Wipe all data from database."""
-    engine = create_engine(connection_string)
     with engine.connect() as connection:
         for table in reversed(Base.metadata.sorted_tables):
             if table.info == "view":