From 5e683274dc66f0bae736ef7d520da1cf7c9f43b3 Mon Sep 17 00:00:00 2001
From: Nargis Sultani <nargis.sultani@cfpb.gov>
Date: Tue, 28 Nov 2023 17:05:10 -0500
Subject: [PATCH] Addressed the comments

---
 db_revisions/env.py | 22 ++++++++++++++++++-
 src/config.py       | 51 ---------------------------------------------
 src/main.py         | 13 +++++++++---
 3 files changed, 31 insertions(+), 55 deletions(-)

diff --git a/db_revisions/env.py b/db_revisions/env.py
index 0f3d5bb..db720db 100644
--- a/db_revisions/env.py
+++ b/db_revisions/env.py
@@ -1,8 +1,28 @@
+import os
 from sqlalchemy import engine_from_config, pool
 from logging.config import fileConfig
 from alembic import context
 from entities.models import Base
-from config import INST_DB_URL, INST_DB_SCHEMA
+from dotenv import load_dotenv
+
+# SBL-specific configuration
+
+ENV = os.getenv("ENV", "LOCAL")
+
+if ENV == "LOCAL":
+    file_dir = os.path.dirname(os.path.realpath(__file__))
+    load_dotenv(f"{file_dir}/../src/.env.local")
+else:
+    load_dotenv()
+
+INST_DB_USER = os.environ.get("INST_DB_USER")
+INST_DB_PWD = os.environ.get("INST_DB_PWD")
+INST_DB_HOST = os.environ.get("INST_DB_HOST")
+INST_DB_NAME = os.environ.get("INST_DB_NAME")
+INST_DB_SCHEMA = os.environ.get("INST_DB_SCHEMA")
+INST_DB_PATH = f"{INST_DB_USER}:{INST_DB_PWD}@{INST_DB_HOST}/{INST_DB_NAME}"
+INST_DB_URL = f"postgresql://{INST_DB_USER}:{INST_DB_PWD}@{INST_DB_HOST}/{INST_DB_NAME}"
+# END SBL-specific configuration
 
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.
diff --git a/src/config.py b/src/config.py
index cd6195b..be326d0 100644
--- a/src/config.py
+++ b/src/config.py
@@ -1,60 +1,9 @@
 import os
-import sys
 from typing import Dict, Any
 from pydantic import TypeAdapter
 from pydantic.networks import HttpUrl, PostgresDsn
 from pydantic.types import SecretStr
 from pydantic_settings import BaseSettings, SettingsConfigDict
-from alembic import command as alembic
-from alembic.config import Config
-from dotenv import load_dotenv
-
-
-# Start Alembic
-if getattr(sys, "frozen", False):
-    # we are running in a bundle
-    basedir = sys._MEIPast
-else:
-    # we are running in a normal Python environment
-    basedir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
-    sys.path.append(basedir)
-
-alembic_dir = os.path.join(basedir, "db_revisions")
-sys.path.append(alembic_dir)
-
-# this specific to SBL configuration
-env_file = os.path.join(basedir, "src/.env.local")
-
-ENV = os.getenv("ENV", "LOCAL")
-
-if ENV == "LOCAL":
-    load_dotenv(env_file)
-else:
-    load_dotenv()
-
-INST_DB_USER = os.environ.get("INST_DB_USER")
-INST_DB_PWD = os.environ.get("INST_DB_PWD")
-INST_DB_HOST = os.environ.get("INST_DB_HOST")
-INST_DB_NAME = os.environ.get("INST_DB_NAME")
-INST_DB_SCHEMA = os.environ.get("INST_DB_SCHEMA")
-INST_DB_PATH = f"{INST_DB_USER}:{INST_DB_PWD}@{INST_DB_HOST}/{INST_DB_NAME}"
-INST_DB_URL = f"postgresql://{INST_DB_USER}:{INST_DB_PWD}@{INST_DB_HOST}/{INST_DB_NAME}"
-# end specific SBL configuration
-
-
-def get_alembic_config(db_url: str = INST_DB_URL) -> Config:
-    alembic_cfg = Config()
-    alembic_cfg.set_main_option("script_location", alembic_dir)
-    alembic_cfg.set_main_option("sqlalchemy.url", str(db_url))
-    return alembic_cfg
-
-
-def upgrade_database(revision: str = "head", db_url: str = INST_DB_URL) -> None:
-    alembic_cfg = get_alembic_config(db_url)
-    alembic.upgrade(alembic_cfg, revision)
-
-
-# End Alembic
 
 JWT_OPTS_PREFIX = "jwt_opts_"
 
diff --git a/src/main.py b/src/main.py
index a85fbdc..0b6fba8 100644
--- a/src/main.py
+++ b/src/main.py
@@ -1,3 +1,4 @@
+import os
 import logging
 from http import HTTPStatus
 from fastapi import FastAPI, HTTPException, Request
@@ -5,12 +6,14 @@
 from fastapi.security import OAuth2AuthorizationCodeBearer
 from fastapi.middleware.cors import CORSMiddleware
 from starlette.middleware.authentication import AuthenticationMiddleware
+from alembic.config import Config
+from alembic import command
 
 from routers import admin_router, institutions_router
 
 from oauth2 import BearerTokenAuthBackend
 
-from config import settings, upgrade_database
+from config import settings
 
 log = logging.getLogger()
 
@@ -18,8 +21,12 @@
 
 
 @app.on_event("startup")
-async def startup() -> None:
-    upgrade_database()
+async def app_start():
+    file_dir = os.path.dirname(os.path.realpath(__file__))
+    alembic_cfg = Config(f"{file_dir}/../alembic.ini")
+    alembic_cfg.set_main_option("script_location", f"{file_dir}/../db_revisions")
+    alembic_cfg.set_main_option("prepend_sys_path", f"{file_dir}/../")
+    command.upgrade(alembic_cfg, "head")
 
 
 @app.exception_handler(HTTPException)