jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Fix formatting issues in database service

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-01 14:43:56 -03:00
commit da82d33493
1 changed files with 30 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

40
src/backend/langflow/services/database/service.py
View file

@ -36,7 +36,10 @@ class DatabaseService(Service):
def _create_engine(self) -> "Engine": def _create_engine(self) -> "Engine":
"""Create the engine for the database.""" """Create the engine for the database."""
settings_service = get_settings_service() settings_service = get_settings_service()
if settings_service.settings.DATABASE_URL and settings_service.settings.DATABASE_URL.startswith("sqlite"): if (
settings_service.settings.DATABASE_URL
and settings_service.settings.DATABASE_URL.startswith("sqlite")
):
connect_args = {"check_same_thread": False} connect_args = {"check_same_thread": False}
else: else:
connect_args = {} connect_args = {}
@ -48,7 +51,9 @@ class DatabaseService(Service):
def __exit__(self, exc_type, exc_value, traceback): def __exit__(self, exc_type, exc_value, traceback):
if exc_type is not None: # If an exception has been raised if exc_type is not None: # If an exception has been raised
logger.error(f"Session rollback because of exception: {exc_type.__name__} {exc_value}") logger.error(
f"Session rollback because of exception: {exc_type.__name__} {exc_value}"
)
self._session.rollback() self._session.rollback()
else: else:
self._session.commit() self._session.commit()
@ -65,7 +70,9 @@ class DatabaseService(Service):
settings_service = get_settings_service() settings_service = get_settings_service()
if settings_service.auth_settings.AUTO_LOGIN: if settings_service.auth_settings.AUTO_LOGIN:
with Session(self.engine) as session: with Session(self.engine) as session:
flows = session.exec(select(models.Flow).where(models.Flow.user_id is None)).all() flows = session.exec(
select(models.Flow).where(models.Flow.user_id is None)
).all()
if flows: if flows:
logger.debug("Migrating flows to default superuser") logger.debug("Migrating flows to default superuser")
username = settings_service.auth_settings.SUPERUSER username = settings_service.auth_settings.SUPERUSER
@ -95,7 +102,9 @@ class DatabaseService(Service):
expected_columns = list(model.model_fields.keys()) expected_columns = list(model.model_fields.keys())
try: try:
available_columns = [col["name"] for col in inspector.get_columns(table)] available_columns = [
col["name"] for col in inspector.get_columns(table)
]
except sa.exc.NoSuchTableError: except sa.exc.NoSuchTableError:
logger.error(f"Missing table: {table}") logger.error(f"Missing table: {table}")
return False return False
@ -152,14 +161,16 @@ class DatabaseService(Service):
try: try:
command.check(alembic_cfg) command.check(alembic_cfg)
except Exception as exc: except Exception as exc:
if isinstance(exc, (util.exc.CommandError, util.exc.AutogenerateDiffsDetected)): if isinstance(
exc, (util.exc.CommandError, util.exc.AutogenerateDiffsDetected)
):
command.upgrade(alembic_cfg, "head") command.upgrade(alembic_cfg, "head")
time.sleep(3) time.sleep(3)
try: try:
command.check(alembic_cfg) command.check(alembic_cfg)
except util.exc.AutogenerateDiffsDetected as e: except util.exc.AutogenerateDiffsDetected as e:
logger.error("AutogenerateDiffsDetected: {exc}") logger.error(f"AutogenerateDiffsDetected: {exc}")
if not fix: if not fix:
raise RuntimeError( raise RuntimeError(
"Something went wrong running migrations. Please, run `langflow migration --fix`" "Something went wrong running migrations. Please, run `langflow migration --fix`"
@ -188,7 +199,10 @@ class DatabaseService(Service):
# We will check that all models are in the database # We will check that all models are in the database
# and that the database is up to date with all columns # and that the database is up to date with all columns
sql_models = [models.Flow, models.User, models.ApiKey] sql_models = [models.Flow, models.User, models.ApiKey]
return [TableResults(sql_model.__tablename__, self.check_table(sql_model)) for sql_model in sql_models] return [
TableResults(sql_model.__tablename__, self.check_table(sql_model))
for sql_model in sql_models
]
def check_table(self, model): def check_table(self, model):
results = [] results = []
@ -197,7 +211,9 @@ class DatabaseService(Service):
expected_columns = list(model.__fields__.keys()) expected_columns = list(model.__fields__.keys())
available_columns = [] available_columns = []
try: try:
available_columns = [col["name"] for col in inspector.get_columns(table_name)] available_columns = [
col["name"] for col in inspector.get_columns(table_name)
]
results.append(Result(name=table_name, type="table", success=True)) results.append(Result(name=table_name, type="table", success=True))
except sa.exc.NoSuchTableError: except sa.exc.NoSuchTableError:
logger.error(f"Missing table: {table_name}") logger.error(f"Missing table: {table_name}")
@ -228,7 +244,9 @@ class DatabaseService(Service):
try: try:
table.create(self.engine, checkfirst=True) table.create(self.engine, checkfirst=True)
except OperationalError as oe: except OperationalError as oe:
logger.warning(f"Table {table} already exists, skipping. Exception: {oe}") logger.warning(
f"Table {table} already exists, skipping. Exception: {oe}"
)
except Exception as exc: except Exception as exc:
logger.error(f"Error creating table {table}: {exc}") logger.error(f"Error creating table {table}: {exc}")
raise RuntimeError(f"Error creating table {table}") from exc raise RuntimeError(f"Error creating table {table}") from exc
@ -240,7 +258,9 @@ class DatabaseService(Service):
if table not in table_names: if table not in table_names:
logger.error("Something went wrong creating the database and tables.") logger.error("Something went wrong creating the database and tables.")
logger.error("Please check your database settings.") logger.error("Please check your database settings.")
raise RuntimeError("Something went wrong creating the database and tables.") raise RuntimeError(
"Something went wrong creating the database and tables."
)
logger.debug("Database and tables created successfully") logger.debug("Database and tables created successfully")