Merge pull request #237 from ArangoDB-Community/support-overload-control

Support overload control

Author: CryptoNinjaGeek

CHANGELOG.md

@@ -1,6 +1,8 @@
 main
 ----
 
+* Added OverloadControlDatabase, enabling the client to react effectively to potential server overloads.
+
 * The db.version() now has a new optional parameter "details" that can be used to return additional information about
   the server version. The default is still false, so the old behavior is preserved.
 

README.md

@@ -14,7 +14,7 @@ database natively supporting documents, graphs and search.
 
 ## Requirements
 
-- ArangoDB version 3.7+
+- ArangoDB version 3.9+
 - Python version 3.8+
 
 ## Installation

arango/database.py

@@ -1,4 +1,10 @@
-__all__ = ["StandardDatabase", "AsyncDatabase", "BatchDatabase", "TransactionDatabase"]
+__all__ = [
+    "StandardDatabase",
+    "AsyncDatabase",
+    "BatchDatabase",
+    "OverloadControlDatabase",
+    "TransactionDatabase",
+]
 
 from datetime import datetime
 from numbers import Number
@@ -75,6 +81,7 @@
     AsyncApiExecutor,
     BatchApiExecutor,
     DefaultApiExecutor,
+    OverloadControlApiExecutor,
     TransactionApiExecutor,
 )
 from arango.formatter import (
@@ -2518,6 +2525,19 @@ def begin_transaction(
             max_size=max_size,
         )
 
+    def begin_controlled_execution(
+        self, max_queue_time_seconds: Optional[float] = None
+    ) -> "OverloadControlDatabase":
+        """Begin a controlled connection, with options to handle server-side overload.
+
+        :param max_queue_time_seconds: Maximum time in seconds a request can be queued
+            on the server-side. If set to 0 or None, the server ignores this setting.
+        :type max_queue_time_seconds: Optional[float]
+        :return: Database API wrapper object specifically for queue bounded execution.
+        :rtype: arango.database.OverloadControlDatabase
+        """
+        return OverloadControlDatabase(self._conn, max_queue_time_seconds)
+
 
 class AsyncDatabase(Database):
     """Database API wrapper tailored specifically for async execution.
@@ -2688,3 +2708,55 @@ def abort_transaction(self) -> bool:
         :raise arango.exceptions.TransactionAbortError: If abort fails.
         """
         return self._executor.abort()
+
+
+class OverloadControlDatabase(Database):
+    """Database API wrapper tailored to gracefully handle server overload scenarios.
+
+    See :func:`arango.database.StandardDatabase.begin_controlled_execution`.
+
+    :param connection: HTTP connection.
+    :param max_queue_time_seconds: Maximum server-side queuing time in seconds.
+        If the server-side queuing time exceeds the client's specified limit,
+        the request will be rejected.
+    :type max_queue_time_seconds: Optional[float]
+    """
+
+    def __init__(
+        self, connection: Connection, max_queue_time_seconds: Optional[float] = None
+    ) -> None:
+        self._executor: OverloadControlApiExecutor
+        super().__init__(
+            connection=connection,
+            executor=OverloadControlApiExecutor(connection, max_queue_time_seconds),
+        )
+
+    def __repr__(self) -> str:  # pragma: no cover
+        return f"<OverloadControlDatabase {self.name}>"
+
+    @property
+    def last_queue_time(self) -> float:
+        """Return the most recently recorded server-side queuing time in seconds.
+
+        :return: Server-side queuing time in seconds.
+        :rtype: float
+        """
+        return self._executor.queue_time_seconds
+
+    @property
+    def max_queue_time(self) -> Optional[float]:
+        """Return the maximum server-side queuing time in seconds.
+
+        :return: Maximum server-side queuing time in seconds.
+        :rtype: Optional[float]
+        """
+        return self._executor.max_queue_time_seconds
+
+    def adjust_max_queue_time(self, max_queue_time_seconds: Optional[float]) -> None:
+        """Adjust the maximum server-side queuing time in seconds.
+
+        :param max_queue_time_seconds: New maximum server-side queuing time
+            in seconds. Setting it to None disables the limit.
+        :type max_queue_time_seconds: Optional[float]
+        """
+        self._executor.max_queue_time_seconds = max_queue_time_seconds

arango/exceptions.py

@@ -232,6 +232,15 @@ class BatchExecuteError(ArangoServerError):
     """Failed to execute batch API request."""
 
 
+#########################################
+# Overload Control Execution Exceptions #
+#########################################
+
+
+class OverloadControlExecutorError(ArangoServerError):
+    """Failed to execute overload controlled API request."""
+
+
 #########################
 # Collection Exceptions #
 #########################

arango/executor.py

@@ -4,6 +4,7 @@
     "AsyncApiExecutor",
     "BatchApiExecutor",
     "TransactionApiExecutor",
+    "OverloadControlApiExecutor",
 ]
 
 from collections import OrderedDict
