diff --git a/README.md b/README.md index c0ff9b1..cc22371 100644 --- a/README.md +++ b/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,6 +61,13 @@ 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 @@ -67,6 +75,9 @@ class Company(DocumentModel): 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") diff --git a/arangodantic/indexes.py b/arangodantic/indexes.py new file mode 100644 index 0000000..8dc8c36 --- /dev/null +++ b/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, + ) diff --git a/arangodantic/models.py b/arangodantic/models.py index 9365472..6b485d0 100644 --- a/arangodantic/models.py +++ b/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): """ diff --git a/arangodantic/tests/conftest.py b/arangodantic/tests/conftest.py index 4fa59d7..3e3ef74 100644 --- a/arangodantic/tests/conftest.py +++ b/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,6 +154,7 @@ 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() @@ -153,6 +162,7 @@ async def extended_identity_collection(configure_db): @pytest.fixture async def link_collection(configure_db): await Link.ensure_collection() + await Link.ensure_indexes() yield await Link.delete_collection() diff --git a/arangodantic/tests/test_model.py b/arangodantic/tests/test_model.py index aca2e70..8abef6b 100644 --- a/arangodantic/tests/test_model.py +++ b/arangodantic/tests/test_model.py @@ -425,3 +425,56 @@ async def test_find_one_with_sort(identity_collection): found = await Identity.find_one(sort=[("name", DESCENDING)]) assert found.name == "Bob" + + +@pytest.mark.asyncio +async def test_identity_indexes(configure_db): + try: + await Identity.ensure_collection() + collection = Identity.get_collection() + + indexes_before = await collection.indexes() + assert len(indexes_before) == 1 + + # Ensure the indexes are created + await Identity.ensure_indexes() + + indexes_after = await collection.indexes() + assert len(indexes_after) == len(indexes_before) + 1 + + created_indexes = [ + index + for index in indexes_after + if index["type"] == "hash" and index["fields"] == ["name"] + ] + assert len(created_indexes) == 1 + finally: + await Identity.delete_collection() + + +@pytest.mark.asyncio +async def test_link_indexes(configure_db): + try: + await Link.ensure_collection() + collection = Link.get_collection() + + indexes_before = await collection.indexes() + # Here we have the primary index, + # and the _from + _to indexes for edges + assert len(indexes_before) == 2 + + # Ensure the indexes are created + await Link.ensure_indexes() + + indexes_after = await collection.indexes() + assert len(indexes_after) == len(indexes_before) + 1 + + created_indexes = [ + index + for index in indexes_after + if index["type"] == "hash" and index["fields"] == ["_from", "_to", "type"] + ] + assert len(created_indexes) == 1 + + finally: + await Link.delete_collection() diff --git a/examples/graph_example.py b/examples/graph_example.py index f98db41..13aecb9 100644 --- a/examples/graph_example.py +++ b/examples/graph_example.py @@ -15,18 +15,25 @@ ModelNotFoundError, configure, ) +from arangodantic.indexes import HashIndex # Define models class Person(DocumentModel): """Documents describing persons.""" + class ArangodanticConfig: + indexes = [HashIndex(fields=["name"])] + name: str class Relation(EdgeModel): """Edge documents describing relation between people.""" + class ArangodanticConfig: + indexes = [HashIndex(fields=["kind"])] + kind: str @@ -72,6 +79,14 @@ async def main(): # Only for demo, you likely want to create the graph in advance. await RelationGraph.ensure_graph() + # Clean up before running. + await Person.truncate_collection() + await Relation.truncate_collection() + + # We need to ensure the indexes separately + await Person.ensure_indexes() + await Relation.ensure_indexes() + # Let's create some example persons alice = Person(name="Alice") bob = Person(name="Bob") diff --git a/examples/readme_example.py b/examples/readme_example.py index e9f000c..8442f32 100644 --- a/examples/readme_example.py +++ b/examples/readme_example.py @@ -8,6 +8,7 @@ from shylock import configure as configure_shylock from arangodantic import ASCENDING, DocumentModel, EdgeModel, configure +from arangodantic.indexes import HashIndex # Define models @@ -21,6 +22,13 @@ 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 @@ -28,6 +36,9 @@ class Company(DocumentModel): class Link(EdgeModel): """Dummy Link Arangodantic model.""" + class ArangodanticConfig: + indexes = [HashIndex(fields=["_from", "_to", "type"], unique=True)] + type: str @@ -51,10 +62,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")