Update database model imports

This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-11-24 10:43:06 -03:00
commit 9d6cafd8a2
24 changed files with 53 additions and 50 deletions

View file

@ -313,7 +313,7 @@ def superuser(
if create_super_user(db=session, username=username, password=password): if create_super_user(db=session, username=username, password=password):
# Verify that the superuser was created # Verify that the superuser was created
from langflow.services.database.models.user.user import User from langflow.services.database.models.user.model import User
user: User = session.query(User).filter(User.username == username).first() user: User = session.query(User).filter(User.username == username).first()
if user is None or not user.is_superuser: if user is None or not user.is_superuser:

View file

@ -2,7 +2,7 @@ from uuid import UUID
from fastapi import APIRouter, HTTPException, Depends from fastapi import APIRouter, HTTPException, Depends
from langflow.api.v1.schemas import ApiKeysResponse, ApiKeyCreateRequest from langflow.api.v1.schemas import ApiKeysResponse, ApiKeyCreateRequest
from langflow.services.auth import utils as auth_utils from langflow.services.auth import utils as auth_utils
from langflow.services.database.models.api_key.api_key import ( from langflow.services.database.models.api_key.model import (
ApiKeyCreate, ApiKeyCreate,
UnmaskedApiKeyRead, UnmaskedApiKeyRead,
) )
@ -13,7 +13,7 @@ from langflow.services.database.models.api_key.crud import (
create_api_key, create_api_key,
delete_api_key, delete_api_key,
) )
from langflow.services.database.models.user.user import User from langflow.services.database.models.user.model import User
from langflow.services.deps import ( from langflow.services.deps import (
get_session, get_session,
get_settings_service, get_settings_service,

View file

@ -18,7 +18,7 @@ from langflow.processing.process import process_graph_cached, process_tweaks
from langflow.services.auth.utils import api_key_security, get_current_active_user from langflow.services.auth.utils import api_key_security, get_current_active_user
from langflow.services.cache.utils import save_uploaded_file from langflow.services.cache.utils import save_uploaded_file
from langflow.services.database.models.flow import Flow from langflow.services.database.models.flow import Flow
from langflow.services.database.models.user.user import User from langflow.services.database.models.user.model import User
from langflow.services.deps import ( from langflow.services.deps import (
get_session, get_session,
get_session_service, get_session_service,

View file

@ -11,7 +11,7 @@ from langflow.api.utils import remove_api_keys
from langflow.api.v1.schemas import FlowListCreate, FlowListRead from langflow.api.v1.schemas import FlowListCreate, FlowListRead
from langflow.services.auth.utils import get_current_active_user from langflow.services.auth.utils import get_current_active_user
from langflow.services.database.models.flow import Flow, FlowCreate, FlowRead, FlowUpdate from langflow.services.database.models.flow import Flow, FlowCreate, FlowRead, FlowUpdate
from langflow.services.database.models.user.user import User from langflow.services.database.models.user.model import User
from langflow.services.deps import get_session, get_settings_service from langflow.services.deps import get_session, get_settings_service
# build router # build router

View file

@ -2,7 +2,7 @@ from enum import Enum
from pathlib import Path from pathlib import Path
from typing import Any, Dict, List, Optional, Union from typing import Any, Dict, List, Optional, Union
from uuid import UUID from uuid import UUID
from langflow.services.database.models.api_key.api_key import ApiKeyRead from langflow.services.database.models.api_key.model import ApiKeyRead
from langflow.services.database.models.flow import FlowCreate, FlowRead from langflow.services.database.models.flow import FlowCreate, FlowRead
from langflow.services.database.models.user import UserRead from langflow.services.database.models.user import UserRead
from langflow.services.database.models.base import orjson_dumps from langflow.services.database.models.base import orjson_dumps

View file

@ -4,7 +4,7 @@ from uuid import UUID
from fastapi import APIRouter, Depends, HTTPException, Query from fastapi import APIRouter, Depends, HTTPException, Query
from langflow.services.auth import utils as auth_utils from langflow.services.auth import utils as auth_utils
from langflow.services.database.models.user.user import User from langflow.services.database.models.user.model import User
from langflow.services.deps import get_settings_service, get_store_service from langflow.services.deps import get_settings_service, get_store_service
from langflow.services.store.exceptions import CustomException from langflow.services.store.exceptions import CustomException
from langflow.services.store.schema import ( from langflow.services.store.schema import (

View file

@ -8,10 +8,10 @@ from fastapi.security import APIKeyHeader, APIKeyQuery, OAuth2PasswordBearer
from jose import JWTError, jwt from jose import JWTError, jwt
from sqlmodel import Session from sqlmodel import Session
from langflow.services.database.models.api_key.api_key import ApiKey from langflow.services.database.models.api_key.model import ApiKey
from langflow.services.database.models.api_key.crud import check_key from langflow.services.database.models.api_key.crud import check_key
from langflow.services.database.models.user.crud import get_user_by_id, get_user_by_username, update_user_last_login_at from langflow.services.database.models.user.crud import get_user_by_id, get_user_by_username, update_user_last_login_at
from langflow.services.database.models.user.user import User from langflow.services.database.models.user.model import User
from langflow.services.deps import get_session, get_settings_service from langflow.services.deps import get_session, get_settings_service
oauth2_login = OAuth2PasswordBearer(tokenUrl="api/v1/login", auto_error=False) oauth2_login = OAuth2PasswordBearer(tokenUrl="api/v1/login", auto_error=False)

View file

@ -1,3 +1,3 @@
from .api_key import ApiKey, ApiKeyCreate, UnmaskedApiKeyRead, ApiKeyRead from .model import ApiKey, ApiKeyCreate, UnmaskedApiKeyRead, ApiKeyRead
__all__ = ["ApiKey", "ApiKeyCreate", "UnmaskedApiKeyRead", "ApiKeyRead"] __all__ = ["ApiKey", "ApiKeyCreate", "UnmaskedApiKeyRead", "ApiKeyRead"]

View file

@ -1,15 +1,15 @@
from pydantic import validator
from sqlmodel import Field, Relationship
from uuid import UUID, uuid4
from typing import Optional, TYPE_CHECKING
from datetime import datetime from datetime import datetime
from langflow.services.database.models.base import SQLModelSerializable from typing import TYPE_CHECKING, Optional
from uuid import UUID, uuid4
from pydantic import validator
from sqlmodel import Field, Relationship, SQLModel
if TYPE_CHECKING: if TYPE_CHECKING:
from langflow.services.database.models.user import User from langflow.services.database.models.user import User
class ApiKeyBase(SQLModelSerializable): class ApiKeyBase(SQLModel):
name: Optional[str] = Field(index=True) name: Optional[str] = Field(index=True)
created_at: datetime = Field(default_factory=datetime.utcnow) created_at: datetime = Field(default_factory=datetime.utcnow)
last_used_at: Optional[datetime] = Field(default=None, nullable=True) last_used_at: Optional[datetime] = Field(default=None, nullable=True)

View file

@ -16,9 +16,3 @@ def orjson_dumps(v, *, default=None, sort_keys=False, indent_2=True):
if default is None: if default is None:
return orjson.dumps(v, option=option).decode() return orjson.dumps(v, option=option).decode()
return orjson.dumps(v, default=default, option=option).decode() return orjson.dumps(v, default=default, option=option).decode()
class SQLModelSerializable(SQLModel):
# TODO[pydantic]: The following keys were removed: `json_loads`, `json_dumps`.
# Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information.
pass

View file

@ -1,3 +1,3 @@
from .component import Component, ComponentModel from .model import Component, ComponentModel
__all__ = ["Component", "ComponentModel"] __all__ = ["Component", "ComponentModel"]

View file

@ -1,11 +1,11 @@
from langflow.services.database.models.base import SQLModelSerializable, SQLModel
from sqlmodel import Field
from typing import Optional
from datetime import datetime
import uuid import uuid
from datetime import datetime
from typing import Optional
from sqlmodel import Field, SQLModel
class Component(SQLModelSerializable, table=True): class Component(SQLModel, table=True):
id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True) id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
frontend_node_id: uuid.UUID = Field(index=True) frontend_node_id: uuid.UUID = Field(index=True)
name: str = Field(index=True) name: str = Field(index=True)

View file

@ -1,3 +1,3 @@
from .flow import Flow, FlowCreate, FlowRead, FlowUpdate from .model import Flow, FlowCreate, FlowRead, FlowUpdate
__all__ = ["Flow", "FlowCreate", "FlowRead", "FlowUpdate"] __all__ = ["Flow", "FlowCreate", "FlowRead", "FlowUpdate"]

View file

@ -5,15 +5,13 @@ from typing import TYPE_CHECKING, Dict, Optional
from uuid import UUID, uuid4 from uuid import UUID, uuid4
from pydantic import field_serializer, field_validator from pydantic import field_serializer, field_validator
from sqlmodel import JSON, Column, Field, Relationship from sqlmodel import JSON, Column, Field, Relationship, SQLModel
from langflow.services.database.models.base import SQLModelSerializable
if TYPE_CHECKING: if TYPE_CHECKING:
from langflow.services.database.models.user import User from langflow.services.database.models.user import User
class FlowBase(SQLModelSerializable): class FlowBase(SQLModel):
name: str = Field(index=True) name: str = Field(index=True)
description: Optional[str] = Field(index=True, nullable=True, default=None) description: Optional[str] = Field(index=True, nullable=True, default=None)
data: Optional[Dict] = Field(default=None, nullable=True) data: Optional[Dict] = Field(default=None, nullable=True)
@ -69,7 +67,7 @@ class FlowRead(FlowBase):
user_id: UUID = Field() user_id: UUID = Field()
class FlowUpdate(SQLModelSerializable): class FlowUpdate(SQLModel):
name: Optional[str] = None name: Optional[str] = None
description: Optional[str] = None description: Optional[str] = None
data: Optional[Dict] = None data: Optional[Dict] = None

View file

@ -1,4 +1,4 @@
from .user import User, UserCreate, UserRead, UserUpdate from .model import User, UserCreate, UserRead, UserUpdate
__all__ = [ __all__ = [
"User", "User",

View file

@ -2,7 +2,7 @@ from datetime import datetime, timezone
from typing import Union from typing import Union
from uuid import UUID from uuid import UUID
from fastapi import Depends, HTTPException, status from fastapi import Depends, HTTPException, status
from langflow.services.database.models.user.user import User, UserUpdate from langflow.services.database.models.user.model import User, UserUpdate
from langflow.services.deps import get_session from langflow.services.deps import get_session
from sqlalchemy.exc import IntegrityError from sqlalchemy.exc import IntegrityError
from sqlmodel import Session from sqlmodel import Session

View file

@ -1,17 +1,18 @@
from langflow.services.database.models.base import SQLModel, SQLModelSerializable from datetime import datetime
from typing import TYPE_CHECKING, Optional
from uuid import UUID, uuid4
from sqlmodel import Field, Relationship from sqlmodel import Field, Relationship
from langflow.services.database.models.base import SQLModel
from datetime import datetime
from typing import Optional, TYPE_CHECKING
from uuid import UUID, uuid4
if TYPE_CHECKING: if TYPE_CHECKING:
from langflow.services.database.models.api_key import ApiKey from langflow.services.database.models.api_key import ApiKey
from langflow.services.database.models.credential import Credential
from langflow.services.database.models.flow import Flow from langflow.services.database.models.flow import Flow
class User(SQLModelSerializable, table=True): class User(SQLModel, table=True):
id: UUID = Field(default_factory=uuid4, primary_key=True, unique=True) id: UUID = Field(default_factory=uuid4, primary_key=True, unique=True)
username: str = Field(index=True, unique=True) username: str = Field(index=True, unique=True)
password: str = Field() password: str = Field()
@ -27,6 +28,10 @@ class User(SQLModelSerializable, table=True):
) )
store_api_key: str = Field(default=None, nullable=True) store_api_key: str = Field(default=None, nullable=True)
flows: list["Flow"] = Relationship(back_populates="user") flows: list["Flow"] = Relationship(back_populates="user")
credentials: list["Credential"] = Relationship(
back_populates="user",
sa_relationship_kwargs={"cascade": "delete"},
)
class UserCreate(SQLModel): class UserCreate(SQLModel):

View file

@ -12,6 +12,11 @@ if TYPE_CHECKING:
from langflow.services.settings.service import SettingsService from langflow.services.settings.service import SettingsService
from langflow.services.store.service import StoreService from langflow.services.store.service import StoreService
from langflow.services.task.service import TaskService from langflow.services.task.service import TaskService
from langflow.services.credentials.service import CredentialService
def get_credential_service() -> "CredentialService":
return service_manager.get(ServiceType.CREDENTIAL_SERVICE)
def get_settings_service() -> "SettingsService": def get_settings_service() -> "SettingsService":

View file

@ -15,3 +15,4 @@ class ServiceType(str, Enum):
SESSION_SERVICE = "session_service" SESSION_SERVICE = "session_service"
TASK_SERVICE = "task_service" TASK_SERVICE = "task_service"
STORE_SERVICE = "store_service" STORE_SERVICE = "store_service"
CREDENTIAL_SERVICE = "credential_service"

View file

@ -44,7 +44,7 @@ def get_factories_and_deps():
def get_or_create_super_user(session: Session, username, password, is_default): def get_or_create_super_user(session: Session, username, password, is_default):
from langflow.services.database.models.user.user import User from langflow.services.database.models.user.model import User
user = session.query(User).filter(User.username == username).first() user = session.query(User).filter(User.username == username).first()
@ -127,7 +127,7 @@ def teardown_superuser(settings_service, session):
try: try:
logger.debug("AUTO_LOGIN is set to False. Removing default superuser if exists.") logger.debug("AUTO_LOGIN is set to False. Removing default superuser if exists.")
username = DEFAULT_SUPERUSER username = DEFAULT_SUPERUSER
from langflow.services.database.models.user.user import User from langflow.services.database.models.user.model import User
user = session.query(User).filter(User.username == username).first() user = session.query(User).filter(User.username == username).first()
if user and user.is_superuser: if user and user.is_superuser:

View file

@ -12,8 +12,8 @@ from fastapi.testclient import TestClient
from httpx import AsyncClient from httpx import AsyncClient
from langflow.graph.graph.base import Graph from langflow.graph.graph.base import Graph
from langflow.services.auth.utils import get_password_hash from langflow.services.auth.utils import get_password_hash
from langflow.services.database.models.flow.flow import Flow, FlowCreate from langflow.services.database.models.flow.model import Flow, FlowCreate
from langflow.services.database.models.user.user import User, UserCreate from langflow.services.database.models.user.model import User, UserCreate
from langflow.services.database.utils import session_getter from langflow.services.database.utils import session_getter
from langflow.services.deps import get_db_service from langflow.services.deps import get_db_service
from sqlmodel import Session, SQLModel, create_engine from sqlmodel import Session, SQLModel, create_engine
@ -260,7 +260,7 @@ def logged_in_headers(client, active_user):
@pytest.fixture @pytest.fixture
def flow(client, json_flow: str, active_user): def flow(client, json_flow: str, active_user):
from langflow.services.database.models.flow.flow import FlowCreate from langflow.services.database.models.flow.model import FlowCreate
loaded_json = json.loads(json_flow) loaded_json = json.loads(json_flow)
flow_data = FlowCreate( flow_data = FlowCreate(

View file

@ -2,7 +2,7 @@ from collections import namedtuple
import uuid import uuid
from langflow.processing.process import Result from langflow.processing.process import Result
from langflow.services.auth.utils import get_password_hash from langflow.services.auth.utils import get_password_hash
from langflow.services.database.models.api_key.api_key import ApiKey from langflow.services.database.models.api_key.model import ApiKey
from langflow.services.deps import get_settings_service from langflow.services.deps import get_settings_service
from langflow.services.database.utils import session_getter from langflow.services.database.utils import session_getter
from langflow.services.deps import get_db_service from langflow.services.deps import get_db_service

View file

@ -1,6 +1,6 @@
from unittest.mock import MagicMock, patch from unittest.mock import MagicMock, patch
from langflow.services.database.models.user.user import User from langflow.services.database.models.user.model import User
from langflow.services.settings.constants import DEFAULT_SUPERUSER, DEFAULT_SUPERUSER_PASSWORD from langflow.services.settings.constants import DEFAULT_SUPERUSER, DEFAULT_SUPERUSER_PASSWORD
from langflow.services.utils import teardown_superuser from langflow.services.utils import teardown_superuser

View file

@ -3,7 +3,7 @@ from datetime import datetime
import pytest import pytest
from langflow.services.auth.utils import create_super_user, get_password_hash from langflow.services.auth.utils import create_super_user, get_password_hash
from langflow.services.database.models.user import UserUpdate from langflow.services.database.models.user import UserUpdate
from langflow.services.database.models.user.user import User from langflow.services.database.models.user.model import User
from langflow.services.database.utils import session_getter from langflow.services.database.utils import session_getter
from langflow.services.deps import get_db_service, get_settings_service from langflow.services.deps import get_db_service, get_settings_service