Add support for defining indexes on models

README.md

@@ -47,6 +47,7 @@ from shylock import ShylockAioArangoDBBackend
 from shylock import configure as configure_shylock
 
 from arangodantic import ASCENDING, DocumentModel, EdgeModel, configure
+from arangodantic.indexes import HashIndex
 
 
 # Define models
@@ -60,13 +61,23 @@ class Owner(BaseModel):
 class Company(DocumentModel):
     """Dummy company Arangodantic model."""
 
+    class ArangodanticConfig:
+        indexes = [
+            HashIndex(fields=["company_id"], unique=True),
+            HashIndex(fields=["owner.first_name"]),
+            HashIndex(fields=["owner.last_name"]),
+        ]
+
     company_id: str
     owner: Owner
 
 
 class Link(EdgeModel):
     """Dummy Link Arangodantic model."""
 
+    class ArangodanticConfig:
+        indexes = [HashIndex(fields=["_from", "_to", "type"], unique=True)]
+
     type: str
 
 
@@ -90,10 +101,16 @@ async def main():
     configure_shylock(await ShylockAioArangoDBBackend.create(db, f"{prefix}shylock"))
     configure(db, prefix=prefix, key_gen=uuid4, lock=Lock)
 
-    # Create collections if they don't yet exist
+    # Create collections and indexes if they don't yet exist
     # Only for demo, you likely want to create the collections in advance.
     await Company.ensure_collection()
+    await Company.ensure_indexes()
     await Link.ensure_collection()
+    await Link.ensure_indexes()
+
+    # Clean up any existing data.
+    await Company.truncate_collection()
+    await Link.truncate_collection()
 
     # Let's create some example entries
     owner = Owner(first_name="John", last_name="Doe")

arangodantic/indexes.py

@@ -0,0 +1,172 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+import pydantic
+from aioarangodb.collection import StandardCollection
+from pydantic import Field
+
+
+class BaseIndex(ABC, pydantic.BaseModel):
+    """
+    Base index model.
+
+    Define fields common to most of the indexes. The index class should be derived
+    from this base class, and add the add_index() function which actually creates
+    the index.
+    """
+
+    fields: List = Field(..., description="List of fields to index")
+    name: Optional[str] = Field(None, description="Optional name for the index.")
+    in_background: Optional[bool] = Field(
+        None, description="Do not hold the collection lock."
+    )
+
+    @abstractmethod
+    async def add_index(self, collection: StandardCollection):
+        """
+        Creates the index on the collection.
+
+        Override this function in the subclass.
+
+        :param collection: The collection to create the index on.
+        :return: Dictionary with the index information.
+        """
+        raise NotImplementedError()
+
+
+class HashIndex(BaseIndex):
+    """
+    Creates a hash index on the collection.
+    """
+
+    unique: Optional[bool] = Field(
+        None, description="Whether the index is unique or not."
+    )
+    sparse: Optional[bool] = Field(
+        None,
+        description="If set to True, documents with None in the field are also indexed,"
+        " otherwise they're skipped",
+    )
+    deduplicate: Optional[bool] = Field(
+        None,
+        description="If set to True, inserting duplicate index values from the same "
+        "document triggers unique constraint errors.",
+    )
+
+    async def add_index(self, collection: StandardCollection):
+        return await collection.add_hash_index(
+            fields=self.fields,
+            unique=self.unique,
+            sparse=self.sparse,
+            deduplicate=self.deduplicate,
+            name=self.name,
+            in_background=self.in_background,
+        )
+
+
+class SkiplistIndex(BaseIndex):
+    """
+    Creates a skiplist index on the collection.
+    """
+
+    unique: Optional[bool] = Field(
+        None, description="Whether the index is unique or not."
+    )
+    sparse: Optional[bool] = Field(
+        None,
+        description="If set to True, documents with None in the field are also indexed,"
+        " otherwise they're skipped",
+    )
+    deduplicate: Optional[bool] = Field(
+        None,
+        description="If set to True, inserting duplicate index values from the same "
+        "document triggers unique constraint errors.",
+    )
+
+    async def add_index(self, collection: StandardCollection):
+        return await collection.add_skiplist_index(
+            fields=self.fields,
+            unique=self.unique,
+            sparse=self.sparse,
+            deduplicate=self.deduplicate,
+            name=self.name,
+            in_background=self.in_background,
+        )
+
+
+class GeoIndex(BaseIndex):
+    """
+    Creates a geo-spatial index on the collection.
+    """
+
+    ordered: Optional[bool] = Field(
+        None, description="Whether the order is longitude, then latitude."
+    )
+
+    async def add_index(self, collection: StandardCollection):
+        return await collection.add_geo_index(
+            fields=self.fields,
+            ordered=self.ordered,
+            name=self.name,
+            in_background=self.in_background,
+        )
+
+
+class FulltextIndex(BaseIndex):
+    """
+    Creates a fulltext index on the collection.
+    """
+
+    min_length: Optional[int] = Field(
+        None, description="Minimum number of characters to index."
+    )
+
+    async def add_index(self, collection: StandardCollection):
+        return await collection.add_fulltext_index(
+            fields=self.fields,
+            min_length=self.min_length,
+            name=self.name,
+            in_background=self.in_background,
+        )
+
+
+class PersistentIndex(BaseIndex):
+    """
+    Creates a persistent index on the collection.
+    """
+
+    unique: Optional[bool] = Field(
+        None, description="Whether the index is unique or not."
+    )
+    sparse: Optional[bool] = Field(
+        None,
+        description="If set to True, documents with None in the field are also indexed,"
+        " otherwise they're skipped",
+    )
+
+    async def add_index(self, collection: StandardCollection):
+        return await collection.add_persistent_index(
+            fields=self.fields,
+            unique=self.unique,
+            sparse=self.sparse,
+            name=self.name,
+            in_background=self.in_background,
+        )
+
+
+class TTLIndex(BaseIndex):
+    """
+    Creates a TTL (time-to-live) index.
+    """
+
+    expiry_time: int = Field(
+        ..., description="Time of expiry in seconds after document creation."
+    )
+
+    async def add_index(self, collection: StandardCollection):
+        return await collection.add_ttl_index(
+            fields=self.fields,
+            expiry_time=self.expiry_time,
+            name=self.name,
+            in_background=self.in_background,
+        )

