qwak-core 0.1.37__py3-none-any.whl → 0.5.19__py3-none-any.whl

qwak/inner/di_configuration/dependency_wiring.py

@@ -0,0 +1,98 @@
+import os
+from pathlib import Path
+from typing import Optional
+
+from qwak.inner.const import QwakConstants
+from qwak.inner.di_configuration import QwakContainer
+from qwak.inner.tool.grpc.grpc_tools import validate_grpc_address
+from qwak.tools.logger import get_qwak_logger
+
+
+logger = get_qwak_logger()
+
+__DEFAULT_CONFIG_FILE_PATH: Path = Path(__file__).parent / "config.yml"
+
+
+def wire_dependencies():
+    container = QwakContainer()
+
+    container.config.from_yaml(__DEFAULT_CONFIG_FILE_PATH)
+    control_plane_grpc_address_override: Optional[str] = os.getenv(
+        QwakConstants.CONTROL_PLANE_GRPC_ADDRESS_ENVAR_NAME
+    )
+
+    if control_plane_grpc_address_override:
+        validate_grpc_address(control_plane_grpc_address_override)
+        __override_control_plane_grpc_address(
+            container, control_plane_grpc_address_override
+        )
+
+    from qwak.clients import (
+        administration,
+        alert_management,
+        alerts_registry,
+        analytics,
+        audience,
+        automation_management,
+        autoscaling,
+        batch_job_management,
+        build_orchestrator,
+        data_versioning,
+        deployment,
+        feature_store,
+        file_versioning,
+        instance_template,
+        integration_management,
+        kube_deployment_captain,
+        logging_client,
+        model_management,
+        project,
+        prompt_manager,
+        system_secret,
+        user_application_instance,
+        vector_store,
+        workspace_manager,
+    )
+
+    container.wire(
+        packages=[
+            administration,
+            alert_management,
+            audience,
+            automation_management,
+            autoscaling,
+            analytics,
+            batch_job_management,
+            build_orchestrator,
+            data_versioning,
+            deployment,
+            file_versioning,
+            instance_template,
+            kube_deployment_captain,
+            logging_client,
+            model_management,
+            project,
+            feature_store,
+            user_application_instance,
+            alerts_registry,
+            workspace_manager,
+            vector_store,
+            integration_management,
+            system_secret,
+            prompt_manager,
+        ]
+    )
+
+    return container
+
+
+def __override_control_plane_grpc_address(
+    container: "QwakContainer", control_plane_grpc_address_override: str
+):
+    logger.debug(
+        "Overriding control plane gRPC address from environment variable to %s.",
+        control_plane_grpc_address_override,
+    )
+    container.config.grpc.core.address.from_value(
+        control_plane_grpc_address_override.strip()
+    )

qwak/inner/di_configuration/session.py

@@ -10,3 +10,8 @@ class Session(metaclass=SingletonMeta):
 
     def get_environment(self):
         return self.__environment
+
+    @classmethod
+    def clear(cls):
+        if cls in cls._instances:
+            del cls._instances[cls]

qwak/inner/instance_template/verify_template_id.py

@@ -0,0 +1,48 @@
+from typing import Optional
+
+from _qwak_proto.qwak.instance_template.instance_template_pb2 import InstanceType
+from qwak.clients.instance_template.client import InstanceTemplateManagementClient
+from qwak.exceptions import QwakException
+from qwak.inner.provider import Provider
+
+INVALID_TEMPLATE_ID_ERROR_FORMAT = (
+    "Invalid instance: {template_id}. Valid instances are: {cpu_templates}. "
+    "For GPU instances you can use {gpu_templates}."
+)
+
+
+def verify_template_id(
+    template_id: str,
+    instance_template_client: InstanceTemplateManagementClient,
+    provider: Optional[Provider] = None,
+) -> None:
+    all_templates = instance_template_client.list_instance_templates()
+    valid_templates = [template for template in all_templates if template.enabled]
+    if provider == Provider.AWS:
+        valid_templates = [
+            template for template in valid_templates if template.aws_supported
+        ]
+    elif provider == Provider.GCP:
+        valid_templates = [
+            template for template in valid_templates if template.gcp_supported
+        ]
+
+    existing_templates = {template.id: template for template in valid_templates}
+    if template_id not in existing_templates.keys():
+        cpu_template_ids = [
+            template.id
+            for template in sorted(valid_templates, key=lambda template: template.order)
+            if template.instance_type == InstanceType.INSTANCE_TYPE_CPU
+        ]
+        gpu_template_ids = [
+            template.id
+            for template in sorted(valid_templates, key=lambda template: template.order)
+            if template.instance_type == InstanceType.INSTANCE_TYPE_GPU
+        ]
+        raise QwakException(
+            INVALID_TEMPLATE_ID_ERROR_FORMAT.format(
+                template_id=template_id,
+                cpu_templates=", ".join(cpu_template_ids),
+                gpu_templates=", ".join(gpu_template_ids),
+            )
+        )

