howler-api 4.1.0.dev1389__tar.gz → 4.1.0.dev1418__tar.gz

{howler_api-4.1.0.dev1389 → howler_api-4.1.0.dev1418}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: howler-api
-Version: 4.1.0.dev1389
+Version: 4.1.0.dev1418
 Summary: Howler - API server
 License: MIT
 Keywords: howler,alerting,gc,canada,cse-cst,cse,cst,cyber,cccs

{howler_api-4.1.0.dev1389 → howler_api-4.1.0.dev1418}/howler/actions/transition.py

@@ -5,6 +5,7 @@ from howler.actions import check_hit_limit
 from howler.common.exceptions import InvalidDataException, NotFoundException
 from howler.common.loader import datastore
 from howler.common.logging import get_logger
+from howler.datastore.exceptions import VersionConflictException
 from howler.helper.workflow import Workflow, WorkflowException
 from howler.odm.models.action import VALID_TRIGGERS
 from howler.odm.models.howler_data import (
@@ -20,11 +21,39 @@ from howler.utils.list_utils import flatten_list
 OPERATION_ID = "transition"
 MAX_HITS_BASIC = 10
 MAX_HITS_ADVANCED = 1000
+MAX_VERSION_CONFLICT_ATTEMPTS = 3
 SKIP_CENTRAL_LIMIT = True  # This operation transforms the query, handles limit check locally
 
 log = get_logger(__file__)
 
 
+def _transition_failure_report(hit_id: str, error: Exception) -> dict[str, str]:
+    if isinstance(error, VersionConflictException):
+        return {
+            "query": f"howler.id:{hit_id}",
+            "outcome": "error",
+            "title": "Version Conflict",
+            "message": "The hit was modified while this transition was running.",
+        }
+
+    return {
+        "query": f"howler.id:{hit_id}",
+        "outcome": "error",
+        "title": "An error occurred while processing.",
+        "message": str(error),
+    }
+
+
+def _transition_hit_with_retry(hit_id: str, transition: HitStatusTransition, user: User, **kwargs) -> None:
+    for attempt in range(MAX_VERSION_CONFLICT_ATTEMPTS):
+        try:
+            hit_service.transition_hit(hit_id, transition, user, **kwargs)
+            return
+        except VersionConflictException:
+            if attempt == MAX_VERSION_CONFLICT_ATTEMPTS - 1:
+                raise
+
+
 def __parse_workflow_actions(workflow: Workflow) -> dict[str, set[str]]:
     """Take in a workflow, and parse the steps and transitions of that workflow into a format understood by the UI"""
     parsed_args: dict[str, set[str]] = {}
@@ -126,22 +155,15 @@ def execute(
     total_processed = 0
     for hit_id in ids:
         try:
-            hit_service.transition_hit(
+            _transition_hit_with_retry(
                 hit_id,
                 cast(HitStatusTransition, HitStatusTransition[transition]),
                 user,
                 **kwargs,
             )
             success_ids.add(hit_id)
-        except (InvalidDataException, NotFoundException, WorkflowException) as e:
-            report.append(
-                {
-                    "query": f"howler.id:{hit_id}",
-                    "outcome": "error",
-                    "title": "An error occurred while processing.",
-                    "message": str(e),
-                }
-            )
+        except (InvalidDataException, NotFoundException, VersionConflictException, WorkflowException) as e:
+            report.append(_transition_failure_report(hit_id, e))
 
         total_processed += 1
         if total_processed % 10 == 0:

{howler_api-4.1.0.dev1389 → howler_api-4.1.0.dev1418}/howler/api/v1/action.py

@@ -143,7 +143,7 @@ def update_action(id: str, user: User, **kwargs) -> Response:
 
     ds = datastore()
 
-    existing_action = ds.action.get(id, as_obj=False)
+    existing_action, server_version = ds.action.get(id, as_obj=False, version=True)
 
     if not existing_action:
         return not_found(err="The specified automation does not exist")
@@ -166,7 +166,7 @@ def update_action(id: str, user: User, **kwargs) -> Response:
         action_obj = Action(updated_action)
         action_obj.action_id = id
 
-        ds.action.save(action_obj.action_id, action_obj, refresh=refresh)
+        ds.action.save(action_obj.action_id, action_obj, version=server_version, refresh=refresh)
     except HowlerException as e:
         return bad_request(err=str(e))
 

{howler_api-4.1.0.dev1389 → howler_api-4.1.0.dev1418}/howler/api/v1/hit.py

@@ -528,7 +528,7 @@ def get_assigned_hits(user, **kwargs):
 @api_login(audit=False, required_priv=["W"])
 @add_etag(getter=hit_service.get_hit, check_if_match=False)
 @parse_parameters(refresh=parse_refresh)
-def add_label(id, label_set, user, **kwargs):
+def add_label(id: str, label_set: str, user: User, server_version: str | None = None, **kwargs):
     """Add labels to a hit.
 
     Variables:
@@ -575,7 +575,8 @@ def add_label(id, label_set, user, **kwargs):
     hit_service.update_hit(
         id,
         [hit_helper.list_add(f"howler.labels.{label_set}", label) for label in labels],
-        user["uname"],
+        user.uname,
+        version=server_version,
         refresh="true" if DEBUG_FORCE_REFRESH else refresh,
     )
 
@@ -591,7 +592,7 @@ def add_label(id, label_set, user, **kwargs):
 @api_login(audit=False, required_priv=["W"])
 @add_etag(getter=hit_service.get_hit, check_if_match=False)
 @parse_parameters(refresh=parse_refresh)
-def remove_labels(id, label_set, user, **kwargs):
+def remove_labels(id: str, label_set: str, user: User, server_version: str | None = None, **kwargs):
     """Remove labels from a hit.
 
     Variables:
@@ -632,7 +633,8 @@ def remove_labels(id, label_set, user, **kwargs):
     hit_service.update_hit(
         id,
         [hit_helper.list_remove(f"howler.labels.{label_set}", label) for label in labels],
-        user["uname"],
+        user.uname,
+        version=server_version,
         refresh="true" if DEBUG_FORCE_REFRESH else refresh,
     )
 
@@ -648,7 +650,7 @@ def remove_labels(id, label_set, user, **kwargs):
 @api_login(audit=False, required_priv=["W"])
 @add_etag(getter=hit_service.get_hit, check_if_match=True)
 @parse_parameters(refresh=parse_refresh)
-def transition(id: str, user: User, **kwargs):
+def transition(id: str, user: User, server_version: str | None = None, **kwargs):
     """Transition a hit
 
     Variables:
@@ -679,11 +681,12 @@ def transition(id: str, user: User, **kwargs):
         return bad_request(err="Invalid data format")
 
     transition = transition_data["transition"]
+    version: str | None = None
     if "If-Match" in request.headers:
         version = request.headers["If-Match"]
     else:
         logger.warning("User is mising version - no If-Match header in request.")
-        version = None
+        version = server_version
 
     try:
         if transition not in HitStatusTransition.list():
@@ -694,7 +697,7 @@ def transition(id: str, user: User, **kwargs):
                 )
             )
 
-        hit_service.transition_hit(
+        hit, version = hit_service.transition_hit(
             id, transition, user, version, refresh=refresh, **kwargs, **transition_data.get("data", {})
         )
     except (WorkflowException, DataStoreException, InvalidDataException) as e:
@@ -704,7 +707,6 @@ def transition(id: str, user: User, **kwargs):
     except HowlerException as e:
         return internal_error(err=str(e))
 
-    hit, version = hit_service.get_hit(id, as_odm=False, version=True)
     return ok(hit), version
 
 
@@ -712,7 +714,7 @@ def transition(id: str, user: User, **kwargs):
 @hit_api.route("/<id>/comments/<comment_id>", methods=["GET"])
 @api_login(audit=False, required_priv=["R"])
 @add_etag(getter=hit_service.get_hit, check_if_match=False)
-def get_comment(id: str, comment_id: str, user: User, server_version: str, **kwargs):
+def get_comment(id: str, comment_id: str, user: User, server_version: str | None = None, **kwargs):
     """Get a comment associated with a particular hit
 
     Variables:
@@ -741,7 +743,7 @@ def get_comment(id: str, comment_id: str, user: User, server_version: str, **kwa
 @hit_api.route("/<id>/comments", methods=["POST"])
 @api_login(audit=False, required_priv=["W"])
 @add_etag(getter=hit_service.get_hit, check_if_match=False)
-def add_comment(id: str, user: dict[str, Any], **kwargs):
+def add_comment(id: str, user: User, server_version: str | None = None, **kwargs):
     """Add a comment
 
     Variables:
@@ -781,12 +783,13 @@ def add_comment(id: str, user: dict[str, Any], **kwargs):
             [
                 hit_helper.list_add(
                     "howler.comment",
-                    Comment({"user": user["uname"], "value": comment_value}),
+                    Comment({"user": user.uname, "value": comment_value}),
                     explanation=f"Added a comment:\n\n{comment_value}",
                     if_missing=True,
                 ),
             ],
-            user["uname"],
+            user.uname,
+            version=server_version,
         )
     except DataStoreException as e:
         return bad_request(err=str(e))
@@ -800,7 +803,7 @@ def add_comment(id: str, user: dict[str, Any], **kwargs):
 @hit_api.route("/<id>/comments/<comment_id>", methods=["PUT"])
 @api_login(audit=False, required_priv=["W"])
 @add_etag(getter=hit_service.get_hit, check_if_match=False)
-def edit_comment(id: str, comment_id: str, user: dict[str, Any], **kwargs):
+def edit_comment(id: str, comment_id: str, user: dict[str, Any], server_version: str | None = None, **kwargs):
     """Edit a comment
 
     Variables:
@@ -865,6 +868,7 @@ def edit_comment(id: str, comment_id: str, user: dict[str, Any], **kwargs):
             ),
         ],
         user["uname"],
+        version=server_version,
     )
 
     return ok(new_hit), version
@@ -874,7 +878,7 @@ def edit_comment(id: str, comment_id: str, user: dict[str, Any], **kwargs):
 @hit_api.route("/<id>/comments", methods=["DELETE"])
 @api_login(audit=False, required_priv=["W"])
 @add_etag(getter=hit_service.get_hit, check_if_match=False)
-def delete_comments(id: str, user: User, **kwargs):
+def delete_comments(id: str, user: User, server_version: str | None = None, **kwargs):
     """Delete a set of comments
 
     Variables:
@@ -923,6 +927,7 @@ def delete_comments(id: str, user: User, **kwargs):
                 for comment in comments
             ],
             user.uname,
+            version=server_version,
         )
 
     except DataStoreException as e:

{howler_api-4.1.0.dev1389 → howler_api-4.1.0.dev1418}/howler/api/v2/ingest.py

@@ -265,7 +265,7 @@ def overwrite(index: str, id: str, **kwargs):
 
     ds = datastore()
 
-    record = ds[index].get(id, as_obj=False, version=False)
+    record, server_version = ds[index].get(id, as_obj=False, version=True)
     if not record:
         return not_found(err="Record %s does not exist" % id)
 
@@ -292,7 +292,7 @@ def overwrite(index: str, id: str, **kwargs):
         ds[index].save(
             id,
             odm(new_record) if odm else new_record,
-            version=kwargs.get("server_version"),
+            version=server_version,
             refresh=refresh,
         )
 

{howler_api-4.1.0.dev1389 → howler_api-4.1.0.dev1418}/howler/datastore/collection.py

@@ -16,6 +16,7 @@ from typing import Any, Callable, Dict, Generic, Literal, Optional, TypeVar, Uni
 import elasticsearch
 from datemath import dm
 from datemath.helpers import DateMathException
+from elastic_transport import ApiResponseMeta
 from opentelemetry import trace
 
 from howler import odm
@@ -1136,13 +1137,46 @@ class ESCollection(Generic[ModelType]):
 
         return data
 
+    def _search_exists(self, key) -> bool:
+        """Check document existence with an alias-safe search query."""
+        result = self.with_retries(
+            self.datastore.client.search,
+            index=self.name,
+            query={"ids": {"values": [key]}},
+            size=0,
+            track_total_hits=True,
+        )
+        total = result["hits"]["total"]
+        return total["value"] > 0 if isinstance(total, dict) else total > 0
+
     def exists(self, key) -> bool:
         """Check if a document exists in the datastore.
 
         :param key: key of the document to get from the datastore
         :return: true/false depending if the document exists or not
         """
-        return cast(bool, self.with_retries(self.datastore.client.exists, index=self.name, id=key, _source=False))
+        if self.ilm_config:
+            return self._search_exists(key)
+
+        try:
+            return cast(bool, self.with_retries(self.datastore.client.exists, index=self.name, id=key, _source=False))
+        except elasticsearch.BadRequestError:
+            logger.exception(
+                "Single-document existence check failed on %s; falling back to an alias-safe search.",
+                self.name,
+            )
+            return self._search_exists(key)
+
+    def _raise_document_not_found(self, key: str, result: Any) -> typing.NoReturn:
+        """Raise the same error shape as an Elasticsearch document lookup for an empty search."""
+        meta = ApiResponseMeta(
+            status=404,
+            http_version="1.1",
+            headers=cast(Any, {}),
+            duration=0.0,
+            node=cast(Any, None),
+        )
+        raise elasticsearch.exceptions.NotFoundError(f"Document with id {key} not found", meta, result)
 
     @overload
     def _get(self, key, retries, version: Literal[False]) -> dict[str, Any]: ...
@@ -1154,7 +1188,7 @@ class ESCollection(Generic[ModelType]):
         """Versioned get-save for atomic update has two paths:
             1. Document doesn't exist at all. Create token will be returned for version.
                This way only the first query to try and create the document will succeed.
-            2. Document exists in hot. A version string with the info needed to do a versioned save is returned.
+            2. Document exists. A version string with the info needed to do a versioned save is returned.
 
         The create token is needed to differentiate between "I'm saving a new
         document non-atomic (version=None)" and "I'm saving a new document
@@ -1173,11 +1207,29 @@ class ESCollection(Generic[ModelType]):
         done = False
         while not done:
             try:
-                doc = self.with_retries(self.datastore.client.get, index=self.name, id=key)
+                if self.ilm_config:
+                    result = self.with_retries(
+                        self.datastore.client.search,
+                        index=self.name,
+                        query={"ids": {"values": [key]}},
+                        size=1,
+                        seq_no_primary_term=True,
+                    )
+                    hits = result["hits"]["hits"]
+                    if not hits:
+                        self._raise_document_not_found(key, result)
+                    doc = hits[0]
+                else:
+                    doc = self.with_retries(self.datastore.client.get, index=self.name, id=key)
+
                 if version:
+                    if self.ilm_config:
+                        version_token = f"{doc['_index']}---{doc['_seq_no']}---{doc['_primary_term']}"
+                    else:
+                        version_token = f"{doc['_seq_no']}---{doc['_primary_term']}"
                     return (
                         normalize_output(doc["_source"]),
-                        f"{doc['_seq_no']}---{doc['_primary_term']}",
+                        version_token,
                     )
                 return normalize_output(doc["_source"])
             except elasticsearch.exceptions.NotFoundError:
@@ -1196,6 +1248,19 @@ class ESCollection(Generic[ModelType]):
 
         return None
 
+    def _get_version_write_target(self, version: str) -> tuple[str, str, str]:
+        """Return the concrete write target and optimistic-concurrency values from a version token."""
+        version_parts = version.split("---")
+        if len(version_parts) == 3:
+            index, seq_no, primary_term = version_parts
+            return index, seq_no, primary_term
+
+        if len(version_parts) == 2:
+            seq_no, primary_term = version_parts
+            return self.name, seq_no, primary_term
+
+        raise DataStoreException(f"Invalid version token for {self.name}: {version!r}")
+
     @overload
     def get(self, key, as_obj: Literal[True], version: Literal[True]) -> tuple[ModelType | None, str]: ...
 
@@ -1325,13 +1390,17 @@ class ESCollection(Generic[ModelType]):
 
         saved_data["id"] = key
         operation = "index"
+        index = self.name
         seq_no = None
         primary_term = None
 
+        if version is None and self.ilm_config:
+            _, version = self.get_if_exists(key, as_obj=False, version=True)
+
         if version == CREATE_TOKEN:
             operation = "create"
         elif version:
-            seq_no, primary_term = version.split("---")
+            index, seq_no, primary_term = self._get_version_write_target(version)
 
         if refresh == "true":
             logger.warning(
@@ -1341,7 +1410,7 @@ class ESCollection(Generic[ModelType]):
         try:
             self.with_retries(
                 self.datastore.client.index,
-                index=self.name,
+                index=index,
                 id=key,
                 document=json.dumps(saved_data),
                 op_type=operation,
@@ -1541,15 +1610,20 @@ class ESCollection(Generic[ModelType]):
         """
         operations = self._validate_operations(operations)
         script = self._create_scripts_from_operations(operations)
+        index = self.name
         seq_no = None
         primary_term = None
+
+        if version is None and self.ilm_config:
+            _, version = self.get_if_exists(key, as_obj=False, version=True)
+
         if version:
-            seq_no, primary_term = version.split("---")
+            index, seq_no, primary_term = self._get_version_write_target(version)
 
         try:
             res = self.with_retries(
                 self.datastore.client.update,
-                index=self.name,
+                index=index,
                 id=key,
                 script=script,
                 if_seq_no=seq_no,
@@ -1559,7 +1633,11 @@ class ESCollection(Generic[ModelType]):
             )
             return (
                 res["result"] == "updated",
-                f"{res['_seq_no']}---{res['_primary_term']}",
+                (
+                    f"{res['_index']}---{res['_seq_no']}---{res['_primary_term']}"
+                    if self.ilm_config
+                    else f"{res['_seq_no']}---{res['_primary_term']}"
+                ),
             )
         except elasticsearch.NotFoundError as e:
             logger.warning("Update - elasticsearch.NotFoundError: %s %s", e.message, e.info)

howler_api-4.1.0.dev1418/howler/external/check_duplicates.py

@@ -0,0 +1,152 @@
+"""Find and optionally remove duplicate documents from an ILM collection."""
+
+import argparse
+import heapq
+import sys
+from collections.abc import Iterator
+from typing import Any
+
+
+def iter_index_ids(collection: Any, index: str, batch_size: int) -> Iterator[str]:
+    """Yield document IDs from *index* in sorted order without retaining prior pages."""
+    search_after = None
+
+    while True:
+        search_args = {
+            "index": index,
+            "query": {"match_all": {}},
+            "sort": [{"id": "asc"}],
+            "size": batch_size,
+            "_source": False,
+            "track_total_hits": False,
+        }
+        if search_after is not None:
+            search_args["search_after"] = search_after
+
+        response = collection.with_retries(collection.datastore.client.search, **search_args)
+        hits = response["hits"]["hits"]
+        if not hits:
+            return
+
+        for hit in hits:
+            yield hit["_id"]
+
+        search_after = hits[-1]["sort"]
+
+
+def _push_next_id(heap: list[tuple[str, str, Iterator[str]]], index: str, document_ids: Iterator[str]) -> None:
+    try:
+        heapq.heappush(heap, (next(document_ids), index, document_ids))
+    except StopIteration:
+        pass
+
+
+def iter_duplicates(collection: Any, indexes: list[str], batch_size: int) -> Iterator[tuple[str, list[str]]]:
+    """Yield each duplicated ID and the physical indexes that contain it.
+
+    This is a k-way merge over paginated, sorted searches. Its memory use is bounded by
+    the number of physical indexes and one Elasticsearch result page per active iterator.
+    """
+    heap: list[tuple[str, str, Iterator[str]]] = []
+    for index in indexes:
+        document_ids = iter_index_ids(collection, index, batch_size)
+        _push_next_id(heap, index, document_ids)
+
+    while heap:
+        document_id, index, document_ids = heapq.heappop(heap)
+        duplicate_indexes = [index]
+        _push_next_id(heap, index, document_ids)
+
+        while heap and heap[0][0] == document_id:
+            _, index, document_ids = heapq.heappop(heap)
+            duplicate_indexes.append(index)
+            _push_next_id(heap, index, document_ids)
+
+        if len(duplicate_indexes) > 1:
+            yield document_id, duplicate_indexes
+
+
+def delete_duplicates(collection: Any, document_id: str, indexes: list[str]) -> int:
+    """Delete all supplied physical-index copies of a document ID."""
+    deleted = 0
+    for index in indexes:
+        response = collection.with_retries(collection.datastore.client.delete, index=index, id=document_id)
+        if response["result"] == "deleted":
+            deleted += 1
+    return deleted
+
+
+def main() -> int:
+    """Run the interactive duplicate scan for an ILM collection."""
+    parser = argparse.ArgumentParser(
+        description="Find duplicate document IDs across an ILM collection's physical indexes."
+    )
+    parser.add_argument("collection", help="Collection name registered in the datastore.")
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=1_000,
+        help="IDs fetched from each physical index per Elasticsearch request (default: 1000).",
+    )
+    args = parser.parse_args()
+
+    if args.batch_size <= 0:
+        parser.error("--batch-size must be a positive integer.")
+
+    from howler.datastore.collection import ESCollection
+
+    # Checking duplicates must not create collections or alter their ILM configuration.
+    ESCollection.IGNORE_ENSURE_COLLECTION = True
+    ESCollection.ENSURE_COLLECTION_WARNED = True
+
+    from howler.common import loader
+
+    datastore = loader.datastore(archive_access=False)
+    collection_names = datastore.ds.get_models()
+    if args.collection not in collection_names:
+        parser.error(
+            f"Unknown collection '{args.collection}'. Valid collections: {', '.join(sorted(collection_names))}"
+        )
+
+    collection = getattr(datastore, args.collection)
+    if collection.ilm_config is None:
+        print(f"Collection '{args.collection}' is not ILM-enabled; nothing to check.")
+        return 0
+
+    collection._refresh_ilm_index_name()
+    indexes = collection.index_list_full
+    if len(indexes) < 2:
+        print(f"Collection '{args.collection}' has fewer than two physical ILM indexes; nothing to check.")
+        return 0
+
+    print(f"Scanning for duplicate IDs in: {', '.join(indexes)}")
+    duplicates_found = 0
+    copies_deleted = 0
+    index_order = {index: position for position, index in enumerate(indexes)}
+
+    try:
+        for document_id, duplicate_indexes in iter_duplicates(collection, indexes, args.batch_size):
+            duplicates_found += 1
+            keep_index = min(duplicate_indexes, key=index_order.__getitem__)
+            stale_indexes = [index for index in duplicate_indexes if index != keep_index]
+            print(
+                f"Duplicate ID '{document_id}' found in {', '.join(duplicate_indexes)}. "
+                f"Keeping the newest copy in '{keep_index}'."
+            )
+            answer = input(f"Delete copies from {', '.join(stale_indexes)}? [yes/NO] ")
+            if answer.lower() in {"y", "yes"}:
+                deleted_count = delete_duplicates(collection, document_id, stale_indexes)
+                copies_deleted += deleted_count
+                print(f"Deleted {deleted_count} copy/copies of '{document_id}'.")
+    except KeyboardInterrupt:
+        print("\nInterrupted; all completed deletions have been retained.", file=sys.stderr)
+        return 130
+    finally:
+        datastore.ds.close()
+
+    print(f"Scan complete. Found {duplicates_found} duplicate ID(s); deleted {copies_deleted} stale copy/copies.")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

{howler_api-4.1.0.dev1389 → howler_api-4.1.0.dev1418}/howler/odm/mixins.py

@@ -15,6 +15,7 @@ from howler.common.exceptions import HowlerRuntimeError
 from howler.common.loader import datastore
 from howler.datastore.collection import ESCollection
 from howler.odm.base import Model
+from howler.utils.compat import Self
 
 ModelType = TypeVar("ModelType", bound=Model)
 
@@ -28,10 +29,10 @@ class _ObjectsDescriptor(Generic[ModelType]):
     """
 
     @overload
-    def __get__(self, obj: None, objtype: type[ModelType]) -> ESCollection[ModelType]: ...
+    def __get__(self: Self, obj: None, objtype: type[ModelType]) -> ESCollection[ModelType]: ...
 
     @overload
-    def __get__(self, obj: ModelType, objtype: type[ModelType]) -> ESCollection[ModelType]: ...
+    def __get__(self: Self, obj: ModelType, objtype: type[ModelType]) -> ESCollection[ModelType]: ...
 
     def __get__(self, obj: ModelType | None, objtype: type[ModelType] | None = None) -> ESCollection[ModelType]:
         """Return the ESCollection for the owner class.
@@ -82,7 +83,11 @@ class DatastoreMixin(Generic[ModelType]):
         """
         return datastore()
 
-    def save(self, refresh: Literal["true", "false", "wait_for"] | None = None) -> bool:
+    def save(
+        self: Self,
+        refresh: Literal["true", "false", "wait_for"] | None = None,
+        version: str | None = None,
+    ) -> bool:
         """Persist the current model instance to the datastore.
 
         Determines the target index from the lowercase class name, extracts the
@@ -96,5 +101,6 @@ class DatastoreMixin(Generic[ModelType]):
         index_name = self.__class__.__name__.lower()
         id_field = self.__class__._Model__id_field  # type: ignore[attr-defined]
         current_id = attrgetter(id_field)(self)
+        collection = self.ds[index_name]
 
-        return self.ds[index_name].save(current_id, self, refresh=refresh)
+        return collection.save(current_id, self, version=version, refresh=refresh)