feat: update database conn envs

pehlicd · Mar 23, 2024 · cef21ae · cef21ae
1 parent 4200eae
commit cef21ae
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 12 deletions.
diff --git a/docker-compose.common.yml b/docker-compose.common.yml
@@ -21,7 +21,11 @@ services:
       - PORT=8080
       - SECRET_MANAGER_TYPE=FILE
       - SECRET_MANAGER_DIRECTORY=/state
-      - DATABASE_CONNECTION_STRING=sqlite:///./db.sqlite3?check_same_thread=False
+      - DATABASE_HOST=""
+      - DATABASE_PORT="3306"
+      - DATABASE_USER="root"
+      - DATABASE_PASSWORD=""
+      - DATABASE_NAME="keep"
       - OPENAI_API_KEY=$OPENAI_API_KEY
       - PUSHER_APP_ID=1
       - PUSHER_APP_KEY=keepappkey

diff --git a/docs/deployment/ecs.mdx b/docs/deployment/ecs.mdx
@@ -38,7 +38,11 @@ sidebarTitle: "AWS ECS"
                 - Protocol: TCP
             ![Container Details](/images/ecs-task-def-frontend3.png)
             - Environment Variables: (This can be static or you can use parameter store or secrets manager)
-                - DATABASE_CONNECTION_STRING
+                - DATABASE_HOST
+                - DATABASE_PORT
+                - DATABASE_USER
+                - DATABASE_PASSWORD
+                - DATABASE_NAME
                 - AUTH_TYPE
                 - KEEP_JWT_SECRET
                 - KEEP_DEFAULT_USERNAME
@@ -79,7 +83,11 @@ sidebarTitle: "AWS ECS"
                 - Protocol: TCP
                 ![Container Details](/images/ecs-task-def-backend3.png)
             - Environment Variables: (This can be static or you can use parameter store or secrets manager)
-                - DATABASE_CONNECTION_STRING
+                - DATABASE_HOST
+                - DATABASE_PORT
+                - DATABASE_USER
+                - DATABASE_PASSWORD
+                - DATABASE_NAME
                 - AUTH_TYPE
                 - KEEP_JWT_SECRET
                 - KEEP_DEFAULT_USERNAME

diff --git a/keep/api/core/db.py b/keep/api/core/db.py
@@ -91,19 +91,21 @@ def __get_conn_impersonate() -> pymysql.connections.Connection:
 # this is a workaround for gunicorn to load the env vars
 #   becuase somehow in gunicorn it doesn't load the .env file
 load_dotenv(find_dotenv())
-db_connection_string = config("DATABASE_CONNECTION_STRING", default=None)
+db_host = config("DATABASE_HOST", default="")
+db_port = config("DATABASE_PORT", default="3306")
+db_user = config("DATABASE_USER", default="root")
+db_password = config("DATABASE_PASSWORD", default="")
+db_name = config("DATABASE_NAME", default="keep")
+db_connection_string = "mysql+pymysql://" + db_user + ":" + db_password + "@" + db_host + ":" + db_port + "/" + db_name
+if db_host == "":
+    db_connection_string = ""
 pool_size = config("DATABASE_POOL_SIZE", default=5, cast=int)
 
 if RUNNING_IN_CLOUD_RUN:
     engine = create_engine(
         "mysql+pymysql://",
         creator=__get_conn,
     )
-elif db_connection_string == "impersonate":
-    engine = create_engine(
-        "mysql+pymysql://",
-        creator=__get_conn_impersonate,
-    )
 elif db_connection_string:
     try:
         logger.info(f"Creating a connection pool with size {pool_size}")

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -139,9 +139,6 @@ def db_session(request, mysql_container):
         )
     SQLModel.metadata.create_all(mock_engine)
 
-    # Mock the environment variables so db.py will use it
-    os.environ["DATABASE_CONNECTION_STRING"] = db_connection_string
-
     # Create a session
     SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=mock_engine)
     session = SessionLocal()