From 74ac238a4e602c3ca49a0c25441a18b14884683a Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 24 Nov 2023 11:18:31 -0300 Subject: [PATCH] Adds Credential table --- .../2ac71eb9c3ae_adds_credential_table.py | 45 ++++++++++++++++++ .../services/database/models/__init__.py | 5 +- .../database/models/credential/__init__.py | 3 ++ .../database/models/credential/model.py | 46 +++++++++++++++++++ .../database/models/credential/schema.py | 8 ++++ 5 files changed, 105 insertions(+), 2 deletions(-) create mode 100644 src/backend/langflow/alembic/versions/2ac71eb9c3ae_adds_credential_table.py create mode 100644 src/backend/langflow/services/database/models/credential/__init__.py create mode 100644 src/backend/langflow/services/database/models/credential/model.py create mode 100644 src/backend/langflow/services/database/models/credential/schema.py diff --git a/src/backend/langflow/alembic/versions/2ac71eb9c3ae_adds_credential_table.py b/src/backend/langflow/alembic/versions/2ac71eb9c3ae_adds_credential_table.py new file mode 100644 index 000000000..3f974dc04 --- /dev/null +++ b/src/backend/langflow/alembic/versions/2ac71eb9c3ae_adds_credential_table.py @@ -0,0 +1,45 @@ +"""Adds Credential table + +Revision ID: c1c8e217a069 +Revises: 7d2162acc8b2 +Create Date: 2023-11-24 10:45:38.465302 + +""" +from typing import Sequence, Union + +import sqlalchemy as sa +import sqlmodel +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = '2ac71eb9c3ae' +down_revision: Union[str, None] = '7d2162acc8b2' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + try: + op.create_table('credential', + sa.Column('name', sqlmodel.sql.sqltypes.AutoString(), nullable=True), + sa.Column('value', sqlmodel.sql.sqltypes.AutoString(), nullable=True), + sa.Column('provider', sqlmodel.sql.sqltypes.AutoString(), nullable=True), + sa.Column('user_id', sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column('id', sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.Column('updated_at', sa.DateTime(), nullable=True), + sa.PrimaryKeyConstraint('id') + ) + except Exception: + pass + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + try: + op.drop_table('credential') + except Exception: + pass + # ### end Alembic commands ### diff --git a/src/backend/langflow/services/database/models/__init__.py b/src/backend/langflow/services/database/models/__init__.py index 3cc4231a3..9dd3e0285 100644 --- a/src/backend/langflow/services/database/models/__init__.py +++ b/src/backend/langflow/services/database/models/__init__.py @@ -1,5 +1,6 @@ +from .api_key import ApiKey +from .credential import Credential from .flow import Flow from .user import User -from .api_key import ApiKey -__all__ = ["Flow", "User", "ApiKey"] +__all__ = ["Flow", "User", "ApiKey", "Credential"] diff --git a/src/backend/langflow/services/database/models/credential/__init__.py b/src/backend/langflow/services/database/models/credential/__init__.py new file mode 100644 index 000000000..2f8b6cd01 --- /dev/null +++ b/src/backend/langflow/services/database/models/credential/__init__.py @@ -0,0 +1,3 @@ +from .model import Credential, CredentialCreate, CredentialRead, CredentialUpdate + +__all__ = ["Credential", "CredentialCreate", "CredentialRead", "CredentialUpdate"] diff --git a/src/backend/langflow/services/database/models/credential/model.py b/src/backend/langflow/services/database/models/credential/model.py new file mode 100644 index 000000000..91c36c2a7 --- /dev/null +++ b/src/backend/langflow/services/database/models/credential/model.py @@ -0,0 +1,46 @@ +from datetime import datetime +from typing import TYPE_CHECKING, Optional +from uuid import UUID, uuid4 + +from sqlmodel import Field, Relationship, SQLModel + +from langflow.services.database.models.credential.schema import AcceptedProviders + +if TYPE_CHECKING: + from langflow.services.database.models.user import User + + +class CredentialBase(SQLModel): + name: Optional[str] = Field(None, description="Name of the credential") + value: Optional[str] = Field(None, description="Encrypted value of the credential") + provider: Optional[str] = Field(None, description="Provider of the credential (e.g OpenAI)") + user_id: UUID = Field(description="User ID associated with this credential") + + +class Credential(CredentialBase, table=True): + id: Optional[UUID] = Field(default_factory=uuid4, primary_key=True, description="Unique ID for the credential") + created_at: datetime = Field(default_factory=datetime.utcnow, description="Creation time of the credential") + updated_at: Optional[datetime] = Field(None, description="Last update time of the credential") + # foreign key to user table + user_id: UUID = Field(description="User ID associated with this credential", foreign_key="user.id") + user: "User" = Relationship(back_populates="credentials") + + if TYPE_CHECKING: + user: "User" = Relationship(back_populates="credentials") + + +class CredentialCreate(CredentialBase): + # AcceptedProviders is a custom Enum + provider: AcceptedProviders = Field(description="Provider of the credential (e.g OpenAI)") + + +class CredentialRead(SQLModel): + id: UUID + name: Optional[str] = Field(None, description="Name of the credential") + provider: Optional[str] = Field(None, description="Provider of the credential (e.g OpenAI)") + + +class CredentialUpdate(SQLModel): + id: UUID # Include the ID for updating + name: Optional[str] = Field(None, description="Name of the credential") + value: Optional[str] = Field(None, description="Encrypted value of the credential") diff --git a/src/backend/langflow/services/database/models/credential/schema.py b/src/backend/langflow/services/database/models/credential/schema.py new file mode 100644 index 000000000..56b6df10a --- /dev/null +++ b/src/backend/langflow/services/database/models/credential/schema.py @@ -0,0 +1,8 @@ +from enum import Enum + + +class AcceptedProviders(str, Enum): + """Accepted providers for credentials.""" + + openai = "openai" + anthropic = "anthropic"