From 8436c66aa767eab3ac78f7080c78a0dab35a0feb Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Mon, 7 Aug 2023 12:00:27 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=93=A6=20chore(manager.py):=20add=20Datab?=
 =?UTF-8?q?aseManager=20class=20to=20handle=20database=20operations=20?=
 =?UTF-8?q?=F0=9F=93=A6=20chore(utils.py):=20add=20initialize=5Fdatabase?=
 =?UTF-8?q?=20function=20and=20session=5Fgetter=20context=20manager=20to?=
 =?UTF-8?q?=20handle=20database=20initialization=20and=20session=20managem?=
 =?UTF-8?q?ent?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../services/database/{base.py => manager.py} | 26 +---------------
 .../langflow/services/database/utils.py       | 31 +++++++++++++++++++
 2 files changed, 32 insertions(+), 25 deletions(-)
 rename src/backend/langflow/services/database/{base.py => manager.py} (77%)
 create mode 100644 src/backend/langflow/services/database/utils.py

diff --git a/src/backend/langflow/services/database/base.py b/src/backend/langflow/services/database/manager.py
similarity index 77%
rename from src/backend/langflow/services/database/base.py
rename to src/backend/langflow/services/database/manager.py
index 9f92c6c25..20db90f66 100644
--- a/src/backend/langflow/services/database/base.py
+++ b/src/backend/langflow/services/database/manager.py
@@ -1,4 +1,3 @@
-from contextlib import contextmanager
 from pathlib import Path
 from langflow.services.base import Service
 from sqlmodel import SQLModel, Session, create_engine
@@ -12,7 +11,7 @@ class DatabaseManager(Service):
 
     def __init__(self, database_url: str):
         self.database_url = database_url
-        # This file is in langflow.services.database.base.py
+        # This file is in langflow.services.database.manager.py
         # the ini is in langflow
         langflow_dir = Path(__file__).parent.parent.parent
         self.script_location = langflow_dir / "alembic"
@@ -65,26 +64,3 @@ class DatabaseManager(Service):
             raise RuntimeError("Something went wrong creating the database and tables.")
         else:
             logger.debug("Database and tables created successfully")
-
-
-@contextmanager
-def session_getter(db_manager: DatabaseManager):
-    try:
-        session = Session(db_manager.engine)
-        yield session
-    except Exception as e:
-        print("Session rollback because of exception:", e)
-        session.rollback()
-        raise
-    finally:
-        session.close()
-
-
-def initialize_database():
-    logger.debug("Initializing database")
-    from langflow.services import service_manager, ServiceType
-
-    database_manager = service_manager.get(ServiceType.DATABASE_MANAGER)
-    database_manager.run_migrations()
-    database_manager.create_db_and_tables()
-    logger.debug("Database initialized")
diff --git a/src/backend/langflow/services/database/utils.py b/src/backend/langflow/services/database/utils.py
new file mode 100644
index 000000000..20b2bbbb4
--- /dev/null
+++ b/src/backend/langflow/services/database/utils.py
@@ -0,0 +1,31 @@
+from typing import TYPE_CHECKING
+from langflow.utils.logger import logger
+from contextlib import contextmanager
+
+from sqlmodel import Session
+
+if TYPE_CHECKING:
+    from langflow.services.database.manager import DatabaseManager
+
+
+def initialize_database():
+    logger.debug("Initializing database")
+    from langflow.services import service_manager, ServiceType
+
+    database_manager = service_manager.get(ServiceType.DATABASE_MANAGER)
+    database_manager.run_migrations()
+    database_manager.create_db_and_tables()
+    logger.debug("Database initialized")
+
+
+@contextmanager
+def session_getter(db_manager: "DatabaseManager"):
+    try:
+        session = Session(db_manager.engine)
+        yield session
+    except Exception as e:
+        print("Session rollback because of exception:", e)
+        session.rollback()
+        raise
+    finally:
+        session.close()