Update datetime handling in ApiKey and Variable models (#1673)

* Update datetime handling in ApiKey and Variable models

* Refactor test_login_successful function in test_login.py

* Fix nullable attribute for created_at and updated_at fields in Variable model
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-04-10 23:04:10 -03:00 • committed by GitHub
commit 836ac08d80
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 24 additions and 18 deletions

View file

@ -9,7 +9,6 @@ from typing import Sequence, Union
import sqlalchemy as sa import sqlalchemy as sa
from alembic import op from alembic import op
from sqlalchemy.dialects import postgresql
from sqlalchemy.engine.reflection import Inspector from sqlalchemy.engine.reflection import Inspector
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
@ -35,7 +34,9 @@ def upgrade() -> None:
"created_at", "created_at",
existing_type=sa.TIMESTAMP(timezone=True), existing_type=sa.TIMESTAMP(timezone=True),
nullable=True, nullable=True,
existing_server_default=sa.text("now()"), # existing_server_default expects str | bool | Identity | Computed | None
# sa.text("now()") is not a valid value for existing_server_default
existing_server_default=False,
) )
# ### end Alembic commands ### # ### end Alembic commands ###
@ -54,7 +55,7 @@ def downgrade() -> None:
"created_at", "created_at",
existing_type=sa.TIMESTAMP(timezone=True), existing_type=sa.TIMESTAMP(timezone=True),
nullable=False, nullable=False,
existing_server_default=sa.text("now()"), existing_server_default=False,
) )
# ### end Alembic commands ### # ### end Alembic commands ###

View file

@ -1,8 +1,8 @@
from datetime import datetime from datetime import datetime, timezone
from typing import TYPE_CHECKING, Optional from typing import TYPE_CHECKING, Optional
from uuid import UUID, uuid4 from uuid import UUID, uuid4
from pydantic import validator from pydantic import field_validator, validator
from sqlmodel import Field, Relationship, SQLModel, Column, func, DateTime from sqlmodel import Field, Relationship, SQLModel, Column, func, DateTime
if TYPE_CHECKING: if TYPE_CHECKING:
@ -11,7 +11,9 @@ if TYPE_CHECKING:
class ApiKeyBase(SQLModel): class ApiKeyBase(SQLModel):
name: Optional[str] = Field(index=True, nullable=True, default=None) name: Optional[str] = Field(index=True, nullable=True, default=None)
created_at: datetime = Field(sa_column=Column(DateTime(timezone=True), server_default=func.now(), nullable=False)) 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) last_used_at: Optional[datetime] = Field(default=None, nullable=True)
total_uses: int = Field(default=0) total_uses: int = Field(default=0)
is_active: bool = Field(default=True) is_active: bool = Field(default=True)
@ -33,6 +35,10 @@ class ApiKeyCreate(ApiKeyBase):
api_key: Optional[str] = None api_key: Optional[str] = None
user_id: Optional[UUID] = None user_id: Optional[UUID] = None
@field_validator("created_at", mode="before")
def set_created_at(cls, v):
return v or datetime.now(timezone.utc)
class UnmaskedApiKeyRead(ApiKeyBase): class UnmaskedApiKeyRead(ApiKeyBase):
id: UUID id: UUID

View file

@ -26,10 +26,12 @@ class Variable(VariableBase, table=True):
) )
# name is unique per user # name is unique per user
created_at: datetime = Field( created_at: datetime = Field(
default=None,
sa_column=Column(DateTime(timezone=True), server_default=func.now(), nullable=True), sa_column=Column(DateTime(timezone=True), server_default=func.now(), nullable=True),
description="Creation time of the variable", description="Creation time of the variable",
) )
updated_at: Optional[datetime] = Field( updated_at: Optional[datetime] = Field(
default=None,
sa_column=Column(DateTime(timezone=True), nullable=True), sa_column=Column(DateTime(timezone=True), nullable=True),
description="Last update time of the variable", description="Last update time of the variable",
) )

View file

@ -1,5 +1,4 @@
import pytest import pytest
from langflow.services.database.models.api_key import ApiKeyCreate from langflow.services.database.models.api_key import ApiKeyCreate
@ -7,7 +6,7 @@ from langflow.services.database.models.api_key import ApiKeyCreate
def api_key(client, logged_in_headers, active_user): def api_key(client, logged_in_headers, active_user):
api_key = ApiKeyCreate(name="test-api-key") api_key = ApiKeyCreate(name="test-api-key")
response = client.post("api/v1/api_key", data=api_key.json(), headers=logged_in_headers) response = client.post("api/v1/api_key", data=api_key.model_dump_json(), headers=logged_in_headers)
assert response.status_code == 200, response.text assert response.status_code == 200, response.text
return response.json() return response.json()

View file

@ -1,17 +1,12 @@
from datetime import datetime from datetime import datetime
from pathlib import Path from pathlib import Path
import pytest
from langflow.graph.graph.base import Graph
from langflow.graph.schema import RunOutputs
from langflow.initial_setup.setup import ( from langflow.initial_setup.setup import (
STARTER_FOLDER_NAME, STARTER_FOLDER_NAME,
create_or_update_starter_projects, create_or_update_starter_projects,
get_project_data, get_project_data,
load_starter_projects, load_starter_projects,
) )
from langflow.memory import delete_messages
from langflow.processing.process import process_tweaks
from langflow.services.database.models.flow.model import Flow from langflow.services.database.models.flow.model import Flow
from langflow.services.deps import session_scope from langflow.services.deps import session_scope
from sqlalchemy import func from sqlalchemy import func

View file

@ -1,8 +1,8 @@
import pytest import pytest
from langflow.services.auth.utils import get_password_hash from langflow.services.auth.utils import get_password_hash
from langflow.services.database.models.user import User from langflow.services.database.models.user import User
from langflow.services.database.utils import session_getter from langflow.services.deps import session_scope
from langflow.services.deps import get_db_service from sqlalchemy.exc import IntegrityError
@pytest.fixture @pytest.fixture
@ -17,9 +17,12 @@ def test_user():
def test_login_successful(client, test_user): def test_login_successful(client, test_user):
# Adding the test user to the database # Adding the test user to the database
with session_getter(get_db_service()) as session: try:
with session_scope() as session:
session.add(test_user) session.add(test_user)
session.commit() session.commit()
except IntegrityError:
pass
response = client.post("api/v1/login", data={"username": "testuser", "password": "testpassword"}) response = client.post("api/v1/login", data={"username": "testuser", "password": "testpassword"})
assert response.status_code == 200 assert response.status_code == 200