genesis-flow 1.0.0__py3-none-any.whl

mlflow/azure/auth_handler.py

@@ -0,0 +1,257 @@
+"""Azure authentication handler with advanced token management."""
+
+from __future__ import annotations
+
+import threading
+import logging
+from datetime import datetime, timedelta
+from typing import Optional
+
+from azure.core.credentials import AccessToken
+from azure.core.exceptions import ClientAuthenticationError
+from azure.identity import (
+    DefaultAzureCredential,
+    ManagedIdentityCredential,
+    ClientSecretCredential,
+)
+
+from mlflow.azure.config import AzureAuthConfig, AuthMethod
+from mlflow.azure.exceptions import TokenAcquisitionError
+from mlflow.azure.utils import is_token_expired
+
+logger = logging.getLogger(__name__)
+
+
+class TokenCache:
+    """Thread-safe token cache with expiration handling."""
+
+    def __init__(self):
+        self._lock = threading.RLock()
+        self._token: Optional[AccessToken] = None
+        self._expiry: Optional[datetime] = None
+
+    def get_token(self, buffer_seconds: int = 300) -> Optional[AccessToken]:
+        """Get cached token if it's still valid.
+
+        Args:
+            buffer_seconds: Buffer time before expiry to consider token invalid
+
+        Returns:
+            Cached token if valid, None otherwise
+        """
+        with self._lock:
+            if self._token and not is_token_expired(self._expiry, buffer_seconds):
+                logger.debug("Using cached token")
+                return self._token
+
+            if self._token:
+                logger.debug("Cached token expired or expiring soon")
+
+            return None
+
+    def set_token(self, token: AccessToken) -> None:
+        """Cache a new token.
+
+        Args:
+            token: Access token to cache
+        """
+        with self._lock:
+            self._token = token
+            # Convert timestamp to datetime
+            self._expiry = datetime.fromtimestamp(token.expires_on)
+
+            logger.debug(
+                "Token cached, expires at %s (in %d seconds)",
+                self._expiry.isoformat(),
+                int((self._expiry - datetime.utcnow()).total_seconds()),
+            )
+
+    def clear(self) -> None:
+        """Clear the cached token."""
+        with self._lock:
+            self._token = None
+            self._expiry = None
+            logger.debug("Token cache cleared")
+
+
+class AzureAuthHandler:
+    """Handles Azure authentication and token management for MLflow PostgreSQL connections."""
+
+    def __init__(self, config: AzureAuthConfig):
+        """Initialize the authentication handler.
+
+        Args:
+            config: Azure authentication configuration
+        """
+        self.config = config
+        self._token_cache = TokenCache()
+        self._credential = None
+
+        if config.enable_debug_logging:
+            logging.getLogger().setLevel(logging.DEBUG)
+
+        logger.info(
+            "Initialized Azure auth handler: auth_method=%s, client_id=%s",
+            config.auth_method.value,
+            config.client_id[:8] + "..." if config.client_id else None,
+        )
+
+    @property
+    def credential(self):
+        """Get or create Azure credential based on configuration."""
+        if self._credential is None:
+            self._credential = self._create_credential()
+        return self._credential
+
+    def _create_credential(self):
+        """Create appropriate Azure credential based on auth method."""
+        auth_method = self.config.auth_method
+
+        logger.info("Creating Azure credential: auth_method=%s", auth_method.value)
+
+        try:
+            if auth_method == AuthMethod.MANAGED_IDENTITY:
+                if self.config.client_id:
+                    logger.info(
+                        "Using user-assigned managed identity: client_id=%s", self.config.client_id
+                    )
+                    return ManagedIdentityCredential(client_id=self.config.client_id)
+                else:
+                    logger.info("Using system-assigned managed identity")
+                    return ManagedIdentityCredential()
+
+            elif auth_method == AuthMethod.SERVICE_PRINCIPAL:
+                if not all(
+                    [self.config.client_id, self.config.client_secret, self.config.tenant_id]
+                ):
+                    raise TokenAcquisitionError(
+                        "Service principal authentication requires client_id, client_secret, and tenant_id"
+                    )
+
+                logger.info(
+                    "Using service principal authentication: client_id=%s, tenant_id=%s",
+                    self.config.client_id,
+                    self.config.tenant_id,
+                )
+                return ClientSecretCredential(
+                    tenant_id=self.config.tenant_id,
+                    client_id=self.config.client_id,
+                    client_secret=self.config.client_secret,
+                )
+
+            elif auth_method == AuthMethod.DEFAULT_AZURE_CREDENTIAL:
+                logger.info("Using DefaultAzureCredential")
+                return DefaultAzureCredential()
+
+            else:
+                raise TokenAcquisitionError(f"Unsupported authentication method: {auth_method}")
+
+        except Exception as e:
+            logger.error(
+                "Failed to create Azure credential: error=%s, auth_method=%s", str(e), auth_method.value
+            )
+            raise TokenAcquisitionError(f"Failed to create credential: {e}") from e
+
+    def _acquire_token_with_retry(self) -> AccessToken:
+        """Acquire token with retry logic."""
+        last_exception = None
+        
+        for attempt in range(self.config.max_retries):
+            try:
+                logger.debug("Acquiring new token from Azure (attempt %d)", attempt + 1)
+                token = self.credential.get_token(self.config.token_scope)
+
+                logger.info(
+                    "Successfully acquired token, expires in %d seconds",
+                    int(token.expires_on - datetime.utcnow().timestamp()),
+                )
+
+                return token
+
+            except ClientAuthenticationError as e:
+                logger.error("Authentication failed: %s", str(e))
+                raise TokenAcquisitionError(f"Authentication failed: {e}") from e
+
+            except Exception as e:
+                last_exception = e
+                if attempt < self.config.max_retries - 1:
+                    delay = self.config.retry_delay * (self.config.retry_backoff ** attempt)
+                    logger.warning(
+                        "Token acquisition failed (attempt %d), retrying in %.1f seconds: %s",
+                        attempt + 1, delay, str(e)
+                    )
+                    import time
+                    time.sleep(delay)
+                else:
+                    logger.error("Token acquisition failed after all retries: %s", str(e))
+
+        raise TokenAcquisitionError(f"Token acquisition failed: {last_exception}") from last_exception
+
+    def get_access_token(self, force_refresh: bool = False) -> str:
+        """Get a valid access token for PostgreSQL authentication.
+
+        Args:
+            force_refresh: Force token refresh even if cached token is valid
+
+        Returns:
+            Valid access token string
+
+        Raises:
+            TokenAcquisitionError: If token acquisition fails
+        """
+        if not self.config.should_use_azure_auth:
+            raise TokenAcquisitionError("Azure authentication is not enabled")
+
+        # Try to get cached token first (unless force refresh)
+        if not force_refresh:
+            cached_token = self._token_cache.get_token(self.config.token_refresh_buffer)
+            if cached_token:
+                return cached_token.token
+
+        # Acquire new token
+        try:
+            new_token = self._acquire_token_with_retry()
+            self._token_cache.set_token(new_token)
+            return new_token.token
+
+        except Exception as e:
+            logger.error("Failed to acquire access token: %s", str(e))
+            raise
+
+    def refresh_token_if_needed(self) -> bool:
+        """Refresh token if it's expiring soon.
+
+        Returns:
+            True if token was refreshed, False otherwise
+        """
+        if not self.config.should_use_azure_auth:
+            return False
+
+        cached_token = self._token_cache.get_token(self.config.token_refresh_buffer)
+        if cached_token:
+            return False  # Token is still valid
+
+        try:
+            self.get_access_token(force_refresh=True)
+            logger.info("Token refreshed proactively")
+            return True
+        except Exception as e:
+            logger.warning("Proactive token refresh failed: %s", str(e))
+            return False
+
+    def clear_token_cache(self) -> None:
+        """Clear the token cache."""
+        self._token_cache.clear()
+        logger.info("Token cache cleared")
+
+    def is_token_valid(self) -> bool:
+        """Check if current cached token is valid.
+
+        Returns:
+            True if token is valid, False otherwise
+        """
+        if not self.config.should_use_azure_auth:
+            return True  # SQL auth doesn't use tokens
+
+        cached_token = self._token_cache.get_token(self.config.token_refresh_buffer)
+        return cached_token is not None

