From 52900503de82d20232a6b307fc871a8d67d95057 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Wed, 18 Oct 2023 18:09:24 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20chore(alembic):=20make=20changes?=
 =?UTF-8?q?=20for=20the=20store?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

🔧 chore(api_key.py): remove unused imports and update save_store_api_key endpoint to use ApiKeyCreateRequest model

🔧 chore(schemas.py): add ApiKeyCreateRequest model

🔧 chore(store.py): remove unused imports and update imports in store.py
---
 .../595c9c2a2ad4_changes_for_the_store.py     | 24 +++++++++++--------
 src/backend/langflow/api/v1/api_key.py        | 10 ++++----
 src/backend/langflow/api/v1/schemas.py        |  4 ++++
 src/backend/langflow/api/v1/store.py          |  5 +---
 4 files changed, 23 insertions(+), 20 deletions(-)

diff --git a/src/backend/langflow/alembic/versions/595c9c2a2ad4_changes_for_the_store.py b/src/backend/langflow/alembic/versions/595c9c2a2ad4_changes_for_the_store.py
index 6ff6079f6..69e8b44d6 100644
--- a/src/backend/langflow/alembic/versions/595c9c2a2ad4_changes_for_the_store.py
+++ b/src/backend/langflow/alembic/versions/595c9c2a2ad4_changes_for_the_store.py
@@ -13,29 +13,33 @@ import sqlmodel
 
 
 # revision identifiers, used by Alembic.
-revision: str = '595c9c2a2ad4'
-down_revision: Union[str, None] = 'eb5866d51fd2'
+revision: str = "595c9c2a2ad4"
+down_revision: Union[str, None] = "eb5866d51fd2"
 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! ###
-    with op.batch_alter_table('flow', schema=None) as batch_op:
-        batch_op.add_column(sa.Column('is_component', sa.Boolean(), nullable=False))
+    with op.batch_alter_table("flow", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("is_component", sa.Boolean(), nullable=False))
 
-    with op.batch_alter_table('user', schema=None) as batch_op:
-        batch_op.add_column(sa.Column('store_api_key', sqlmodel.sql.sqltypes.AutoString(), nullable=True))
+    with op.batch_alter_table("user", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "store_api_key", sqlmodel.sql.sqltypes.AutoString(), nullable=True
+            )
+        )
 
     # ### end Alembic commands ###
 
 
 def downgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    with op.batch_alter_table('user', schema=None) as batch_op:
-        batch_op.drop_column('store_api_key')
+    with op.batch_alter_table("user", schema=None) as batch_op:
+        batch_op.drop_column("store_api_key")
 
-    with op.batch_alter_table('flow', schema=None) as batch_op:
-        batch_op.drop_column('is_component')
+    with op.batch_alter_table("flow", schema=None) as batch_op:
+        batch_op.drop_column("is_component")
 
     # ### end Alembic commands ###
diff --git a/src/backend/langflow/api/v1/api_key.py b/src/backend/langflow/api/v1/api_key.py
index d98da1c66..4ec6d8175 100644
--- a/src/backend/langflow/api/v1/api_key.py
+++ b/src/backend/langflow/api/v1/api_key.py
@@ -1,6 +1,6 @@
 from uuid import UUID
-from fastapi import APIRouter, HTTPException, Depends, Body
-from langflow.api.v1.schemas import ApiKeysResponse
+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 (
     ApiKeyCreate,
@@ -15,7 +15,6 @@ from langflow.services.database.models.api_key.crud import (
 )
 from langflow.services.database.models.user.user import User
 from langflow.services.deps import (
-    get_store_service,
     get_session,
     get_settings_service,
 )
@@ -25,8 +24,7 @@ from typing import TYPE_CHECKING
 from sqlmodel import Session
 
 if TYPE_CHECKING:
-    from langflow.services.store.service import StoreService
-    from langflow.services.settings.service import SettingsService
+    pass
 
 router = APIRouter(tags=["APIKey"], prefix="/api_key")
 
@@ -73,7 +71,7 @@ def delete_api_key_route(
 
 @router.post("/store")
 def save_store_api_key(
-    api_key: Body(str, embed=True),
+    api_key: ApiKeyCreateRequest,
     current_user: User = Depends(auth_utils.get_current_active_user),
     db: Session = Depends(get_session),
     settings_service=Depends(get_settings_service),
diff --git a/src/backend/langflow/api/v1/schemas.py b/src/backend/langflow/api/v1/schemas.py
index 37e7d712d..73e45d673 100644
--- a/src/backend/langflow/api/v1/schemas.py
+++ b/src/backend/langflow/api/v1/schemas.py
@@ -198,3 +198,7 @@ class Token(BaseModel):
     access_token: str
     refresh_token: str
     token_type: str
+
+
+class ApiKeyCreateRequest(BaseModel):
+    api_key: str
diff --git a/src/backend/langflow/api/v1/store.py b/src/backend/langflow/api/v1/store.py
index 0ef3a7362..fe7f6712d 100644
--- a/src/backend/langflow/api/v1/store.py
+++ b/src/backend/langflow/api/v1/store.py
@@ -1,5 +1,4 @@
-from datetime import timezone
-from typing import List, TYPE_CHECKING, Optional
+from typing import List, Optional
 from uuid import UUID
 from langflow.services.auth import utils as auth_utils
 from langflow.services.database.models.flow.flow import Flow
@@ -11,9 +10,7 @@ from langflow.services.deps import (
 )
 from langflow.services.store.schema import ComponentResponse
 
-from sqlmodel import Session, select
 from fastapi import APIRouter, Depends, HTTPException, Query
-from sqlalchemy.exc import IntegrityError
 from datetime import datetime
 
 from langflow.services.store.service import StoreService