Merge branch 'authentication' into login

This commit is contained in:
Cristhian Zanforlin Lousa 2023-08-10 15:27:33 -03:00
commit b6794c76fd
4 changed files with 45 additions and 8 deletions

View file

@ -6,18 +6,19 @@ from fastapi.security import OAuth2PasswordBearer
from fastapi import Depends, HTTPException, status from fastapi import Depends, HTTPException, status
from datetime import datetime, timedelta, timezone from datetime import datetime, timedelta, timezone
from langflow.services.utils import get_session
from langflow.database.models.token import TokenData from langflow.database.models.token import TokenData
from langflow.database.models.user import get_user, User from langflow.database.models.user import get_user, User
from langflow.services.utils import get_session
# TODO: Move to env - Test propose!!!!! # TODO: Move to env - Test propose!!!!!
SECRET_KEY = "698619adad2d916f1f32d264540976964b3c0d3828e0870a65add5800a8cc6b9" SECRET_KEY = "698619adad2d916f1f32d264540976964b3c0d3828e0870a65add5800a8cc6b9"
ALGORITHM = "HS256" ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 30 ACCESS_TOKEN_EXPIRE_MINUTES = 60
REFRESH_TOKEN_EXPIRE_MINUTES = 180
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto") pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token") oauth2_scheme = OAuth2PasswordBearer(tokenUrl="login")
def verify_password(plain_password, hashed_password): 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: if expires_delta:
expire = datetime.now(timezone.utc) + expires_delta expire = datetime.now(timezone.utc) + expires_delta
else: 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 to_encode["exp"] = expire
return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM) 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( async def get_current_user(
token: Annotated[str, Depends(oauth2_scheme)], db: Session = Depends(get_session) token: Annotated[str, Depends(oauth2_scheme)], db: Session = Depends(get_session)
): ) -> User:
credentials_exception = HTTPException( credentials_exception = HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED, status_code=status.HTTP_401_UNAUTHORIZED,
detail="Could not validate credentials", detail="Could not validate credentials",

View file

@ -3,6 +3,7 @@ from pydantic import BaseModel
class Token(BaseModel): class Token(BaseModel):
access_token: str access_token: str
refresh_token: str
token_type: str token_type: str

View file

@ -7,6 +7,7 @@ from langflow.auth.auth import (
ACCESS_TOKEN_EXPIRE_MINUTES, ACCESS_TOKEN_EXPIRE_MINUTES,
authenticate_user, authenticate_user,
create_access_token, create_access_token,
create_refresh_token,
) )
from sqlalchemy.orm import Session from sqlalchemy.orm import Session
@ -24,10 +25,16 @@ def create_user_token(user: User) -> dict:
expires_delta=access_token_expires, 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( async def login_to_get_access_token(
form_data: OAuth2PasswordRequestForm = Depends(), db: Session = Depends(get_session) form_data: OAuth2PasswordRequestForm = Depends(), db: Session = Depends(get_session)
): ):

View file

@ -1,11 +1,16 @@
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING
from langflow.services.base import Service from langflow.services.base import Service
from langflow.services.utils import get_settings_manager
from sqlmodel import SQLModel, Session, create_engine from sqlmodel import SQLModel, Session, create_engine
from langflow.utils.logger import logger from langflow.utils.logger import logger
from alembic.config import Config from alembic.config import Config
from alembic import command from alembic import command
from langflow.services.database import models # noqa from langflow.services.database import models # noqa
if TYPE_CHECKING:
from sqlalchemy.engine import Engine
class DatabaseManager(Service): class DatabaseManager(Service):
name = "database_manager" name = "database_manager"
@ -17,7 +22,19 @@ class DatabaseManager(Service):
langflow_dir = Path(__file__).parent.parent.parent langflow_dir = Path(__file__).parent.parent.parent
self.script_location = langflow_dir / "alembic" self.script_location = langflow_dir / "alembic"
self.alembic_cfg_path = langflow_dir / "alembic.ini" 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): def __enter__(self):
self._session = Session(self.engine) self._session = Session(self.engine)