From 0a7acb4e665d38999612f56bfb263a8d6e3cf970 Mon Sep 17 00:00:00 2001
From: Ritchie Martori <ritchie@voxel51.com>
Date: Wed, 2 Oct 2024 08:37:03 -0700
Subject: [PATCH] refactor: remove pydantic from execution_store

---
 fiftyone/factory/repos/execution_store.py | 21 +++++++++++++--------
 fiftyone/operators/store/models.py        | 21 +++++++++++++--------
 fiftyone/operators/store/store.py         |  1 -
 3 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/fiftyone/factory/repos/execution_store.py b/fiftyone/factory/repos/execution_store.py
index 69fa78b077..b8fe6b786a 100644
--- a/fiftyone/factory/repos/execution_store.py
+++ b/fiftyone/factory/repos/execution_store.py
@@ -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]:
@@ -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,
diff --git a/fiftyone/operators/store/models.py b/fiftyone/operators/store/models.py
index 5fd6c89e21..f6cba9f841 100644
--- a/fiftyone/operators/store/models.py
+++ b/fiftyone/operators/store/models.py
@@ -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
 
@@ -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."""
 
diff --git a/fiftyone/operators/store/store.py b/fiftyone/operators/store/store.py
index 15e0ab3a01..cdf761dd46 100644
--- a/fiftyone/operators/store/store.py
+++ b/fiftyone/operators/store/store.py
@@ -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