From 5f2a949ec5ec42ed7ea67098abeabbac18c41121 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 28 Aug 2023 16:21:50 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20fix(manager.py):=20import=20miss?= =?UTF-8?q?ing=20dependencies=20to=20resolve=20NameError=20and=20improve?= =?UTF-8?q?=20code=20readability=20=E2=9C=A8=20feat(manager.py):=20add=20m?= =?UTF-8?q?ethod=20run=5Fmigrations=5Ftest()=20to=20check=20if=20all=20mod?= =?UTF-8?q?els=20are=20in=20the=20database=20and=20up=20to=20date=20with?= =?UTF-8?q?=20all=20columns=20=E2=9C=A8=20feat(manager.py):=20add=20method?= =?UTF-8?q?=20check=5Ftable()=20to=20check=20if=20a=20table=20exists=20and?= =?UTF-8?q?=20if=20all=20expected=20columns=20are=20present=20=F0=9F=94=A7?= =?UTF-8?q?=20fix(utils.py):=20import=20missing=20dependencies=20to=20reso?= =?UTF-8?q?lve=20NameError=20and=20improve=20code=20readability=20?= =?UTF-8?q?=E2=9C=A8=20feat(utils.py):=20add=20dataclasses=20Result=20and?= =?UTF-8?q?=20TableResults=20to=20store=20check=20results=20for=20tables?= =?UTF-8?q?=20and=20columns?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../langflow/services/database/manager.py | 35 +++++++++++++++++++ .../langflow/services/database/utils.py | 14 ++++++++ 2 files changed, 49 insertions(+) diff --git a/src/backend/langflow/services/database/manager.py b/src/backend/langflow/services/database/manager.py index a95010bc2..2b599a0ba 100644 --- a/src/backend/langflow/services/database/manager.py +++ b/src/backend/langflow/services/database/manager.py @@ -1,6 +1,7 @@ from pathlib import Path from typing import TYPE_CHECKING from langflow.services.base import Service +from langflow.services.database.utils import Result, TableResults from langflow.services.utils import get_settings_manager from sqlalchemy import inspect import sqlalchemy as sa @@ -100,6 +101,40 @@ class DatabaseManager(Service): alembic_cfg.set_main_option("sqlalchemy.url", self.database_url) command.upgrade(alembic_cfg, "head") + def run_migrations_test(self): + # This method is used for testing purposes only + # We will check that all models are in the database + # and that the database is up to date with all columns + sql_models = [models.Flow, models.User, models.ApiKey] + results = [] + for sql_model in sql_models: + results.append( + TableResults(sql_model.__tablename__, self.check_table(sql_model)) + ) + return results + + def check_table(self, model): + results = [] + inspector = inspect(self.engine) + table_name = model.__tablename__ + expected_columns = list(model.__fields__.keys()) + try: + available_columns = [ + col["name"] for col in inspector.get_columns(table_name) + ] + results.append(Result(name=table_name, type="table", success=True)) + except sa.exc.NoSuchTableError: + logger.error(f"Missing table: {table_name}") + results.append(Result(name=table_name, type="table", success=False)) + + for column in expected_columns: + if column not in available_columns: + logger.error(f"Missing column: {column} in table {table_name}") + results.append(Result(name=column, type="column", success=False)) + else: + results.append(Result(name=column, type="column", success=True)) + return results + def create_db_and_tables(self): logger.debug("Creating database and tables") try: diff --git a/src/backend/langflow/services/database/utils.py b/src/backend/langflow/services/database/utils.py index 47204e14d..e6afae184 100644 --- a/src/backend/langflow/services/database/utils.py +++ b/src/backend/langflow/services/database/utils.py @@ -1,3 +1,4 @@ +from dataclasses import dataclass from typing import TYPE_CHECKING from langflow.utils.logger import logger from contextlib import contextmanager @@ -53,3 +54,16 @@ def session_getter(db_manager: "DatabaseManager"): raise finally: session.close() + + +@dataclass +class Result: + name: str + type: str + success: bool + + +@dataclass +class TableResults: + table_name: str + results: list[Result]