Adds teardown logic for database and other services (#852)

poetry.lock

@@ -7041,6 +7041,17 @@ files = [
     {file = "types_pytz-2023.3.0.1-py3-none-any.whl", hash = "sha256:65152e872137926bb67a8fe6cc9cfd794365df86650c5d5fdc7b167b0f38892e"},
 ]
 
+[[package]]
+name = "types-pywin32"
+version = "306.0.0.4"
+description = "Typing stubs for pywin32"
+optional = false
+python-versions = "*"
+files = [
+    {file = "types-pywin32-306.0.0.4.tar.gz", hash = "sha256:ae4bbec80d535053236d4bebedf55f58dee89cf5883d277f0fa89e857f3ff337"},
+    {file = "types_pywin32-306.0.0.4-py3-none-any.whl", hash = "sha256:f76a343ed6933008af85e158063963f923e54f2f461e697b2929b4178c7b77a1"},
+]
+
 [[package]]
 name = "types-pyyaml"
 version = "6.0.12.11"
@@ -7773,4 +7784,4 @@ local = ["ctransformers", "llama-cpp-python", "sentence-transformers"]
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.9,<3.11"
-content-hash = "c877b4d713eef71815d858d30976ab21c42e5eadcc2df8159e940e03323681ee"
+content-hash = "a3a506d483c2db7169a9790090095d1764aa5be223d135c6fc3fc2768dfef36c"

pyproject.toml

@@ -103,6 +103,7 @@ types-python-jose = "^3.3.4.8"
 types-passlib = "^1.7.7.13"
 pytest-mock = "^3.11.1"
 pytest-xdist = "^3.3.1"
+types-pywin32 = "^306.0.0.4"
 
 
 [tool.poetry.extras]

src/backend/langflow/__main__.py

@@ -356,7 +356,7 @@ def superuser(
     with session_getter(db_manager) as session:
         from langflow.services.auth.utils import create_super_user
 
-        if create_super_user(session, username, password):
+        if create_super_user(db=session, username=username, password=password):
             # Verify that the superuser was created
             from langflow.services.database.models.user.user import User
 

src/backend/langflow/api/v1/chat.py

@@ -61,14 +61,13 @@ async def chat(
         await websocket.close(code=status.WS_1011_INTERNAL_ERROR, reason=str(exc))
     except Exception as exc:
         logger.error(f"Error in chat websocket: {exc}")
-        if isinstance(exc, HTTPException):
+        messsage = exc.detail if isinstance(exc, HTTPException) else str(exc)
-            exc = exc.detail
         if "Could not validate credentials" in str(exc):
             await websocket.close(
                 code=status.WS_1008_POLICY_VIOLATION, reason="Unauthorized"
             )
         else:
-            await websocket.close(code=status.WS_1011_INTERNAL_ERROR, reason=str(exc))
+            await websocket.close(code=status.WS_1011_INTERNAL_ERROR, reason=messsage)
 
 
 @router.post("/build/init/{flow_id}", response_model=InitResponse, status_code=201)

src/backend/langflow/main.py

@@ -10,7 +10,7 @@ from langflow.api import router
 
 from langflow.interface.utils import setup_llm_caching
 from langflow.services.database.utils import initialize_database
-from langflow.services.manager import initialize_services
+from langflow.services.manager import initialize_services, teardown_services
 from langflow.utils.logger import configure
 
 
@@ -40,6 +40,7 @@ def create_app():
     app.on_event("startup")(initialize_services)
     app.on_event("startup")(initialize_database)
     app.on_event("startup")(setup_llm_caching)
+    app.on_event("shutdown")(teardown_services)
     return app
 
 

src/backend/langflow/services/auth/utils.py

@@ -37,7 +37,12 @@ async def api_key_security(
     result: Optional[Union[ApiKey, User]] = None
     if settings_manager.auth_settings.AUTO_LOGIN:
         # Get the first user
-        settings_manager.auth_settings.FIRST_SUPERUSER
+        if not settings_manager.auth_settings.FIRST_SUPERUSER:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail="Missing first superuser credentials",
+            )
+
         result = get_user_by_username(
             db, settings_manager.auth_settings.FIRST_SUPERUSER
         )
@@ -80,6 +85,9 @@ async def get_current_user(
     if isinstance(token, Coroutine):
         token = await token
 
+    if settings_manager.auth_settings.SECRET_KEY is None:
+        raise credentials_exception
+
     try:
         payload = jwt.decode(
             token,
@@ -150,22 +158,16 @@ def create_token(data: dict, expires_delta: timedelta):
 
 
 def create_super_user(
+    username: str,
+    password: str,
     db: Session = Depends(get_session),
-    username: Optional[str] = None,
-    password: Optional[str] = None,
 ) -> User:
-    settings_manager = get_settings_manager()
+    super_user = get_user_by_username(db, username)
-
-    super_user = get_user_by_username(
-        db, username or settings_manager.auth_settings.FIRST_SUPERUSER
-    )
 
     if not super_user:
         super_user = User(
-            username=username or settings_manager.auth_settings.FIRST_SUPERUSER,
+            username=username,
-            password=get_password_hash(
+            password=get_password_hash(password),
-                password or settings_manager.auth_settings.FIRST_SUPERUSER_PASSWORD
-            ),
             is_superuser=True,
             is_active=True,
             last_login_at=None,
@@ -179,7 +181,15 @@ def create_super_user(
 
 
 def create_user_longterm_token(db: Session = Depends(get_session)) -> dict:
-    super_user = create_super_user(db)
+    settings_manager = get_settings_manager()
+    username = settings_manager.auth_settings.FIRST_SUPERUSER
+    password = settings_manager.auth_settings.FIRST_SUPERUSER_PASSWORD
+    if not username or not password:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Missing first superuser credentials",
+        )
+    super_user = create_super_user(db=db, username=username, password=password)
 
     access_token_expires_longterm = timedelta(days=365)
     access_token = create_token(

src/backend/langflow/services/base.py

@@ -1,2 +1,8 @@
-class Service:
+from abc import ABC
+
+
+class Service(ABC):
     name: str
+
+    def teardown(self):
+        pass

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()

src/backend/langflow/services/manager.py

@@ -1,5 +1,6 @@
 from langflow.services.schema import ServiceType
 from typing import TYPE_CHECKING, List, Optional
+from langflow.utils.logger import logger
 
 if TYPE_CHECKING:
     from langflow.services.factory import ServiceFactory
@@ -42,6 +43,7 @@ class ServiceManager:
         """
         Create a new service given its name, handling dependencies.
         """
+        logger.debug(f"Create service {service_name}")
         self._validate_service_creation(service_name)
 
         # Create dependencies first
@@ -74,9 +76,21 @@ class ServiceManager:
         Update a service by its name.
         """
         if service_name in self.services:
+            logger.debug(f"Update service {service_name}")
             self.services.pop(service_name, None)
             self.get(service_name)
 
+    def teardown(self):
+        """
+        Teardown all the services.
+        """
+        for service in self.services.values():
+            logger.debug(f"Teardown service {service.name}")
+            service.teardown()
+        self.services = {}
+        self.factories = {}
+        self.dependencies = {}
+
 
 service_manager = ServiceManager()
 
@@ -121,7 +135,7 @@ def initialize_session_manager():
     """
     Initialize the session manager.
     """
-    from langflow.services.session import factory as session_manager_factory
+    from langflow.services.session import factory as session_manager_factory  # type: ignore
     from langflow.services.cache import factory as cache_factory
 
     initialize_settings_manager()
@@ -134,3 +148,10 @@ def initialize_session_manager():
         session_manager_factory.SessionManagerFactory(),
         dependencies=[ServiceType.CACHE_MANAGER],
     )
+
+
+def teardown_services():
+    """
+    Teardown all the services.
+    """
+    service_manager.teardown()

src/backend/langflow/services/settings/auth.py

@@ -11,10 +11,11 @@ from langflow.utils.logger import logger
 class AuthSettings(BaseSettings):
     # Login settings
     CONFIG_DIR: str
-    SECRET_KEY: Optional[str] = Field(
+    SECRET_KEY: str = Field(
-        None,
+        default="",
         description="Secret key for JWT. If not provided, a random one will be generated.",
         env="LANGFLOW_SECRET_KEY",
+        allow_mutation=False,
     )
     ALGORITHM: str = "HS256"
     ACCESS_TOKEN_EXPIRE_MINUTES: int = 60

src/backend/langflow/services/settings/manager.py

@@ -35,5 +35,10 @@ class SettingsManager(Service):
                 )
 
         settings = Settings(**settings_dict)
-        auth_settings = AuthSettings(CONFIG_DIR=settings.CONFIG_DIR)
+        if not settings.CONFIG_DIR:
+            raise ValueError("CONFIG_DIR must be set in settings")
+
+        auth_settings = AuthSettings(
+            CONFIG_DIR=settings.CONFIG_DIR,
+        )
         return cls(settings, auth_settings)

src/backend/langflow/services/settings/utils.py

@@ -43,5 +43,5 @@ def write_secret_to_file(path: Path, value: str) -> None:
 
 
 def read_secret_from_file(path: Path) -> str:
-    with path.open("rb") as f:
+    with path.open("r") as f:
         return f.read()

tests/test_user.py

@@ -9,7 +9,13 @@ from langflow.services.database.models.user import UserUpdate
 
 @pytest.fixture
 def super_user(client, session):
-    return create_super_user(session)
+    settings_manager = get_settings_manager()
+    auth_settings = settings_manager.auth_settings
+    return create_super_user(
+        db=session,
+        username=auth_settings.FIRST_SUPERUSER,
+        password=auth_settings.FIRST_SUPERUSER_PASSWORD,
+    )
 
 
 @pytest.fixture