@@ -16,6 +17,7 @@
     AsyncExecuteError,
     BatchExecuteError,
     BatchStateError,
+    OverloadControlExecutorError,
     TransactionAbortError,
     TransactionCommitError,
     TransactionInitError,
@@ -32,6 +34,7 @@
     "AsyncApiExecutor",
     "BatchApiExecutor",
     "TransactionApiExecutor",
+    "OverloadControlApiExecutor",
 ]
 
 T = TypeVar("T")
@@ -428,3 +431,84 @@ def abort(self) -> bool:
         if resp.is_success:
             return True
         raise TransactionAbortError(resp, request)
+
+
+class OverloadControlApiExecutor:
+    """Allows setting the maximum acceptable server-side queuing time
+        for client requests.
+
+    :param connection: HTTP connection.
+    :type connection: arango.connection.BasicConnection |
+        arango.connection.JwtConnection | arango.connection.JwtSuperuserConnection
+    :param max_queue_time_seconds: Maximum server-side queuing time in seconds.
+    :type max_queue_time_seconds: float
+    """
+
+    def __init__(
+        self, connection: Connection, max_queue_time_seconds: Optional[float] = None
+    ) -> None:
+        self._conn = connection
+        self._max_queue_time_seconds = max_queue_time_seconds
+        self._queue_time_seconds = 0.0
+
+    @property
+    def context(self) -> str:  # pragma: no cover
+        return "overload-control"
+
+    @property
+    def queue_time_seconds(self) -> float:
+        """Return the most recent request queuing/de-queuing time.
+            Defaults to 0 if no request has been sent.
+
+        :return: Server-side queuing time in seconds.
+        :rtype: float
+        """
+        return self._queue_time_seconds
+
+    @property
+    def max_queue_time_seconds(self) -> Optional[float]:
+        """Return the maximum server-side queuing time.
+
+        :return: Maximum server-side queuing time in seconds.
+        :rtype: Optional[float]
+        """
+        return self._max_queue_time_seconds
+
+    @max_queue_time_seconds.setter
+    def max_queue_time_seconds(self, value: Optional[float]) -> None:
+        """Set the maximum server-side queuing time.
+            Setting it to None disables the feature.
+
+        :param value: Maximum server-side queuing time in seconds.
+        :type value: Optional[float]
+        """
+        self._max_queue_time_seconds = value
+
+    def execute(
+        self,
+        request: Request,
+        response_handler: Callable[[Response], T],
+    ) -> T:
+        """Execute an API request and return the result.
+
+        :param request: HTTP request.
+        :type request: arango.request.Request
+        :param response_handler: HTTP response handler.
+        :type response_handler: callable
+        :return: API execution result.
+        """
+        if self._max_queue_time_seconds is not None:
+            request.headers["x-arango-queue-time-seconds"] = str(
+                self._max_queue_time_seconds
+            )
+        resp = self._conn.send_request(request)
+
+        if not resp.is_success:
+            raise OverloadControlExecutorError(resp, request)
+
+        if "X-Arango-Queue-Time-Seconds" in resp.headers:
+            self._queue_time_seconds = float(
+                resp.headers["X-Arango-Queue-Time-Seconds"]
+            )
+
+        return response_handler(resp)

docs/contributing.rst

@@ -40,7 +40,7 @@ To ensure PEP8_ compliance, run flake8_:
 .. code-block:: bash
 
     ~$ pip install flake8
-    ~$ git clone https://github.com/joowani/python-arango.git
+    ~$ git clone https://github.com/ArangoDB-Community/python-arango.git
     ~$ cd python-arango
     ~$ flake8
 
@@ -57,7 +57,7 @@ To run the test suite (use your own host, port and root password):
 .. code-block:: bash
 
     ~$ pip install pytest
-    ~$ git clone https://github.com/joowani/python-arango.git
+    ~$ git clone https://github.com/ArangoDB-Community/python-arango.git
     ~$ cd python-arango
     ~$ py.test --complete --host=127.0.0.1 --port=8529 --passwd=passwd
 
