diff --git a/arangoasync/collection.py b/arangoasync/collection.py index ae36fa8..ed9b7e4 100644 --- a/arangoasync/collection.py +++ b/arangoasync/collection.py @@ -1,11 +1,16 @@ __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, ) @@ -13,7 +18,7 @@ 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 `__ + """ # 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 `__ + """ # 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) diff --git a/arangoasync/exceptions.py b/arangoasync/exceptions.py index e7b2cbf..1bdcdef 100644 --- a/arangoasync/exceptions.py +++ b/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.""" diff --git a/arangoasync/executor.py b/arangoasync/executor.py index 7d9c888..48fafe5 100644 --- a/arangoasync/executor.py +++ b/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: diff --git a/tests/test_database.py b/tests/test_database.py index 671ef59..8531ecc 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -31,7 +31,7 @@ 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 @@ -39,6 +39,13 @@ async def test_create_drop_database(url, sys_db_name, root, password): 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()