Document insertion and retrieval

arangoasync/collection.py

@@ -1,19 +1,24 @@
 __all__ = ["Collection", "StandardCollection"]
 
 
-from typing import Generic, Optional, Tuple, TypeVar
+from typing import Generic, Optional, Tuple, TypeVar, cast
 
-from arangoasync.errno import HTTP_NOT_FOUND, HTTP_PRECONDITION_FAILED
+from arangoasync.errno import (
+    HTTP_BAD_PARAMETER,
+    HTTP_NOT_FOUND,
+    HTTP_PRECONDITION_FAILED,
+)
 from arangoasync.exceptions import (
     DocumentGetError,
+    DocumentInsertError,
     DocumentParseError,
     DocumentRevisionError,
 )
 from arangoasync.executor import ApiExecutor
 from arangoasync.request import Method, Request
 from arangoasync.response import Response
 from arangoasync.serialization import Deserializer, Serializer
-from arangoasync.typings import Json, Result
+from arangoasync.typings import Json, Params, Result
 
 T = TypeVar("T")
 U = TypeVar("U")
@@ -83,6 +88,21 @@ def _extract_id(self, body: Json) -> str:
         except KeyError:
             raise DocumentParseError('Field "_key" or "_id" required')
 
+    def _ensure_key_from_id(self, body: Json) -> Json:
+        """Return the body with "_key" field if it has "_id" field.
+
+        Args:
+            body (dict): Document body.
+
+        Returns:
+            dict: Document body with "_key" field if it has "_id" field.
+        """
+        if "_id" in body and "_key" not in body:
+            doc_id = self._validate_id(body["_id"])
+            body = body.copy()
+            body["_key"] = doc_id[len(self._id_prefix) :]
+        return body
+
     def _prep_from_doc(
         self,
         document: str | Json,
@@ -172,7 +192,10 @@ async def get(
         Raises:
             DocumentRevisionError: If the revision is incorrect.
             DocumentGetError: If retrieval fails.
-        """
+
+        References:
+            - `get-a-document <https://docs.arangodb.com/stable/develop/http-api/documents/#get-a-document>`__
+        """  # noqa: E501
         handle, headers = self._prep_from_doc(document, rev, check_rev)
 
         if allow_dirty_read:
@@ -195,3 +218,102 @@ def response_handler(resp: Response) -> Optional[U]:
                 raise DocumentGetError(resp, request)
 
         return await self._executor.execute(request, response_handler)
+
+    async def insert(
+        self,
+        document: T,
+        wait_for_sync: Optional[bool] = None,
+        return_new: Optional[bool] = None,
+        return_old: Optional[bool] = None,
+        silent: Optional[bool] = None,
+        overwrite: Optional[bool] = None,
+        overwrite_mode: Optional[str] = None,
+        keep_null: Optional[bool] = None,
+        merge_objects: Optional[bool] = None,
+        refill_index_caches: Optional[bool] = None,
+        version_attribute: Optional[str] = None,
+    ) -> Result[bool | Json]:
+        """Insert a new document.
+
+        Args:
+            document (dict): Document to insert. If it contains the "_key" or "_id"
+                field, the value is used as the key of the new document (otherwise
+                it is auto-generated). Any "_rev" field is ignored.
+            wait_for_sync (bool | None): Wait until document has been synced to disk.
+            return_new (bool | None): Additionally return the complete new document
+                under the attribute `new` in the result.
+            return_old (bool | None): Additionally return the complete old document
+                under the attribute `old` in the result. Only available if the
+                `overwrite` option is used.
+            silent (bool | None): If set to `True`, no document metadata is returned.
+                This can be used to save resources.
+            overwrite (bool | None): If set to `True`, operation does not fail on
+                duplicate key and existing document is overwritten (replace-insert).
+            overwrite_mode (str | None): Overwrite mode. Supersedes **overwrite**
+                option. May be one of "ignore", "replace", "update" or "conflict".
+            keep_null (bool | None): If set to `True`, fields with value None are
+                retained in the document. Otherwise, they are removed completely.
+                Applies only when **overwrite_mode** is set to "update"
+                (update-insert).
+            merge_objects (bool | None): If set to True, sub-dictionaries are merged
+                instead of the new one overwriting the old one. Applies only when
+                **overwrite_mode** is set to "update" (update-insert).
+            refill_index_caches (bool | None): Whether to add new entries to
+                in-memory index caches if document insertions affect the edge index
+                or cache-enabled persistent indexes.
+            version_attribute (str | None): Support for simple external versioning to
+                document operations. Only applicable if **overwrite** is set to true
+                or **overwriteMode** is set to "update" or "replace".
+
+        References:
+            - `create-a-document <https://docs.arangodb.com/stable/develop/http-api/documents/#create-a-document>`__
+        """  # noqa: E501
+        if isinstance(document, dict):
+            # We assume that the document deserializer works with dictionaries.
+            document = cast(T, self._ensure_key_from_id(document))
+
+        params: Params = {}
+        if wait_for_sync is not None:
+            params["waitForSync"] = wait_for_sync
+        if return_new is not None:
+            params["returnNew"] = return_new
+        if return_old is not None:
+            params["returnOld"] = return_old
+        if silent is not None:
+            params["silent"] = silent
+        if overwrite is not None:
+            params["overwrite"] = overwrite
+        if overwrite_mode is not None:
+            params["overwriteMode"] = overwrite_mode
+        if keep_null is not None:
+            params["keepNull"] = keep_null
+        if merge_objects is not None:
+            params["mergeObjects"] = merge_objects
+        if refill_index_caches is not None:
+            params["refillIndexCaches"] = refill_index_caches
+        if version_attribute is not None:
+            params["versionAttribute"] = version_attribute
+
+        request = Request(
+            method=Method.POST,
+            endpoint=f"/_api/document/{self._name}",
+            params=params,
+            data=self._doc_serializer.dumps(document),
+        )
+
+        def response_handler(resp: Response) -> bool | Json:
+            if resp.is_success:
+                if silent is True:
+                    return True
+                return self._executor.deserialize(resp.raw_body)
+            msg: Optional[str] = None
+            if resp.status_code == HTTP_BAD_PARAMETER:
+                msg = (
+                    "Body does not contain a valid JSON representation of "
+                    "one document."
+                )
+            elif resp.status_code == HTTP_NOT_FOUND:
+                msg = "Collection not found."
+            raise DocumentInsertError(resp, request, msg)
+
+        return await self._executor.execute(request, response_handler)

arangoasync/exceptions.py

@@ -50,7 +50,10 @@ class ArangoServerError(ArangoError):
     def __init__(
         self, resp: Response, request: Request, msg: Optional[str] = None
     ) -> None:
-        msg = msg or resp.error_message or resp.status_text
+        if msg is None:
+            msg = resp.error_message or resp.status_text
+        else:
+            msg = f"{msg} ({resp.error_message or resp.status_text})"
         self.error_message = resp.error_message
         self.error_code = resp.error_code
         if self.error_code is not None:
@@ -112,6 +115,10 @@ class DocumentGetError(ArangoServerError):
     """Failed to retrieve document."""
 
 
+class DocumentInsertError(ArangoServerError):
+    """Failed to insert document."""
+
+
 class DocumentParseError(ArangoClientError):
     """Failed to parse document input."""
 

arangoasync/executor.py

@@ -37,6 +37,12 @@ def serializer(self) -> Serializer[Json]:
     def deserializer(self) -> Deserializer[Json, Jsons]:
         return self._conn.deserializer
 
+    def serialize(self, data: Json) -> str:
+        return self.serializer.dumps(data)
+
+    def deserialize(self, data: bytes) -> Json:
+        return self.deserializer.loads(data)
+
     async def execute(
         self, request: Request, response_handler: Callable[[Response], T]
     ) -> T:

tests/test_database.py

@@ -31,14 +31,21 @@ async def test_create_drop_database(url, sys_db_name, root, password):
         # Create a new database
         db_name = generate_db_name()
         assert await sys_db.create_database(db_name) is True
-        await client.db(db_name, auth_method="basic", auth=auth, verify=True)
+        new_db = await client.db(db_name, auth_method="basic", auth=auth, verify=True)
         assert await sys_db.has_database(db_name) is True
 
         # List available databases
         dbs = await sys_db.databases()
         assert db_name in dbs
         assert "_system" in dbs
 
+        # TODO move this to a separate test
+        col_name = generate_col_name()
+        col = await new_db.create_collection(col_name)
+        await col.insert({"_key": "1", "a": 1})
+        doc = await col.get("1")
+        assert doc["_key"] == "1"
+
         # Drop the newly created database
         assert await sys_db.delete_database(db_name) is True
         non_existent_db = generate_db_name()