peak-sdk 1.0.0__py3-none-any.whl

peak/session.py

@@ -0,0 +1,259 @@
+#
+# # Copyright © 2023 Peak AI Limited. or its affiliates. All Rights Reserved.
+# #
+# # Licensed under the Apache License, Version 2.0 (the "License"). You
+# # may not use this file except in compliance with the License. A copy of
+# # the License is located at:
+# #
+# # https://github.com/PeakBI/peak-sdk/blob/main/LICENSE
+# #
+# # or in the "license" file accompanying this file. This file is
+# # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# # ANY KIND, either express or implied. See the License for the specific
+# # language governing permissions and limitations under the License.
+# #
+# # This file is part of the peak-sdk.
+# # see (https://github.com/PeakBI/peak-sdk)
+# #
+# # You should have received a copy of the APACHE LICENSE, VERSION 2.0
+# # along with this program. If not, see <https://apache.org/licenses/LICENSE-2.0>
+#
+"""Session module for Peak API."""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any, Dict, Iterator, List, Optional
+
+from peak import exceptions
+from peak.constants import DOWNLOAD_CHUNK_SIZE, ContentType, HttpMethods, Stage
+from peak.handler import Handler
+from peak.helpers import get_base_domain
+from peak.logger import logger
+
+DEFAULT_SESSION = None
+
+
+def _get_default_session() -> Session:
+    """Get the global default session object. Creates one if not already created and re-uses it.
+
+    Returns:
+        Session: The default session object
+    """
+    global DEFAULT_SESSION  # noqa: PLW0603
+    if DEFAULT_SESSION is None:
+        logger.debug("Creating DEFAULT_SESSION object")
+        DEFAULT_SESSION = Session()
+
+    logger.debug("DEFAULT_SESSION already present, reusing the object")
+    return DEFAULT_SESSION
+
+
+class Session:
+    """A session stores credentials which are used to authenticate the requests.
+
+    By default, a DEFAULT_SESSION is created which reads the credentials from the env variables.
+    Custom Session objects can be created and used if you want to work with multiple tenants.
+    """
+
+    auth_token: str
+    stage: Stage
+    base_domain: str
+    handler: Handler
+
+    def __init__(
+        self,
+        auth_token: Optional[str] = None,
+        stage: Optional[str] = None,
+    ) -> None:
+        """Initialize a session for the Peak API.
+
+        Args:
+            auth_token (str | None): Authentication token. Both API Key and Bearer tokens are supported.
+                Picks up from `API_KEY` environment variable if not provided.
+            stage (str | None): Name of the stage where tenant is created. Default is `prod`.
+        """
+        self.base_domain: str = ""
+        self._set_auth_token(auth_token)
+        self._set_stage(stage)
+        self._set_base_domain()
+        self.handler = Handler()
+
+    def create_request(
+        self,
+        endpoint: str,
+        method: HttpMethods,
+        content_type: ContentType,
+        *,
+        params: Optional[Dict[str, Any]] = None,
+        body: Optional[Dict[str, Any]] = None,
+        path: Optional[str] = None,
+        ignore_files: Optional[list[str]] = None,
+        subdomain: Optional[str] = "service",
+    ) -> Any:
+        """Prepares a request to be sent over the network.
+
+        Adds auth_token to the headers and creates URL using STAGE.
+        To be used with endpoints which returns JSON parsable response.
+
+        Args:
+            endpoint (str): The endpoint to send the request to.
+            method (HttpMethods): The HTTP method to use.
+            content_type (ContentType): The content type of the request.
+            params (Dict[str, Any], optional): params to send to the request, defaults to None
+            body (Dict[str, Any], optional): body to send to the request, defaults to None
+            path (Optional[str] optional): path to the file or folder that will be compressed and used as artifact, required for multipart requests.
+            ignore_files(Optional[list[str]]): Ignore files to be used when creating artifact, used only for multipart requests.
+            subdomain (Optional[str]): Subdomain for the endpoint. Defaults to `service`.
+
+        Returns:
+            Any: response dict object.
+        """
+        headers: Dict[str, str] = {"Authorization": self.auth_token}
+        base_domain: str = get_base_domain(stage=self.stage.value, subdomain=subdomain)
+        url: str = f"{base_domain}/{endpoint}"
+        return self.handler.make_request(
+            url,
+            method,
+            content_type=content_type,
+            headers=headers,
+            params=params or {},
+            body=body or {},
+            path=path,
+            ignore_files=ignore_files,
+            session_meta={
+                "stage": self.stage,
+            },
+        ).json()
+
+    def create_generator_request(
+        self,
+        endpoint: str,
+        method: HttpMethods,
+        content_type: ContentType,
+        response_key: str,
+        *,
+        params: Optional[Dict[str, Any]] = None,
+        body: Optional[Dict[str, Any]] = None,
+        path: Optional[str] = None,
+        subdomain: Optional[str] = "service",
+    ) -> Iterator[Dict[str, Any]]:
+        """Prepares a request to be sent over the network.
+
+        Adds auth_token to the headers and creates URL using STAGE.
+        Returns an iterator which automatically handles pagination and returns a new page at each iteration.
+        To be used with list endpoints only, which returns `pageNumber`, `pageCount` keys in response.
+
+        # noqa: DAR201
+
+        Args:
+            endpoint (str): The endpoint to send the request to.
+            method (HttpMethods): The HTTP method to use.
+            content_type (ContentType): The content type of the request.
+            response_key (str): key in the response dict which contains actual list data.
+            params (Optional[Dict[str, Any]]): params to send to the request.
+            body (Optional[Dict[str, Any]]): body to send to the request.
+            path (Optional[str]): path to the file or folder that will be compressed and used as artifact.
+            subdomain (Optional[str]): Subdomain for the endpoint. Defaults to `service`.
+
+        Yields:
+            Iterator[Dict[str, Any]]: paginated response json, element wise.
+
+        Raises:
+            StopIteration: There are no more pages to list
+        """
+        page_number: int = 1
+        page_count: int = 1
+        params = params or {}
+        while page_number <= page_count:
+            params = {**params, "pageNumber": page_number}
+            response = self.create_request(
+                endpoint,
+                method,
+                content_type,
+                params=params,
+                body=body,
+                path=path,
+                subdomain=subdomain,
+            )
+            page_count = response["pageCount"]
+            yield from response[response_key]
+            page_number += 1
+        return f"No more {response_key} to list"
+
+    def create_download_request(
+        self,
+        endpoint: str,
+        method: HttpMethods,
+        content_type: ContentType,
+        download_path: str,
+        params: Optional[Dict[str, Any]] = None,
+        body: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Prepares a request to be sent over the network.
+
+        Adds auth_token to the headers and creates URL using STAGE.
+        To be used for file download requests.
+
+        Args:
+            endpoint (str): The endpoint to send the request to.
+            method (HttpMethods): The HTTP method to use.
+            content_type (ContentType): The content type of the request.
+            download_path (str): Path where the downloaded file will be stored.
+            params (Dict[str, Any], optional): params to send to the request, defaults to None
+            body (Dict[str, Any], optional): body to send to the request, only used in multipart requests, defaults to None
+
+        Raises:
+            InvalidPathException: The download_path is invalid.
+        """
+        headers: Dict[str, str] = {"Authorization": self.auth_token}
+        url: str = f"{self.base_domain}/{endpoint}"
+        response: Any = self.handler.make_request(
+            url,
+            method,
+            content_type=content_type,
+            headers=headers,
+            params=params or {},
+            body=body or {},
+            request_kwargs={
+                "stream": True,
+                "allow_redirects": True,
+            },
+        )
+        try:
+            with Path(download_path).open("wb") as fd:
+                for chunk in response.iter_content(chunk_size=DOWNLOAD_CHUNK_SIZE):
+                    fd.write(chunk)
+        except IsADirectoryError:
+            raise exceptions.InvalidPathException(
+                download_path,
+                "Path should include the name with which the downloaded file should be stored.",
+            ) from None
+
+    def _set_base_domain(self) -> None:
+        self.base_domain = get_base_domain(stage=self.stage.value)
+
+    def _set_auth_token(self, auth_token: Optional[str]) -> None:
+        if auth_token is not None:
+            self.auth_token = auth_token
+            return
+
+        logger.info("auth_token not given, searching for API_KEY in env variables")
+        if not os.environ.get("API_KEY"):
+            raise exceptions.MissingEnvironmentVariableException(env_var="API_KEY")
+        self.auth_token = os.environ["API_KEY"]
+
+    def _set_stage(self, stage: Optional[str]) -> None:
+        if stage is not None:
+            self.stage = Stage(stage)
+            return
+
+        logger.info("stage not given, searching for STAGE in env variables")
+        if not os.environ.get("STAGE"):
+            logger.info("STAGE environment variable is not set, defaulting to PROD")
+            self.stage = Stage.PROD
+            return
+        self.stage = Stage(os.environ["STAGE"])
+
+
+__all__: List[str] = ["Session", "_get_default_session"]

peak/telemetry.py

@@ -0,0 +1,201 @@
+#
+# # Copyright © 2023 Peak AI Limited. or its affiliates. All Rights Reserved.
+# #
+# # Licensed under the Apache License, Version 2.0 (the "License"). You
+# # may not use this file except in compliance with the License. A copy of
+# # the License is located at:
+# #
+# # https://github.com/PeakBI/peak-sdk/blob/main/LICENSE
+# #
+# # or in the "license" file accompanying this file. This file is
+# # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# # ANY KIND, either express or implied. See the License for the specific
+# # language governing permissions and limitations under the License.
+# #
+# # This file is part of the peak-sdk.
+# # see (https://github.com/PeakBI/peak-sdk)
+# #
+# # You should have received a copy of the APACHE LICENSE, VERSION 2.0
+# # along with this program. If not, see <https://apache.org/licenses/LICENSE-2.0>
+#
+"""Decorator for sending telemetry data for each request."""
+from __future__ import annotations
+
+import platform
+import threading
+import uuid
+from contextlib import suppress
+from functools import wraps
+from typing import Any, Callable, Dict, Optional
+
+import requests
+
+import peak.config
+from peak.constants import ContentType, HttpMethods, Stage
+from peak.exceptions import BaseHttpException
+from peak.helpers import get_base_domain
+
+from ._version import __version__
+
+F = Callable[..., requests.Response]
+session_id = str(uuid.uuid4())
+
+
+def get_status_code(error: Optional[Exception]) -> int | None:
+    """It takes an exception object and returns the status code associated with it.
+
+    Args:
+        error (Optional[Exception]): The exception object to check
+
+    Returns:
+        int | None: The status code related to the error or None if no status code found
+    """
+    if not error:
+        return 200
+
+    if isinstance(error, BaseHttpException):
+        return error.STATUS_CODE
+
+    return None
+
+
+def telemetry(make_request: F) -> F:
+    """A decorator that wraps over the make_request function to send telemetry requests as required.
+
+    Args:
+        make_request (F): The make_request function to wrap in the decorator
+
+    Returns:
+        F: the wrapped function that sends telemetry data for each request
+    """
+
+    def get_telemetry_url(session_meta: Optional[Dict[str, Any]] = None) -> str:
+        """Returns the telemetry url for the given stage.
+
+        Args:
+            session_meta (Optional[Dict[str, Any]]): Session metadata object that contains information like stage
+
+        Returns:
+            str: The telemetry URL
+        """
+        stage = Stage.PROD
+        if session_meta:
+            stage = session_meta["stage"] if "stage" in session_meta else Stage.PROD
+        base_domain = get_base_domain(stage.value, "service")
+        return f"{base_domain}/resource-usage/api/v1/telemetry"
+
+    def get_telemetry_data() -> Dict[str, Any]:
+        return {
+            "sdkVersion": __version__,
+            "os": platform.platform(),
+            "hostname": platform.uname().node,
+            "pythonVersion": platform.python_version(),
+            "sessionId": session_id,
+            "requestId": str(uuid.uuid4()),
+        }
+
+    @wraps(make_request)
+    def wrapper(
+        self: Any,
+        url: str,
+        method: HttpMethods,
+        content_type: ContentType,
+        headers: Optional[Dict[str, str]] = None,
+        params: Optional[Dict[str, Any]] = None,
+        body: Optional[Dict[str, Any]] = None,
+        path: Optional[str] = None,
+        request_kwargs: Optional[Dict[str, int | bool | str | float]] = None,
+        ignore_files: Optional[list[str]] = None,
+        session_meta: Optional[Dict[str, Any]] = None,
+    ) -> requests.Response:
+        """A decorator that wraps over the make_request function to send telemetry requests as required.
+
+        Args:
+            self (Any): the object instance of Handler class on which the make_request call is being made
+            url (str): url to send the request to
+            method (HttpMethods): The HTTP method to use, e.g. get, post, put, delete
+            content_type (ContentType): content type of the request
+            headers (Dict[str, str]): headers to send with the request
+            params (Dict[str, Any]): params to send to the request
+            body (Dict[str, Any]): body to send to the request
+            path (str): path to the file or folder that will be compressed and used as artifact, defaults to None
+            request_kwargs(Dict[str, int | bool | str | float] | None): extra arguments to be passed when making the request.
+            ignore_files(Optional[list[str]]): Ignore files to be used when creating artifact
+            session_meta(Dict[str, Any]): Metadata about the session object, like - stage
+
+        Returns:
+            requests.Response: response json
+
+        Raises:
+            BaseHttpException: The http request failed.
+            Exception: Some other error occurred.
+        """
+
+        def make_telemetry_request(
+            res: Optional[requests.Response] = None,
+            error: Optional[Exception] = None,
+        ) -> None:
+            telemetry_url = get_telemetry_url(session_meta)
+
+            telemetry_body = {
+                "response": res.json() if method != HttpMethods.GET and res else None,
+                "error": str(error) if error else None,
+                "url": url,
+                "requestMethod": method.value,
+                "statusCode": get_status_code(error),
+                "source": peak.config.SOURCE.value,
+                **get_telemetry_data(),
+            }
+
+            with suppress(Exception):
+                make_request(
+                    self,
+                    telemetry_url,
+                    HttpMethods.POST,
+                    ContentType.APPLICATION_JSON,
+                    headers=headers,
+                    body=telemetry_body,
+                )
+
+        def trigger_usage_collection(
+            res: Optional[Any] = None,
+            error: Optional[Exception] = None,
+        ) -> None:
+            thr = threading.Thread(
+                target=make_telemetry_request,
+                kwargs={
+                    "res": res,
+                    "error": error,
+                },
+            )
+
+            thr.start()
+
+        try:
+            custom_headers = {f"x-peak-{key}": value for (key, value) in get_telemetry_data().items()}
+            custom_headers = {
+                **custom_headers,
+                **(headers or {}),
+            }
+
+            res = make_request(
+                self,
+                url,
+                method,
+                content_type=content_type,
+                headers=custom_headers,
+                params=params or {},
+                body=body or {},
+                path=path,
+                ignore_files=ignore_files,
+                request_kwargs=request_kwargs,
+            )
+
+            trigger_usage_collection(res=res)
+        except Exception as e:
+            trigger_usage_collection(error=e)
+            raise
+        else:
+            return res
+
+    return wrapper

peak/template.py

@@ -0,0 +1,231 @@
+#
+# # Copyright © 2023 Peak AI Limited. or its affiliates. All Rights Reserved.
+# #
+# # Licensed under the Apache License, Version 2.0 (the "License"). You
+# # may not use this file except in compliance with the License. A copy of
+# # the License is located at:
+# #
+# # https://github.com/PeakBI/peak-sdk/blob/main/LICENSE
+# #
+# # or in the "license" file accompanying this file. This file is
+# # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# # ANY KIND, either express or implied. See the License for the specific
+# # language governing permissions and limitations under the License.
+# #
+# # This file is part of the peak-sdk.
+# # see (https://github.com/PeakBI/peak-sdk)
+# #
+# # You should have received a copy of the APACHE LICENSE, VERSION 2.0
+# # along with this program. If not, see <https://apache.org/licenses/LICENSE-2.0>
+#
+"""Template module which handles all things related to templates."""
+from __future__ import annotations
+
+import os
+import re
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
+
+import jinja2
+import yaml
+from jinja2 import Environment
+from jinja2.ext import Extension
+
+from peak import exceptions
+
+
+def _parse_jinja_template(template_path: Path, params: Dict[str, Any]) -> str:
+    """Read, parse and render the Jinja template text."""
+    jinja_loader = _CustomJinjaLoader()
+    jinja_env = jinja2.Environment(  # TODO: show warning if variable not found in params  # noqa: TD002, TD003, RUF100
+        loader=jinja_loader,
+        autoescape=False,  # noqa: S701
+        extensions=[_IncludeWithIndentation],
+    )
+    jinja_template: jinja2.Template = jinja_env.get_template(str(template_path))
+    return jinja_template.render(params, env=os.environ)
+
+
+def load_template(file: Union[Path, str], params: Optional[Dict[str, Any]]) -> Dict[str, Any]:
+    """Load a template file through `Jinja` into a dictionary.
+
+    This function performs the following steps:
+        * Passes the `YAML` file to be loaded and parsed through **`Jinja`**
+        * **`Jinja`** substitutes the variables with their values as they are found in `params`
+        * Loads any other files that are referenced using the Jinja `{% include %}` directive, if it is present.
+        * Updates the `context` key path within `image` definitions with respect to its relative parent file path.
+        * Loads the rendered YAML file into a `dictionary`.
+
+    Args:
+        file (Union[Path, str]): Path to the templated `YAML` file to be loaded.
+        params (Dict[str, Any] | None, optional): Named parameters to be passed to Jinja. Defaults to `{}`.
+
+    Returns:
+        Dict[str, Any]: Dictionary containing the rendered YAML file
+    """
+    params = {} if params is None else params
+    file = Path(file)
+    template: str = _parse_jinja_template(file, params)
+    return yaml.safe_load(template)  # type: ignore[no-any-return]
+
+
+class _CustomJinjaLoader(jinja2.BaseLoader):
+    """Custom Jinja loader class which handles the include directive.
+
+    Inspired from the jinja2.FileSystemLoader class.
+    """
+
+    def __init__(
+        self,
+        search_path: Optional[List[str]] = None,
+        encoding: str = "utf-8",
+    ) -> None:
+        """Initialize all variables.
+
+        Args:
+            search_path (List[str] | None, optional): Path(s) of the directory to search file in.
+            encoding (str): Encoding to use when reading files. Defaults to "utf-8".
+        """
+        if search_path is None:
+            search_path = ["."]
+
+        self.search_path: List[str] = [os.fspath(p) for p in search_path]
+        self.encoding: str = encoding
+        self.root_file_path: Optional[str] = None
+        self.build_context_regex: re.Pattern[str] = re.compile(r"^(\s*context\s*:\s+)(.+)$", flags=re.MULTILINE)
+        self.seen_files: Set[Path] = set()
+
+    def _update_image_build_context(self, source: str, file_parent_dir: str) -> str:
+        """Updates the context key in image definition to the relative path of the file where it is being imported.
+
+        Args:
+            source (str): Content of the file
+            file_parent_dir (str): Directory where the file is located
+
+        Returns:
+            str: Content of the file with updated context value if present.
+        """
+        if self.root_file_path is None:
+            return source
+
+        def substitute(match: re.Match[str]) -> str:
+            """Substitute the context key with the relative path of the file where it is being imported."""
+            context_path: str = match.group(2)
+            context_path = context_path.strip().strip('"').strip("'")
+            final_path: str = str(Path(file_parent_dir) / Path(context_path))
+            final_relative_path: str = os.path.relpath(final_path, self.root_file_path)
+            return f"{match.group(1)}{final_relative_path}"
+
+        return self.build_context_regex.sub(substitute, source)
+
+    def get_source(self, _: jinja2.Environment, template_path: str) -> Tuple[str, str, Callable[[], bool]]:
+        """Searches and reads the template file.
+
+        Args:
+            _ (jinja2.Environment): Jinja environment variable.
+            template_path (str): Path of the template file.
+
+        # noqa: DAR401
+        Raises:
+            jinja2.TemplateNotFound: The template file is not found on given path.
+
+        Returns:
+            Tuple[str, str, Callable[[], bool]]: Tuple containing 3 variables
+                1. Content of the template file
+                2. Normalized path where the file was found
+                3. Function which checks if the file was modified (not useful in this case)
+        """
+        template_path = template_path.strip()
+        for search_path in self.search_path:
+            # Use posixpath even on Windows to avoid "drive:" or UNC
+            # segments breaking out of the search directory.
+            file_path: str = str(Path(search_path) / Path(template_path))
+
+            if Path(file_path).is_file():
+                break
+        else:
+            error_msg: str = f"File does not exist at path: {template_path!r}"
+            raise jinja2.TemplateNotFound(error_msg)
+
+        file_path_obj = Path(file_path)
+        absolute_file_path: Path = file_path_obj.resolve()
+
+        if absolute_file_path in self.seen_files:
+            error_msg = f"Failed to render template, circular include directive found at {absolute_file_path!r}"
+            raise exceptions.InvalidTemplateException(error_msg)
+
+        file_parent_dir = str(file_path_obj.parent)
+        with Path(file_path).open(encoding=self.encoding) as f:
+            contents: str = f.read()
+            contents = self._update_image_build_context(contents, file_parent_dir)
+
+        self.seen_files.add(absolute_file_path)
+        self.search_path.append(file_parent_dir)
+        if self.root_file_path is None:
+            self.root_file_path = file_parent_dir
+
+        return contents, os.path.normpath(file_path), lambda: True
+
+
+class _IncludeWithIndentation(Extension):
+    """Override Jinja include directive to preserve indentation.
+
+    Inspired from: https://github.com/stereobutter/jinja2_workarounds
+    """
+
+    @staticmethod
+    def _include_statement_regex(block_start: str, block_end: str) -> re.Pattern[str]:
+        """Get the compiled regex for finding the include directives in template."""
+        return re.compile(
+            rf"""
+            (^.*)
+            (?=
+                (
+                    {re.escape(block_start)}
+                    (?P<block_start_modifier> [\+|-]?)
+                    (?P<statement>
+                        \s*include
+                        \s+.*?
+                    )
+                    (?P<block_end_modifier> [\+|-]?)
+                    {re.escape(block_end)}
+                )
+            )
+            .*$
+            """,
+            flags=re.MULTILINE | re.VERBOSE,
+        )
+
+    def preprocess(self, source: str, _: Optional[str], __: Optional[str] = None) -> str:
+        """Enclose all include directives.
+
+        For all the regex matches in the text, enclose the include block
+        with indent filter blocks and return updated text.
+        """
+        env: Environment = self.environment
+
+        block_start: str = env.block_start_string
+        block_end: str = env.block_end_string
+        pattern: re.Pattern[str] = self._include_statement_regex(block_start=block_start, block_end=block_end)
+        re.compile("\n")
+
+        def add_indentation_filter(match: re.Match[str]) -> str:
+            """Add indent filter to include blocks."""
+            content_before_include: str | Any = match.group(1)
+            include_statement: str | Any = match.group("statement").replace("indent content", "")
+
+            block_start_modifier: str | Any = match.group("block_start_modifier") or ""
+            block_end_modifier: str | Any = match.group("block_end_modifier") or ""
+
+            start_filter: str = (
+                f"{block_start + block_start_modifier} filter indent({len(content_before_include)}) {block_end}"
+            )
+            include_block: str = f"{block_start} {include_statement} {block_end}"
+            end_filter: str = f"{block_start} endfilter {block_end_modifier + block_end}"
+
+            return f"{content_before_include}{start_filter}{include_block}{end_filter}"
+
+        return pattern.sub(add_indentation_filter, source)
+
+
+__all__: List[str] = ["load_template"]