Database creation and deletion (#22)

apetenchea · web-flow · commit aad94da1fe75 · 2024-10-05T22:02:30.000+03:00
* Database creation and deletion

* Fixing docs
diff --git a/arangoasync/collection.py b/arangoasync/collection.py
@@ -195,9 +195,9 @@ async def get(
         def response_handler(resp: Response) -> Optional[U]:
             if resp.is_success:
                 return self._doc_deserializer.loads(resp.raw_body)
-            elif resp.error_code == HTTP_NOT_FOUND:
+            elif resp.status_code == HTTP_NOT_FOUND:
                 return None
-            elif resp.error_code == HTTP_PRECONDITION_FAILED:
+            elif resp.status_code == HTTP_PRECONDITION_FAILED:
                 raise DocumentRevisionError(resp, request)
             else:
                 raise DocumentGetError(resp, request)
diff --git a/arangoasync/database.py b/arangoasync/database.py
@@ -8,19 +8,22 @@
 
 from arangoasync.collection import CollectionType, StandardCollection
 from arangoasync.connection import Connection
-from arangoasync.errno import HTTP_NOT_FOUND
+from arangoasync.errno import HTTP_FORBIDDEN, HTTP_NOT_FOUND
 from arangoasync.exceptions import (
     CollectionCreateError,
     CollectionDeleteError,
     CollectionListError,
+    DatabaseCreateError,
+    DatabaseDeleteError,
+    DatabaseListError,
     ServerStatusError,
 )
 from arangoasync.executor import ApiExecutor, DefaultApiExecutor
 from arangoasync.request import Method, Request
 from arangoasync.response import Response
 from arangoasync.serialization import Deserializer, Serializer
 from arangoasync.typings import Json, Jsons, Params, Result
-from arangoasync.wrapper import KeyOptions, ServerStatusInformation
+from arangoasync.wrapper import KeyOptions, ServerStatusInformation, User
 
 T = TypeVar("T")
 U = TypeVar("U")
@@ -76,6 +79,137 @@ def response_handler(resp: Response) -> ServerStatusInformation:
 
         return await self._executor.execute(request, response_handler)
 
+    async def has_database(self, name: str) -> Result[bool]:
+        """Check if a database exists.
+
+        Args:
+            name (str): Database name.
+
+        Returns:
+            bool: `True` if the database exists, `False` otherwise.
+
+        Raises:
+            DatabaseListError: If failed to retrieve the list of databases.
+        """
+        request = Request(method=Method.GET, endpoint="/_api/database")
+
+        def response_handler(resp: Response) -> bool:
+            if not resp.is_success:
+                raise DatabaseListError(resp, request)
+            body = self.deserializer.loads(resp.raw_body)
+            return name in body["result"]
+
+        return await self._executor.execute(request, response_handler)
+
+    async def create_database(
+        self,
+        name: str,
+        users: Optional[Sequence[Json | User]] = None,
+        replication_factor: Optional[int | str] = None,
+        write_concern: Optional[int] = None,
+        sharding: Optional[bool] = None,
+    ) -> Result[bool]:
+        """Create a new database.
+
+        Args:
+            name (str): Database name.
+            users (list | None): Optional list of users with access to the new
+                database, where each user is of :class:`User
+                <arangoasync.wrapper.User>` type, or a dictionary with fields
+                "username", "password" and "active". If not set, the default user
+                **root** will be used to ensure that the new database will be
+                accessible after it is created.
+            replication_factor (int | str | None): Default replication factor for new
+                collections created in this database. Special values include
+                “satellite”, which will replicate the collection to every DB-Server
+                (Enterprise Edition only), and 1, which disables replication. Used
+                for clusters only.
+            write_concern (int | None): Default write concern for collections created
+            in this database. Determines how many copies of each shard are required
+            to be in sync on different DB-Servers. If there are less than these many
+            copies in the cluster a shard will refuse to write. Writes to shards with
+            enough up-to-date copies will succeed at the same time, however. Value of
+            this parameter can not be larger than the value of **replication_factor**.
+            Used for clusters only.
+            sharding (str | None): Sharding method used for new collections in this
+                database. Allowed values are: "", "flexible" and "single". The first
+                two are equivalent. Used for clusters only.
+
+        Returns:
+            bool: True if the database was created successfully.
+
+        Raises:
+            DatabaseCreateError: If creation fails.
+        """
+        data: Json = {"name": name}
+
+        options: Json = {}
+        if replication_factor is not None:
+            options["replicationFactor"] = replication_factor
+        if write_concern is not None:
+            options["writeConcern"] = write_concern
+        if sharding is not None:
+            options["sharding"] = sharding
+        if options:
+            data["options"] = options
+
+        if users is not None:
+            data["users"] = [
+                {
+                    "username": user["username"],
+                    "passwd": user["password"],
+                    "active": user.get("active", True),
+                    "extra": user.get("extra", {}),
+                }
+                for user in users
+            ]
+
+        request = Request(
+            method=Method.POST,
+            endpoint="/_api/database",
+            data=self.serializer.dumps(data),
+        )
+
+        def response_handler(resp: Response) -> bool:
+            if resp.is_success:
+                return True
+            raise DatabaseCreateError(resp, request)
+
+        return await self._executor.execute(request, response_handler)
+
+    async def delete_database(
+        self, name: str, ignore_missing: bool = False
+    ) -> Result[bool]:
+        """Delete a database.
+
+        Args:
+            name (str): Database name.
+            ignore_missing (bool): Do not raise an exception on missing database.
+
+        Returns:
+            bool: True if the database was deleted successfully, `False` if the
+                database was not found but **ignore_missing** was set to `True`.
+
+        Raises:
+            DatabaseDeleteError: If deletion fails.
+        """
+        request = Request(method=Method.DELETE, endpoint=f"/_api/database/{name}")
+
+        def response_handler(resp: Response) -> bool:
+            if resp.is_success:
+                return True
+            if resp.status_code == HTTP_NOT_FOUND and ignore_missing:
+                return False
+            if resp.status_code == HTTP_FORBIDDEN:
+                raise DatabaseDeleteError(
+                    resp,
+                    request,
+                    "This request can only be executed in the _system database.",
+                )
+            raise DatabaseDeleteError(resp, request)
+
+        return await self._executor.execute(request, response_handler)
+
     def collection(
         self,
         name: str,
@@ -231,7 +365,7 @@ async def create_collection(
             data["isSystem"] = is_system
         if key_options is not None:
             if isinstance(key_options, dict):
-                key_options = KeyOptions(key_options)
+                key_options = KeyOptions(data=key_options)
             key_options.validate()
             data["keyOptions"] = key_options.to_dict()
         if schema is not None:
@@ -311,7 +445,7 @@ def response_handler(resp: Response) -> bool:
             nonlocal ignore_missing
             if resp.is_success:
                 return True
-            if resp.error_code == HTTP_NOT_FOUND and ignore_missing:
+            if resp.status_code == HTTP_NOT_FOUND and ignore_missing:
                 return False
             raise CollectionDeleteError(resp, request)
 
diff --git a/arangoasync/exceptions.py b/arangoasync/exceptions.py
@@ -92,6 +92,18 @@ class ClientConnectionError(ArangoClientError):
     """The request was unable to reach the server."""
 
 
+class DatabaseCreateError(ArangoServerError):
+    """Failed to create database."""
+
+
+class DatabaseDeleteError(ArangoServerError):
+    """Failed to delete database."""
+
+
+class DatabaseListError(ArangoServerError):
+    """Failed to retrieve databases."""
+
+
 class DeserializationError(ArangoClientError):
     """Failed to deserialize the server response."""
 
diff --git a/arangoasync/wrapper.py b/arangoasync/wrapper.py
@@ -54,18 +54,7 @@ class KeyOptions(JsonWrapper):
 
     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
-            }
-
     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
@@ -78,15 +67,26 @@ class KeyOptions(JsonWrapper):
             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.
+        data (dict | None): Key options. If this parameter is specified, the
+            other parameters are ignored.
+
+    Example:
+        .. code-block:: json
+
+            {
+                "type": "autoincrement",
+                "increment": 5,
+                "allowUserKeys": true
+            }
     """
 
     def __init__(
         self,
-        data: Optional[Json] = None,
         allow_user_keys: bool = True,
         generator_type: str = "traditional",
         increment: Optional[int] = None,
         offset: Optional[int] = None,
+        data: Optional[Json] = None,
     ) -> None:
         if data is None:
             data = {
@@ -123,6 +123,61 @@ def validate(self) -> None:
             )
 
 
+class User(JsonWrapper):
+    """User information.
+
+    https://docs.arangodb.com/stable/develop/http-api/users/#get-a-user
+
+    Args:
+        username (str): The name of the user.
+        password (str | None): The user password as a string. Note that user
+            password is not returned back by the server.
+        active (bool): `True` if user is active, `False` otherwise.
+        extra (dict | None): Additional user information. For internal use only.
+            Should not be set or modified by end users.
+
+    Example:
+        .. code-block:: json
+
+            {
+                "username": "john",
+                "password": "secret",
+                "active": true,
+                "extra": {}
+            }
+    """
+
+    def __init__(
+        self,
+        username: str,
+        password: Optional[str] = None,
+        active: bool = True,
+        extra: Optional[Json] = None,
+    ) -> None:
+        data = {"username": username, "active": active}
+        if password is not None:
+            data["password"] = password
+        if extra is not None:
+            data["extra"] = extra
+        super().__init__(data)
+
+    @property
+    def username(self) -> str:
+        return self._data.get("username")  # type: ignore[return-value]
+
+    @property
+    def password(self) -> Optional[str]:
+        return self._data.get("password")
+
+    @property
+    def active(self) -> bool:
+        return self._data.get("active")  # type: ignore[return-value]
+
+    @property
+    def extra(self) -> Optional[Json]:
+        return self._data.get("extra")
+
+
 class ServerStatusInformation(JsonWrapper):
     """Status information about the server.
 
diff --git a/tests/helpers.py b/tests/helpers.py
@@ -1,6 +1,15 @@
 from uuid import uuid4
 
 
+def generate_db_name():
+    """Generate and return a random database name.
+
+    Returns:
+        str: Random database name.
+    """
+    return f"test_database_{uuid4().hex}"
+
+
 def generate_col_name():
     """Generate and return a random collection name.
 
diff --git a/tests/test_connection.py b/tests/test_connection.py
@@ -109,10 +109,10 @@ async def test_BasicConnection_prep_response_bad_response(
     with pytest.raises(ServerConnectionError):
         connection.raise_for_status(request, response)
 
-    error = b'{"error": true, "errorMessage": "msg", "errorNum": 404}'
+    error = b'{"error": true, "errorMessage": "msg", "errorNum": 1234}'
     response = Response(Method.GET, url, {}, 0, "ERROR", error)
     connection.prep_response(request, response)
-    assert response.error_code == 404
+    assert response.error_code == 1234
     assert response.error_message == "msg"
 
 
diff --git a/tests/test_database.py b/tests/test_database.py
@@ -3,7 +3,7 @@
 from arangoasync.auth import Auth
 from arangoasync.client import ArangoClient
 from arangoasync.collection import StandardCollection
-from tests.helpers import generate_col_name
+from tests.helpers import generate_col_name, generate_db_name
 
 
 @pytest.mark.asyncio
@@ -17,6 +17,28 @@ async def test_database_misc_methods(url, sys_db_name, root, password):
         assert status["server"] == "arango"
 
 
+@pytest.mark.asyncio
+async def test_create_drop_database(url, sys_db_name, root, password):
+    auth = Auth(username=root, password=password)
+
+    # TODO also handle exceptions
+    # TODO use more options (cluster must be enabled for that)
+    async with ArangoClient(hosts=url) as client:
+        sys_db = await client.db(
+            sys_db_name, auth_method="basic", auth=auth, verify=True
+        )
+        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)
+        assert await sys_db.has_database(db_name) is True
+        assert await sys_db.delete_database(db_name) is True
+        non_existent_db = generate_db_name()
+        assert await sys_db.has_database(non_existent_db) is False
+        assert (
+            await sys_db.delete_database(non_existent_db, ignore_missing=True) is False
+        )
+
+
 @pytest.mark.asyncio
 async def test_create_drop_collection(url, sys_db_name, root, password):
     auth = Auth(username=root, password=password)
@@ -28,7 +50,10 @@ async def test_create_drop_collection(url, sys_db_name, root, password):
         col = await db.create_collection(col_name)
         assert isinstance(col, StandardCollection)
         assert await db.has_collection(col_name)
-        await db.delete_collection(col_name)
+        assert await db.delete_collection(col_name) is True
         assert not await db.has_collection(col_name)
         non_existent_col = generate_col_name()
         assert await db.has_collection(non_existent_col) is False
+        assert (
+            await db.delete_collection(non_existent_col, ignore_missing=True) is False
+        )