@@ -66,7 +66,7 @@ To run the test suite with coverage report:
 .. code-block:: bash
 
     ~$ pip install coverage pytest pytest-cov
-    ~$ git clone https://github.com/joowani/python-arango.git
+    ~$ git clone https://github.com/ArangoDB-Community/python-arango.git
     ~$ cd python-arango
     ~$ py.test --complete --host=127.0.0.1 --port=8529 --passwd=passwd --cov=kq
 
@@ -82,9 +82,9 @@ Sphinx_. To build an HTML version on your local machine:
 .. code-block:: bash
 
     ~$ pip install sphinx sphinx_rtd_theme
-    ~$ git clone https://github.com/joowani/python-arango.git
-    ~$ cd python-arango/docs
-    ~$ sphinx-build . build  # Open build/index.html in a browser
+    ~$ git clone https://github.com/ArangoDB-Community/python-arango.git
+    ~$ cd python-arango
+    ~$ python -m sphinx -b html -W docs docs/_build/  # Open build/index.html in a browser
 
 As always, thank you for your contribution!
 

docs/index.rst

@@ -10,7 +10,7 @@ Welcome to the documentation for **python-arango**, a Python driver for ArangoDB
 Requirements
 =============
 
-- ArangoDB version 3.7+
+- ArangoDB version 3.9+
 - Python version 3.8+
 
 Installation
@@ -38,6 +38,7 @@ Contents
     cursor
     async
     batch
+    overload
     transaction
     admin
     user

docs/overload.rst

@@ -0,0 +1,58 @@
+Overload API Execution
+----------------------
+:ref:`OverloadControlDatabase` is designed to handle time-bound requests. It allows setting a maximum server-side
+queuing time for client requests via the *max_queue_time_seconds* parameter. If the server's queueing time for a
+request surpasses this defined limit, the request will be rejected. This mechanism provides you with more control over
+request handling, enabling your application to react effectively to potential server overloads.
+
+Additionally, the response from ArangoDB will always include the most recent request queuing/dequeuing time from the
+server's perspective. This can be accessed via the :attr:`~.OverloadControlDatabase.last_queue_time` property.
+
+**Example:**
+
+.. testcode::
+
+    from arango import errno
+    from arango import ArangoClient
+    from arango.exceptions import OverloadControlExecutorError
+
+    # Initialize the ArangoDB client.
+    client = ArangoClient()
+
+    # Connect to "test" database as root user.
+    db = client.db('test', username='root', password='passwd')
+
+    # Begin controlled execution.
+    controlled_db = db.begin_controlled_execution(max_queue_time_seconds=7.5)
+
+    # All requests surpassing the specified limit will be rejected.
+    controlled_aql = controlled_db.aql
+    controlled_col = controlled_db.collection('students')
+
+    # On API execution, the last_queue_time property gets updated.
+    controlled_col.insert({'_key': 'Neal'})
+
+    # Retrieve the last recorded queue time.
+    assert controlled_db.last_queue_time >= 0
+
+    try:
+        controlled_aql.execute('RETURN 100000')
+    except OverloadControlExecutorError as err:
+        assert err.http_code == errno.HTTP_PRECONDITION_FAILED
+        assert err.error_code == errno.QUEUE_TIME_REQUIREMENT_VIOLATED
+
+    # Retrieve the maximum allowed queue time.
+    assert controlled_db.max_queue_time == 7.5
+
+    # Adjust the maximum allowed queue time.
+    controlled_db.adjust_max_queue_time(0.0001)
+
+    # Disable the maximum allowed queue time.
+    controlled_db.adjust_max_queue_time(None)
+
+.. note::
+    Setting *max_queue_time_seconds* to 0 or a non-numeric value will cause ArangoDB to ignore the header.
+
+See :ref:`OverloadControlDatabase` for API specification.
+See the `official documentation <https://www.arangodb.com/docs/stable/http/general.html#overload-control>`_ for
+details on ArangoDB's overload control options.

docs/specs.rst

@@ -135,6 +135,15 @@ HTTPClient
 .. autoclass:: arango.http.HTTPClient
     :members:
 
+.. _OverloadControlDatabase:
+
+OverloadControlDatabase
+=======================
+
+.. autoclass:: arango.database.OverloadControlDatabase
+    :inherited-members:
+    :members:
+
 .. _Pregel:
 
 Pregel