add support for get_by_ids and count

README.md

@@ -92,11 +92,49 @@ Product.find({"stock": {">=": 3}})
 Product.find({"stock": {op.GTE: 3}})
 Product.find({"stock": {">=": 1}}, order_by=[('unit_value', Query.ASCENDING)], limit=25, offset=50)
 Product.find(order_by=[('unit_value', Query.ASCENDING), ('stock', Query.DESCENDING)], limit=2)
+
+# Count documents, with or without filters
+Product.count()
+Product.count({"stock": {op.GTE: 3}})
+
+# Fetch multiple documents by Firestore document ID in one batch request
+Product.get_by_ids(["firestore-doc-id-1", "firestore-doc-id-2"])
 ```
 
 The query operators are found at
 [https://firebase.google.com/docs/firestore/query-data/queries#query_operators](https://firebase.google.com/docs/firestore/query-data/queries#query_operators).
 
+If your model customizes the document ID field by extending `BareModel`, use
+`get_by_doc_ids()` instead of `get_by_ids()`.
+
+### Atomic increments
+
+Numeric fields can be incremented atomically without first reloading and saving the
+whole model. Pass a negative amount to decrement the value.
+
+```python
+from firedantic import Model
+
+
+class Product(Model):
+    __collection__ = "products"
+    product_id: str
+    stock: int
+
+
+product = Product(product_id="abc-123", stock=10)
+product.save()
+
+product.increment("stock", 5)
+assert product.stock == 15
+
+product.increment("stock", -3)
+assert product.stock == 12
+```
+
+This uses Firestore's atomic increment operation. If the field does not exist or the
+current field value is not numeric, Firestore sets the field to the given amount.
+
 ### Async usage
 
 Firedantic can also be used in an async way, like this:
@@ -126,12 +164,16 @@ configure(client, prefix="firedantic-test-")
 class Person(AsyncModel):
     __collection__ = "persons"
     name: str
+    login_count: int = 0
 
 
 async def main():
     alice = Person(name="Alice")
     await alice.save()
     print(f"Saved Alice as {alice.id}")
+    await alice.increment("login_count")
+    print(f"Alice has logged in {alice.login_count} time(s)")
+
     bob = Person(name="Bob")
     await bob.save()
     print(f"Saved Bob as {bob.id}")
@@ -144,6 +186,12 @@ async def main():
     assert bob.id == found_bob.id
     print(f"Found Bob: {found_bob.id}")
 
+    people = await Person.get_by_ids([alice.id, bob.id])
+    print(f"Found {len(people)} people")
+
+    total_people = await Person.count()
+    print(f"Total people: {total_people}")
+
     await alice.delete()
     print("Deleted Alice")
     await bob.delete()
@@ -327,13 +375,19 @@ Firedantic has basic support for
 The following methods can be used in a transaction:
 
 - `Model.delete(transaction=transaction)`
+- `Model.count(transaction=transaction)`
 - `Model.find_one(transaction=transaction)`
 - `Model.find(transaction=transaction)`
 - `Model.get_by_doc_id(transaction=transaction)`
+- `Model.get_by_doc_ids(transaction=transaction)`
 - `Model.get_by_id(transaction=transaction)`
+- `Model.get_by_ids(transaction=transaction)`
+- `Model.increment("field", amount, transaction=transaction)`
 - `Model.reload(transaction=transaction)`
 - `Model.save(transaction=transaction)`
 - `SubModel.get_by_id(transaction=transaction)`
+- `SubModel.get_by_ids(transaction=transaction)`
+- `SubModel.increment("field", amount, transaction=transaction)`
 
 When using transactions, note that read operations must come before write operations.
 

firebase.json

@@ -5,7 +5,7 @@
     },
     "ui": {
       "enabled": true,
-      "port": 4000
+      "port": 4343
     }
   }
 }

firedantic/_async/model.py

@@ -8,6 +8,7 @@
     AsyncDocumentReference,
     DocumentSnapshot,
     FieldFilter,
+    Increment,
 )
 from google.cloud.firestore_v1.async_query import AsyncQuery
 from google.cloud.firestore_v1.async_transaction import AsyncTransaction
@@ -24,6 +25,7 @@
 
 TAsyncBareModel = TypeVar("TAsyncBareModel", bound="AsyncBareModel")
 TAsyncBareSubModel = TypeVar("TAsyncBareSubModel", bound="AsyncBareSubModel")
+Number = Union[int, float]
 logger = getLogger("firedantic")
 
 # https://firebase.google.com/docs/firestore/query-data/queries#query_operators
@@ -100,6 +102,49 @@ async def save(
             await doc_ref.set(data)
         setattr(self, self.__document_id__, doc_ref.id)
 
+    async def increment(
+        self,
+        field: str,
+        amount: Number = 1,
+        transaction: Optional[AsyncTransaction] = None,
+    ) -> None:
+        """
+        Atomically increments or decrements a numeric field in the database.
+
+        :param field: Firestore field path to update.
+        :param amount: Amount to increment by. Use a negative value to decrement.
+        :param transaction: Optional transaction to use.
+        """
+        doc_ref = self._get_doc_ref()
+        data = {field: Increment(amount)}
+        if transaction is not None:
+            transaction.update(doc_ref, data)
+        else:
+            await doc_ref.update(data)
+        self._increment_field_locally(field, amount)
+
+    def _increment_field_locally(self, field: str, amount: Number) -> None:
+        if "." in field:
+            return
+
+        model_field = field
+        if model_field not in self.__class__.model_fields:
+            for field_name, field_info in self.__class__.model_fields.items():
+                if field_info.alias == field:
+                    model_field = field_name
+                    break
+            else:
+                return
+
+        current_value = getattr(self, model_field, None)
+        if isinstance(current_value, (int, float)) and not isinstance(
+            current_value, bool
+        ):
+            value = current_value + amount
+        else:
+            value = amount
+        setattr(self, model_field, value)
+
     async def delete(self, transaction: Optional[AsyncTransaction] = None) -> None:
         """
         Deletes this model from the database.
@@ -181,28 +226,59 @@ async def find(  # pylint: disable=too-many-arguments
         if offset is not None:
             query = query.offset(offset)  # type: ignore
 
-        def _cls(doc_id: str, data: Dict[str, Any]) -> TAsyncBareModel:
-            if cls.__document_id__ in data:
-                logger.warning(
-                    "%s document ID %s contains conflicting %s in data with value %s",
-                    cls.__name__,
-                    doc_id,
-                    cls.__document_id__,
-                    data[cls.__document_id__],
-                )
-            data[cls.__document_id__] = doc_id
-            model = cls(**data)
-            setattr(model, cls.__document_id__, doc_id)
-            return model
-
         return [
-            _cls(doc_id, doc_dict)
+            cls._create_from_doc_data(doc_id, doc_dict)
             async for doc_id, doc_dict in (
                 (doc.id, doc.to_dict()) async for doc in query.stream(transaction=transaction)  # type: ignore
             )
             if doc_dict is not None
         ]
 
+    @classmethod
+    async def count(
+        cls: Type[TAsyncBareModel],
+        filter_: Optional[Dict[str, Union[str, dict]]] = None,
+        transaction: Optional[AsyncTransaction] = None,
+    ) -> int:
+        """
+        Returns the number of models matching the given filter.
+
+        Example: `Company.count()`.
+        Example: `Company.count({"owner.first_name": "John"})`.
+
+        :param filter_: The filter criteria.
+        :param transaction: Optional transaction to use.
+        :return: Number of matching models.
+        """
+        query: Union[AsyncQuery, AsyncCollectionReference] = cls._get_col_ref()
+        if filter_:
+            for key, value in filter_.items():
+                query = cls._add_filter(query, key, value)
+
+        aggregation_results = await query.count().get(transaction=transaction)  # type: ignore
+        for aggregation_result in aggregation_results:
+            if len(aggregation_result) == 0:
+                continue
+            return int(aggregation_result[0].value)
+        return 0
+
+    @classmethod
+    def _create_from_doc_data(
+        cls: Type[TAsyncBareModel], doc_id: str, data: Dict[str, Any]
+    ) -> TAsyncBareModel:
+        if cls.__document_id__ in data:
+            logger.warning(
+                "%s document ID %s contains conflicting %s in data with value %s",
+                cls.__name__,
+                doc_id,
+                cls.__document_id__,
+                data[cls.__document_id__],
+            )
+        data[cls.__document_id__] = doc_id
+        model = cls(**data)
+        setattr(model, cls.__document_id__, doc_id)
+        return model
+
     @classmethod
     def _add_filter(
         cls, query: Union[AsyncQuery, AsyncCollectionReference], field: str, value: Any
@@ -279,10 +355,47 @@ async def get_by_doc_id(
             raise ModelNotFoundError(
                 f"No '{cls.__name__}' found with {cls.__document_id__} '{doc_id}'"
             )
-        data[cls.__document_id__] = doc_id
-        model = cls(**data)
-        setattr(model, cls.__document_id__, doc_id)
-        return model
+        return cls._create_from_doc_data(doc_id, data)
+
+    @classmethod
+    async def get_by_doc_ids(
+        cls: Type[TAsyncBareModel],
+        doc_ids: Iterable[str],
+        transaction: Optional[AsyncTransaction] = None,
+    ) -> List[TAsyncBareModel]:
+        """
+        Returns models matching the given document IDs.
+
+        Missing documents are omitted from the returned list.
+
+        :param doc_ids: Document IDs to fetch.
+        :param transaction: Optional transaction to use.
+        :return: List of found models.
+        :raise ModelNotFoundError: Raised if a document ID is invalid.
+        """
+        doc_refs = []
+        for doc_id in doc_ids:
+            try:
+                cls._validate_document_id(doc_id)
+            except InvalidDocumentID as e:
+                raise ModelNotFoundError(
+                    f"No '{cls.__name__}' found with {cls.__document_id__} '{doc_id}'"
+                ) from e
+            doc_refs.append(cls._get_col_ref().document(doc_id))
+        if len(doc_refs) == 0:
+            return []
+
+        return [
+            cls._create_from_doc_data(doc_id, doc_dict)
+            async for doc_id, doc_dict in (
+                (doc.id, doc.to_dict())
+                async for doc in CONFIGURATIONS["db"].get_all(
+                    doc_refs,
+                    transaction=transaction,
+                )
+            )
+            if doc_dict is not None
+        ]
 
     @classmethod
     async def truncate_collection(cls, batch_size: int = 128) -> int:
@@ -372,6 +485,20 @@ async def get_by_id(
         """
         return await cls.get_by_doc_id(id_, transaction=transaction)
 
+    @classmethod
+    async def get_by_ids(
+        cls: Type[TAsyncBareModel],
+        ids: Iterable[str],
+        transaction: Optional[AsyncTransaction] = None,
+    ) -> List[TAsyncBareModel]:
+        """
+        Get models by document IDs.
+
+        :param ids: Document IDs.
+        :param transaction: Optional transaction to use.
+        """
+        return await cls.get_by_doc_ids(ids, transaction=transaction)
+
 
 class AsyncBareSubCollection(ABC):
     __collection_tpl__: Optional[str] = None
@@ -445,6 +572,20 @@ async def get_by_id(
         """
         return await cls.get_by_doc_id(id_, transaction=transaction)
 
+    @classmethod
+    async def get_by_ids(
+        cls: Type[TAsyncBareModel],
+        ids: Iterable[str],
+        transaction: Optional[AsyncTransaction] = None,
+    ) -> List[TAsyncBareModel]:
+        """
+        Get items by document IDs.
+
+        :param ids: Document IDs.
+        :param transaction: Optional transaction to use.
+        """
+        return await cls.get_by_doc_ids(ids, transaction=transaction)
+
 
 class AsyncSubCollection(AsyncBareSubCollection, ABC):
     __document_id__ = "id"