Source code for steamship.base.configuration

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Optional

import inflection
from pydantic import AnyHttpUrl, SecretStr

from steamship.base.model import CamelModel, to_camel
from steamship.cli.login import login
from steamship.utils.utils import format_uri

DEFAULT_WEB_BASE = "https://steamship.com/"
DEFAULT_APP_BASE = "https://steamship.run/"
DEFAULT_API_BASE = "https://api.steamship.com/api/v1/"

ENVIRONMENT_VARIABLES_TO_PROPERTY = {
    "STEAMSHIP_API_KEY": "api_key",
    "STEAMSHIP_API_BASE": "api_base",
    "STEAMSHIP_APP_BASE": "app_base",
    "STEAMSHIP_WEB_BASE": "web_base",
    "STEAMSHIP_WORKSPACE_ID": "workspace_id",
    "STEAMSHIP_WORKSPACE_HANDLE": "workspace_handle",
}
DEFAULT_CONFIG_FILE = Path.home() / ".steamship.json"

# This stops us from including the `client` object in the dict() output, which is fine in a dict()
# but explodes if that dict() is turned into JSON. Sadly the `exclude` option in Pydantic doesn't
# cascade down nested objects, so we have to use this structure to catch all the possible combinations
EXCLUDE_FROM_DICT = {
    "client": True,
    "blocks": {"__all__": {"client": True, "tags": {"__all__": {"client": True}}}},
    "tags": {"__all__": {"client": True}},
}




class Configuration(CamelModel):
    api_key: SecretStr
    api_base: AnyHttpUrl = DEFAULT_API_BASE
    app_base: AnyHttpUrl = DEFAULT_APP_BASE
    web_base: AnyHttpUrl = DEFAULT_WEB_BASE
    workspace_id: str = None
    workspace_handle: str = None
    profile: Optional[str] = None

    # For use in deployed packages and plugins for tracing. Do not set manually
    request_id: Optional[str] = None

    def __init__(
        self,
        config_file: Optional[Path] = None,
        **kwargs,
    ):
        # First set the profile
        kwargs["profile"] = profile = kwargs.get("profile") or os.getenv("STEAMSHIP_PROFILE")

        # Then load configuration from a file if provided
        config_dict = self._load_from_file(
            config_file or DEFAULT_CONFIG_FILE,
            profile,
            raise_on_exception=config_file is not None,
        )
        config_dict.update(self._get_config_dict_from_environment())
        kwargs.update({k: v for k, v in config_dict.items() if kwargs.get(k) is None})
        kwargs = {to_camel(k): v for k, v in kwargs.items()}

        kwargs["apiBase"] = format_uri(kwargs.get("apiBase"))
        kwargs["appBase"] = format_uri(kwargs.get("appBase"))
        kwargs["webBase"] = format_uri(kwargs.get("webBase"))

        if not kwargs.get("apiKey"):
            api_key = login(
                kwargs.get("apiBase") or DEFAULT_API_BASE,
                kwargs.get("webBase") or DEFAULT_WEB_BASE,
            )
            Configuration._save_api_key_to_file(
                api_key, profile, config_file or DEFAULT_CONFIG_FILE
            )
            kwargs["apiKey"] = api_key

        super().__init__(**kwargs)

    @staticmethod
    def _load_from_file(
        file: Path, profile: str = None, raise_on_exception: bool = False
    ) -> Optional[dict]:
        try:
            with file.open() as f:
                config_file = json.load(f)
                if profile:
                    if "profiles" not in config_file or profile not in config_file["profiles"]:
                        raise RuntimeError(f"Profile {profile} requested but not found in {file}")
                    config = config_file["profiles"][profile]
                else:
                    config = config_file
                return {inflection.underscore(k): v for k, v in config.items()}
        except FileNotFoundError:
            if raise_on_exception:
                raise Exception(f"Tried to load configuration file at {file} but it did not exist.")
        except Exception as err:
            if raise_on_exception:
                raise err
        return {}

    @staticmethod
    def _get_config_dict_from_environment():
        """Overrides configuration with environment variables."""
        return {
            property_name: os.getenv(environment_variable_name, None)
            for environment_variable_name, property_name in ENVIRONMENT_VARIABLES_TO_PROPERTY.items()
            if environment_variable_name in os.environ
        }

    @staticmethod
    def _save_api_key_to_file(new_api_key: Optional[str], profile: Optional[str], file_path: Path):
        # Minimally rewrite config file, adding api key
        try:
            with file_path.open() as f:
                config_file = json.load(f)
            if profile:
                if "profiles" not in config_file or profile not in config_file["profiles"]:
                    raise RuntimeError(f"Could not update API key for {profile} in {file_path}")
                config = config_file["profiles"][profile]
            else:
                config = config_file
        except FileNotFoundError:
            config_file = {}
            config = config_file

        config["apiKey"] = new_api_key

        with file_path.open("w") as f:
            json.dump(config_file, f, indent="\t")



    @staticmethod
    def default_config_file_has_api_key() -> bool:
        return Configuration._load_from_file(DEFAULT_CONFIG_FILE).get("api_key") is not None




    @staticmethod
    def remove_api_key_from_default_config():
        Configuration._save_api_key_to_file(None, None, DEFAULT_CONFIG_FILE)