Introducing BasicConnection

Author: apetenchea

.pre-commit-config.yaml

@@ -38,4 +38,4 @@ repos:
       hooks:
         -   id: mypy
             files: ^arangoasync/
-            additional_dependencies: ['types-requests', "types-setuptools"]
+            additional_dependencies: ["types-requests", "types-setuptools"]

arangoasync/auth.py

@@ -0,0 +1,77 @@
+__all__ = [
+    "Auth",
+    "JwtToken",
+]
+
+from dataclasses import dataclass
+
+import jwt
+
+from arangoasync.exceptions import JWTExpiredError
+
+
+@dataclass
+class Auth:
+    """Authentication details for the ArangoDB instance.
+
+    Attributes:
+        username (str): Username.
+        password (str): Password.
+        encoding (str): Encoding for the password (default: utf-8)
+    """
+
+    username: str
+    password: str
+    encoding: str = "utf-8"
+
+
+class JwtToken:
+    """JWT token.
+
+    Args:
+        token (str | bytes): JWT token.
+
+    Raises:
+        JWTExpiredError: If the token expired.
+    """
+
+    def __init__(self, token: str | bytes) -> None:
+        self._token = token
+        self._validate()
+
+    @property
+    def token(self) -> str | bytes:
+        """Get token."""
+        return self._token
+
+    @token.setter
+    def token(self, token: str | bytes) -> None:
+        """Set token.
+
+        Raises:
+            JWTExpiredError: If the token expired.
+        """
+        self._token = token
+        self._validate()
+
+    def _validate(self) -> None:
+        """Validate the token."""
+        if type(self._token) is not str:
+            raise TypeError("Token must be a string")
+        try:
+            jwt_payload = jwt.decode(
+                self._token,
+                issuer="arangodb",
+                algorithms=["HS256"],
+                options={
+                    "require_exp": True,
+                    "require_iat": True,
+                    "verify_iat": True,
+                    "verify_exp": True,
+                    "verify_signature": False,
+                },
+            )
+        except jwt.ExpiredSignatureError:
+            raise JWTExpiredError("JWT token has expired")
+
+        self._token_exp = jwt_payload["exp"]

arangoasync/connection.py

@@ -0,0 +1,104 @@
+__all__ = [
+    "BaseConnection",
+    "BasicConnection",
+]
+
+from abc import ABC, abstractmethod
+from typing import Any, List
+
+from arangoasync.http import HTTPClient
+from arangoasync.request import Method, Request
+from arangoasync.resolver import HostResolver
+from arangoasync.response import Response
+
+
+class BaseConnection(ABC):
+    """Blueprint for connection to a specific ArangoDB database.
+
+    Args:
+        sessions (list): List of client sessions.
+        host_resolver (HostResolver): Host resolver.
+        http_client (HTTPClient): HTTP client.
+        db_name (str): Database name.
+    """
+
+    def __init__(
+        self,
+        sessions: List[Any],
+        host_resolver: HostResolver,
+        http_client: HTTPClient,
+        db_name: str,
+    ) -> None:
+        self._sessions = sessions
+        self._db_endpoint = f"/_db/{db_name}"
+        self._host_resolver = host_resolver
+        self._http_client = http_client
+        self._db_name = db_name
+
+    @property
+    def db_name(self) -> str:
+        """Return the database name."""
+        return self._db_name
+
+    def prep_response(selfs, resp: Response) -> None:
+        """Prepare response for return."""
+        # TODO: Populate response fields
+
+    async def process_request(self, request: Request) -> Response:
+        """Process request."""
+        # TODO add accept-encoding header option
+        # TODO regulate number of tries
+        # TODO error handling
+        host_index = self._host_resolver.get_host_index()
+        return await self._http_client.send_request(self._sessions[host_index], request)
+
+    async def ping(self) -> int:
+        """Ping host to check if connection is established.
+
+        Returns:
+            int: Response status code.
+        """
+        request = Request(method=Method.GET, endpoint="/_api/collection")
+        resp = await self.send_request(request)
+        # TODO check raise ServerConnectionError
+        return resp.status_code
+
+    @abstractmethod
+    async def send_request(self, request: Request) -> Response:  # pragma: no cover
+        """Send an HTTP request to the ArangoDB server.
+
+        Args:
+            request (Request): HTTP request.
+
+        Returns:
+            Response: HTTP response.
+        """
+        raise NotImplementedError
+
+
+class BasicConnection(BaseConnection):
+    """Connection to a specific ArangoDB database.
+
+    Allows for basic authentication to be used (username and password).
+
+    Args:
+        sessions (list): List of client sessions.
+        host_resolver (HostResolver): Host resolver.
+        http_client (HTTPClient): HTTP client.
+        db_name (str): Database name.
+    """
+
+    def __init__(
+        self,
+        sessions: List[Any],
+        host_resolver: HostResolver,
+        http_client: HTTPClient,
+        db_name: str,
+    ) -> None:
+        super().__init__(sessions, host_resolver, http_client, db_name)
+
+    async def send_request(self, request: Request) -> Response:
+        """Send an HTTP request to the ArangoDB server."""
+        response = await self.process_request(request)
+        self.prep_response(response)
+        return response

arangoasync/exceptions.py

@@ -0,0 +1,24 @@
+class ArangoError(Exception):
+    """Base class for all exceptions in python-arango-async."""
+
+
+class ArangoClientError(ArangoError):
+    """Base class for all client-related exceptions.
+
+    Args:
+        msg (str): Error message.
+
+    Attributes:
+        source (str): Source of the error (always set to "client")
+        message (str): Error message.
+    """
+
+    source = "client"
+
+    def __init__(self, msg: str) -> None:
+        super().__init__(msg)
+        self.message = msg
+
+
+class JWTExpiredError(ArangoClientError):
+    """JWT token has expired."""

