From 2dd2b31f4152225b0f379f50f5098523c08d323b Mon Sep 17 00:00:00 2001 From: Aliaksandr Kuzmik Date: Tue, 17 Sep 2024 00:58:47 +0200 Subject: [PATCH] Draft fix --- sdks/python/src/opik/config.py | 18 +++- sdks/python/src/opik/opik_login.py | 120 ++++++++++++++-------- sdks/python/tests/unit/test_opik_login.py | 100 ++++++++++++++---- 3 files changed, 169 insertions(+), 69 deletions(-) diff --git a/sdks/python/src/opik/config.py b/sdks/python/src/opik/config.py index f7e961b51..33f70ad59 100644 --- a/sdks/python/src/opik/config.py +++ b/sdks/python/src/opik/config.py @@ -9,7 +9,12 @@ from . import dict_utils -PathType = Union[pathlib.Path, str, List[Union[pathlib.Path, str]], Tuple[Union[pathlib.Path, str], ...]] +PathType = Union[ + pathlib.Path, + str, + List[Union[pathlib.Path, str]], + Tuple[Union[pathlib.Path, str], ...], +] _SESSION_CACHE_DICT: Dict[str, Any] = {} @@ -19,7 +24,7 @@ OPIK_PROJECT_DEFAULT_NAME: Final[str] = "Default Project" OPIK_WORKSPACE_DEFAULT_NAME: Final[str] = "default" -CONFIG_FILE_PATH_DEFAULT: Final[str] = '~/.opik.config' +CONFIG_FILE_PATH_DEFAULT: Final[str] = "~/.opik.config" LOGGER = logging.getLogger(__name__) @@ -39,7 +44,9 @@ def __init__( def _read_file(self, file_path: pathlib.Path) -> dict[str, Any]: config = configparser.ConfigParser() config.read(file_path) - config_values = {section: dict(config.items(section)) for section in config.sections()} + config_values = { + section: dict(config.items(section)) for section in config.sections() + } if "opik" in config_values: return config_values["opik"] @@ -151,13 +158,14 @@ def save_to_file(self) -> None: if self.api_key is not None: config_file_content["opik"]["api_key"] = self.api_key - with open(self.config_file_fullpath, mode="w+", encoding="utf-8") as config_file: + with open( + self.config_file_fullpath, mode="w+", encoding="utf-8" + ) as config_file: config_file_content.write(config_file) LOGGER.info(f"Saved configuration to a file: {self.config_file_fullpath}") - def update_session_config(key: str, value: Any) -> None: _SESSION_CACHE_DICT[key] = value diff --git a/sdks/python/src/opik/opik_login.py b/sdks/python/src/opik/opik_login.py index 378a21a46..c16d63142 100644 --- a/sdks/python/src/opik/opik_login.py +++ b/sdks/python/src/opik/opik_login.py @@ -1,9 +1,6 @@ -""" - -""" +""" """ import logging -import sys from getpass import getpass from typing import Final, Optional @@ -11,12 +8,16 @@ import opik.config from opik import httpx_client -from opik.config import OPIK_BASE_URL_LOCAL, OPIK_BASE_URL_CLOUD, OPIK_WORKSPACE_DEFAULT_NAME +from opik.config import ( + OPIK_BASE_URL_LOCAL, + OPIK_BASE_URL_CLOUD, + OPIK_WORKSPACE_DEFAULT_NAME, +) from opik.exceptions import ConfigurationError LOGGER = logging.getLogger(__name__) -HEALTH_CHECK_URL_POSTFIX: Final[str] = '/is-alive/ping' +HEALTH_CHECK_URL_POSTFIX: Final[str] = "/is-alive/ping" HEALTH_CHECK_TIMEOUT: Final[float] = 1.0 @@ -24,7 +25,7 @@ def is_interactive() -> bool: """ Returns True if in interactive mode """ - #return bool(getattr(sys, "ps1", sys.flags.interactive)) + # return bool(getattr(sys, "ps1", sys.flags.interactive)) return True @@ -58,9 +59,11 @@ def is_workspace_name_correct(api_key: str, workspace: str) -> bool: url = "https://www.comet.com/api/rest/v2/workspaces" client = httpx.Client() - client.headers.update({ - "Authorization": f"{api_key}", - }) + client.headers.update( + { + "Authorization": f"{api_key}", + } + ) try: response = client.get(url=url) @@ -88,9 +91,11 @@ def is_api_key_correct(api_key: str) -> bool: url = "https://www.comet.com/api/rest/v2/account-details" client = httpx.Client() - client.headers.update({ - "Authorization": f"{api_key}", - }) + client.headers.update( + { + "Authorization": f"{api_key}", + } + ) try: response = client.get(url=url) @@ -116,9 +121,11 @@ def get_default_workspace(api_key: str) -> str: url = "https://www.comet.com/api/rest/v2/account-details" client = httpx.Client() - client.headers.update({ - "Authorization": f"{api_key}", - }) + client.headers.update( + { + "Authorization": f"{api_key}", + } + ) try: response = client.get(url=url) @@ -126,15 +133,17 @@ def get_default_workspace(api_key: str) -> str: raise ConnectionError(f"Error while getting default workspace name: {str(e)}") if response.status_code != 200: - raise ConnectionError(f"Error while getting default workspace name: {response.text}") + raise ConnectionError( + f"Error while getting default workspace name: {response.text}" + ) return response.json()["defaultWorkspaceName"] def _update_config( - api_key: Optional[str], - url: str, - workspace: str, + api_key: Optional[str], + url: str, + workspace: str, ) -> None: """ Save changes to config file and update current session config @@ -183,7 +192,9 @@ def _ask_for_url() -> str: LOGGER.error("Wrong URL. Please try again.") retries -= 1 - raise ConfigurationError("Can't use URL provided by user. Opik instance is not active or not found.") + raise ConfigurationError( + "Can't use URL provided by user. Opik instance is not active or not found." + ) def _ask_for_api_key() -> str: @@ -211,7 +222,9 @@ def _ask_for_workspace(api_key: str) -> str: retries = 3 while retries > 0: - user_input_workspace = input("Please enter your cloud Opik instance workspace name: ") + user_input_workspace = input( + "Please enter your cloud Opik instance workspace name: " + ) if is_workspace_name_correct(api_key, user_input_workspace): return user_input_workspace @@ -237,10 +250,10 @@ def ask_user_for_approval(message: str) -> bool: def login( - api_key: Optional[str] = None, - url: Optional[str] = None, - workspace: Optional[str] = None, - use_local: bool = False + api_key: Optional[str] = None, + url: Optional[str] = None, + workspace: Optional[str] = None, + use_local: bool = False, ) -> None: """ Args: @@ -267,8 +280,8 @@ def login( def _login_cloud( - api_key: Optional[str], - workspace: Optional[str], + api_key: Optional[str], + workspace: Optional[str], ) -> None: """ Login to cloud Opik instance @@ -280,35 +293,51 @@ def _login_cloud( if is_interactive() is False and api_key is None and current_config.api_key is None: raise ConfigurationError("No API key provided for cloud Opik instance.") - if is_interactive() is False and workspace is None and current_config.workspace is None: + if ( + is_interactive() is False + and workspace is None + and current_config.workspace is None + ): raise ConfigurationError("No workspace name provided for cloud Opik instance.") # Ask for API key if api_key is None and current_config.api_key is None: - LOGGER.info("You can find your API key here: https://www.comet.com/api/my/settings/") + LOGGER.info( + "You can find your API key here: https://www.comet.com/api/my/settings/" + ) api_key = _ask_for_api_key() config_file_needs_updating = True elif api_key is None and current_config.api_key is not None: api_key = current_config.api_key - # if workspace already configured - will use this value - if "workspace" in current_config.model_fields_set and current_config.workspace != OPIK_WORKSPACE_DEFAULT_NAME: - workspace = current_config.workspace - - # Check what their default workspace is, and we ask them if they want to use the default workspace - if workspace is None: - default_workspace = get_default_workspace(api_key) - use_default_workspace = ask_user_for_approval(f"Do you want to use \"{default_workspace}\" workspace? Y/n: ") - if use_default_workspace: + # Check passed workspace (if it was passed) + if workspace is not None: + if is_workspace_name_correct(api_key, workspace): config_file_needs_updating = True - workspace = default_workspace + else: + raise ConfigurationError("Workspace `%s` is incorrect for the given API key.", workspace) + else: + # Workspace was not passed, we check if there is already configured value + # if workspace already configured - will use this value + if ( + "workspace" in current_config.model_fields_set + and current_config.workspace != OPIK_WORKSPACE_DEFAULT_NAME + ): + workspace = current_config.workspace + + # Check what their default workspace is, and we ask them if they want to use the default workspace + if workspace is None: + default_workspace = get_default_workspace(api_key) + use_default_workspace = ask_user_for_approval( + f'Do you want to use "{default_workspace}" workspace? Y/n: ' + ) - if not is_workspace_name_correct(api_key, workspace): - LOGGER.warning("Workspace name is incorrect.") + if use_default_workspace: + workspace = default_workspace + else: workspace = _ask_for_workspace(api_key=api_key) - else: - workspace = _ask_for_workspace(api_key=api_key) + config_file_needs_updating = True if config_file_needs_updating: _update_config( @@ -351,7 +380,8 @@ def _login_local(url: Optional[str]) -> None: return use_url = ask_user_for_approval( - f"Found local Opik instance on: {OPIK_BASE_URL_LOCAL}\nDo you want to use it? Y/n: ") + f"Found local Opik instance on: {OPIK_BASE_URL_LOCAL}\nDo you want to use it? Y/n: " + ) if use_url: _update_config( diff --git a/sdks/python/tests/unit/test_opik_login.py b/sdks/python/tests/unit/test_opik_login.py index 8ee115918..8fb1f23b8 100644 --- a/sdks/python/tests/unit/test_opik_login.py +++ b/sdks/python/tests/unit/test_opik_login.py @@ -7,28 +7,90 @@ @pytest.mark.parametrize( "api_key, url, workspace, local, should_raise", [ - (None, "http://example.com", "workspace1", True, False), # Missing api_key, local=True - (None, "http://example.com", "workspace1", False, True), # Missing api_key, local=False - ("apikey123", None, "workspace1", True, True), # Missing url, local=True - ("apikey123", None, "workspace1", False, True), # Missing url, local=False - ("apikey123", "http://example.com", None, True, True), # Missing workspace, local=True - ("apikey123", "http://example.com", None, False, True), # Missing workspace, local=False - (None, None, "workspace1", True, True), # Missing api_key and url, local=True - (None, None, "workspace1", False, True), # Missing api_key and url, local=False - (None, "http://example.com", None, True, True), # Missing api_key and workspace, local=True - (None, "http://example.com", None, False, True), # Missing api_key and workspace, local=False - ("apikey123", None, None, True, True), # Missing url and workspace, local=True - ("apikey123", None, None, False, True), # Missing url and workspace, local=False - (None, None, None, True, True), # All missing, local=True - (None, None, None, False, True), # All missing, local=False - ("apikey123", "http://example.com", "workspace1", True, False), # All present, local=True - ("apikey123", "http://example.com", "workspace1", False, False), # All present, local=False - ] + ( + None, + "http://example.com", + "workspace1", + True, + False, + ), # Missing api_key, local=True + ( + None, + "http://example.com", + "workspace1", + False, + True, + ), # Missing api_key, local=False + ("apikey123", None, "workspace1", True, True), # Missing url, local=True + ("apikey123", None, "workspace1", False, True), # Missing url, local=False + ( + "apikey123", + "http://example.com", + None, + True, + True, + ), # Missing workspace, local=True + ( + "apikey123", + "http://example.com", + None, + False, + True, + ), # Missing workspace, local=False + (None, None, "workspace1", True, True), # Missing api_key and url, local=True + (None, None, "workspace1", False, True), # Missing api_key and url, local=False + ( + None, + "http://example.com", + None, + True, + True, + ), # Missing api_key and workspace, local=True + ( + None, + "http://example.com", + None, + False, + True, + ), # Missing api_key and workspace, local=False + ("apikey123", None, None, True, True), # Missing url and workspace, local=True + ( + "apikey123", + None, + None, + False, + True, + ), # Missing url and workspace, local=False + (None, None, None, True, True), # All missing, local=True + (None, None, None, False, True), # All missing, local=False + ( + "apikey123", + "http://example.com", + "workspace1", + True, + False, + ), # All present, local=True + ( + "apikey123", + "http://example.com", + "workspace1", + False, + False, + ), # All present, local=False + ], ) def test_login__force_new_settings__fail(api_key, url, workspace, local, should_raise): if should_raise: with pytest.raises(ConfigurationError): - login(api_key=api_key, url=url, workspace=workspace, force=True, use_local=local) + login( + api_key=api_key, + url=url, + workspace=workspace, + force=True, + use_local=local, + ) else: # No exception should be raised - login(api_key=api_key, url=url, workspace=workspace, force=True, use_local=local) + login( + api_key=api_key, url=url, workspace=workspace, force=True, use_local=local + )