qwak-core 0.4.272__py3-none-any.whl → 0.4.274__py3-none-any.whl

frogml_storage/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.4.274"

frogml_storage/artifactory/__init__.py

@@ -0,0 +1 @@
+from ._artifactory_api import ArtifactoryApi, StartTransactionResponse

frogml_storage/artifactory/_artifactory_api.py

@@ -0,0 +1,315 @@
+import json
+import os
+from typing import Optional
+
+from requests import Response
+from tqdm.auto import tqdm
+from tqdm.utils import CallbackIOWrapper
+from urllib3 import Retry
+
+from frogml_storage.logging import logger
+from frogml_storage.utils import join_url
+from frogml_storage.constants import CHECKSUM_SHA2_HEADER
+from frogml_storage.models.entity_manifest import Checksums, EntityManifest
+from frogml_storage.models.frogml_entity_type_info import FrogMLEntityTypeInfo
+from frogml_storage.http import HTTPClient
+from frogml_storage.models import DownloadContext
+
+
+class StartTransactionResponse:
+    files_upload_path: str
+    lead_upload_path: str
+    dependencies_upload_path: str
+    code_upload_path: str
+    transaction_id: str
+
+    def __init__(
+        self,
+        files_upload_path,
+        lead_upload_path,
+        dependencies_upload_path,
+        code_upload_path,
+        transaction_id,
+    ):
+        self.files_upload_path = files_upload_path
+        self.lead_upload_path = lead_upload_path
+        self.dependencies_upload_path = dependencies_upload_path
+        self.code_upload_path = code_upload_path
+        self.transaction_id = transaction_id
+
+
+class ArtifactoryApi:
+    def __init__(self, uri, auth=None, http_client=None):
+        self.uri = uri
+        if http_client is not None:
+            self.http_client = http_client
+        else:
+            self.auth = auth
+            self.http_client = HTTPClient(auth=auth)
+
+    def start_transaction(
+        self,
+        entity_type_info: FrogMLEntityTypeInfo,
+        repository: str,
+        entity_name: str,
+        version: Optional[str],
+    ) -> StartTransactionResponse:
+        """
+        Initializes an upload. Returns transaction ID and upload path
+        """
+        if version is None:
+            start_transaction_url = (
+                f"{self.uri}/api/machinelearning/{repository}/"
+                f"{entity_type_info.entity_type}/{entity_name}/start-transaction"
+            )
+        else:
+            start_transaction_url = (
+                f"{self.uri}/api/machinelearning/{repository}/{entity_type_info.entity_type}"
+                f"/{entity_name}/start-transaction/{version}"
+            )
+        try:
+            response = self.http_client.post(start_transaction_url)
+            response.raise_for_status()
+            files_upload_path = response.json()["filesUploadPath"]
+            lead_upload_path = response.json()["leadUploadPath"]
+            dependencies_upload_path = response.json()["dependenciesUploadPath"]
+            code_upload_path = response.json()["codeUploadPath"]
+            transaction_id = response.json()["transactionId"]
+        except Exception as exception:
+            err = (
+                f"Error occurred while trying to start an upload transaction for "
+                f"{entity_type_info.entity_type}: '{entity_name}'"
+                f" Error: '{exception}'"
+            )
+            logger.error(err, exc_info=False)
+            raise exception
+        return StartTransactionResponse(
+            files_upload_path=files_upload_path,
+            lead_upload_path=lead_upload_path,
+            dependencies_upload_path=dependencies_upload_path,
+            code_upload_path=code_upload_path,
+            transaction_id=transaction_id,
+        )
+
+    def end_transaction(
+        self,
+        entity_type: FrogMLEntityTypeInfo,
+        repository: str,
+        entity_name: str,
+        entity_manifest: EntityManifest,
+        transaction_id: str,
+        version: str,
+        properties: Optional[dict[str, str]],
+    ) -> None:
+        """
+        Upload model-manifest.json | dataset-manifest.json file, makes the model | dataset available in the repository
+        """
+        filename = entity_type.metadata_file_name
+
+        url = join_url(
+            self.uri,
+            "api",
+            "machinelearning",
+            repository,
+            entity_type.entity_type,
+            "entity-manifest",
+            entity_name,
+            version,
+            transaction_id,
+        )
+
+        json_entity_manifest = entity_manifest.to_json()
+        self.upload_entity_manifest(
+            entity_type=entity_type,
+            filename=filename,
+            payload=json_entity_manifest,
+            url=url,
+            properties=properties,
+        )
+
+    def download_file(self, args: DownloadContext) -> None:
+        filename = os.path.basename(args.target_path)
+        try:
+            url = f"{self.uri}/{args.repo_key}/{args.source_url}"
+            with self.http_client.get(url=url, stream=True) as response:
+                response.raise_for_status()
+                total_size = int(response.headers.get("content-length", 0))
+                with open(args.target_path, "wb") as file:
+                    with self.__initialize_progress_bar(total_size, filename) as pbar:
+                        for chunk in response.iter_content(chunk_size=8192):
+                            if chunk:
+                                file.write(chunk)
+                                pbar.update(len(chunk))
+
+        except Exception as exception:
+            self.__handle_download_exception(exception, args.target_path, filename)
+
+    def get_entity_manifest(
+        self,
+        entity_type_info: FrogMLEntityTypeInfo,
+        repository: str,
+        entity_name: str,
+        namespace: Optional[str],
+        version: Optional[str],
+    ) -> dict:
+        url = join_url(
+            self.uri,
+            "api",
+            "machinelearning",
+            repository,
+            entity_type_info.entity_type,
+            "entity-manifest",
+            namespace,
+            entity_name,
+            version,
+        )
+        try:
+            with self.http_client.get(url=url) as r:
+                r.raise_for_status()
+                return r.json()
+        except Exception as exception:
+            err = f"Error occurred while trying to get {entity_type_info.entity_type} info file. Error: '{exception}'"
+            logger.error(err, exc_info=False)
+            raise exception
+
+    @staticmethod
+    def __handle_download_exception(
+        exception: Exception, target_path: str, filename: str
+    ) -> None:
+        if os.path.exists(target_path):
+            os.remove(target_path)
+        err = f"Error occurred while trying to download file: '{filename}' Error: '{exception}'"
+        logger.error(err, exc_info=False)
+        raise exception
+
+    def get_artifact_checksum(self, download_context: DownloadContext) -> str:
+        url = f"{self.uri}/{download_context.repo_key}/{download_context.source_url}"
+        try:
+            with self.http_client.head(url=url, stream=True) as response:
+                response.raise_for_status()
+                return response.headers.get(CHECKSUM_SHA2_HEADER)
+
+        except Exception as exception:
+            logger.error(exception.__cause__, exc_info=False)
+            raise exception
+
+    def upload_entity_manifest(
+        self,
+        entity_type: FrogMLEntityTypeInfo,
+        filename: str,
+        payload: str,
+        url: str,
+        properties: Optional[dict[str, str]],
+        stream: bool = False,
+    ) -> None:
+        body_part_name = f"{entity_type.body_part_stream}"
+
+        try:
+            files = {
+                f"{body_part_name}": (
+                    f"{body_part_name}",
+                    payload,
+                    "application/octet-stream",
+                ),  # Include the InputStream
+                "additionalData": (
+                    "additionalData",
+                    json.dumps(properties),
+                    "application/octet-stream",
+                ),  # Include the object
+            }
+            with self.http_client.put(url=url, files=files, stream=stream) as response:
+                response.raise_for_status()
+        except Exception as exception:
+            err = f"Error occurred while trying to upload file: '{filename}' Error: '{exception}'"
+            logger.error(err, exc_info=False)
+            raise exception
+
+    def upload_file(self, file_path: str, url: str) -> None:
+        wrapped_file = None
+        try:
+            file_size = os.stat(file_path).st_size
+            with (
+                self.__initialize_progress_bar(file_size, file_path) as pbar,
+                open(file_path, "rb") as file,
+            ):
+                wrapped_file = CallbackIOWrapper(pbar.update, file, "read")
+                with self.http_client.put(url=url, payload=wrapped_file) as response:
+                    response.raise_for_status()
+        except Exception as exception:
+            err = f"Error occurred while trying to upload file: '{file_path}' Error: '{exception}'"
+            logger.error(err, exc_info=False)
+            raise type(exception)(f"{err} File: {file_path}") from exception
+        finally:
+            if wrapped_file is not None:
+                wrapped_file.close()
+
+    def checksum_deployment(
+        self,
+        checksum: Checksums,
+        url: str,
+        full_path: str,
+        stream: bool = False,
+    ) -> bool:
+        response = self.http_client.put(
+            url=url,
+            headers={"X-Checksum-Sha256": checksum.sha2, "X-Checksum-Deploy": "true"},
+            stream=stream,
+        )
+        if response.status_code != 200 and response.status_code != 201:
+            return False
+        else:
+            file_size = os.path.getsize(full_path)
+            pbar = self.__initialize_progress_bar(file_size, full_path)
+            pbar.update(file_size)
+            pbar.close()
+            return True
+
+    @staticmethod
+    def __initialize_progress_bar(total_size: int, filename: str) -> tqdm:
+        return tqdm(
+            total=total_size, unit="B", unit_scale=True, desc=filename, initial=0
+        )
+
+    def encrypt_password(self) -> Response:
+        """
+        returns encrypted password as text
+        """
+        return self.http_client.get(
+            url=join_url(self.uri, "/api/security/encryptedPassword")
+        )
+
+    def ping(self) -> Response:
+        """
+        Sends a ping to Artifactory to validate login status
+        """
+        url = join_url(self.uri, "api/system/ping")
+        return self.http_client.get(url=url)
+
+    def get_artifactory_version(self) -> Response:
+        return self.http_client.get(url=join_url(self.uri, "/api/system/version"))
+
+    def create_machinelearning_local_repo(self, repo_name: str) -> Response:
+        data = {
+            "rclass": "local",
+            "packageType": "machinelearning",
+        }
+        return self.http_client.put(
+            url=join_url(self.uri, "/api/repositories/" + repo_name), json=data
+        )
+
+    def delete_frogml_local_repo(self, repo_name: str) -> Response:
+        return self.http_client.delete(
+            url=join_url(self.uri, "/api/repositories/" + repo_name)
+        )
+
+
+class RetryWithLog(Retry):
+    """
+    Adding extra logs before making a retry request
+    """
+
+    def __init__(self, *args, **kwargs):
+        history = kwargs.get("history")
+        if history is not None:
+            logger.debug(f"Error: ${history[-1].error}\nretrying...")
+        super().__init__(*args, **kwargs)

