Skip to content

Replace NullPool with QueuePool and consolidate database engines #621

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Dagonite wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Replace NullPool with QueuePool and consolidate database engines #621

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9cc621f
Replace NullPool with QueuePool and consolidate database engines
Dagonite Mar 10, 2026
5eb95e7
Formatting and linting commit
invalid-email-address Mar 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 2 additions & 20 deletions fia_api/core/repositories.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,12 @@
"""Provides a generic repository class for performing database operations."""

import logging
import os
from collections.abc import Sequence
from typing import Generic, TypeVar

from sqlalchemy import NullPool, create_engine, func, select
from sqlalchemy import func, select
from sqlalchemy.exc import MultipleResultsFound, NoResultFound
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.orm import Session

from fia_api.core.exceptions import NonUniqueRecordError
from fia_api.core.models import Base
Expand All @@ -17,23 +16,6 @@

logger = logging.getLogger(__name__)

DB_USERNAME = os.environ.get("DB_USERNAME", "postgres")
DB_PASSWORD = os.environ.get("DB_PASSWORD", "password")
DB_IP = os.environ.get("DB_IP", "localhost")

ENGINE = create_engine(
f"postgresql+psycopg2://{DB_USERNAME}:{DB_PASSWORD}@{DB_IP}:5432/fia",
poolclass=NullPool,
)

SESSION = sessionmaker(ENGINE)


def ensure_db_connection(db: Session) -> None:
"""Test connection to database."""
with db as session:
session.execute(select(1))


class Repo(Generic[T]):
"""
Expand Down
14 changes: 12 additions & 2 deletions fia_api/core/session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import os
from collections.abc import Generator

from sqlalchemy import NullPool, create_engine
from sqlalchemy import create_engine, select
from sqlalchemy.orm import Session, sessionmaker

DB_USERNAME = os.environ.get("DB_USERNAME", "postgres")
Expand All @@ -11,10 +11,14 @@

ENGINE = create_engine(
DB_URL,
poolclass=NullPool,
pool_size=5,
max_overflow=10,
pool_pre_ping=True,
pool_recycle=300,
)
Comment on lines 12 to 18
Copy link

Copilot AI Mar 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The connection pool sizing/tuning is hard-coded (pool_size/max_overflow/pool_recycle). This can cause unexpected connection counts in production (e.g., per-worker pools) and makes it hard to tune without a redeploy. Consider sourcing these values from environment variables (with sensible defaults) so deployments can adjust based on DB limits and worker count.

Copilot uses AI. Check for mistakes.

SessionLocal = sessionmaker(bind=ENGINE, autoflush=False, autocommit=False)
SESSION = sessionmaker(ENGINE)
Copy link

Copilot AI Mar 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are now two session factories bound to the same ENGINE (SessionLocal and SESSION) with different configuration (SessionLocal sets autoflush/autocommit, SESSION does not). This duplication can lead to inconsistent behavior depending on which one is imported. Consider exporting a single configured sessionmaker (or making SESSION an alias of SessionLocal) and using it consistently.

Suggested change
SESSION = sessionmaker(ENGINE)
SESSION = SessionLocal

Copilot uses AI. Check for mistakes.


def get_db_session() -> Generator[Session, None, None]:
Expand All @@ -23,3 +27,9 @@ def get_db_session() -> Generator[Session, None, None]:
yield db
finally:
db.close()


def ensure_db_connection(db: Session) -> None:
"""Test connection to database."""
with db as session:
session.execute(select(1))
Comment on lines +34 to +35
Copy link

Copilot AI Mar 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ensure_db_connection uses with db as session: on a Session instance that is typically owned/closed by the caller (e.g., FastAPI dependency). Entering the context manager will close the passed-in session on exit, which is surprising and can cause the caller to operate on a closed session (and also results in double-close with get_db_session()). Prefer executing directly on the provided Session without wrapping it in a context manager, or have this helper create/own its own Session instead.

Suggested change
with db as session:
session.execute(select(1))
db.execute(select(1))

Copilot uses AI. Check for mistakes.
3 changes: 1 addition & 2 deletions fia_api/routers/health.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,7 @@
from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session

from fia_api.core.repositories import ensure_db_connection
from fia_api.core.session import get_db_session
from fia_api.core.session import ensure_db_connection, get_db_session

health_router = APIRouter()

Expand Down
5 changes: 3 additions & 2 deletions test/core/test_repositories_integration.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,8 @@
from sqlalchemy.orm import Session, sessionmaker

from fia_api.core.models import Base, Instrument, Job, JobOwner, JobType, Run, Script, State
from fia_api.core.repositories import ENGINE, SESSION, Repo, ensure_db_connection
from fia_api.core.repositories import Repo
from fia_api.core.session import ENGINE, SESSION, ensure_db_connection
from fia_api.core.specifications.job import JobSpecification

DB_USERNAME = os.environ.get("DB_USERNAME", "postgres")
Expand Down Expand Up @@ -204,7 +205,7 @@ def test_jobs_by_instrument_sort_by_job_field(job_repo):
assert result == expected


@mock.patch("fia_api.core.repositories.select")
@mock.patch("fia_api.core.session.select")
def test_ensure_db_connection_raises_httpexception(mock_select):
"""Test exception raised when runtime error occurs"""
mock_session_object = MagicMock()
Expand Down
2 changes: 1 addition & 1 deletion test/e2e/conftest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
import pytest
from sqlalchemy.orm import make_transient

from fia_api.core.repositories import SESSION
from fia_api.core.session import SESSION
from test.e2e.constants import TEST_INSTRUMENT, TEST_JOB, TEST_RUN, TEST_SCRIPT
from test.utils import setup_database

Expand Down
2 changes: 1 addition & 1 deletion test/e2e/test_autoreduction_routes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@
from sqlalchemy import delete, select

from fia_api.core.models import Job, Run
from fia_api.core.repositories import SESSION
from fia_api.core.responses import JobResponse
from fia_api.core.session import SESSION
from test.e2e.constants import API_KEY_HEADER, STAFF_HEADER, TEST_JOB, USER_HEADER
from test.e2e.test_core import client

Expand Down
2 changes: 1 addition & 1 deletion test/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from faker.providers import BaseProvider

from fia_api.core.models import Base, Instrument, Job, JobOwner, JobType, Run, Script, State
from fia_api.core.repositories import ENGINE, SESSION
from fia_api.core.session import ENGINE, SESSION

random.seed(0)

Expand Down