🔧 fix(__main__.py): import display_results function from utils.util to fix NameError

 feat(__main__.py): add migration command to run database migrations with an option to run in test mode
🔧 fix(util.py): import console from __main__ module to fix NameError
 feat(util.py): add display_results function to display migration results in a table format
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-08-28 16:22:45 -03:00
commit a154955154
2 changed files with 38 additions and 2 deletions

View file

@ -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()

View file

@ -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