From d7297960ee7fd8443ec173e83c2b2284bec15603 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Justin=20K=C3=B6stinger?= <jkoestinger@gmail.com>
Date: Tue, 4 May 2021 01:13:24 +0200
Subject: [PATCH] Added basic typesense document handling

---
 type_sense/models.py               | 55 +++++++++++++++++++++++++++---
 type_sense/tests/test_typesense.py |  2 +-
 2 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/type_sense/models.py b/type_sense/models.py
index 2c4d84d..2a5abb6 100644
--- a/type_sense/models.py
+++ b/type_sense/models.py
@@ -9,30 +9,77 @@ class TypeSenseMixin(models.Model):
 
     typesense_fields = []
     typesense_schema_name = None
+    typesense_default_sorting_field = None
 
     def get_typesense_schema_name(self) -> str:
         if self.typesense_schema_name:
             return self.typesense_schema_name
         return self.__class__.__name__.lower()
 
-    def _get_schema(self):
+    def get_typesense_fields(self):
+        fields = deepcopy(self.typesense_fields)
+
+        # Auto adds the pk as Id if absent
+        if not any(field["name"] == "id" for field in fields):
+            fields.append({"name": "id", "type": "string", "attribute": "pk"})
+
+        return fields
+
+    def _get_typesense_schema(self):
         schema = {
             "name": self.get_typesense_schema_name(),
             "fields": deepcopy(self.typesense_fields),
         }
 
+        if self.typesense_default_sorting_field:
+            schema["default_sorting_field"] = self.typesense_default_sorting_field
+
+        existing_fields = []
+
         for field in schema["fields"]:
+            existing_fields.append(field["name"])
             if "attribute" in field:
                 del field["attribute"]
 
         return schema
 
-    def create_collection(self):
-        client.collections.create(self._get_schema())
+    def _get_typesense_field_value(self, field):
+        if "attribute" in field:
+            attribute = field["attribute"]
+        else:
+            attribute = field["name"]
+        return getattr(self, attribute)
+
+    def create_typesense_collection(self):
+        client.collections.create(self._get_typesense_schema())
 
-    def drop_collection(self):
+    def drop_typesense_collection(self):
         client.collections[self.get_typesense_schema_name()].delete()
 
+    def retrieve_typesense_collection(self):
+        return client.collections[self.get_typesense_schema_name()].retrieve()
+
+    def upsert_typesense_document(self):
+        fields = {}
+
+        for field in self.get_typesense_fields():
+            fields[field["name"]] = self._get_typesense_field_value(field)
+
+        return fields
+
+    def delete_typesense_document(self):
+        client.collections[self.get_typesense_schema_name()].documents[
+            str(self.pk)
+        ].delete()
+
+    def save(self, *args, **kwargs):
+        super().save(*args, **kwargs)
+        self.upsert_typesense_document()
+
+    def delete(self, **kwargs):
+        self.delete_typesense_document()
+        super().delete(**kwargs)
+
     @classmethod
     def check(cls, **kwargs):
         errors = super().check(**kwargs)
diff --git a/type_sense/tests/test_typesense.py b/type_sense/tests/test_typesense.py
index 67aadbd..a361dc3 100644
--- a/type_sense/tests/test_typesense.py
+++ b/type_sense/tests/test_typesense.py
@@ -47,4 +47,4 @@ def test_typesense_schema_custom_fields(self):
             ],
             "name": "customfieldsmockmodel",
         }
-        self.assertDictEqual(m._get_schema(), expected)
+        self.assertDictEqual(m._get_typesense_schema(), expected)