steamship/base/client.py

from __future__ import annotations

import logging
import typing
from abc import ABC
from inspect import isclass
from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union

import inflection
from pydantic import BaseModel, PrivateAttr
from requests import Session

from steamship.base.configuration import Configuration
from steamship.base.error import SteamshipError
from steamship.base.mime_types import MimeTypes
from steamship.base.model import CamelModel, to_camel
from steamship.base.request import Request
from steamship.base.tasks import Task, TaskState
from steamship.utils.url import Verb, is_local

_logger = logging.getLogger(__name__)

T = TypeVar("T")  # TODO (enias): Do we need this?


def _multipart_name(path: str, val: Any) -> List[Tuple[Optional[str], str, Optional[str]]]:
    """Decode any object into a series of HTTP Multi-part segments that Vapor will consume.

    https://github.com/vapor/multipart-kit
    When sending a JSON object in a MultiPart request, Vapor wishes to see multi part segments as follows:

    single_key
    array_key[idx]
    obj_key[prop]

    So a File with a list of one tag with kind=Foo would be transmitted as setting the part:
    [tags][0][kind]
    """
    ret = []
    if isinstance(val, dict):
        for key, subval in val.items():
            ret.extend(_multipart_name(f"{path}[{key}]", subval))
    elif isinstance(val, list):
        for idx, subval in enumerate(val):
            ret.extend(_multipart_name(f"{path}[{idx}]", subval))
    elif val is not None:
        ret.append((path, val, None))
    return ret


