Adding KeyOptions object

Author: apetenchea

arangoasync/collection.py

@@ -1,4 +1,4 @@
-__all__ = ["Collection", "Collection", "StandardCollection"]
+__all__ = ["Collection", "CollectionType", "StandardCollection"]
 
 
 from enum import Enum

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
+                <arangoasync.wrapper.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)

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")

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,16 +14,16 @@ 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
 
     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()