From cc14c5123870872731895ef9c1569c0f0231c719 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Sat, 13 Apr 2024 11:21:24 -0300 Subject: [PATCH] Update created_at field in api_key/crud.py and api_key/model.py (#1698) * Update created_at field in api_key/crud.py and api_key/model.py --- .../services/database/models/api_key/crud.py | 1 + .../services/database/models/api_key/model.py | 16 ++++++++++------ 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/backend/base/langflow/services/database/models/api_key/crud.py b/src/backend/base/langflow/services/database/models/api_key/crud.py index 33bdc757995d..37fe08cd15fc 100644 --- a/src/backend/base/langflow/services/database/models/api_key/crud.py +++ b/src/backend/base/langflow/services/database/models/api_key/crud.py @@ -24,6 +24,7 @@ def create_api_key(session: Session, api_key_create: ApiKeyCreate, user_id: UUID api_key=generated_api_key, name=api_key_create.name, user_id=user_id, + created_at=api_key_create.created_at or datetime.datetime.now(datetime.timezone.utc), ) session.add(api_key) diff --git a/src/backend/base/langflow/services/database/models/api_key/model.py b/src/backend/base/langflow/services/database/models/api_key/model.py index e110258b92d1..ee0d1e71d00e 100644 --- a/src/backend/base/langflow/services/database/models/api_key/model.py +++ b/src/backend/base/langflow/services/database/models/api_key/model.py @@ -3,17 +3,18 @@ from uuid import UUID, uuid4 from pydantic import field_validator, validator -from sqlmodel import Field, Relationship, SQLModel, Column, func, DateTime +from sqlmodel import Column, DateTime, Field, Relationship, SQLModel, func if TYPE_CHECKING: from langflow.services.database.models.user import User +def utc_now(): + return datetime.now(timezone.utc) + + class ApiKeyBase(SQLModel): name: Optional[str] = Field(index=True, nullable=True, default=None) - created_at: datetime = Field( - default=None, sa_column=Column(DateTime(timezone=True), server_default=func.now(), nullable=False) - ) last_used_at: Optional[datetime] = Field(default=None, nullable=True) total_uses: int = Field(default=0) is_active: bool = Field(default=True) @@ -21,7 +22,9 @@ class ApiKeyBase(SQLModel): class ApiKey(ApiKeyBase, table=True): id: UUID = Field(default_factory=uuid4, primary_key=True, unique=True) - + created_at: datetime = Field( + default=None, sa_column=Column(DateTime(timezone=True), server_default=func.now(), nullable=False) + ) api_key: str = Field(index=True, unique=True) # User relationship # Delete API keys when user is deleted @@ -34,10 +37,11 @@ class ApiKey(ApiKeyBase, table=True): class ApiKeyCreate(ApiKeyBase): api_key: Optional[str] = None user_id: Optional[UUID] = None + created_at: Optional[datetime] = Field(default_factory=utc_now) @field_validator("created_at", mode="before") def set_created_at(cls, v): - return v or datetime.now(timezone.utc) + return v or utc_now() class UnmaskedApiKeyRead(ApiKeyBase):