Add support for TTL policies (#44)

- Support for TTL policies and specifying a `__ttl_field__` in the model classes.
- Update dependencies.
- Update pre-commit hooks.
- Remove support for Python 3.7, which is EOL, now requires 3.8.1 or newer.
- Run tests on Python 3.11 in addition to older versions.
- Update authors (company was renamed to IOXIO Ltd).

.github/workflows/tests.yaml

@@ -10,7 +10,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.7", "3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
     steps:
       - name: Checkout 🔁
         uses: actions/checkout@v3

.pre-commit-config.yaml

@@ -2,7 +2,7 @@
 # See https://pre-commit.com/hooks.html for more hooks
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.3.0
+    rev: v4.4.0
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
@@ -28,22 +28,22 @@ repos:
         language: system
         pass_filenames: false
   - repo: https://github.com/pycqa/isort
-    rev: 5.10.1
+    rev: 5.12.0
     hooks:
       - id: isort
   - repo: https://github.com/psf/black
-    rev: 22.6.0
+    rev: 23.9.1
     hooks:
       - id: black
-  - repo: https://gitlab.com/pycqa/flake8
-    rev: 4.0.1
+  - repo: https://github.com/pycqa/flake8
+    rev: 6.1.0
     hooks:
       - id: flake8
   - repo: https://github.com/twu/skjold
-    rev: v0.5.0
+    rev: v0.6.1
     hooks:
       - id: skjold
   - repo: https://github.com/pre-commit/mirrors-prettier
-    rev: v2.7.1
+    rev: v3.0.3
     hooks:
       - id: prettier

CHANGELOG.md

@@ -7,9 +7,20 @@ this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased]
 
+## [0.4.0] - 2023-10-03
+
+### Added
+
+- Support for TTL policies and specifying a `__ttl_field__` in the model classes.
+
+### Removed
+
+- Remove support for Python 3.7, which is EOL, require 3.8.1 or newer.
+
 ### Changed
 
 - Switch tests to use new firestore emulator, improve documentation about running it.
+- Update authors (company was renamed to IOXIO Ltd).
 
 ## [0.3.0] - 2022-08-04
 
@@ -181,7 +192,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.3.0...HEAD
+[unreleased]: https://github.com/ioxiocom/firedantic/compare/0.4.0...HEAD
+[0.4.0]: https://github.com/ioxiocom/firedantic/compare/0.3.0...0.4.0
 [0.3.0]: https://github.com/ioxiocom/firedantic/compare/0.2.8...0.3.0
 [0.2.8]: https://github.com/ioxiocom/firedantic/compare/0.2.7...0.2.8
 [0.2.7]: https://github.com/ioxiocom/firedantic/compare/0.2.6...0.2.7

README.md

@@ -144,12 +144,7 @@ async def main():
 
 
 if __name__ == "__main__":
-    # Starting from Python 3.7 ->
-    # asyncio.run(main())
-
-    # Compatible with Python 3.6 ->
-    loop = asyncio.get_event_loop()
-    result = loop.run_until_complete(main())
+    asyncio.run(main())
 ```
 
 ## Subcollections
@@ -191,6 +186,89 @@ async def get_user_purchases(user_id: str, period: str = "2021") -> int:
 
 ```
 
