Fix formatting issues in database service

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-01 14:43:56 -03:00
commit da82d33493

View file

@ -36,7 +36,10 @@ class DatabaseService(Service):
def _create_engine(self) -> "Engine":
"""Create the engine for the database."""
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}
else:
connect_args = {}
@ -48,7 +51,9 @@ class DatabaseService(Service):
def __exit__(self, exc_type, exc_value, traceback):
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()
else:
self._session.commit()
@ -65,7 +70,9 @@ class DatabaseService(Service):
settings_service = get_settings_service()
if settings_service.auth_settings.AUTO_LOGIN:
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:
logger.debug("Migrating flows to default superuser")
username = settings_service.auth_settings.SUPERUSER
@ -95,7 +102,9 @@ class DatabaseService(Service):
expected_columns = list(model.model_fields.keys())
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:
logger.error(f"Missing table: {table}")
return False
@ -152,14 +161,16 @@ class DatabaseService(Service):
try:
command.check(alembic_cfg)
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")
time.sleep(3)
try:
command.check(alembic_cfg)
except util.exc.AutogenerateDiffsDetected as e:
logger.error("AutogenerateDiffsDetected: {exc}")
logger.error(f"AutogenerateDiffsDetected: {exc}")
if not fix:
raise RuntimeError(
"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
# and that the database is up to date with all columns
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):
results = []
@ -197,7 +211,9 @@ class DatabaseService(Service):
expected_columns = list(model.__fields__.keys())
available_columns = []
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))
except sa.exc.NoSuchTableError:
logger.error(f"Missing table: {table_name}")
@ -228,7 +244,9 @@ class DatabaseService(Service):
try:
table.create(self.engine, checkfirst=True)
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:
logger.error(f"Error creating table {table}: {exc}")
raise RuntimeError(f"Error creating table {table}") from exc
@ -240,7 +258,9 @@ class DatabaseService(Service):
if table not in table_names:
logger.error("Something went wrong creating the database and tables.")
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")