frogml_storage/authentication/login/__init__.py

@@ -0,0 +1 @@
+from ._login_cli import login as frogml_login

frogml_storage/authentication/login/_login_cli.py

@@ -0,0 +1,239 @@
+from typing import List, Optional, Union
+
+import typer
+
+from frogml_storage.utils import assemble_artifact_url, join_url
+from frogml_storage.authentication.login._login_command import run as run_login
+from frogml_storage.authentication.models import AuthConfig, LoginArguments
+from frogml_storage.authentication.utils import (
+    get_frogml_configuration,
+    parse_cli_config_server,
+    read_jfrog_cli_config,
+)
+from frogml_storage.authentication.utils import (
+    get_list_of_servers_from_config,
+    login_input_checks,
+    is_login_without_params,
+)
+from frogml_storage.constants import CONFIG_FILE_PATH, JFROG_CLI_CONFIG_FILE_PATH
+
+
+def login(
+    url: Optional[str],
+    username: Optional[str],
+    password: Optional[str],
+    token: Optional[str],
+    anonymous: Optional[bool],
+    is_interactive: Optional[bool],
+) -> bool:
+    if anonymous is None:
+        anonymous = False
+    if is_login_without_params(url, username, password, token, anonymous):
+        if not is_interactive:
+            return __login_by_frogml_configuration_file_flow()
+        else:
+            return __interactive_flow()
+    elif login_input_checks(url, username, password, token, anonymous):
+        url = assemble_artifact_url(url)
+        return __login_by_command_line_params(url, username, password, token, anonymous)
+    else:
+        typer.echo(
+            f"{__get_failed_reason(url, username, password, token, anonymous)}Login failed"
+        )
+        return False
+
+
+def __get_failed_reason(
+    url: Optional[str],
+    username: Optional[str],
+    password: Optional[str],
+    token: Optional[str],
+    anonymous: Optional[bool],
+) -> str:
+    details = ""
+    if url is None:
+        details += "Url is required.\n"
+    if token is not None and (
+        any(params is not None for params in [username, password])
+    ):
+        details += "Token is specified when username/password is specified.\n"
+    if anonymous is True and any(
+        params is not None for params in [username, password, token]
+    ):
+        details += "Anonymous is specified with authentication details.\n"
+    if username is None and password is not None:
+        details += "Username is required when password is specified.\n"
+    elif username is not None and password is None:
+        details += "Password is required when username is specified.\n"
+    return details
+
+
+def __run_interactive_mode() -> Union[LoginArguments, None]:
+    jfrog_cli_config = read_jfrog_cli_config()
+
+    if jfrog_cli_config is not None:
+        jfrog_cli_servers = get_list_of_servers_from_config(jfrog_cli_config)
+        if jfrog_cli_servers is not None and len(jfrog_cli_servers) > 0:
+            login_method_id = typer.prompt(
+                f"Please select from the following options:\n"  # nosec B608
+                f"1.Login by jfrog-cli configuration file: {JFROG_CLI_CONFIG_FILE_PATH}\n"
+                f"2.Connecting to a new server\n"
+            )
+
+            while (
+                not login_method_id.isdigit()
+                or int(login_method_id) > 2
+                or int(login_method_id) <= 0
+            ):
+                login_method_id = typer.prompt(
+                    "Bad Input. Choose your preferred login option"
+                )
+
+            login_method_id = int(login_method_id)
+            if login_method_id == 1:
+                return __prompt_jfrog_cli_configuration_list(
+                    jfrog_cli_config, jfrog_cli_servers
+                )
+
+        return __prompt_manual_details()
+
+    return None
+
+
+def __prompt_manual_details() -> Union[LoginArguments, None]:
+    login_args = LoginArguments()
+    login_args.artifactory_url = typer.prompt("Enter artifactory base url")
+
+    if login_args.artifactory_url is not None:
+        login_args.artifactory_url = join_url(login_args.artifactory_url, "artifactory")
+
+    auth_options = ["Username and Password", "Access Token", "Anonymous Access"]
+    authentication_types = ""
+    for index, item in enumerate(auth_options):
+        authentication_types += f"{index}: {item} \n"
+    selected_auth_type = typer.prompt(
+        f"Choose your preferred authentication option:\n{authentication_types}"
+    )
+    while (
+        not selected_auth_type.isdigit()
+        or int(selected_auth_type) >= len(auth_options)
+        or int(selected_auth_type) < 0
+    ):
+        selected_auth_type = typer.prompt(
+            "Bad Input. Choose your preferred authentication"
+        )
+    selected_auth_type = int(selected_auth_type)
+    if selected_auth_type == 0:
+        return __prompt_username_password(login_args)
+    elif selected_auth_type == 1:
+        return __prompt_access_token(login_args)
+    elif selected_auth_type == 2:
+        login_args.isAnonymous = True
+        return login_args
+    return None
+
+
+def __prompt_username_password(login_args: LoginArguments) -> LoginArguments:
+    username = typer.prompt("Enter JFrog user name")
+    password = typer.prompt("Enter JFrog password", hide_input=True)
+    login_args.username = username
+    login_args.password = password
+    return login_args
+
+
+def __prompt_access_token(login_args: LoginArguments) -> LoginArguments:
+    token = typer.prompt("Enter JFrog access token", hide_input=True)
+    login_args.access_token = token
+    return login_args
+
+
+def __prompt_jfrog_cli_configuration_list(
+    jfrog_cli_config: dict, jfrog_cli_servers: List[str]
+) -> Union[LoginArguments, None]:
+    list_server_options = ""
+    for index, item in enumerate(jfrog_cli_servers):
+        list_server_options += f"{index}: {item}\n"
+    server_index_cli_conf = typer.prompt(
+        f"Found the following servers in your JFrog CLI configuration, "
+        f"choose one of the following:\n{list_server_options}"
+    )
+    while (
+        not server_index_cli_conf.isdigit()
+        or int(server_index_cli_conf) < 0
+        or int(server_index_cli_conf) >= len(jfrog_cli_servers)
+    ):
+        server_index_cli_conf = typer.prompt(
+            "Invalid choice. Please choose a number from the list"
+        )
+    server_index_cli_conf = int(server_index_cli_conf)
+    typer.echo(f"{jfrog_cli_servers[server_index_cli_conf]} was chosen")
+    servers = jfrog_cli_config.get("servers")
+    if servers is not None:
+        return parse_cli_config_server(servers[server_index_cli_conf])
+    else:
+        raise ValueError("No servers found in the JFrog CLI configuration.")
+
+
+def __execute_login(auth_config: AuthConfig, anonymous: bool) -> bool:
+    if run_login(auth_config, anonymous):
+        success_message = "Logged in successfully"
+        if auth_config.artifactory_url is not None:
+            success_message += (
+                f" to: {auth_config.artifactory_url.replace('/artifactory', '')}"
+            )
+        typer.echo(success_message)
+        return True
+    else:
+        typer.echo("Failed to login, bad authentication")
+        return False
+
+
+def __login_by_frogml_configuration_file_flow() -> bool:
+    login_args = get_frogml_configuration()
+    auth_config = None
+    if login_args is not None:
+        typer.echo(
+            f"Using existing frogml authentication config file: {CONFIG_FILE_PATH}"
+        )
+        auth_config = AuthConfig(
+            artifactory_url=login_args.artifactory_url,
+            user=login_args.username,
+            password=login_args.password,
+            access_token=login_args.access_token,
+        )
+
+    if (
+        auth_config is not None
+        and login_args is not None
+        and __execute_login(auth_config, login_args.isAnonymous)
+    ):
+        return True
+
+    return __interactive_flow()
+
+
+def __interactive_flow() -> bool:
+    login_args = __run_interactive_mode()
+    if login_args is not None:
+        auth_config = AuthConfig(
+            artifactory_url=login_args.artifactory_url,
+            user=login_args.username,
+            password=login_args.password,
+            access_token=login_args.access_token,
+        )
+        return __execute_login(auth_config, login_args.isAnonymous)
+    return False
+
+
+def __login_by_command_line_params(
+    url: str,
+    username: Optional[str],
+    password: Optional[str],
+    token: Optional[str],
+    anonymous: Optional[bool],
+) -> bool:
+    anonymous_value = anonymous if anonymous is not None else False
+    auth_config = AuthConfig(
+        artifactory_url=url, user=username, password=password, access_token=token
+    )
+    return __execute_login(auth_config, anonymous_value)