mlflow/azure/client.py

@@ -0,0 +1,319 @@
+"""
+This module provides utilities for performing Azure Blob Storage operations without requiring
+the heavyweight azure-storage-blob library dependency
+"""
+
+import logging
+import urllib
+from copy import deepcopy
+
+from mlflow.utils import rest_utils
+from mlflow.utils.file_utils import read_chunk
+
+_logger = logging.getLogger(__name__)
+_PUT_BLOCK_HEADERS = {
+    "x-ms-blob-type": "BlockBlob",
+}
+
+
+def put_adls_file_creation(sas_url, headers):
+    """Performs an ADLS Azure file create `Put` operation
+    (https://docs.microsoft.com/en-us/rest/api/storageservices/datalakestoragegen2/path/create)
+
+    Args:
+        sas_url: A shared access signature URL referring to the Azure ADLS server
+            to which the file creation command should be issued.
+        headers: Additional headers to include in the Put request body.
+    """
+    request_url = _append_query_parameters(sas_url, {"resource": "file"})
+
+    request_headers = {}
+    for name, value in headers.items():
+        if _is_valid_adls_put_header(name):
+            request_headers[name] = value
+        else:
+            _logger.debug("Removed unsupported '%s' header for ADLS Gen2 Put operation", name)
+
+    with rest_utils.cloud_storage_http_request(
+        "put", request_url, headers=request_headers
+    ) as response:
+        rest_utils.augmented_raise_for_status(response)
+
+
+def patch_adls_file_upload(sas_url, local_file, start_byte, size, position, headers, is_single):
+    """
+    Performs an ADLS Azure file create `Patch` operation
+    (https://docs.microsoft.com/en-us/rest/api/storageservices/datalakestoragegen2/path/update)
+
+    Args:
+        sas_url: A shared access signature URL referring to the Azure ADLS server
+            to which the file update command should be issued.
+        local_file: The local file to upload
+        start_byte: The starting byte of the local file to upload
+        size: The number of bytes to upload
+        position: Positional offset of the data in the Patch request
+        headers: Additional headers to include in the Patch request body
+        is_single: Whether this is the only patch operation for this file
+    """
+    new_params = {"action": "append", "position": str(position)}
+    if is_single:
+        new_params["flush"] = "true"
+    request_url = _append_query_parameters(sas_url, new_params)
+
+    request_headers = {}
+    for name, value in headers.items():
+        if _is_valid_adls_patch_header(name):
+            request_headers[name] = value
+        else:
+            _logger.debug("Removed unsupported '%s' header for ADLS Gen2 Patch operation", name)
+
+    data = read_chunk(local_file, size, start_byte)
+    with rest_utils.cloud_storage_http_request(
+        "patch", request_url, data=data, headers=request_headers
+    ) as response:
+        rest_utils.augmented_raise_for_status(response)
+
+
+def patch_adls_flush(sas_url, position, headers):
+    """Performs an ADLS Azure file flush `Patch` operation
+    (https://docs.microsoft.com/en-us/rest/api/storageservices/datalakestoragegen2/path/update)
+
+    Args:
+        sas_url: A shared access signature URL referring to the Azure ADLS server
+            to which the file update command should be issued.
+        position: The final size of the file to flush.
+        headers: Additional headers to include in the Patch request body.
+
+    """
+    request_url = _append_query_parameters(sas_url, {"action": "flush", "position": str(position)})
+
+    request_headers = {}
+    for name, value in headers.items():
+        if _is_valid_adls_put_header(name):
+            request_headers[name] = value
+        else:
+            _logger.debug("Removed unsupported '%s' header for ADLS Gen2 Patch operation", name)
+
+    with rest_utils.cloud_storage_http_request(
+        "patch", request_url, headers=request_headers
+    ) as response:
+        rest_utils.augmented_raise_for_status(response)
+
+
+def put_block(sas_url, block_id, data, headers):
+    """
+    Performs an Azure `Put Block` operation
+    (https://docs.microsoft.com/en-us/rest/api/storageservices/put-block)
+
+    Args:
+        sas_url: A shared access signature URL referring to the Azure Block Blob
+            to which the specified data should be staged.
+        block_id: A base64-encoded string identifying the block.
+        data: Data to include in the Put Block request body.
+        headers: Additional headers to include in the Put Block request body
+            (the `x-ms-blob-type` header is always included automatically).
+    """
+    request_url = _append_query_parameters(sas_url, {"comp": "block", "blockid": block_id})
+
+    request_headers = deepcopy(_PUT_BLOCK_HEADERS)
+    for name, value in headers.items():
+        if _is_valid_put_block_header(name):
+            request_headers[name] = value
+        else:
+            _logger.debug("Removed unsupported '%s' header for Put Block operation", name)
+
+    with rest_utils.cloud_storage_http_request(
+        "put", request_url, data=data, headers=request_headers
+    ) as response:
+        rest_utils.augmented_raise_for_status(response)
+
+
+def put_block_list(sas_url, block_list, headers):
+    """Performs an Azure `Put Block List` operation
+    (https://docs.microsoft.com/en-us/rest/api/storageservices/put-block-list)
+
+    Args:
+        sas_url: A shared access signature URL referring to the Azure Block Blob
+            to which the specified data should be staged.
+        block_list: A list of uncommitted base64-encoded string block IDs to commit. For
+            more information, see
+            https://docs.microsoft.com/en-us/rest/api/storageservices/put-block-list.
+        headers: Headers to include in the Put Block request body.
+
+    """
+    request_url = _append_query_parameters(sas_url, {"comp": "blocklist"})
+    data = _build_block_list_xml(block_list)
+
+    request_headers = {}
+    for name, value in headers.items():
+        if _is_valid_put_block_list_header(name):
+            request_headers[name] = value
+        else:
+            _logger.debug("Removed unsupported '%s' header for Put Block List operation", name)
+
+    with rest_utils.cloud_storage_http_request(
+        "put", request_url, data=data, headers=request_headers
+    ) as response:
+        rest_utils.augmented_raise_for_status(response)
+
+
+def _append_query_parameters(url, parameters):
+    parsed_url = urllib.parse.urlparse(url)
+    query_dict = dict(urllib.parse.parse_qsl(parsed_url.query))
+    query_dict.update(parameters)
+    new_query = urllib.parse.urlencode(query_dict)
+    new_url_components = parsed_url._replace(query=new_query)
+    return urllib.parse.urlunparse(new_url_components)
+
+
+def _build_block_list_xml(block_list):
+    xml = '<?xml version="1.0" encoding="utf-8"?>\n<BlockList>\n'
+    for block_id in block_list:
+        # Because block IDs are base64-encoded and base64 strings do not contain
+        # XML special characters, we can safely insert the block ID directly into
+        # the XML document
+        xml += f"<Uncommitted>{block_id}</Uncommitted>\n"
+    xml += "</BlockList>"
+    return xml
+
+
+def _is_valid_put_block_list_header(header_name):
+    """
+    Returns:
+        True if the specified header name is a valid header for the Put Block List operation,
+        False otherwise. For a list of valid headers, see https://docs.microsoft.com/en-us/
+        rest/api/storageservices/put-block-list#request-headers and https://docs.microsoft.com/
+        en-us/rest/api/storageservices/
+        specifying-conditional-headers-for-blob-service-operations#Subheading1.
+    """
+    return header_name.startswith("x-ms-meta-") or header_name in {
+        "Authorization",
+        "Date",
+        "x-ms-date",
+        "x-ms-version",
+        "Content-Length",
+        "Content-MD5",
+        "x-ms-content-crc64",
+        "x-ms-blob-cache-control",
+        "x-ms-blob-content-type",
+        "x-ms-blob-content-encoding",
+        "x-ms-blob-content-language",
+        "x-ms-blob-content-md5",
+        "x-ms-encryption-scope",
+        "x-ms-tags",
+        "x-ms-lease-id",
+        "x-ms-client-request-id",
+        "x-ms-blob-content-disposition",
+        "x-ms-access-tier",
+        "If-Modified-Since",
+        "If-Unmodified-Since",
+        "If-Match",
+        "If-None-Match",
+    }
+
+
+def _is_valid_put_block_header(header_name):
+    """
+    Returns:
+        True if the specified header name is a valid header for the Put Block operation, False
+        otherwise. For a list of valid headers, see
+        https://docs.microsoft.com/en-us/rest/api/storageservices/put-block#request-headers and
+        https://docs.microsoft.com/en-us/rest/api/storageservices/put-block#
+        request-headers-customer-provided-encryption-keys.
+    """
+    return header_name in {
+        "Authorization",
+        "x-ms-date",
+        "x-ms-version",
+        "Content-Length",
+        "Content-MD5",
+        "x-ms-content-crc64",
+        "x-ms-encryption-scope",
+        "x-ms-lease-id",
+        "x-ms-client-request-id",
+        "x-ms-encryption-key",
+        "x-ms-encryption-key-sha256",
+        "x-ms-encryption-algorithm",
+    }
+
+
+def _is_valid_adls_put_header(header_name):
+    """
+    Returns:
+        True if the specified header name is a valid header for the ADLS Put operation, False
+        otherwise. For a list of valid headers, see
+        https://docs.microsoft.com/en-us/rest/api/storageservices/datalakestoragegen2/path/create
+    """
+    return header_name in {
+        "Cache-Control",
+        "Content-Encoding",
+        "Content-Language",
+        "Content-Disposition",
+        "x-ms-cache-control",
+        "x-ms-content-type",
+        "x-ms-content-encoding",
+        "x-ms-content-language",
+        "x-ms-content-disposition",
+        "x-ms-rename-source",
+        "x-ms-lease-id",
+        "x-ms-properties",
+        "x-ms-permissions",
+        "x-ms-umask",
+        "x-ms-owner",
+        "x-ms-group",
+        "x-ms-acl",
+        "x-ms-proposed-lease-id",
+        "x-ms-expiry-option",
+        "x-ms-expiry-time",
+        "If-Match",
+        "If-None-Match",
+        "If-Modified-Since",
+        "If-Unmodified-Since",
+        "x-ms-source-if-match",
+        "x-ms-source-if-none-match",
+        "x-ms-source-if-modified-since",
+        "x-ms-source-if-unmodified-since",
+        "x-ms-encryption-key",
+        "x-ms-encryption-key-sha256",
+        "x-ms-encryption-algorithm",
+        "x-ms-encryption-context",
+        "x-ms-client-request-id",
+        "x-ms-date",
+        "x-ms-version",
+    }
+
+
+def _is_valid_adls_patch_header(header_name):
+    """
+    Returns:
+        True if the specified header name is a valid header for the ADLS Patch operation, False
+        otherwise. For a list of valid headers, see
+        https://docs.microsoft.com/en-us/rest/api/storageservices/datalakestoragegen2/path/update
+    """
+    return header_name in {
+        "Content-Length",
+        "Content-MD5",
+        "x-ms-lease-id",
+        "x-ms-cache-control",
+        "x-ms-content-type",
+        "x-ms-content-disposition",
+        "x-ms-content-encoding",
+        "x-ms-content-language",
+        "x-ms-content-md5",
+        "x-ms-properties",
+        "x-ms-owner",
+        "x-ms-group",
+        "x-ms-permissions",
+        "x-ms-acl",
+        "If-Match",
+        "If-None-Match",
+        "If-Modified-Since",
+        "If-Unmodified-Since",
+        "x-ms-encryption-key",
+        "x-ms-encryption-key-sha256",
+        "x-ms-encryption-algorithm",
+        "x-ms-encryption-context",
+        "x-ms-client-request-id",
+        "x-ms-date",
+        "x-ms-version",
+    }