+## TTL Policies
+
+Firedantic has support for defining TTL policies and creating the 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)
+
+```python
+from datetime import datetime
+
+from firedantic import Model, configure, get_all_subclasses, set_up_ttl_policies
+from google.cloud.firestore import Client
+from google.cloud.firestore_admin_v1 import FirestoreAdminClient
+
+
+class ExpiringModel(Model):
+    __collection__ = "expiringModel"
+    __ttl_field__ = "expire"
+
+    expire: datetime
+    content: str
+
+
+def main():
+    configure(Client(), prefix="firedantic-test-")
+    set_up_ttl_policies(
+        gcloud_project="my-project",
+        models=get_all_subclasses(Model),
+        client=FirestoreAdminClient(),
+    )
+
+
+if __name__ == "__main__":
+    main()
+```
+
+### TTL Policy Example (async)
+
+```python
+import asyncio
+from datetime import datetime
+
+from firedantic import (
+    AsyncModel,
+    async_set_up_ttl_policies,
+    configure,
+    get_all_subclasses,
+)
+from google.cloud.firestore import AsyncClient
+from google.cloud.firestore_admin_v1.services.firestore_admin import (
+    FirestoreAdminAsyncClient,
+)
+
+
+class ExpiringModel(AsyncModel):
+    __collection__ = "expiringModel"
+    __ttl_field__ = "expire"
+
+    expire: datetime
+    content: str
+
+
+async def main():
+    configure(AsyncClient(), prefix="firedantic-test-")
+    await async_set_up_ttl_policies(
+        gcloud_project="my-project",
+        models=get_all_subclasses(AsyncModel),
+        client=FirestoreAdminAsyncClient(),
+    )
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
+```
+
 ## Development
 
 PRs are welcome!

firedantic/__init__.py

@@ -8,6 +8,9 @@
     AsyncSubCollection,
     AsyncSubModel,
 )
+from firedantic._async.ttl_policy import (
+    set_up_ttl_policies as async_set_up_ttl_policies,
+)
 from firedantic._sync.helpers import truncate_collection
 from firedantic._sync.model import (
     BareModel,
@@ -17,5 +20,7 @@
     SubCollection,
     SubModel,
 )
+from firedantic._sync.ttl_policy import set_up_ttl_policies
 from firedantic.configurations import CONFIGURATIONS, configure
 from firedantic.exceptions import *
+from firedantic.utils import get_all_subclasses

firedantic/_async/helpers.py

@@ -6,6 +6,7 @@ async def truncate_collection(
 ) -> int:
     """Removes all documents inside a collection.
 
+    :param col_ref: A collection reference to the collection to be truncated.
     :param batch_size: Batch size for listing documents.
     :return: Number of removed documents.
     """

firedantic/_async/model.py

@@ -38,12 +38,16 @@
 }
 
 
-def _get_col_ref(cls, name) -> AsyncCollectionReference:
-    if name is None:
+def get_collection_name(cls, name: Optional[str]) -> str:
+    if not name:
         raise CollectionNotDefined(f"Missing collection name for {cls.__name__}")
 
+    return f"{CONFIGURATIONS['prefix']}{name}"
+
+
+def _get_col_ref(cls, name: Optional[str]) -> AsyncCollectionReference:
     collection: AsyncCollectionReference = CONFIGURATIONS["db"].collection(
-        CONFIGURATIONS["prefix"] + name
+        get_collection_name(cls, name)
     )
     return collection
 
@@ -56,6 +60,7 @@ class AsyncBareModel(pydantic.BaseModel, ABC):
 
     __collection__: Optional[str] = None
     __document_id__: str
+    __ttl_field__: Optional[str] = None
 
     async def save(self) -> None:
         """
@@ -213,6 +218,10 @@ def _get_col_ref(cls) -> AsyncCollectionReference:
         """Returns the collection reference."""
         return _get_col_ref(cls, cls.__collection__)
 
+    @classmethod
+    def get_collection_name(cls) -> str:
+        return get_collection_name(cls, cls.__collection__)
+
     def _get_doc_ref(self) -> AsyncDocumentReference:
         """
         Returns the document reference.

firedantic/_async/ttl_policy.py

