diff --git a/fia_api/core/repositories.py b/fia_api/core/repositories.py index 63d5b27a..ca5dc9db 100644 --- a/fia_api/core/repositories.py +++ b/fia_api/core/repositories.py @@ -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 @@ -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]): """ diff --git a/fia_api/core/session.py b/fia_api/core/session.py index bf1048a8..704066e3 100644 --- a/fia_api/core/session.py +++ b/fia_api/core/session.py @@ -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") @@ -11,10 +11,14 @@ ENGINE = create_engine( DB_URL, - poolclass=NullPool, + pool_size=5, + max_overflow=10, + pool_pre_ping=True, + pool_recycle=300, ) SessionLocal = sessionmaker(bind=ENGINE, autoflush=False, autocommit=False) +SESSION = sessionmaker(ENGINE) def get_db_session() -> Generator[Session, None, None]: @@ -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)) diff --git a/fia_api/routers/health.py b/fia_api/routers/health.py index d26b9fc9..29d2c67d 100644 --- a/fia_api/routers/health.py +++ b/fia_api/routers/health.py @@ -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() diff --git a/test/core/test_repositories_integration.py b/test/core/test_repositories_integration.py index 59dac022..cc5b7e1e 100644 --- a/test/core/test_repositories_integration.py +++ b/test/core/test_repositories_integration.py @@ -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") @@ -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() diff --git a/test/e2e/conftest.py b/test/e2e/conftest.py index d1e6e678..83ea5f0b 100644 --- a/test/e2e/conftest.py +++ b/test/e2e/conftest.py @@ -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 diff --git a/test/e2e/test_autoreduction_routes.py b/test/e2e/test_autoreduction_routes.py index fcd9dd84..c5c74efe 100644 --- a/test/e2e/test_autoreduction_routes.py +++ b/test/e2e/test_autoreduction_routes.py @@ -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 diff --git a/test/utils.py b/test/utils.py index df53dc4e..ef2e59ae 100644 --- a/test/utils.py +++ b/test/utils.py @@ -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)