diff --git a/src/backend/langflow/auth/auth.py b/src/backend/langflow/auth/auth.py index 3a7e324c2..067184053 100644 --- a/src/backend/langflow/auth/auth.py +++ b/src/backend/langflow/auth/auth.py @@ -6,18 +6,19 @@ from fastapi.security import OAuth2PasswordBearer from fastapi import Depends, HTTPException, status from datetime import datetime, timedelta, timezone +from langflow.services.utils import get_session from langflow.database.models.token import TokenData from langflow.database.models.user import get_user, User -from langflow.services.utils import get_session # TODO: Move to env - Test propose!!!!! SECRET_KEY = "698619adad2d916f1f32d264540976964b3c0d3828e0870a65add5800a8cc6b9" ALGORITHM = "HS256" -ACCESS_TOKEN_EXPIRE_MINUTES = 30 +ACCESS_TOKEN_EXPIRE_MINUTES = 60 +REFRESH_TOKEN_EXPIRE_MINUTES = 180 pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto") -oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token") +oauth2_scheme = OAuth2PasswordBearer(tokenUrl="login") def verify_password(plain_password, hashed_password): @@ -33,7 +34,18 @@ def create_access_token(data: dict, expires_delta: timedelta = None): # type: i if expires_delta: expire = datetime.now(timezone.utc) + expires_delta else: - expire = datetime.now(timezone.utc) + timedelta(minutes=15) + expire = datetime.now(timezone.utc) + timedelta( + minutes=ACCESS_TOKEN_EXPIRE_MINUTES + ) + to_encode["exp"] = expire + return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM) + + +def create_refresh_token(data: dict): + to_encode = data.copy() + expire = datetime.now(timezone.utc) + timedelta( + minutes=REFRESH_TOKEN_EXPIRE_MINUTES + ) to_encode["exp"] = expire return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM) @@ -47,7 +59,7 @@ def authenticate_user(db: Session, username: str, password: str): async def get_current_user( token: Annotated[str, Depends(oauth2_scheme)], db: Session = Depends(get_session) -): +) -> User: credentials_exception = HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Could not validate credentials", diff --git a/src/backend/langflow/database/models/token.py b/src/backend/langflow/database/models/token.py index 080286787..e78743877 100644 --- a/src/backend/langflow/database/models/token.py +++ b/src/backend/langflow/database/models/token.py @@ -3,6 +3,7 @@ from pydantic import BaseModel class Token(BaseModel): access_token: str + refresh_token: str token_type: str diff --git a/src/backend/langflow/routers/login.py b/src/backend/langflow/routers/login.py index 3cb71a82e..8108a2d18 100644 --- a/src/backend/langflow/routers/login.py +++ b/src/backend/langflow/routers/login.py @@ -7,6 +7,7 @@ from langflow.auth.auth import ( ACCESS_TOKEN_EXPIRE_MINUTES, authenticate_user, create_access_token, + create_refresh_token, ) from sqlalchemy.orm import Session @@ -24,10 +25,16 @@ def create_user_token(user: User) -> dict: expires_delta=access_token_expires, ) - return {"access_token": access_token, "token_type": "bearer"} + refresh_token = create_refresh_token(data={"sub": user.username}) + + return { + "access_token": access_token, + "refresh_token": refresh_token, + "token_type": "bearer", + } -@router.post("/token", response_model=Token) +@router.post("/login", response_model=Token) async def login_to_get_access_token( form_data: OAuth2PasswordRequestForm = Depends(), db: Session = Depends(get_session) ): diff --git a/src/backend/langflow/services/database/manager.py b/src/backend/langflow/services/database/manager.py index 92385a457..bf90703e2 100644 --- a/src/backend/langflow/services/database/manager.py +++ b/src/backend/langflow/services/database/manager.py @@ -1,11 +1,16 @@ from pathlib import Path +from typing import TYPE_CHECKING from langflow.services.base import Service +from langflow.services.utils import get_settings_manager from sqlmodel import SQLModel, Session, create_engine from langflow.utils.logger import logger from alembic.config import Config from alembic import command from langflow.services.database import models # noqa +if TYPE_CHECKING: + from sqlalchemy.engine import Engine + class DatabaseManager(Service): name = "database_manager" @@ -17,7 +22,19 @@ class DatabaseManager(Service): langflow_dir = Path(__file__).parent.parent.parent self.script_location = langflow_dir / "alembic" self.alembic_cfg_path = langflow_dir / "alembic.ini" - self.engine = create_engine(database_url) + self.engine = self._create_engine() + + def _create_engine(self) -> "Engine": + """Create the engine for the database.""" + settings_manager = get_settings_manager() + if ( + settings_manager.settings.DATABASE_URL + and settings_manager.settings.DATABASE_URL.startswith("sqlite") + ): + connect_args = {"check_same_thread": False} + else: + connect_args = {} + return create_engine(self.database_url, connect_args=connect_args) def __enter__(self): self._session = Session(self.engine)