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):
# 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()
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 langflow.api.v1.schemas import ApiKeysResponse, ApiKeyCreateRequest
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,
UnmaskedApiKeyRead,
)
@ -13,7 +13,7 @@ from langflow.services.database.models.api_key.crud import (
create_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 (
get_session,
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.cache.utils import save_uploaded_file
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 (
get_session,
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.services.auth.utils import get_current_active_user
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
# build router

View file

@ -2,7 +2,7 @@ from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional, Union
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.user import UserRead
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 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.store.exceptions import CustomException
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 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.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
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"]

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 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:
from langflow.services.database.models.user import User
class ApiKeyBase(SQLModelSerializable):
class ApiKeyBase(SQLModel):
name: Optional[str] = Field(index=True)
created_at: datetime = Field(default_factory=datetime.utcnow)
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:
return orjson.dumps(v, 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"]

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
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)
frontend_node_id: uuid.UUID = 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"]

View file

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

View file

@ -2,7 +2,7 @@ from datetime import datetime, timezone
from typing import Union
from uuid import UUID
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 sqlalchemy.exc import IntegrityError
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 datetime import datetime
from typing import Optional, TYPE_CHECKING
from uuid import UUID, uuid4
from langflow.services.database.models.base import SQLModel
if TYPE_CHECKING:
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
class User(SQLModelSerializable, table=True):
class User(SQLModel, table=True):
id: UUID = Field(default_factory=uuid4, primary_key=True, unique=True)
username: str = Field(index=True, unique=True)
password: str = Field()
@ -27,6 +28,10 @@ class User(SQLModelSerializable, table=True):
)
store_api_key: str = Field(default=None, nullable=True)
flows: list["Flow"] = Relationship(back_populates="user")
credentials: list["Credential"] = Relationship(
back_populates="user",
sa_relationship_kwargs={"cascade": "delete"},
)
class UserCreate(SQLModel):

View file

@ -12,6 +12,11 @@ if TYPE_CHECKING:
from langflow.services.settings.service import SettingsService
from langflow.services.store.service import StoreService
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":

View file

@ -15,3 +15,4 @@ class ServiceType(str, Enum):
SESSION_SERVICE = "session_service"
TASK_SERVICE = "task_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):
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()
@ -127,7 +127,7 @@ def teardown_superuser(settings_service, session):
try:
logger.debug("AUTO_LOGIN is set to False. Removing default superuser if exists.")
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()
if user and user.is_superuser:

View file

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

View file

@ -2,7 +2,7 @@ from collections import namedtuple
import uuid
from langflow.processing.process import Result
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.database.utils import session_getter
from langflow.services.deps import get_db_service

View file

@ -1,6 +1,6 @@
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.utils import teardown_superuser

View file

@ -3,7 +3,7 @@ from datetime import datetime
import pytest
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.user import User
from langflow.services.database.models.user.model import User
from langflow.services.database.utils import session_getter
from langflow.services.deps import get_db_service, get_settings_service