Add composite indexes support (#64)

.github/workflows/tests.yaml

@@ -22,7 +22,8 @@ jobs:
 
       - name: Install Firestore emulator 🚧
         run: |
-          curl -sL https://firebase.tools | bash
+          # with at least 13.5.1 and 13.6.0 tests can't connect to the emulator
+          npm install -g [email protected]
 
       - name: Start Firestore emulator 🚦
         run: |

CHANGELOG.md

@@ -7,6 +7,12 @@ this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased]
 
+## [0.7.0] - 2024-03-27
+
+### Added
+
+- Support for composite indexes via `__composite_indexes__` property in model classes.
+
 ## [0.6.0] - 2024-02-26
 
 ### Added
@@ -212,7 +218,8 @@ this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 - Update README.md
 - Update .gitignore
 
-[unreleased]: https://github.com/ioxiocom/firedantic/compare/0.6.0...HEAD
+[unreleased]: https://github.com/ioxiocom/firedantic/compare/0.7.0...HEAD
+[0.7.0]: https://github.com/ioxiocom/firedantic/compare/0.6.0...0.7.0
 [0.6.0]: https://github.com/ioxiocom/firedantic/compare/0.5.1...0.6.0
 [0.5.1]: https://github.com/ioxiocom/firedantic/compare/0.5.0...0.5.1
 [0.5.0]: https://github.com/ioxiocom/firedantic/compare/0.4.0...0.5.0

README.md

@@ -190,63 +190,102 @@ async def get_user_purchases(user_id: str, period: str = "2021") -> int:
 
 ```
 
-## TTL Policies
+## Composite Indexes and TTL Policies
 
-Firedantic has support for defining TTL policies and creating the policies.
+Firedantic has support for defining composite indexes and TTL policies as well as
+creating them.
+
+### Composite indexes
+
+Composite indexes of a collection are defined in `__composite_indexes__`, which is a
+list of all indexes to be created.
+
+To define an index, you can use `collection_index` or `collection_group_index`,
+depending on the query scope of the index. Each of these takes in an arbitrary amount of
+tuples, where the first element is the field name and the second is the order
+(`ASCENDING`/`DESCENDING`).
+
+The `set_up_composite_indexes` and `async_set_up_composite_indexes` functions are used
+to create indexes.
+
+For more details, see the example further down.
+
+### TTL Policies
 
 The field used for the TTL policy should be a datetime field and the name of the field
 should be defined in `__ttl_field__`. The `set_up_ttl_policies` and
 `async_set_up_ttl_policies` functions are used to set up the policies.
 
-Below are examples (both sync and async) to show how to use Firedantic to set up the TTL
-policies.
-
 Note: The TTL policies can not be set up in the Firestore emulator.
 
-### TTL Policy Example (sync)
+### Examples
+
+Below are examples (both sync and async) to show how to use Firedantic to set up
+composite indexes and TTL policies.
+
+The examples use `async_set_up_composite_indexes_and_ttl_policies` and
+`set_up_composite_indexes_and_ttl_policies` functions to set up both composite indexes
+and TTL policies. However, you can use separate functions to set up only either one of
+them.
+
+#### Composite Index and TTL Policy Example (sync)
 
 ```python
 from datetime import datetime
 
-from firedantic import Model, configure, get_all_subclasses, set_up_ttl_policies
-from google.cloud.firestore import Client
+from firedantic import (
+    collection_index,
+    collection_group_index,
+    configure,
+    get_all_subclasses,
+    Model,
+    set_up_composite_indexes_and_ttl_policies,
+)
+from google.cloud.firestore import Client, Query
 from google.cloud.firestore_admin_v1 import FirestoreAdminClient
 
 
 class ExpiringModel(Model):
     __collection__ = "expiringModel"
     __ttl_field__ = "expire"
+    __composite_indexes__ = [
+        collection_index(("content", Query.ASCENDING), ("expire", Query.DESCENDING)),
+        collection_group_index(("content", Query.DESCENDING), ("expire", Query.ASCENDING)),
+    ]
 
     expire: datetime
     content: str
 
 
 def main():
     configure(Client(), prefix="firedantic-test-")
