<filename>redbrick/common/client.py
"""Graphql Client responsible for make API requests."""
from typing import Dict
import requests

import aiohttp
import tenacity

from redbrick.version_check import __version__
from redbrick.utils.logging import print_error
from redbrick.common.constants import MAX_RETRY_ATTEMPTS


class RBClient:
    """Client to communicate with RedBrick AI GraphQL Server."""

    def __init__(
        self,
        api_key: str,
        url: str,
    ) -> None:
        """Construct RBClient."""
        self.session = requests.Session()
        assert (
            len(api_key) == 43
        ), "Invalid Api Key length, make sure you've copied it correctly"
        self.api_key = api_key
        self.url = url.rstrip("/") + "/graphql/"

    def __del__(self) -> None:
        """Garbage collect and close session."""
        self.session.close()

    @property
    def headers(self) -> Dict:
        """Get request headers."""
        return {  # Default: Accept-Encoding = gzip, deflate, br
            "ApiKey": self.api_key,
            "RB-SDK-Version": __version__,
        }

    @tenacity.retry(
        stop=tenacity.stop_after_attempt(MAX_RETRY_ATTEMPTS),
        wait=tenacity.wait_exponential(multiplier=1, min=1, max=10),
        retry=tenacity.retry_if_not_exception_type(
            (KeyboardInterrupt, PermissionError, ValueError)
        ),
    )
    def execute_query(
        self, query: str, variables: Dict, raise_for_error: bool = True
    ) -> Dict:
        """Execute a graphql query."""
        response = self.session.post(
            self.url,
            headers=self.headers,
            json={"query": query, "variables": variables},
        )
        self._check_status_msg(response.status_code)
        return self._process_json_response(response.json(), raise_for_error)

    @tenacity.retry(
        stop=tenacity.stop_after_attempt(MAX_RETRY_ATTEMPTS),
        wait=tenacity.wait_exponential(multiplier=1, min=1, max=10),
        retry=tenacity.retry_if_not_exception_type(
            (KeyboardInterrupt, PermissionError, ValueError)
        ),
    )
    async def execute_query_async(
        self,
        aio_session: aiohttp.ClientSession,
        query: str,
        variables: Dict,
        raise_for_error: bool = True,
    ) -> Dict:
        """Execute a graphql query using asyncio."""
        async with aio_session.post(
            self.url,
            headers=self.headers,
            json={"query": query, "variables": variables},
        ) as response:
            self._check_status_msg(response.status)
            return self._process_json_response(await response.json(), raise_for_error)

    @staticmethod
    def _check_status_msg(response_status: int) -> None:
        if response_status >= 500:
            raise ConnectionError(
                "Internal Server Error: You are probably using an invalid API key"
            )
        if response_status == 403:
            raise PermissionError("Problem authenticating with Api Key")

    @staticmethod
    def _process_json_response(
        response_data: Dict, raise_for_error: bool = True
    ) -> Dict:
        """Process JSON resonse."""
        if "errors" in response_data:
            errors = []
            for error in response_data["errors"]:
                errors.append(error["message"])
                print_error(error["message"])

            if raise_for_error:
                raise ValueError("\n".join(errors))

            del response_data["errors"]

        res = {}
        if "data" in response_data:
            res = response_data["data"]
        else:
            res = response_data
        return res

<filename>redbrick/common/client.py
"""Graphql Client responsible for make API requests."""
from typing import Dict
import requests

import aiohttp
import tenacity

from redbrick.version_check import __version__
from redbrick.utils.logging import print_error
from redbrick.common.constants import MAX_RETRY_ATTEMPTS


class RBClient:
    """Client to communicate with RedBrick AI GraphQL Server."""

    def __init__(
        self,
        api_key: str,
        url: str,
    ) -> None:
        """Construct RBClient."""
        self.session = requests.Session()
        assert (
            len(api_key) == 43
        ), "Invalid Api Key length, make sure you've copied it correctly"
        self.api_key = api_key
        self.url = url.rstrip("/") + "/graphql/"

    def __del__(self) -> None:
        """Garbage collect and close session."""
        self.session.close()

    @property
    def headers(self) -> Dict:
        """Get request headers."""
        return {  # Default: Accept-Encoding = gzip, deflate, br
            "ApiKey": self.api_key,
            "RB-SDK-Version": __version__,
        }

    @tenacity.retry(
        stop=tenacity.stop_after_attempt(MAX_RETRY_ATTEMPTS),
        wait=tenacity.wait_exponential(multiplier=1, min=1, max=10),
        retry=tenacity.retry_if_not_exception_type(
            (KeyboardInterrupt, PermissionError, ValueError)
        ),
    )
    def execute_query(
        self, query: str, variables: Dict, raise_for_error: bool = True
    ) -> Dict:
        """Execute a graphql query."""
        response = self.session.post(
            self.url,
            headers=self.headers,
            json={"query": query, "variables": variables},
        )
        self._check_status_msg(response.status_code)
        return self._process_json_response(response.json(), raise_for_error)

    @tenacity.retry(
        stop=tenacity.stop_after_attempt(MAX_RETRY_ATTEMPTS),
        wait=tenacity.wait_exponential(multiplier=1, min=1, max=10),
        retry=tenacity.retry_if_not_exception_type(
            (KeyboardInterrupt, PermissionError, ValueError)
        ),
    )
    async def execute_query_async(
        self,
        aio_session: aiohttp.ClientSession,
        query: str,
        variables: Dict,
        raise_for_error: bool = True,
    ) -> Dict:
        """Execute a graphql query using asyncio."""
        async with aio_session.post(
            self.url,
            headers=self.headers,
            json={"query": query, "variables": variables},
        ) as response:
            self._check_status_msg(response.status)
            return self._process_json_response(await response.json(), raise_for_error)

    @staticmethod
    def _check_status_msg(response_status: int) -> None:
        if response_status >= 500:
            raise ConnectionError(
                "Internal Server Error: You are probably using an invalid API key"
            )
        if response_status == 403:
            raise PermissionError("Problem authenticating with Api Key")

    @staticmethod
    def _process_json_response(
        response_data: Dict, raise_for_error: bool = True
    ) -> Dict:
        """Process JSON resonse."""
        if "errors" in response_data:
            errors = []
            for error in response_data["errors"]:
                errors.append(error["message"])
                print_error(error["message"])

            if raise_for_error:
                raise ValueError("\n".join(errors))

            del response_data["errors"]

        res = {}
        if "data" in response_data:
            res = response_data["data"]
        else:
            res = response_data
        return res

Graphql Client responsible for make API requests. Client to communicate with RedBrick AI GraphQL Server. Construct RBClient. Garbage collect and close session. Get request headers. # Default: Accept-Encoding = gzip, deflate, br Execute a graphql query. Execute a graphql query using asyncio. Process JSON resonse.