qwak/inner/model_loggers_utils.py

@@ -3,6 +3,7 @@ import re
 from typing import Optional
 
 import requests
+
 from qwak.clients.model_management.client import ModelsManagementClient
 from qwak.exceptions import QwakException
 
@@ -67,23 +68,25 @@ def fetch_build_id() -> Optional[str]:
 def upload_data(
     upload_url: str,
     data: bytes,
+    headers: dict,
     content_type: str = "text/plain",
-) -> None:
+):
     """
     Upload data
-    Args:
-        upload_url: the url to upload to.
-        data: the data to upload
-        content_type: Uploaded content-type
+    :param upload_url: the url to upload to.
+    :param data: the data to upload
+    :param headers: authentication details for upload data
+    :param content_type: Uploaded content-type
     """
     try:
+        headers["Content-Type"] = content_type
         http_response = requests.put(  # nosec B113
             upload_url,
             data=data,
-            headers={"content-type": content_type},
+            headers=headers,
         )
 
-        if http_response.status_code != 200:
+        if http_response.status_code not in [200, 201]:
             raise QwakException(
                 f"Failed to upload data. "
                 f"Status: [{http_response.status_code}], "

qwak/inner/provider.py

@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class Provider(Enum):
+    AWS = 1
+    GCP = 2

qwak/inner/tool/auth.py

@@ -1,6 +1,9 @@
 import warnings
+from filelock import FileLock
 
 from qwak.inner.di_configuration.session import Session
+from abc import ABC, abstractmethod
+from typing import Optional
 
 warnings.filterwarnings(action="ignore", module=".*jose.*")
 
@@ -13,12 +16,29 @@ from qwak.exceptions import QwakLoginException  # noqa E402
 from qwak.inner.const import QwakConstants  # noqa E402
 
 
-class Auth0ClientBase:
+class BaseAuthClient(ABC):
+    @abstractmethod
+    def get_token(self) -> Optional[str]:
+        pass
+
+    @abstractmethod
+    def login(self) -> None:
+        pass
+
+    @abstractmethod
+    def get_base_url(self) -> str:
+        """
+        Returns the base URL for the authentication service.
+        """
+        return ""
+
+
+class Auth0ClientBase(BaseAuthClient):
     _TOKENS_FIELD = "TOKENS"
 
     def __init__(
         self,
-        api_key,
+        api_key=None,
         auth_file=QwakConstants.QWAK_AUTHORIZATION_FILE,
         audience=QwakConstants.TOKEN_AUDIENCE,
     ):
@@ -30,8 +50,15 @@ class Auth0ClientBase:
         self.audience = audience
         self.api_key = api_key
 
+    def get_base_url(self) -> str:
+        return QwakConstants.QWAK_APP_URL
+
     # Returns Non if token is expired
     def get_token(self):
+        if self._environment != Session().get_environment():
+            self.token = None
+            self.api_key = None
+            self._environment = Session().get_environment()
         try:
             if not self.token:
                 self._config.read(self._auth_file)
@@ -54,6 +81,12 @@ class Auth0ClientBase:
     def login(self):
         from qwak.clients.administration import AuthenticationClient
 
+        if not self.api_key:
+            from qwak.inner.di_configuration import UserAccountConfiguration
+
+            user_account = UserAccountConfiguration().get_user_config()
+            self.api_key = user_account.api_key
+
         self.token = AuthenticationClient().authenticate(self.api_key).access_token
 
         from pathlib import Path
@@ -61,12 +94,14 @@ class Auth0ClientBase:
         Path(self._auth_file).parent.mkdir(parents=True, exist_ok=True)
         self._config.read(self._auth_file)
 
-        with open(self._auth_file, "w") as configfile:
-            self._config[self._environment] = {
-                self._TOKENS_FIELD: json.dumps(self.token)
-            }
+        lock_path = f"{self._auth_file}.lock"  # Create a lock file
+        with FileLock(lock_path):  # Use file lock
+            with open(self._auth_file, "w") as configfile:
+                self._config[self._environment] = {
+                    self._TOKENS_FIELD: json.dumps(self.token)
+                }
 
-            self._config.write(configfile)
+                self._config.write(configfile)
 
     def get_claims(self):
         try:

qwak/inner/tool/grpc/grpc_auth.py

@@ -1,17 +1,36 @@
+from typing import Callable, Tuple
+
 import grpc
-from qwak.inner.tool.auth import Auth0ClientBase
 
 _SIGNATURE_HEADER_KEY = "authorization"
 
 
-class Auth0Client(grpc.AuthMetadataPlugin, Auth0ClientBase):
+class Auth0Client(grpc.AuthMetadataPlugin):
     def __init__(self):
-        from qwak.inner.di_configuration import UserAccountConfiguration
+        self._auth_client = None
+
+    def __call__(
+        self,
+        context: grpc.AuthMetadataContext,
+        callback: Callable[[Tuple[Tuple[str, str]], None], None],
+    ):
+        """Implements authentication by passing metadata to a callback.
 
-        user_account = UserAccountConfiguration().get_user_config()
-        api_key = user_account.api_key
+        Args:
+            context: An AuthMetadataContext providing information on the RPC that
+                the plugin is being called to authenticate.
+            callback: A callback that accepts a tuple of metadata key/value pairs and a None
+                parameter.
+        """
+        # Get token from Auth0 client
+        token = self.get_token()
+        metadata = ((_SIGNATURE_HEADER_KEY, f"Bearer {token}"),)
+        callback(metadata, None)
 
-        Auth0ClientBase.__init__(self, api_key=api_key)
+    def get_token(self) -> str:
+        """Get the authentication token."""
+        if not self._auth_client:
+            from qwak.inner.tool.auth import Auth0ClientBase
 
-    def __call__(self, context, callback):
-        callback(((_SIGNATURE_HEADER_KEY, "Bearer {}".format(self.get_token())),), None)
+            self._auth_client = Auth0ClientBase()
+        return self._auth_client.get_token()

qwak/inner/tool/grpc/grpc_tools.py

@@ -1,15 +1,18 @@
 import logging
+import re
 import time
 from abc import ABC, abstractmethod
 from random import randint
-from typing import Optional, Tuple, Callable
+from typing import Callable, Optional, Tuple
+from urllib.parse import urlparse, ParseResult
 
 import grpc
-from qwak.exceptions import QwakException
 
+from qwak.exceptions import QwakException, QwakGrpcAddressException
 from .grpc_auth import Auth0Client
 
 logger = logging.getLogger()
+HOSTNAME_REGEX: str = r"^(?!-)(?:[A-Za-z0-9-]{1,63}\.)*[A-Za-z0-9-]{1,63}(?<!-)$"
 
 
 def create_grpc_channel(
@@ -19,7 +22,7 @@ def create_grpc_channel(
     auth_metadata_plugin: grpc.AuthMetadataPlugin = None,
     timeout: int = 100,
     options=None,
-    backoff_options={},
+    backoff_options=None,
     max_attempts=4,
     status_for_retry=(grpc.StatusCode.UNAVAILABLE,),
     attempt=0,
@@ -40,6 +43,9 @@ def create_grpc_channel(
         status_for_retry: grpc statuses to retry upon
     Returns: Returns a grpc.Channel
     """
+    if backoff_options is None:
+        backoff_options = {}
+
     if not url:
         raise QwakException("Unable to create gRPC channel. URL has not been defined.")
 
@@ -101,11 +107,14 @@ def create_grpc_channel_or_none(
     auth_metadata_plugin: grpc.AuthMetadataPlugin = None,
     timeout: int = 30,
     options=None,
-    backoff_options={},
+    backoff_options=None,
     max_attempts=2,
     status_for_retry=(grpc.StatusCode.UNAVAILABLE,),
     attempt=0,
 ) -> Callable[[Optional[str], Optional[bool]], Optional[grpc.Channel]]:
+    if backoff_options is None:
+        backoff_options = {}
+
     def deferred_channel(
         url_overwrite: Optional[str] = None, ssl_overwrite: Optional[bool] = None
     ):
@@ -129,6 +138,117 @@ def create_grpc_channel_or_none(
     return deferred_channel
 
 
+def validate_grpc_address(
+    grpc_address: str,
+    is_port_specification_allowed: bool = False,
+    is_url_scheme_allowed: bool = False,
+):
+    """
+    Validate gRPC address format
+    Args:
+        grpc_address (str): gRPC address to validate
+        is_port_specification_allowed (bool): Whether to allow port specification in the address
+        is_url_scheme_allowed (bool): Whether to allow URL scheme in the address
+    Raises:
+        QwakGrpcAddressException: If the gRPC address is invalid
+    """
+    parsed_grpc_address: ParseResult = parse_address(grpc_address)
+    hostname: str = get_hostname_from_address(parsed_grpc_address)
+    validate_paths_are_not_included_in_address(parsed_grpc_address)
+
+    if not is_url_scheme_allowed:
+        __validate_url_scheme_not_included_in_address(parsed_grpc_address)
+
+    if not is_port_specification_allowed:
+        __validate_port_not_included_in_address(parsed_grpc_address)
+
+    if not is_valid_hostname(hostname):
+        raise QwakGrpcAddressException(
+            "gRPC address must be a simple hostname or fully qualified domain name.",
+            parsed_grpc_address,
+        )
+
+
+def validate_paths_are_not_included_in_address(
+    parsed_grpc_address: ParseResult,
+) -> None:
+    has_invalid_path: bool = (
+        parsed_grpc_address.path not in {"", "/"}
+        or parsed_grpc_address.query
+        or parsed_grpc_address.fragment
+    )
+
+    if has_invalid_path:
+        raise QwakGrpcAddressException(
+            "gRPC address must not contain paths, queries, or fragments.",
+            parsed_grpc_address,
+        )
+
+
+def get_hostname_from_address(parsed_grpc_address: ParseResult) -> str:
+    hostname: Optional[str] = parsed_grpc_address.hostname
+    if not hostname:
+        raise QwakGrpcAddressException(
+            "gRPC address must contain a valid hostname.", parsed_grpc_address
+        )
+
+    return hostname
+
+
+def __validate_url_scheme_not_included_in_address(
+    parsed_grpc_address: ParseResult,
+) -> None:
+    if parsed_grpc_address.scheme:
+        raise QwakGrpcAddressException(
+            "URL scheme is not allowed in the gRPC address.", parsed_grpc_address
+        )
+
+
+def __validate_port_not_included_in_address(parsed_grpc_address: ParseResult):
+    try:
+        port: Optional[int] = parsed_grpc_address.port
+    except ValueError as exc:
+        raise QwakGrpcAddressException(
+            "Invalid port specification in the gRPC address.", parsed_grpc_address
+        ) from exc
+
+    if port:
+        raise QwakGrpcAddressException(
+            "Port specification is not allowed in the gRPC address.",
+            parsed_grpc_address,
+        )
+
+
+def parse_address(grpc_address: str) -> ParseResult:
+    if not grpc_address or not grpc_address.strip():
+        raise QwakGrpcAddressException(
+            "gRPC address must not be empty or whitespace.", grpc_address
+        )
+
+    trimmed_address: str = grpc_address.strip()
+    parsed_address: ParseResult = urlparse(
+        trimmed_address if "://" in trimmed_address else f"//{trimmed_address}"
+    )
+
+    return parsed_address
+
+
+def is_valid_hostname(hostname: str) -> bool:
+    """
+    Validate that the supplied hostname conforms to RFC-style label rules:
+    anchored pattern enforces full-string validation, negative lookahead/lookbehind block
+    leading or trailing hyphens per label, and each dot-separated label must be 1-63
+    alphanumeric/hyphen characters.
+
+    Args:
+        hostname (str): The hostname to validate.
+    Returns:
+        bool: True if the hostname is valid, False otherwise.
+    """
+    hostname_pattern: re.Pattern = re.compile(HOSTNAME_REGEX)
+    return bool(hostname_pattern.fullmatch(hostname))
+
+
 class SleepingPolicy(ABC):
     @abstractmethod
     def sleep(self, try_i: int):

qwak/inner/tool/grpc/grpc_try_wrapping.py

@@ -1,17 +1,152 @@
+import logging
+import functools
+import inspect
+from inspect import BoundArguments, Signature
+from typing import Any, Callable, Optional
+
 import grpc
 from qwak.exceptions import QwakException
 
+logger = logging.getLogger()
+
+
+def grpc_try_catch_wrapper(
+    exception_message: str,
+    reraise_non_grpc_error_original_exception: Optional[bool] = False,
+) -> Callable:
+    """A decorator for handling exceptions in client methods.
+
+    This decorator wraps a function, catching any `grpc.RpcError` or other
+    `Exception`. it re-raises them as a `QwakException`. The error message can be a
+    static string or a template formatted with the decorated function's
+    runtime arguments.
+
+    Parameters
+    ----------
+    exception_message : str
+        The error message template. This can be a static string or a format
+        string using placeholders that match the decorated function's
+        parameter names (e.g., "Failed on item {item_id}").
+    reraise_non_grpc_error_original_exception : bool, optional
+        If `True`, any exception that is *not* a `grpc.RpcError` will be
+        re-raised in its original form. If `False` (the default), all
+        exceptions are wrapped in `QwakException`.
+
+    Returns
+    -------
+    Callable
+        A decorator that wraps the function with the error handling logic.
+
+    Raises
+    ------
+    QwakException
+        Raised when the decorated function encounters any exception, unless
+        `reraise_non_grpc_error_original_exception` is `True` and the error
+        is not a `grpc.RpcError`. The new exception's message is formatted
+        and includes details from the original error.
+    Exception
+        The original non-gRPC exception is re-raised if
+        `reraise_non_grpc_error_original_exception` is set to `True`.
+
+    Usage
+    -----
+    The decorator can be used with static strings and formatted strings
+
+    **1. Formatted Error Messages**
+
+    Use placeholders matching the function's parameter names to create
+    context-specific error messages.
+
+        @grpc_try_catch_wrapper("Failed to delete version {version_number} for FeatureSet {featureset_name}")
+        def delete_featureset_version(self, featureset_name: str, version_number: int):
+            # ... gRPC call ...
+
+    If this function fails with a gRPC error, a `QwakException` is raised with a message like:
+    `"Failed to delete version 2 for FeatureSet customer_churn - <original gRPC error code> - <original gRPC error details>."`
+
+    **2. Static Error Messages**
+
+    When no dynamic context is needed, a simple string is sufficient.
+
+        @grpc_try_catch_wrapper("Failed to retrieve featuresets mapping")
+        def get_featuresets_mapping(self):
+            # ... gRPC call ...
 
-def grpc_try_catch_wrapper(exception_message):
-    def decorator(function):
+    If this fails, the `QwakException` message will be:
+    `"Failed to retrieve featuresets mapping - <original gRPC error code> - <original gRPC error details>."`
+    """
+
+    def decorator(function: Callable):
+        @functools.wraps(function)
         def _inner_wrapper(*args, **kwargs):
             try:
                 return function(*args, **kwargs)
             except grpc.RpcError as e:
-                raise QwakException(
-                    exception_message + f"error is: {e.details()} (code {e.code})"
+                error_message: str = __get_error_msg(
+                    exception_message=exception_message, e=e
+                )
+
+                formatted_message = __get_formatted_error_message(
+                    exception_message=error_message,
+                    function=function,
+                    args=args,
+                    kwargs=kwargs,
+                )
+
+                raise QwakException(formatted_message) from e
+
+            except Exception as e:
+                if reraise_non_grpc_error_original_exception:
+                    raise e
+
+                formatted_message = __get_formatted_error_message(
+                    exception_message=exception_message,
+                    function=function,
+                    args=args,
+                    kwargs=kwargs,
                 )
+                raise QwakException(f"{formatted_message}. Error is: {e}.") from e
 
         return _inner_wrapper
 
     return decorator
+
+
+def __get_error_msg(exception_message: str, e: grpc.RpcError) -> str:
+    err_msg: str = (
+        f"{exception_message} - {e.code() if hasattr(e, 'code') else 'UNKNOWN'}"
+    )
+    if hasattr(e, "details") and e.details() is not None:
+        err_msg: str = f"{err_msg} - {e.details()}"
+    elif hasattr(e, "debug_error_string"):
+        err_msg: str = f"{err_msg} - {e.debug_error_string()}"
+    logger.debug(f"{err_msg}: {e}")
+    return err_msg
+
+
+def __get_formatted_error_message(
+    exception_message: str, function: Callable, args: tuple[Any], kwargs: dict[str, Any]
+) -> str:
+    """Formats an error message string with the runtime arguments of a function.
+
+    It inspects the function's signature and binds the passed args and kwargs
+    to their parameter names, then uses them to format the message string.
+
+    Args:
+        exception_message (str): The error message template with placeholders.
+        function (Callable): The decorated function.
+        args (tuple[Any]): The positional arguments passed to the function.
+        kwargs (dict[str, Any]): The keyword arguments passed to the function.
+
+    Returns:
+        str: The error message with placeholders filled in.
+    """
+    try:
+        sig: Signature = inspect.signature(function)
+        bound_args: BoundArguments = sig.bind(*args, **kwargs)
+        bound_args.apply_defaults()
+
+        return exception_message.format(**bound_args.arguments)
+    except (ValueError, KeyError, TypeError):
+        # Fallback if formatting fails (e.g., missing key in template)
+        return exception_message