-    set_up_ttl_policies(
+    set_up_composite_indexes_and_ttl_policies(
         gcloud_project="my-project",
         models=get_all_subclasses(Model),
         client=FirestoreAdminClient(),
     )
+    # or use set_up_composite_indexes / set_up_ttl_policies functions separately
 
 
 if __name__ == "__main__":
     main()
 ```
 
-### TTL Policy Example (async)
+#### Composite Index and TTL Policy Example (async)
 
 ```python
 import asyncio
 from datetime import datetime
 
 from firedantic import (
     AsyncModel,
-    async_set_up_ttl_policies,
+    async_set_up_composite_indexes_and_ttl_policies,
+    collection_index,
+    collection_group_index,
     configure,
     get_all_subclasses,
 )
-from google.cloud.firestore import AsyncClient
+from google.cloud.firestore import AsyncClient, Query
 from google.cloud.firestore_admin_v1.services.firestore_admin import (
     FirestoreAdminAsyncClient,
 )
@@ -255,18 +294,23 @@ from google.cloud.firestore_admin_v1.services.firestore_admin import (
 class ExpiringModel(AsyncModel):
     __collection__ = "expiringModel"
     __ttl_field__ = "expire"
+    __composite_indexes__ = [
+        collection_index(("content", Query.ASCENDING), ("expire", Query.DESCENDING)),
+        collection_group_index(("content", Query.DESCENDING), ("expire", Query.ASCENDING)),
+    ]
 
     expire: datetime
     content: str
 
 
 async def main():
     configure(AsyncClient(), prefix="firedantic-test-")
-    await async_set_up_ttl_policies(
+    await async_set_up_composite_indexes_and_ttl_policies(
         gcloud_project="my-project",
         models=get_all_subclasses(AsyncModel),
         client=FirestoreAdminAsyncClient(),
     )
+    # or await async_set_up_composite_indexes / async_set_up_ttl_policies separately
 
 
 if __name__ == "__main__":

firedantic/__init__.py

@@ -1,5 +1,11 @@
 # flake8: noqa
 from firedantic._async.helpers import truncate_collection as async_truncate_collection
+from firedantic._async.indexes import (
+    set_up_composite_indexes as async_set_up_composite_indexes,
+)
+from firedantic._async.indexes import (
+    set_up_composite_indexes_and_ttl_policies as async_set_up_composite_indexes_and_ttl_policies,
+)
 from firedantic._async.model import (
     AsyncBareModel,
     AsyncBareSubCollection,
@@ -12,6 +18,10 @@
     set_up_ttl_policies as async_set_up_ttl_policies,
 )
 from firedantic._sync.helpers import truncate_collection
+from firedantic._sync.indexes import (
+    set_up_composite_indexes,
+    set_up_composite_indexes_and_ttl_policies,
+)
 from firedantic._sync.model import (
     BareModel,
     BareSubCollection,
@@ -21,6 +31,7 @@
     SubModel,
 )
 from firedantic._sync.ttl_policy import set_up_ttl_policies
+from firedantic.common import collection_group_index, collection_index
 from firedantic.configurations import CONFIGURATIONS, configure
 from firedantic.exceptions import *
 from firedantic.utils import get_all_subclasses

firedantic/_async/indexes.py

@@ -0,0 +1,144 @@
+from logging import getLogger
+from typing import Iterable, List, Optional, Set, Type
+
+from google.api_core.operation_async import AsyncOperation
+from google.cloud.firestore_admin_v1 import (
+    CreateIndexRequest,
+    Index,
+    ListIndexesRequest,
+)
+from google.cloud.firestore_admin_v1.services.firestore_admin import (
+    FirestoreAdminAsyncClient,
+)
+
+from firedantic._async.model import AsyncBareModel
+from firedantic._async.ttl_policy import set_up_ttl_policies
+from firedantic.common import IndexDefinition, IndexField
+
+logger = getLogger("firedantic")
+
+
+async def get_existing_indexes(
+    client: FirestoreAdminAsyncClient, path: str
+) -> Set[IndexDefinition]:
+    """
+    Get existing database indexes and return a set of them
+    for easy comparison with other indexes
+
+    :param client: The Firestore admin client.
+    :param path: Index path in Firestore.
+    :return: Set of IndexDef tuples
+    """
+    raw_indexes = []
+    request = ListIndexesRequest({"parent": path})
+    operation = await client.list_indexes(request=request)
+    async for page in operation.pages:
+        raw_indexes.extend(list(page.indexes))
+
+    indexes = set()
+    for raw_index in raw_indexes:
+        # apparently `list_indexes` returns all indexes in all collections
+        if not raw_index.name.startswith(path):
+            continue
+        query_scope = raw_index.query_scope.name
+        fields = tuple(
+            IndexField(name=f.field_path, order=f.order.name)  # noqa
+            for f in raw_index.fields
+            if f.field_path != "__name__"
+        )
+        indexes.add(IndexDefinition(query_scope=query_scope, fields=fields))
+    return indexes
+
+
+async def create_composite_index(
+    client: FirestoreAdminAsyncClient,
+    index: IndexDefinition,
+    path: str,
+) -> AsyncOperation:
+    """
+    Create a composite index in Firestore
+
+    :param client: The Firestore admin client.
+    :param index: Index definition.
+    :param path: Index path in Firestore.
+    :return: Operation that was launched to create the index.
+    """
+    request = CreateIndexRequest(
+        {
+            "parent": path,
+            "index": Index(
+                {
+                    "query_scope": index.query_scope,
+                    "fields": [
+                        {"field_path": field[0], "order": field[1]}
+                        for field in list(index.fields)
+                    ],
+                }
+            ),
+        }
+    )
+    return await client.create_index(request=request)
+
+
+async def set_up_composite_indexes(
+    gcloud_project: str,
+    models: Iterable[Type[AsyncBareModel]],
+    database: str = "(default)",
+    client: Optional[FirestoreAdminAsyncClient] = None,
+) -> List[AsyncOperation]:
+    """
+    Set up composite indexes for models.
+
+    :param gcloud_project: The technical name of the project in Google Cloud.
+    :param models: Models for which to set up composite indexes.
+    :param database: The Firestore database instance (it now supports multiple).
+    :param client: The Firestore admin client.
+    :return: List of operations that were launched to create indexes.
+    """
+    if not client:
+        client = FirestoreAdminAsyncClient()
+
+    operations = []
+    for model in models:
+        if not model.__composite_indexes__:
+            continue
+        path = (
+            f"projects/{gcloud_project}/databases/{database}/"
+            f"collectionGroups/{model.__collection__}"
+        )
+        indexes_in_db = await get_existing_indexes(client, path=path)
+        model_indexes = set(model.__composite_indexes__)
+        existing_indexes = indexes_in_db.intersection(model_indexes)
+        new_indexes = model_indexes.difference(indexes_in_db)
+
+        for index in existing_indexes:
+            log_str = "Composite index already exists in DB: %s, collection: %s"
+            logger.debug(log_str, index, model.get_collection_name())
+
+        for index in new_indexes:
+            log_str = "Creating new composite index: %s, collection: %s"
+            logger.info(log_str, index, model.get_collection_name())
+            operation = await create_composite_index(client, index, path)
+            operations.append(operation)
+
+    return operations
+
+
+async def set_up_composite_indexes_and_ttl_policies(
+    gcloud_project: str,
+    models: Iterable[Type[AsyncBareModel]],
+    database: str = "(default)",
+    client: Optional[FirestoreAdminAsyncClient] = None,
+) -> List[AsyncOperation]:
+    """
+    Set up indexes and TTL policies that are defined in the model
+
+    :param gcloud_project: The technical name of the project in Google Cloud.
+    :param models: Models for which to set up composite indexes and TTL policies.
+    :param database: The Firestore database instance (it now supports multiple).
+    :param client: The Firestore admin client.
+    :return: List of operations that were launched.
+    """
+    ops = await set_up_composite_indexes(gcloud_project, models, database, client)
+    ops.extend(await set_up_ttl_policies(gcloud_project, models, database, client))
+    return ops

firedantic/_async/model.py

@@ -1,6 +1,6 @@
 from abc import ABC
 from logging import getLogger
-from typing import Any, Dict, List, Literal, Optional, Tuple, Type, TypeVar, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, TypeVar, Union
 
 import pydantic
 from google.cloud.firestore_v1 import (
@@ -13,6 +13,7 @@
 
 import firedantic.operators as op
 from firedantic import async_truncate_collection
+from firedantic.common import IndexDefinition, OrderDirection
 from firedantic.configurations import CONFIGURATIONS
 from firedantic.exceptions import (
     CollectionNotDefined,
@@ -62,6 +63,7 @@ class AsyncBareModel(pydantic.BaseModel, ABC):
     __collection__: Optional[str] = None
     __document_id__: str
     __ttl_field__: Optional[str] = None
+    __composite_indexes__: Optional[Iterable[IndexDefinition]]
 
     async def save(self) -> None:
         """
@@ -96,8 +98,7 @@ def get_document_id(self):
             self._validate_document_id(doc_id)
         return getattr(self, self.__document_id__, None)
 
-    _OrderDirection = Union[Literal["ASCENDING"], Literal["DESCENDING"]]
-    _OrderBy = List[Tuple[str, _OrderDirection]]
+    _OrderBy = List[Tuple[str, OrderDirection]]
 
     @classmethod
     async def find(