mlflow/azure/config.py

@@ -0,0 +1,120 @@
+"""Configuration management for Azure authentication."""
+
+from __future__ import annotations
+
+from enum import Enum
+import os
+from typing import Optional
+
+
+class AuthMethod(str, Enum):
+    """Supported authentication methods."""
+
+    MANAGED_IDENTITY = "managed_identity"
+    SERVICE_PRINCIPAL = "service_principal"
+    SQL_AUTH = "sql_auth"
+    DEFAULT_AZURE_CREDENTIAL = "default_azure_credential"
+
+
+class AzureAuthConfig:
+    """Configuration for Azure authentication.
+
+    This configuration can be loaded from environment variables or
+    provided programmatically.
+    """
+
+    def __init__(
+        self,
+        auth_enabled: Optional[bool] = None,
+        auth_method: Optional[str] = None,
+        client_id: Optional[str] = None,
+        client_secret: Optional[str] = None,
+        tenant_id: Optional[str] = None,
+        token_scope: str = "https://ossrdbms-aad.database.windows.net/.default",
+        connection_timeout: int = 30,
+        token_refresh_buffer: int = 300,
+        pool_size: int = 20,
+        max_overflow: int = 10,
+        pool_recycle: int = 3600,
+        pool_pre_ping: bool = True,
+        max_retries: int = 3,
+        retry_delay: float = 1.0,
+        retry_backoff: float = 2.0,
+        enable_debug_logging: bool = False,
+    ):
+        """Initialize Azure authentication configuration."""
+        
+        # Core settings - use provided values or environment variables
+        self.auth_enabled = (
+            auth_enabled 
+            if auth_enabled is not None 
+            else os.getenv("MLFLOW_AZURE_AUTH_ENABLED", "false").lower() == "true"
+        )
+        
+        self.auth_method = AuthMethod(
+            auth_method.lower() 
+            if auth_method 
+            else os.getenv("MLFLOW_AZURE_AUTH_METHOD", "sql_auth").lower()
+        )
+
+        # Azure specific settings - check both provided values and environment
+        self.client_id = client_id or os.getenv("AZURE_CLIENT_ID")
+        self.client_secret = client_secret or os.getenv("AZURE_CLIENT_SECRET") 
+        self.tenant_id = tenant_id or os.getenv("AZURE_TENANT_ID")
+
+        # Connection settings
+        self.token_scope = token_scope
+        self.connection_timeout = connection_timeout
+        self.token_refresh_buffer = token_refresh_buffer
+
+        # Pool settings
+        self.pool_size = pool_size
+        self.max_overflow = max_overflow
+        self.pool_recycle = pool_recycle
+        self.pool_pre_ping = pool_pre_ping
+
+        # Retry settings
+        self.max_retries = max_retries
+        self.retry_delay = retry_delay
+        self.retry_backoff = retry_backoff
+
+        # Logging
+        self.enable_debug_logging = (
+            enable_debug_logging
+            or os.getenv("MLFLOW_AZURE_DEBUG", "false").lower() == "true"
+        )
+
+        # Additional environment variable checks
+        # Only override auth method if not explicitly set
+        if (os.getenv("MLFLOW_POSTGRES_USE_MANAGED_IDENTITY", "false").lower() == "true" and 
+            auth_method is None and 
+            os.getenv("MLFLOW_AZURE_AUTH_METHOD") is None):
+            self.auth_enabled = True
+            self.auth_method = AuthMethod.MANAGED_IDENTITY
+
+        # Validate configuration
+        self._validate()
+
+    def _validate(self):
+        """Validate configuration settings."""
+        if self.auth_method == AuthMethod.SERVICE_PRINCIPAL:
+            if not self.client_secret:
+                raise ValueError("client_secret is required for service principal authentication")
+            if not self.client_id:
+                raise ValueError("client_id is required for service principal authentication")
+            if not self.tenant_id:
+                raise ValueError("tenant_id is required for service principal authentication")
+
+    @property
+    def should_use_azure_auth(self) -> bool:
+        """Check if Azure authentication should be used."""
+        return self.auth_enabled and self.auth_method != AuthMethod.SQL_AUTH
+
+    @property
+    def requires_token_refresh(self) -> bool:
+        """Check if the auth method requires token refresh."""
+        return self.auth_method in [
+            AuthMethod.MANAGED_IDENTITY,
+            AuthMethod.SERVICE_PRINCIPAL,
+            AuthMethod.DEFAULT_AZURE_CREDENTIAL,
+        ]