From 4d08d511e9155f7a0deab8a5e223d0f11aecde9b Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Mon, 28 Aug 2023 11:18:01 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(manager.py):=20fix=20typo=20?=
 =?UTF-8?q?in=20logger.warn=20to=20logger.warning=20=E2=9C=A8=20feat(manag?=
 =?UTF-8?q?er.py):=20add=20check=5Fschema=5Fhealth=20method=20to=20Databas?=
 =?UTF-8?q?eManager=20to=20verify=20the=20integrity=20of=20the=20database?=
 =?UTF-8?q?=20schema=20=F0=9F=90=9B=20fix(utils.py):=20add=20error=20handl?=
 =?UTF-8?q?ing=20when=20checking=20schema=20health=20in=20initialize=5Fdat?=
 =?UTF-8?q?abase=20method?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../langflow/services/database/manager.py     | 54 ++++++++++++++++---
 .../langflow/services/database/utils.py       |  5 ++
 2 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/src/backend/langflow/services/database/manager.py b/src/backend/langflow/services/database/manager.py
index bf90703e2..a95010bc2 100644
--- a/src/backend/langflow/services/database/manager.py
+++ b/src/backend/langflow/services/database/manager.py
@@ -2,6 +2,8 @@ from pathlib import Path
 from typing import TYPE_CHECKING
 from langflow.services.base import Service
 from langflow.services.utils import get_settings_manager
+from sqlalchemy import inspect
+import sqlalchemy as sa
 from sqlmodel import SQLModel, Session, create_engine
 from langflow.utils.logger import logger
 from alembic.config import Config
@@ -54,6 +56,41 @@ class DatabaseManager(Service):
         with Session(self.engine) as session:
             yield session
 
+    def check_schema_health(self) -> bool:
+        inspector = inspect(self.engine)
+
+        model_mapping = {
+            "flow": models.Flow,
+            "user": models.User,
+            "apikey": models.ApiKey,
+            # Add other SQLModel classes here
+        }
+
+        # To account for tables that existed in older versions
+        legacy_tables = ["flowstyle"]
+
+        for table, model in model_mapping.items():
+            expected_columns = list(model.__fields__.keys())
+
+            try:
+                available_columns = [
+                    col["name"] for col in inspector.get_columns(table)
+                ]
+            except sa.exc.NoSuchTableError:
+                logger.error(f"Missing table: {table}")
+                return False
+
+            for column in expected_columns:
+                if column not in available_columns:
+                    logger.error(f"Missing column: {column} in table {table}")
+                    return False
+
+        for table in legacy_tables:
+            if table in inspector.get_table_names():
+                logger.warn(f"Legacy table exists: {table}")
+
+        return True
+
     def run_migrations(self):
         logger.info(
             f"Running DB migrations in {self.script_location} on {self.database_url}"
@@ -76,9 +113,14 @@ class DatabaseManager(Service):
         from sqlalchemy import inspect
 
         inspector = inspect(self.engine)
-        if "flow" not in inspector.get_table_names():
-            logger.error("Something went wrong creating the database and tables.")
-            logger.error("Please check your database settings.")
-            raise RuntimeError("Something went wrong creating the database and tables.")
-        else:
-            logger.debug("Database and tables created successfully")
+        current_tables = ["flow", "user", "apikey"]
+        table_names = inspector.get_table_names()
+        for table in current_tables:
+            if table not in table_names:
+                logger.error("Something went wrong creating the database and tables.")
+                logger.error("Please check your database settings.")
+                raise RuntimeError(
+                    "Something went wrong creating the database and tables."
+                )
+
+        logger.debug("Database and tables created successfully")
diff --git a/src/backend/langflow/services/database/utils.py b/src/backend/langflow/services/database/utils.py
index 94bcd6651..efbf2b2c8 100644
--- a/src/backend/langflow/services/database/utils.py
+++ b/src/backend/langflow/services/database/utils.py
@@ -13,6 +13,11 @@ def initialize_database():
     from langflow.services import service_manager, ServiceType
 
     database_manager = service_manager.get(ServiceType.DATABASE_MANAGER)
+    try:
+        database_manager.check_schema_health()
+    except Exception as exc:
+        logger.error(f"Error checking schema health: {exc}")
+        raise RuntimeError("Error checking schema health") from exc
     try:
         database_manager.run_migrations()
     except CommandError as exc: