🔧 chore(manager.py): add teardown method to DatabaseManager class for cleaning up the database
🔧 chore(manager.py): remove default superuser if auto_login is enabled during teardown to ensure clean database state
This commit is contained in:
parent
29a616bc77
commit
8d96c32c2b
1 changed files with 21 additions and 0 deletions
|
|
@ -1,6 +1,7 @@
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
from langflow.services.base import Service
|
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.database.utils import Result, TableResults
|
||||||
from langflow.services.utils import get_settings_manager
|
from langflow.services.utils import get_settings_manager
|
||||||
from sqlalchemy import inspect
|
from sqlalchemy import inspect
|
||||||
|
|
@ -159,3 +160,23 @@ class DatabaseManager(Service):
|
||||||
)
|
)
|
||||||
|
|
||||||
logger.debug("Database and tables created successfully")
|
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()
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue