From a15495515468e84a93236a0af847c00be882dda4 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Mon, 28 Aug 2023 16:22:45 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20fix(=5F=5Fmain=5F=5F.py):=20impo?=
 =?UTF-8?q?rt=20display=5Fresults=20function=20from=20utils.util=20to=20fi?=
 =?UTF-8?q?x=20NameError=20=E2=9C=A8=20feat(=5F=5Fmain=5F=5F.py):=20add=20?=
 =?UTF-8?q?migration=20command=20to=20run=20database=20migrations=20with?=
 =?UTF-8?q?=20an=20option=20to=20run=20in=20test=20mode=20=F0=9F=94=A7=20f?=
 =?UTF-8?q?ix(util.py):=20import=20console=20from=20=5F=5Fmain=5F=5F=20mod?=
 =?UTF-8?q?ule=20to=20fix=20NameError=20=E2=9C=A8=20feat(util.py):=20add?=
 =?UTF-8?q?=20display=5Fresults=20function=20to=20display=20migration=20re?=
 =?UTF-8?q?sults=20in=20a=20table=20format?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/backend/langflow/__main__.py   | 15 +++++++++++++++
 src/backend/langflow/utils/util.py | 25 +++++++++++++++++++++++--
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/src/backend/langflow/__main__.py b/src/backend/langflow/__main__.py
index c2b7b0fd6..5c5582fdd 100644
--- a/src/backend/langflow/__main__.py
+++ b/src/backend/langflow/__main__.py
@@ -5,6 +5,7 @@ from langflow.services.database.utils import session_getter
 from langflow.services.manager import initialize_services, initialize_settings_manager
 from langflow.services.utils import get_db_manager, get_settings_manager
 from langflow.utils.util import get_number_of_workers
+from langflow.utils.util import display_results
 from multiprocess import Process  # type: ignore
 import platform
 from pathlib import Path
@@ -19,6 +20,10 @@ from langflow.utils.logger import configure, logger
 import webbrowser
 from dotenv import load_dotenv
 
+from rich.console import Console
+
+console = Console()
+
 app = typer.Typer()
 
 
@@ -340,6 +345,16 @@ def superuser(
             typer.echo("Superuser creation failed.")
 
 
+@app.command()
+def migration(test: bool = typer.Option(False, help="Run migrations in test mode.")):
+    initialize_services()
+    db_manager = get_db_manager()
+    if not test:
+        db_manager.run_migrations()
+    results = db_manager.run_migrations_test()
+    display_results(results)
+
+
 def main():
     app()
 
diff --git a/src/backend/langflow/utils/util.py b/src/backend/langflow/utils/util.py
index 90c9cd015..8f0ff216a 100644
--- a/src/backend/langflow/utils/util.py
+++ b/src/backend/langflow/utils/util.py
@@ -4,12 +4,14 @@ import importlib
 from functools import wraps
 from typing import Optional, Dict, Any, Union
 
-from docstring_parser import parse  # type: ignore
+from docstring_parser import parse
+from langflow.__main__ import console  # type: ignore
 
 from langflow.template.frontend_node.constants import FORCE_SHOW_FIELDS
 from langflow.utils import constants
 from langflow.utils.logger import logger
-from multiprocess import cpu_count  # type: ignore
+from multiprocess import cpu_count
+from rich.table import Table  # type: ignore
 
 
 def build_template_from_function(
@@ -465,3 +467,22 @@ def get_number_of_workers(workers=None):
         workers = (cpu_count() * 2) + 1
     logger.debug(f"Number of workers: {workers}")
     return workers
+
+
+def display_results(results):
+    """
+    Display the results of the migration.
+    """
+    for table_results in results:
+        table = Table(title=f"Migration {table_results.table_name}")
+        table.add_column("Name")
+        table.add_column("Type")
+        table.add_column("Status")
+
+        for result in table_results.results:
+            status = "Success" if result.success else "Failure"
+            color = "green" if result.success else "red"
+            table.add_row(result.name, result.type, f"[{color}]{status}[/{color}]")
+
+        console.print(table)
+        console.print()  # Print a new line