arangodantic/models.py

@@ -1,7 +1,7 @@
 import textwrap
 from abc import ABC
 from functools import lru_cache
-from typing import Optional, Type, TypeVar, Union
+from typing import List, Optional, Type, TypeVar, Union
 
 import aioarangodb.exceptions
 import pydantic
@@ -23,6 +23,7 @@
     MultipleModelsFoundError,
     UniqueConstraintError,
 )
+from arangodantic.indexes import BaseIndex
 from arangodantic.utils import (
     FilterTypes,
     SortTypes,
@@ -43,6 +44,9 @@ class ArangodanticCollectionConfig(pydantic.BaseModel):
     collection_name: Optional[str] = Field(
         None, description="Override the name of the collection to use"
     )
+    indexes: Optional[List[BaseIndex]] = Field(
+        None, description="Override the collection indexes."
+    )
 
 
 class Model(pydantic.BaseModel, ABC):
@@ -239,6 +243,45 @@ async def ensure_collection(cls, *args, **kwargs):
         if not await db.has_collection(name):
             await db.create_collection(name, *args, **kwargs)
 
+    @classmethod
+    async def ensure_indexes(cls):
+        """
+        Ensure the collection's indexes are created.
+
+        The configuration for the indexes are defined in the ArangodanticConfig class
+        for a model. The class should define an `indexes` list, which contains
+        the desired indexing class to be run, and it's index definitions.
+        Supported index classes are:
+        - HashIndex (add_hash_index)
+        - GeoIndex (add_geo_index)
+        - TTLIndex (add_ttl_index)
+        - FulltextIndex (add_fulltext_index)
+        - PersistentIndex (add_persistent_index)
+        - SkiplistIndex (add_skiplist_index)
+
+        All field definition key-values are passed to the index function as keyword
+        arguments.
+
+        Example:
+            class Node(DocumentModel):
+                class ArangodanticConfig:
+                    indexes = [
+                        HashIndex(fields=["field1", "field2"]),
+                        HashIndex(fields=["field3"], unique=True),
+                        HashIndex(fields=["field4"], sparse=True),
+                        HashIndex(fields=["field5"], deduplicate=True),
+                    ]
+        """
+        cls_config: ArangodanticCollectionConfig = getattr(
+            cls, "ArangodanticConfig", ArangodanticCollectionConfig()
+        )
+        indexes = getattr(cls_config, "indexes", None)
+
+        if indexes:
+            collection = cls.get_collection()
+            for index in indexes:
+                await index.add_index(collection)
+
     @classmethod
     async def delete_collection(cls, ignore_missing: bool = True):
         """

arangodantic/tests/conftest.py

@@ -11,6 +11,7 @@
 from shylock import configure as configure_shylock
 
 from arangodantic import DocumentModel, EdgeDefinition, EdgeModel, Graph, configure
+from arangodantic.indexes import HashIndex
 
 HOSTS = "http://localhost:8529"
 USERNAME = "root"
@@ -51,6 +52,9 @@ def rand_str(length: int) -> str:
 class Identity(DocumentModel):
     """Dummy identity Arangodantic model."""
 
+    class ArangodanticConfig:
+        indexes = [HashIndex(fields=["name"])]
+
     name: str = ""
 
 
@@ -79,6 +83,9 @@ async def before_save(
 class Link(EdgeModel):
     """Dummy Arangodantic edge model."""
 
+    class ArangodanticConfig:
+        indexes = [HashIndex(fields=["_from", "_to", "type"], unique=True)]
+
     type: str
 
 
@@ -125,6 +132,7 @@ class ArangodanticConfig:
 @pytest.fixture
 async def identity_collection(configure_db):
     await Identity.ensure_collection()
+    await Identity.ensure_indexes()
     yield
     await Identity.delete_collection()
 
@@ -146,13 +154,15 @@ async def identity_bob(identity_collection):
 @pytest.fixture
 async def extended_identity_collection(configure_db):
     await ExtendedIdentity.ensure_collection()
+    await ExtendedIdentity.ensure_indexes()
     yield
     await ExtendedIdentity.delete_collection()
 
 
 @pytest.fixture
 async def link_collection(configure_db):
     await Link.ensure_collection()
+    await Link.ensure_indexes()
     yield
     await Link.delete_collection()