frogml_storage/authentication/login/_login_command.py

@@ -0,0 +1,74 @@
+import http
+from typing import Optional
+
+import requests
+from requests.auth import AuthBase, HTTPBasicAuth
+
+from frogml_storage.artifactory import ArtifactoryApi
+from frogml_storage.logging import logger
+from frogml_storage.authentication.models import BearerAuth, EmptyAuth
+from frogml_storage.authentication.utils import (
+    get_encrypted_password,
+    save_auth_config,
+)
+from frogml_storage.authentication.models import AuthConfig
+from frogml_storage.authentication.utils import (
+    is_username_password_login,
+    is_access_token_login,
+    is_anonymous_login,
+)
+
+
+def run(auth_config: AuthConfig, anonymous: Optional[bool] = False) -> bool:
+    if is_username_password_login(auth_config, anonymous):
+        connection_validation_result = __login_by_username_password(auth_config)
+    elif is_access_token_login(auth_config, anonymous):
+        connection_validation_result = __validate_server_connection(
+            auth_config, BearerAuth(auth_config.access_token)
+        )
+    elif is_anonymous_login(auth_config, anonymous):
+        connection_validation_result = __validate_server_connection(
+            auth_config, EmptyAuth()
+        )
+    else:
+        connection_validation_result = False
+
+    if connection_validation_result:
+        save_auth_config(auth_config)
+    return connection_validation_result
+
+
+def __login_by_username_password(auth_config: AuthConfig) -> bool:
+    if auth_config.user is not None and auth_config.password is not None:
+        auth_token = HTTPBasicAuth(auth_config.user, auth_config.password)
+        connection_validation_result = __validate_server_connection(
+            auth_config, auth_token
+        )
+        if connection_validation_result:
+            encrypted_password = get_encrypted_password(auth_config, auth_token)
+            if encrypted_password is not None:
+                auth_config.password = encrypted_password
+                return True
+    return False
+
+
+def __validate_server_connection(auth_config: AuthConfig, auth_token: AuthBase) -> bool:
+    success = False
+    try:
+        logger.debug("Attempting to ping artifactory")
+        response = ArtifactoryApi(auth_config.artifactory_url, auth_token).ping()
+        if response.status_code == http.HTTPStatus.OK:
+            success = True
+        else:
+            logger.debug(
+                f"Expected {http.HTTPStatus.OK} status but got {response.status_code} "
+                f"when using url {auth_config.artifactory_url}"
+            )
+    except requests.exceptions.ConnectionError as e:
+        logger.debug(f"Unable to connect to the provided url :{e}.")
+    except requests.exceptions.MissingSchema as e:
+        logger.debug(f"Invalid Artifactory URL provided: {e}.")
+    except requests.exceptions.RequestException as e:
+        logger.debug(f"Unexpected request exception: {e}.")
+
+    return success

frogml_storage/authentication/models/__init__.py

@@ -0,0 +1,3 @@
+from ._auth_config import AuthConfig
+from ._auth import BearerAuth, EmptyAuth
+from ._login import LoginArguments

frogml_storage/authentication/models/_auth.py

@@ -0,0 +1,24 @@
+from typing import Any
+
+from requests import PreparedRequest
+from requests.auth import AuthBase
+
+
+class BearerAuth(AuthBase):
+    def __init__(self, token: str):
+        self.token: str = token
+
+    def __call__(self, r: PreparedRequest) -> PreparedRequest:
+        r.headers["Authorization"] = f"Bearer {self.token}"
+        return r
+
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, BearerAuth):
+            return False
+
+        return self.token == other.token
+
+
+class EmptyAuth(AuthBase):
+    def __call__(self, r: PreparedRequest) -> PreparedRequest:
+        return r