From 6b743c603a76ccc9ee71ab544783842ab78076cc Mon Sep 17 00:00:00 2001 From: Alex Petenchea Date: Sat, 5 Oct 2024 19:10:47 +0300 Subject: [PATCH] Adding KeyOptions object --- arangoasync/collection.py | 2 +- arangoasync/database.py | 20 +++++--- arangoasync/wrapper.py | 99 +++++++++++++++++++++++++++++++++++---- tests/test_wrapper.py | 28 +++++++++-- 4 files changed, 127 insertions(+), 22 deletions(-) diff --git a/arangoasync/collection.py b/arangoasync/collection.py index b0606a8..9e826f4 100644 --- a/arangoasync/collection.py +++ b/arangoasync/collection.py @@ -1,4 +1,4 @@ -__all__ = ["Collection", "Collection", "StandardCollection"] +__all__ = ["Collection", "CollectionType", "StandardCollection"] from enum import Enum diff --git a/arangoasync/database.py b/arangoasync/database.py index 7cba762..aaed5d6 100644 --- a/arangoasync/database.py +++ b/arangoasync/database.py @@ -20,7 +20,7 @@ from arangoasync.response import Response from arangoasync.serialization import Deserializer, Serializer from arangoasync.typings import Json, Jsons, Params, Result -from arangoasync.wrapper import ServerStatusInformation +from arangoasync.wrapper import KeyOptions, ServerStatusInformation T = TypeVar("T") U = TypeVar("U") @@ -140,7 +140,7 @@ async def create_collection( computed_values: Optional[Jsons] = None, distribute_shards_like: Optional[str] = None, is_system: Optional[bool] = False, - key_options: Optional[Json] = None, + key_options: Optional[KeyOptions | Json] = None, schema: Optional[Json] = None, shard_keys: Optional[Sequence[str]] = None, sharding_strategy: Optional[str] = None, @@ -179,7 +179,10 @@ async def create_collection( way as the shards of the other collection. is_system (bool | None): If `True`, create a system collection. In this case, the collection name should start with an underscore. - key_options (dict | None): Additional options for key generation. + key_options (KeyOptions | dict | None): Additional options for key + generation. You may use a :class:`KeyOptions + ` object for easier configuration, + or pass a dictionary directly. schema (dict | None): Optional object that specifies the collection level schema for documents. shard_keys (list | None): In a cluster, this attribute determines which @@ -204,6 +207,7 @@ async def create_collection( StandardCollection: Collection API wrapper. Raises: + ValueError: If parameters are invalid. CollectionCreateError: If the operation fails. """ data: Json = {"name": name} @@ -226,7 +230,10 @@ async def create_collection( if is_system is not None: data["isSystem"] = is_system if key_options is not None: - data["keyOptions"] = key_options + if isinstance(key_options, dict): + key_options = KeyOptions(key_options) + key_options.validate() + data["keyOptions"] = key_options.to_dict() if schema is not None: data["schema"] = schema if shard_keys is not None: @@ -304,9 +311,8 @@ def response_handler(resp: Response) -> bool: nonlocal ignore_missing if resp.is_success: return True - if resp.error_code == HTTP_NOT_FOUND: - if ignore_missing: - return False + if resp.error_code == HTTP_NOT_FOUND and ignore_missing: + return False raise CollectionDeleteError(resp, request) return await self._executor.execute(request, response_handler) diff --git a/arangoasync/wrapper.py b/arangoasync/wrapper.py index da2e974..61c289d 100644 --- a/arangoasync/wrapper.py +++ b/arangoasync/wrapper.py @@ -1,10 +1,12 @@ -from typing import Any, Dict, Iterator, Optional, Tuple +from typing import Any, Iterator, Optional, Tuple +from arangoasync.typings import Json -class Wrapper: - """Wrapper over server response objects.""" - def __init__(self, data: Dict[str, Any]) -> None: +class JsonWrapper: + """Wrapper over server request/response objects.""" + + def __init__(self, data: Json) -> None: self._data = data def __getitem__(self, key: str) -> Any: @@ -42,9 +44,88 @@ def items(self) -> Iterator[Tuple[str, Any]]: """Return an iterator over the dictionary’s key-value pairs.""" return iter(self._data.items()) + def to_dict(self) -> Json: + """Return the dictionary.""" + return self._data + + +class KeyOptions(JsonWrapper): + """Additional options for key generation, used on collections. + + https://docs.arangodb.com/stable/develop/http-api/collections/#create-a-collection_body_keyOptions + + Example: + .. code-block:: json + + "keyOptions": { + "type": "autoincrement", + "increment": 5, + "allowUserKeys": true + } -class ServerStatusInformation(Wrapper): + Args: + data (dict | None): Key options. If this parameter is specified, the + other parameters are ignored. + allow_user_keys (bool): If set to `True`, then you are allowed to supply own + key values in the `_key` attribute of documents. If set to `False`, then + the key generator is solely responsible for generating keys and an error + is raised if you supply own key values in the `_key` attribute of + documents. + generator_type (str): Specifies the type of the key generator. The currently + available generators are "traditional", "autoincrement", "uuid" and + "padded". + increment (int | None): The increment value for the "autoincrement" key + generator. Not allowed for other key generator types. + offset (int | None): The initial offset value for the "autoincrement" key + generator. Not allowed for other key generator types. """ + + def __init__( + self, + data: Optional[Json] = None, + allow_user_keys: bool = True, + generator_type: str = "traditional", + increment: Optional[int] = None, + offset: Optional[int] = None, + ) -> None: + if data is None: + data = { + "allowUserKeys": allow_user_keys, + "type": generator_type, + } + if increment is not None: + data["increment"] = increment + if offset is not None: + data["offset"] = offset + super().__init__(data) + + def validate(self) -> None: + """Validate key options.""" + if "type" not in self: + raise ValueError('"type" value is required for key options') + if "allowUserKeys" not in self: + raise ValueError('"allowUserKeys" value is required for key options') + + allowed_types = {"autoincrement", "uuid", "padded", "traditional"} + if self["type"] not in allowed_types: + raise ValueError( + f"Invalid key generator type '{self['type']}', " + f"expected one of {allowed_types}" + ) + + if self.get("increment") is not None and self["type"] != "autoincrement": + raise ValueError( + '"increment" value is only allowed for "autoincrement" ' "key generator" + ) + if self.get("offset") is not None and self["type"] != "autoincrement": + raise ValueError( + '"offset" value is only allowed for "autoincrement" ' "key generator" + ) + + +class ServerStatusInformation(JsonWrapper): + """Status information about the server. + https://docs.arangodb.com/stable/develop/http-api/administration/#get-server-status-information Example: @@ -92,7 +173,7 @@ class ServerStatusInformation(Wrapper): } """ - def __init__(self, data: Dict[str, Any]) -> None: + def __init__(self, data: Json) -> None: super().__init__(data) @property @@ -132,13 +213,13 @@ def hostname(self) -> Optional[str]: return self._data.get("hostname") @property - def server_info(self) -> Optional[Dict[str, Any]]: + def server_info(self) -> Optional[Json]: return self._data.get("serverInfo") @property - def coordinator(self) -> Optional[Dict[str, Any]]: + def coordinator(self) -> Optional[Json]: return self._data.get("coordinator") @property - def agency(self) -> Optional[Dict[str, Any]]: + def agency(self) -> Optional[Json]: return self._data.get("agency") diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index d2396b4..8e14e36 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -1,8 +1,10 @@ -from arangoasync.wrapper import Wrapper +import pytest + +from arangoasync.wrapper import JsonWrapper, KeyOptions def test_basic_wrapper(): - wrapper = Wrapper({"a": 1, "b": 2}) + wrapper = JsonWrapper({"a": 1, "b": 2}) assert wrapper["a"] == 1 assert wrapper["b"] == 2 @@ -12,7 +14,7 @@ def test_basic_wrapper(): del wrapper["a"] assert "a" not in wrapper - wrapper = Wrapper({"a": 1, "b": 2}) + wrapper = JsonWrapper({"a": 1, "b": 2}) keys = list(iter(wrapper)) assert keys == ["a", "b"] assert len(wrapper) == 2 @@ -20,8 +22,8 @@ def test_basic_wrapper(): assert "a" in wrapper assert "c" not in wrapper - assert repr(wrapper) == "Wrapper({'a': 1, 'b': 2})" - wrapper = Wrapper({"a": 1, "b": 2}) + assert repr(wrapper) == "JsonWrapper({'a': 1, 'b': 2})" + wrapper = JsonWrapper({"a": 1, "b": 2}) assert str(wrapper) == "{'a': 1, 'b': 2}" assert wrapper == {"a": 1, "b": 2} @@ -30,3 +32,19 @@ def test_basic_wrapper(): items = list(wrapper.items()) assert items == [("a", 1), ("b", 2)] + assert wrapper.to_dict() == {"a": 1, "b": 2} + + +def test_KeyOptions(): + options = KeyOptions(generator_type="autoincrement") + options.validate() + with pytest.raises(ValueError, match="Invalid key generator type 'invalid_type'"): + KeyOptions(generator_type="invalid_type").validate() + with pytest.raises(ValueError, match='"increment" value'): + KeyOptions(generator_type="uuid", increment=5).validate() + with pytest.raises(ValueError, match='"offset" value'): + KeyOptions(generator_type="uuid", offset=5).validate() + with pytest.raises(ValueError, match='"type" value'): + KeyOptions(data={"allowUserKeys": True}).validate() + with pytest.raises(ValueError, match='"allowUserKeys" value'): + KeyOptions(data={"type": "autoincrement"}).validate()