class Client(CamelModel, ABC):
    """Client model.py class.

    Separated primarily as a hack to prevent circular imports.
    """

    config: Configuration
    _session: Session = PrivateAttr()

    def __init__(
        self,
        api_key: str = None,
        api_base: str = None,
        app_base: str = None,
        web_base: str = None,
        workspace: str = None,
        fail_if_workspace_exists: bool = False,
        profile: str = None,
        config_file: str = None,
        config: Configuration = None,
        trust_workspace_config: bool = False,  # For use by lambda_handler; don't fetch the workspace
        **kwargs,
    ):
        """Create a new client.

        If `workspace` is provided, it will anchor the client in a workspace by that name, creating it if necessary.
        Otherwise the `default` workspace will be used.
        """
        if config is not None and not isinstance(config, Configuration):
            config = Configuration.parse_obj(config)

        self._session = Session()
        config = config or Configuration(
            api_key=api_key,
            api_base=api_base,
            app_base=app_base,
            web_base=web_base,
            workspace_handle=workspace,
            profile=profile,
            config_file=config_file,
        )

        super().__init__(config=config)
        # The lambda_handler will pass in the workspace via the workspace_id, so we need to plumb this through to make sure
        # that the workspace switch performed doesn't mistake `workspace=None` as a request for the default workspace
        self.switch_workspace(
            workspace_handle=workspace or config.workspace_handle,
            workspace_id=config.workspace_id,
            fail_if_workspace_exists=fail_if_workspace_exists,
            trust_workspace_config=trust_workspace_config,
        )

    def switch_workspace(  # noqa: C901
        self,
        workspace_handle: str = None,
        workspace_id: str = None,
        fail_if_workspace_exists: bool = False,
        trust_workspace_config: bool = False,
        # For use by lambda_handler; don't fetch the workspacetrust_workspace_config: bool = False, # For use by lambda_handler; don't fetch the workspace
    ):
        """Switches this client to the requested workspace, possibly creating it. If all arguments are None, the client
        actively switches into the default workspace.

        - API calls are performed manually to not result in circular imports.
        - Note that the default workspace is technically not necessary for API usage; it will be assumed by the Engine
          in the absense of a Workspace ID or Handle being manually specified in request headers.
        """
        workspace = None

        if workspace_handle is None and workspace_id is None:
            # Switch to the default workspace since no named or ID'ed workspace was provided
            workspace_handle = "default"

        if fail_if_workspace_exists:
            logging.info(
                f"[Client] Creating workspace with handle/id: {workspace_handle}/{workspace_id}."
            )
        else:
            logging.info(
                f"[Client] Creating/Fetching workspace with handle/id: {workspace_handle}/{workspace_id}."
            )

        # Zero out the workspace_handle on the config block in case we're being invoked from
        # `init` (otherwise we'll attempt to create the space IN that non-existant workspace)
        old_workspace_handle = self.config.workspace_handle
        self.config.workspace_handle = None

        if trust_workspace_config:
            if workspace_handle is None or workspace_id is None:
                raise SteamshipError(
                    message="Attempted a trusted workspace switch without providing both workspace handle and workspace id."
                )
            return_id = workspace_id
            return_handle = workspace_handle
        else:
            try:
                if workspace_handle is not None and workspace_id is not None:
                    get_params = {
                        "handle": workspace_handle,
                        "id": workspace_id,
                        "fetchIfExists": False,
                    }
                    workspace = self.post("workspace/get", get_params)
                elif workspace_handle is not None:
                    get_params = {
                        "handle": workspace_handle,
                        "fetchIfExists": not fail_if_workspace_exists,
                    }
                    workspace = self.post("workspace/create", get_params)
                elif workspace_id is not None:
                    get_params = {"id": workspace_id}
                    workspace = self.post("workspace/get", get_params)

            except SteamshipError as e:
                self.config.workspace_handle = old_workspace_handle
                raise e

            if workspace is None:
                raise SteamshipError(
                    message="Was unable to switch to new workspace: server returned empty Workspace."
                )

            return_id = workspace.get("workspace", {}).get("id")
            return_handle = workspace.get("workspace", {}).get("handle")

        if return_id is None or return_handle is None:
            raise SteamshipError(
                message="Was unable to switch to new workspace: server returned empty ID and Handle."
            )

        # Finally, set the new workspace.
        self.config.workspace_id = return_id
        self.config.workspace_handle = return_handle
        logging.info(f"[Client] Switched to workspace {return_handle}/{return_id}")

    def dict(self, **kwargs) -> dict:
        # Because of the trick we do to hack these in as both static and member methods (with different
        # implementations), Pydantic will try to include them by default. So we have to suppress that otherwise
        # downstream serialization into JSON will fail.
        if "exclude" not in kwargs:
            kwargs["exclude"] = {
                "use": True,
                "use_plugin": True,
                "_instance_use": True,
                "_instance_use_plugin": True,
                "config": {"api_key"},
            }
        elif isinstance(kwargs["exclude"], set):
            kwargs["exclude"].add("use")
            kwargs["exclude"].add("use_plugin")
            kwargs["exclude"].add("_instance_use")
            kwargs["exclude"].add("_instance_use_plugin")
            kwargs["exclude"].add(
                "config"
            )  # the set version cannot exclude subcomponents; we must remove all config
        elif isinstance(kwargs["exclude"], dict):
            kwargs["exclude"]["use"] = True
            kwargs["exclude"]["use_plugin"] = True
            kwargs["exclude"]["_instance_use"] = True
            kwargs["exclude"]["_instance_use_plugin"] = True
            kwargs["exclude"]["config"] = {"api_key"}

        return super().dict(**kwargs)

    def _url(
        self,
        is_package_call: bool = False,
        package_owner: str = None,
        operation: str = None,
    ):
        if not is_package_call:
            # Regular API call
            base = self.config.api_base
        else:
            # Do the invocable version
            if package_owner is None:
                return SteamshipError(
                    code="UserMissing",
                    message="Cannot invoke a package endpoint without the package owner's user handle.",
                    suggestion="Provide the package_owner option, or initialize your package with a lookup.",
                )

            base = self.config.app_base
            if not is_local(base):
                # We want to prepend the user handle
                parts = base.split("//")
                base = f"{parts[0]}//{package_owner}.{'//'.join(parts[1:])}"

        # Clean leading and trailing "/"
        if base[len(base) - 1] == "/":
            base = base[:-1]
        if operation[0] == "/":
            operation = operation[1:]

        return f"{base}/{operation}"

    def _headers(  # noqa: C901
        self,
        is_package_call: bool = False,
        package_owner: str = None,
        package_id: str = None,
        package_instance_id: str = None,
        as_background_task: bool = False,
        wait_on_tasks: List[Union[str, Task]] = None,
    ):
        headers = {"Authorization": f"Bearer {self.config.api_key.get_secret_value()}"}

        if self.config.workspace_id:
            headers["X-Workspace-Id"] = self.config.workspace_id
        elif self.config.workspace_handle:
            headers["X-Workspace-Handle"] = self.config.workspace_handle

        if is_package_call:
            if package_owner:
                headers["X-Package-Owner-Handle"] = package_owner
            if package_id:
                headers["X-Package-Id"] = package_id
            if package_instance_id:
                headers["X-Package-Instance-Id"] = package_instance_id

        if wait_on_tasks:
            # Will result in the engine persisting the inbound HTTP request as a Task for deferred
            # execution. Additionally, the task will be scheduled to first wait on the other tasks
            # provided in the list of IDs. Accepts a list of EITHER Task objects OR task_id strings.
            as_background_task = True
            task_ids = []
            for task_or_id in wait_on_tasks:
                if isinstance(task_or_id, str):
                    task_ids.append(task_or_id)
                elif isinstance(task_or_id, Task):
                    task_ids.append(task_or_id.task_id)
                else:
                    raise SteamshipError(
                        message=f"`wait_on_tasks` should only contain Task or str objects. Got a {type(task_or_id)}."
                    )

            headers["X-Task-Dependency"] = ",".join(task_ids)

        if as_background_task:
            # Will result in the engine persisting the inbound HTTP request as a Task for deferred
            # execution. The client will receive task information back instead of the synchronous API response.
            # That task can be polled for eventual completion.
            headers["X-Task-Background"] = "true"

        return headers

    @staticmethod
    def _prepare_data(payload: Union[Request, dict]):
        if payload is None:
            data = {}
        elif isinstance(payload, dict):
            data = payload
        elif isinstance(payload, BaseModel):
            data = payload.dict(by_alias=True)
        else:
            raise RuntimeError(f"Unable to parse payload of type {type(payload)}")

        return data

    @staticmethod
    def _response_data(resp, raw_response: bool = False):
        if resp is None:
            return None

        if raw_response:
            return resp.content

        if resp.headers:
            ct = None
            if "Content-Type" in resp.headers:
                ct = resp.headers["Content-Type"]
            if "content-type" in resp.headers:
                ct = resp.headers["content-type"]
            if ct is not None:
                ct = ct.split(";")[0]  # application/json; charset=utf-8
                if ct in [MimeTypes.TXT, MimeTypes.MKD, MimeTypes.HTML]:
                    return resp.text
                elif ct == MimeTypes.JSON:
                    return resp.json()
                else:
                    return resp.content

    @staticmethod
    def _prepare_multipart_data(data, file):
        # Note: requests seems to have a bug passing boolean (and maybe numeric?)
        # values in the midst of multipart form data. You need to manually convert
        # it to a string; otherwise it will pass as False or True (with the capital),
        # which is not standard notation outside of Python.
        for key in data:
            if data[key] is False:
                data[key] = "false"
            elif data[key] is True:
                data[key] = "true"

        result = {}
        for key, val in data.items():
            for t in _multipart_name(key, val):
                result[t[0]] = t
        result["file"] = file
        return result

    def _add_client_to_response(self, expect: Type, response_data: Any):
        if isinstance(response_data, dict):
            self._add_client_to_object(expect, response_data)
        elif isinstance(response_data, list):
            for el in response_data:
                typing_parameters = typing.get_args(expect)
                self._add_client_to_response(
                    typing_parameters[0] if typing_parameters else None, el
                )

        return response_data

    def _add_client_to_object(self, expect, response_data):
        if expect and isclass(expect):
            if len(response_data.keys()) == 1 and list(response_data.keys())[0] in (
                to_camel(expect.__name__),
                to_camel(expect.__name__).replace("package", "invocable"),
                # Hack since engine uses "App" instead of "Package"
                "index",
                "pluginInstance",  # Inlined here since `expect` may be a subclass of pluginInstance
            ):
                # TODO (enias): Hack since the engine responds with incosistent formats e.g. {"plugin" : {plugin_fields}}
                for _, v in response_data.items():
                    self._add_client_to_response(expect, v)
            elif issubclass(expect, BaseModel):
                response_data["client"] = self
                try:
                    key_to_type = typing.get_type_hints(expect)
                    for k, v in response_data.items():
                        self._add_client_to_response(key_to_type.get(inflection.underscore(k)), v)
                except NameError:
                    # typing.get_type_hints fails for Workspace
                    pass

    def call(  # noqa: C901
        self,
        verb: Verb,
        operation: str,
        payload: Union[Request, dict] = None,
        file: Any = None,
        expect: Type[T] = None,
        debug: bool = False,
        raw_response: bool = False,
        is_package_call: bool = False,
        package_owner: str = None,
        package_id: str = None,
        package_instance_id: str = None,
        as_background_task: bool = False,
        wait_on_tasks: List[Union[str, Task]] = None,
        timeout_s: Optional[float] = None,
    ) -> Union[
        Any, Task
    ]:  # TODO (enias): I would like to list all possible return types using interfaces instead of Any
        """Post to the Steamship API.

        All responses have the format::

        .. code-block:: json

           {
               "data": "<actual response>",
               "error": {"reason": "<message>"}
           } # noqa: RST203

        For the Python client we return the contents of the `data` field if present, and we raise an exception
        if the `error` field is filled in.
        """
        # TODO (enias): Review this codebase
        url = self._url(
            is_package_call=is_package_call,
            package_owner=package_owner,
            operation=operation,
        )

        headers = self._headers(
            is_package_call=is_package_call,
            package_owner=package_owner,
            package_id=package_id,
            package_instance_id=package_instance_id,
            as_background_task=as_background_task,
            wait_on_tasks=wait_on_tasks,
        )

        data = self._prepare_data(payload=payload)

        logging.debug(
            f"Making {verb} to {url} in workspace {self.config.workspace_handle}/{self.config.workspace_id}"
        )
        if verb == Verb.POST:
            if file is not None:
                files = self._prepare_multipart_data(data, file)
                resp = self._session.post(url, files=files, headers=headers, timeout=timeout_s)
            else:
                resp = self._session.post(url, json=data, headers=headers, timeout=timeout_s)
        elif verb == Verb.GET:
            resp = self._session.get(url, params=data, headers=headers, timeout=timeout_s)
        else:
            raise Exception(f"Unsupported verb: {verb}")

        logging.debug(f"From {verb} to {url} got HTTP {resp.status_code}")

        if debug is True:
            logging.debug(f"Got response {resp}")

        response_data = self._response_data(resp, raw_response=raw_response)

        logging.debug(f"Response JSON {response_data}")

        task = None
        error = None

        if isinstance(response_data, dict):
            if "status" in response_data:
                try:
                    task = Task.parse_obj(
                        {**response_data["status"], "client": self, "expect": expect}
                    )
                    if "state" in response_data["status"]:
                        if response_data["status"]["state"] == "failed":
                            error = SteamshipError.from_dict(response_data["status"])
                            logging.warning(f"Client received error from server: {error}")
                except TypeError as e:
                    # There's an edge case here -- if a Steamship package returns the JSON dictionary
                    #
                    # { "status": "status string" }
                    #
                    # Then the above handler will attempt to parse it and throw... But we don't actually want to throw
                    # since we don't take a strong opinion on what the response type of a package endpoint ought to be.
                    # It *may* choose to conform to the SteamshipResponse<T> type, but it doesn't have to.
                    if not is_package_call:
                        raise e

                if task is not None and task.state == TaskState.failed:
                    error = task.as_error()

            if "data" in response_data:
                if expect is not None:
                    if issubclass(expect, SteamshipError):
                        data = expect.from_dict({**response_data["data"], "client": self})
                    elif issubclass(expect, BaseModel):
                        data = expect.parse_obj(
                            self._add_client_to_response(expect, response_data["data"])
                        )
                    else:
                        raise RuntimeError(f"obj of type {expect} does not have a from_dict method")
                else:
                    data = response_data["data"]

                if task:
                    task.output = data
            else:
                data = response_data

        else:
            data = response_data

        if error is not None:
            logging.warning(f"Client received error from server: {error}", exc_info=error)
            raise error

        if not resp.ok:
            raise SteamshipError(
                f"API call did not complete successfully.  Server returned: {response_data}"
            )

        elif task is not None:
            return task
        elif data is not None and expect is not None:
            # if we have data AND we expect it to be of a certain type,
            # we should probably make sure that expectation is met.
            if not isinstance(data, expect):
                raise SteamshipError(
                    message=f"Inconsistent response from server (data does not match expected type: {expect}.)",
                    suggestion="Please contact support via [email protected] and report what caused this error.",
                )
            return data
        elif data is not None:
            return data
        else:
            raise SteamshipError("Inconsistent response from server. Please contact support.")

    def post(
        self,
        operation: str,
        payload: Union[Request, dict, BaseModel] = None,
        file: Any = None,
        expect: Any = None,
        debug: bool = False,
        raw_response: bool = False,
        is_package_call: bool = False,
        package_owner: str = None,
        package_id: str = None,
        package_instance_id: str = None,
        as_background_task: bool = False,
        wait_on_tasks: List[Union[str, Task]] = None,
        timeout_s: Optional[float] = None,
    ) -> Union[
        Any, Task
    ]:  # TODO (enias): I would like to list all possible return types using interfaces instead of Any
        return self.call(
            verb=Verb.POST,
            operation=operation,
            payload=payload,
            file=file,
            expect=expect,
            debug=debug,
            raw_response=raw_response,
            is_package_call=is_package_call,
            package_owner=package_owner,
            package_id=package_id,
            package_instance_id=package_instance_id,
            as_background_task=as_background_task,
            wait_on_tasks=wait_on_tasks,
            timeout_s=timeout_s,
        )

    def get(
        self,
        operation: str,
        payload: Union[Request, dict] = None,
        file: Any = None,
        expect: Any = None,
        debug: bool = False,
        raw_response: bool = False,
        is_package_call: bool = False,
        package_owner: str = None,
        package_id: str = None,
        package_instance_id: str = None,
        as_background_task: bool = False,
        wait_on_tasks: List[Union[str, Task]] = None,
        timeout_s: Optional[float] = None,
    ) -> Union[
        Any, Task
    ]:  # TODO (enias): I would like to list all possible return types using interfaces instead of Any
        return self.call(
            verb=Verb.GET,
            operation=operation,
            payload=payload,
            file=file,
            expect=expect,
            debug=debug,
            raw_response=raw_response,
            is_package_call=is_package_call,
            package_owner=package_owner,
            package_id=package_id,
            package_instance_id=package_instance_id,
            as_background_task=as_background_task,
            wait_on_tasks=wait_on_tasks,
            timeout_s=timeout_s,
        )

    def logs(
        self,
        offset: int = 0,
        number: int = 50,
        invocable_handle: Optional[str] = None,
        instance_handle: Optional[str] = None,
        invocable_version_handle: Optional[str] = None,
        path: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Return generated logs for a client.

        The logs will be workspace-scoped. You will only receive logs
        for packages and plugins that you own.

        :param offset: The index of the first log entry to return. This can be used with `number` to page through logs.
        :param number: The number of log entries to return. This can be used with `offset` to page through logs.
        :param invocable_handle: Enables optional filtering based on the handle of package or plugin. Example: `my-package`
        :param instance_handle: Enables optional filtering based on the handle of package instance or plugin instance. Example: `my-instance`
        :param invocable_version_handle: Enables optional filtering based on the version handle of package or plugin. Example: `0.0.2`
        :param path: Enables optional filtering based on request path. Example: `/generate`.
        :return: Returns a dictionary containing the offset and number of log entries as well as a list of `entries` that match the specificed filters.
        """
        args = {"from": offset, "size": number}
        if invocable_handle:
            args["invocableHandle"] = invocable_handle
        if instance_handle:
            args["invocableInstanceHandle"] = instance_handle
        if invocable_version_handle:
            args["invocableVersionHandle"] = invocable_version_handle
        if path:
            args["invocablePath"] = path

        return self.post("logs/list", args)