arangoasync/http.py

@@ -9,6 +9,7 @@
 
 from aiohttp import BaseConnector, BasicAuth, ClientSession, ClientTimeout, TCPConnector
 
+from arangoasync.auth import Auth
 from arangoasync.request import Request
 from arangoasync.response import Response
 
@@ -74,7 +75,7 @@ class AioHTTPClient(HTTPClient):
         timeout (aiohttp.ClientTimeout | None): Client timeout settings.
             300s total timeout by default for a complete request/response operation.
         read_bufsize (int): Size of read buffer (64KB default).
-        auth (aiohttp.BasicAuth | None): HTTP authentication helper.
+        auth (Auth | None): HTTP authentication helper.
             Should be used for specifying authorization data in client API.
         compression_threshold (int): Will compress requests to the server if the size
             of the request body (in bytes) is at least the value of this option.
@@ -88,7 +89,7 @@ def __init__(
         connector: Optional[BaseConnector] = None,
         timeout: Optional[ClientTimeout] = None,
         read_bufsize: int = 2**16,
-        auth: Optional[BasicAuth] = None,
+        auth: Optional[Auth] = None,
         compression_threshold: int = 1024,
     ) -> None:
         self._connector = connector or TCPConnector(
@@ -100,7 +101,13 @@ def __init__(
             connect=60,  # max number of seconds for acquiring a pool connection
         )
         self._read_bufsize = read_bufsize
-        self._auth = auth
+        self._auth = (
+            BasicAuth(
+                login=auth.username, password=auth.password, encoding=auth.encoding
+            )
+            if auth
+            else None
+        )
         self._compression_threshold = compression_threshold
 
     def create_session(self, host: str) -> ClientSession:

arangoasync/resolver.py

@@ -0,0 +1,115 @@
+__all__ = [
+    "HostResolver",
+    "SingleHostResolver",
+    "RoundRobinHostResolver",
+    "DefaultHostResolver",
+    "get_resolver",
+]
+
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+
+class HostResolver(ABC):
+    """Abstract base class for host resolvers.
+
+    Args:
+        host_count (int): Number of hosts.
+        max_tries (int): Maximum number of attempts to try a host.
+
+    Raises:
+        ValueError: If max_tries is less than host_count.
+    """
+
+    def __init__(self, host_count: int = 1, max_tries: Optional[int] = None) -> None:
+        max_tries = max_tries or host_count * 3
+        if max_tries < host_count:
+            raise ValueError(
+                "The maximum number of attempts cannot be "
+                "lower than the number of hosts."
+            )
+        self._host_count = host_count
+        self._max_tries = max_tries
+        self._index = 0
+
+    @abstractmethod
+    def get_host_index(self) -> int:  # pragma: no cover
+        """Return the index of the host to use.
+
+        Returns:
+            int: Index of the host.
+        """
+        raise NotImplementedError
+
+    def change_host(self) -> None:
+        """If there aer multiple hosts available, switch to the next one."""
+        self._index = (self._index + 1) % self.host_count
+
+    @property
+    def host_count(self) -> int:
+        """Return the number of hosts."""
+        return self._host_count
+
+    @property
+    def max_tries(self) -> int:
+        """Return the maximum number of attempts."""
+        return self._max_tries
+
+
+class SingleHostResolver(HostResolver):
+    """Single host resolver. Always returns the same host index."""
+
+    def __init__(self, host_count: int, max_tries: Optional[int] = None) -> None:
+        super().__init__(host_count, max_tries)
+
+    def get_host_index(self) -> int:
+        return self._index
+
+
+class RoundRobinHostResolver(HostResolver):
+    """Round-robin host resolver. Changes host every time.
+
+    Useful for bulk inserts or updates.
+
+    Note:
+        Do not use this resolver for stream transactions.
+        Transaction IDs cannot be shared across different coordinators.
+    """
+
+    def __init__(self, host_count: int, max_tries: Optional[int] = None) -> None:
+        super().__init__(host_count, max_tries)
+        self._index = -1
+
+    def get_host_index(self, indexes_to_filter: Optional[List[int]] = None) -> int:
+        self.change_host()
+        return self._index
+
+
+DefaultHostResolver = SingleHostResolver
+
+
+def get_resolver(
+    strategy: str,
+    host_count: int,
+    max_tries: Optional[int] = None,
+) -> HostResolver:
+    """Return a host resolver based on the strategy.
+
+    Args:
+        strategy (str): Resolver strategy.
+        host_count (int): Number of hosts.
+        max_tries (int): Maximum number of attempts to try a host.
+
+    Returns:
+        HostResolver: Host resolver.
+
+    Raises:
+        ValueError: If the strategy is not supported.
+    """
+    if strategy == "roundrobin":
+        return RoundRobinHostResolver(host_count, max_tries)
+    if strategy == "single":
+        return SingleHostResolver(host_count, max_tries)
+    if strategy == "default":
+        return DefaultHostResolver(host_count, max_tries)
+    raise ValueError(f"Unsupported host resolver strategy: {strategy}")

docs/specs.rst

@@ -4,11 +4,23 @@ API Specification
 This page contains the specification for all classes and methods available in
 python-arango-async.
 
+.. automodule:: arangoasync.auth
+    :members:
+
+.. automodule:: arangoasync.connection
+    :members:
+
+.. automodule:: arangoasync.exceptions
+    :members: ArangoError, ArangoClientError
+
 .. automodule:: arangoasync.http
     :members:
 
 .. automodule:: arangoasync.request
     :members:
 
+.. automodule:: arangoasync.resolver
+    :members:
+
 .. automodule:: arangoasync.response
     :members: