From 82fbeace068a09a83e1c8d02f958a1e40643c39b Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Fri, 25 Aug 2023 12:03:51 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=80=20chore(api=5Fkey):=20update=20imp?=
 =?UTF-8?q?ort=20and=20export=20names=20in=20=5F=5Finit=5F=5F.py=20for=20b?=
 =?UTF-8?q?etter=20clarity=20and=20consistency?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

🔀 chore(api_key): add UnmaskedApiKeyRead model to represent an unmasked API key

🔀 chore(api_key): add user relationship to ApiKey model for easier access to associated user

🔀 chore(api_key): add user_id field to ApiKeyBase and ApiKeyCreate models for easier user association

🔀 chore(api_key): add mask_api_key validator to ApiKeyRead model to mask the API key for security reasons
---
 .../database/models/api_key/__init__.py       |  4 +--
 .../database/models/api_key/api_key.py        | 28 ++++++++++++++++---
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/src/backend/langflow/services/database/models/api_key/__init__.py b/src/backend/langflow/services/database/models/api_key/__init__.py
index c97425ee8..fbb8265b9 100644
--- a/src/backend/langflow/services/database/models/api_key/__init__.py
+++ b/src/backend/langflow/services/database/models/api_key/__init__.py
@@ -1,3 +1,3 @@
-from .api_key import ApiKey, ApiKeyCreate, ApiKeyRead
+from .api_key import ApiKey, ApiKeyCreate, UnmaskedApiKeyRead, ApiKeyRead
 
-__all__ = ["ApiKey", "ApiKeyCreate", "ApiKeyRead"]
+__all__ = ["ApiKey", "ApiKeyCreate", "UnmaskedApiKeyRead", "ApiKeyRead"]
diff --git a/src/backend/langflow/services/database/models/api_key/api_key.py b/src/backend/langflow/services/database/models/api_key/api_key.py
index 784b25229..1006b1c0f 100644
--- a/src/backend/langflow/services/database/models/api_key/api_key.py
+++ b/src/backend/langflow/services/database/models/api_key/api_key.py
@@ -1,24 +1,44 @@
-from sqlmodel import Field
+from pydantic import validator
+from sqlmodel import Field, Relationship
 from uuid import UUID, uuid4
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 from datetime import datetime
 from langflow.services.database.models.base import SQLModelSerializable
 
+if TYPE_CHECKING:
+    from langflow.services.database.models.user import User
+
 
 class ApiKeyBase(SQLModelSerializable):
     api_key: str = Field(index=True, unique=True)
     name: Optional[str] = Field(index=True)
-    create_at: datetime = Field(default_factory=datetime.utcnow)
+    created_at: datetime = Field(default_factory=datetime.utcnow)
     last_used_at: Optional[datetime] = Field(default=None)
+    user_id: UUID = Field()
 
 
 class ApiKey(ApiKeyBase, table=True):
     id: UUID = Field(default_factory=uuid4, primary_key=True, unique=True)
+    # User relationship
+    user_id: UUID = Field(index=True, foreign_key="user.id")
+    user: "User" = Relationship(back_populates="api_keys")
 
 
 class ApiKeyCreate(ApiKeyBase):
-    pass
+    api_key: Optional[str] = None
+    user_id: Optional[UUID] = None
+
+
+class UnmaskedApiKeyRead(ApiKeyBase):
+    id: UUID
 
 
 class ApiKeyRead(ApiKeyBase):
     id: UUID
+    api_key: Optional[str] = None
+    user_id: Optional[UUID] = None
+
+    @validator("api_key", always=True)
+    def mask_api_key(cls, v):
+        # This validator will always run, and will mask the API key
+        return f"{'*' * 8}{v[-4:]}"