Skip to content

Commit

Permalink
refactor: remove pydantic from execution_store
Browse files Browse the repository at this point in the history
  • Loading branch information
ritch committed Oct 3, 2024
1 parent 9f9182b commit 0a7acb4
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 17 deletions.
21 changes: 13 additions & 8 deletions fiftyone/factory/repos/execution_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,8 @@ def __init__(self, collection: Collection):

def create_store(self, store_name, permissions=None) -> StoreDocument:
"""Creates a store in the execution store."""
store_doc = StoreDocument(
store_name=store_name, permissions=permissions
)
self._collection.insert_one(store_doc.dict())
store_doc = StoreDocument(store_name=store_name, value=permissions)
self._collection.insert_one(store_doc.to_mongo_dict())
return store_doc

def list_stores(self) -> list[str]:
Expand All @@ -40,14 +38,21 @@ def set_key(self, store_name, key, value, ttl=None) -> KeyDocument:
now = datetime.datetime.now()
expiration = KeyDocument.get_expiration(ttl)
key_doc = KeyDocument(
store_name=store_name, key=key, value=value, updated_at=now
store_name=store_name,
key=key,
value=value,
updated_at=now,
expires_at=expiration,
)

# Prepare the update operations
update_fields = {
"$set": key_doc.dict(
exclude={"created_at", "expires_at", "store_name", "key"}
),
"$set": {
k: v
for k, v in key_doc.to_mongo_dict().items()
if k
not in {"_id", "created_at", "expires_at", "store_name", "key"}
},
"$setOnInsert": {
"store_name": store_name,
"key": key,
Expand Down
21 changes: 13 additions & 8 deletions fiftyone/operators/store/models.py
Original file line number Diff line number Diff line change
@@ -1,21 +1,19 @@
"""
Store and key models for the execution store.
"""

from pydantic import BaseModel, Field
from dataclasses import dataclass, field, asdict
from typing import Optional, Dict, Any
import datetime


class KeyDocument(BaseModel):
@dataclass
class KeyDocument:
"""Model representing a key in the store."""

store_name: str
key: str
value: Any
created_at: datetime.datetime = Field(
created_at: datetime.datetime = field(
default_factory=datetime.datetime.now
)
_id: Optional[Any] = None
updated_at: Optional[datetime.datetime] = None
expires_at: Optional[datetime.datetime] = None

Expand All @@ -24,10 +22,17 @@ def get_expiration(ttl: Optional[int]) -> Optional[datetime.datetime]:
"""Gets the expiration date for a key with the given TTL."""
if ttl is None:
return None

return datetime.datetime.now() + datetime.timedelta(seconds=ttl)

def to_mongo_dict(self, exclude_id: bool = True) -> Dict[str, Any]:
"""Serializes the document to a MongoDB dictionary."""
data = asdict(self)
if exclude_id:
data.pop("_id", None)
return data


@dataclass
class StoreDocument(KeyDocument):
"""Model representing a Store."""

Expand Down
1 change: 0 additions & 1 deletion fiftyone/operators/store/store.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
"""

import logging
from fiftyone.factory.repo_factory import RepositoryFactory
from fiftyone.operators.store.service import ExecutionStoreService
from typing import Any, Optional

Expand Down

0 comments on commit 0a7acb4

Please sign in to comment.