From 8d96c32c2b4a42dae66afaa562e6816a50cffefd Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Thu, 31 Aug 2023 10:21:03 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20chore(manager.py):=20add=20teard?= =?UTF-8?q?own=20method=20to=20DatabaseManager=20class=20for=20cleaning=20?= =?UTF-8?q?up=20the=20database=20=F0=9F=94=A7=20chore(manager.py):=20remov?= =?UTF-8?q?e=20default=20superuser=20if=20auto=5Flogin=20is=20enabled=20du?= =?UTF-8?q?ring=20teardown=20to=20ensure=20clean=20database=20state?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../langflow/services/database/manager.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/src/backend/langflow/services/database/manager.py b/src/backend/langflow/services/database/manager.py index 2b599a0ba..1159cbf7a 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.models.user.crud import get_user_by_username from langflow.services.database.utils import Result, TableResults from langflow.services.utils import get_settings_manager from sqlalchemy import inspect @@ -159,3 +160,23 @@ class DatabaseManager(Service): ) logger.debug("Database and tables created successfully") + + def teardown(self): + logger.debug("Tearing down database") + try: + settings_manager = get_settings_manager() + # remove the default superuser if auto_login is enabled + # using the FIRST_SUPERUSER to get the user + if settings_manager.auth_settings.AUTO_LOGIN: + logger.debug("Removing default superuser") + username = settings_manager.auth_settings.FIRST_SUPERUSER + with Session(self.engine) as session: + user = get_user_by_username(session, username) + session.delete(user) + session.commit() + logger.debug("Default superuser removed") + + except Exception as exc: + logger.error(f"Error tearing down database: {exc}") + + self.engine.dispose()