@@ -0,0 +1,72 @@
+from logging import getLogger
+from typing import Iterable, List, Optional, Type
+
+from google.api_core.operation_async import AsyncOperation
+from google.cloud.firestore_admin_v1 import Field
+from google.cloud.firestore_admin_v1.services.firestore_admin import (
+    FirestoreAdminAsyncClient,
+)
+
+from firedantic._async.model import AsyncBareModel
+from firedantic.utils import remove_prefix
+
+logger = getLogger("firedantic")
+
+
+async def set_up_ttl_policies(
+    gcloud_project: str,
+    models: Iterable[Type[AsyncBareModel]],
+    database: str = "(default)",
+    client: Optional[FirestoreAdminAsyncClient] = None,
+) -> List[AsyncOperation]:
+    """
+    Set up TTL policies for models.
+
+    :param gcloud_project: The technical name of the project in Google Cloud.
+    :param models: Models for which to set up the TTL policy.
+    :param database: The Firestore database instance (it now supports multiple).
+    :param client: The Firestore admin client.
+    :return: List of operations that were launched to enable the policies.
+    """
+    if not client:
+        client = FirestoreAdminAsyncClient()
+
+    operations = []
+    for model in models:
+        if not model.__ttl_field__:
+            continue
+
+        # Get current details of the field
+        path = client.field_path(
+            project=gcloud_project,
+            database=database,
+            collection=model.get_collection_name(),
+            field=model.__ttl_field__,
+        )
+        field_obj = await client.get_field({"name": path})
+
+        # Variables for logging
+        readable_state = remove_prefix(str(field_obj.ttl_config.state), "State.")
+        log_str = '"%s", collection: "%s", field: "%s", state: "%s"'
+        log_params = [
+            model.__class__.__name__,
+            model.get_collection_name(),
+            model.__ttl_field__,
+            readable_state,
+        ]
+
+        if field_obj.ttl_config.state == Field.TtlConfig.State.STATE_UNSPECIFIED:
+            logger.info("Setting up new TTL config: " + log_str, *log_params)
+            field_obj.ttl_config = Field.TtlConfig(
+                {"state": Field.TtlConfig.State.CREATING}
+            )
+            operation = await client.update_field({"field": field_obj})
+            operations.append(operation)
+        elif field_obj.ttl_config.state == Field.TtlConfig.State.CREATING:
+            logger.info("TTL config is still being created: " + log_str, *log_params)
+        elif field_obj.ttl_config.state == Field.TtlConfig.State.NEEDS_REPAIR:
+            logger.error("TTL config needs repair: " + log_str, *log_params)
+        elif field_obj.ttl_config.state == Field.TtlConfig.State.ACTIVE:
+            logger.debug("TTL config is active: " + log_str, *log_params)
+
+    return operations

firedantic/_sync/helpers.py

@@ -4,6 +4,7 @@
 def truncate_collection(col_ref: CollectionReference, batch_size: int = 128) -> int:
     """Removes all documents inside a collection.
 
+    :param col_ref: A collection reference to the collection to be truncated.
     :param batch_size: Batch size for listing documents.
     :return: Number of removed documents.
     """

firedantic/_sync/model.py

@@ -38,12 +38,16 @@
 }
 
 
-def _get_col_ref(cls, name) -> CollectionReference:
-    if name is None:
+def get_collection_name(cls, name: Optional[str]) -> str:
+    if not name:
         raise CollectionNotDefined(f"Missing collection name for {cls.__name__}")
 
+    return f"{CONFIGURATIONS['prefix']}{name}"
+
+
+def _get_col_ref(cls, name: Optional[str]) -> CollectionReference:
     collection: CollectionReference = CONFIGURATIONS["db"].collection(
-        CONFIGURATIONS["prefix"] + name
+        get_collection_name(cls, name)
     )
     return collection
 
@@ -56,6 +60,7 @@ class BareModel(pydantic.BaseModel, ABC):
 
     __collection__: Optional[str] = None
     __document_id__: str
+    __ttl_field__: Optional[str] = None
 
     def save(self) -> None:
         """
@@ -209,6 +214,10 @@ def _get_col_ref(cls) -> CollectionReference:
         """Returns the collection reference."""
         return _get_col_ref(cls, cls.__collection__)
 
+    @classmethod
+    def get_collection_name(cls) -> str:
+        return get_collection_name(cls, cls.__collection__)
+
     def _get_doc_ref(self) -> DocumentReference:
         """
         Returns the document reference.