Unix socket support about python-arango HOT 5 CLOSED

from typing import MutableMapping, Optional, Tuple, Union
from urllib.parse import urlparse
​
import requests_unixsocket
from requests import Session
from requests_toolbelt import MultipartEncoder
​
from arango.http import DefaultHTTPClient
from arango.response import Response
from arango.typings import Headers
​
DEFAULT_REQUEST_TIMEOUT = 60
​
​
class UnixSocketsHTTPClient(DefaultHTTPClient):
    """Unix Socket HTTP client implementation.
​
    :param unix_socket_url: Unix socket URL.
    :type unix_socket_url: str
    :param request_timeout: Timeout in seconds for each individual connection.
    :type request_timeout: int | float
    :param retry_attempts: Number of retry attempts.
    :type retry_attempts: int
    :param backoff_factor: Backoff factor for retry attempts.
    :type backoff_factor: float
    :param pool_connections: The number of urllib3 connection pools to cache.
    :type pool_connections: int
    :param pool_maxsize: The maximum number of connections to save in the pool.
    :type pool_maxsize: int
    :param pool_timeout: If set, then the pool will be set to block=True,
        and requests will block for pool_timeout seconds and raise
        EmptyPoolError if no connection is available within the time period.
    :type pool_timeout: int | float | None
    """
​
    def __init__(
        self,
        unix_socket_url: str,
        request_timeout: Union[int, float, None] = DEFAULT_REQUEST_TIMEOUT,
        retry_attempts: int = 3,
        backoff_factor: float = 1.0,
        pool_connections: int = 10,
        pool_maxsize: int = 10,
        pool_timeout: Union[int, float, None] = None,
    ) -> None:
        self.unix_socket_url = unix_socket_url
        self.request_timeout = request_timeout
        self._retry_attempts = retry_attempts
        self._backoff_factor = backoff_factor
        self._pool_connections = pool_connections
        self._pool_maxsize = pool_maxsize
        self._pool_timeout = pool_timeout
​
    def create_session(self, host: str) -> Session:
        """Create and return a new session/connection.
​
        :param host: ArangoDB host URL.
        :type host: str
        :returns: requests session object
        :rtype: requests.Session
        """
        # retry_strategy = Retry(
        #     total=self._retry_attempts,
        #     backoff_factor=self._backoff_factor,
        #     status_forcelist=[429, 500, 502, 503, 504],
        #     allowed_methods=["HEAD", "GET", "OPTIONS"],
        # )
​
        unix_adapter = requests_unixsocket.UnixAdapter(
            self.request_timeout,
            self._pool_connections,
            # connection_timeout=self.request_timeout,
            # pool_maxsize=self._pool_maxsize,
            # pool_timeout=self._pool_timeout,
            # max_retries=retry_strategy,
        )
​
        session = Session()
        session.mount("http+unix://", unix_adapter)
​
        return session
​
    def send_request(
        self,
        session: Session,
        method: str,
        url: str,
        headers: Optional[Headers] = None,
        params: Optional[MutableMapping[str, str]] = None,
        data: Union[str, MultipartEncoder, None] = None,
        auth: Optional[Tuple[str, str]] = None,
    ) -> Response:
        """Send an HTTP request.
​
        :param session: Requests session object.
        :type session: requests.Session
        :param method: HTTP method in lowercase (e.g. "post").
        :type method: str
        :param url: Request URL.
        :type url: str
        :param headers: Request headers.
        :type headers: dict
        :param params: URL (query) parameters.
        :type params: dict
        :param data: Request payload.
        :type data: str | MultipartEncoder | None
        :param auth: Username and password.
        :type auth: tuple
        :returns: HTTP response.
        :rtype: arango.response.Response
        """
        url_path = urlparse(url).path
        unix_url = f"http+unix://{self.unix_socket_url}/{url_path.lstrip('/')}"
​
        response = session.request(
            method=method,
            url=unix_url,
            params=params,
            data=data,
            headers=headers,
            auth=auth,
            timeout=self.request_timeout,
        )
        return Response(
            method=method,
            url=response.url,
            headers=response.headers,
            status_code=response.status_code,
            status_text=response.reason,
            raw_body=response.text,
        )
​
from arango import ArangoClient
​
# Instantiate the custom Unix socket HTTP client
unix_socket_http_client = UnixSocketsHTTPClient('path/to/unix/socket')
​
# Create an ArangoDB client instance with the custom HTTP client
db = ArangoClient(http_client=unix_socket_http_client).db(verify=True)