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