atlan-application-sdk 2.1.1__py3-none-any.whl → 2.3.0__py3-none-any.whl

application_sdk/application/__init__.py

@@ -1,13 +1,20 @@
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Dict, List, Optional, Tuple, Type
 
+from typing_extensions import deprecated
+
 from application_sdk.activities import ActivitiesInterface
 from application_sdk.clients.base import BaseClient
 from application_sdk.clients.utils import get_workflow_client
-from application_sdk.constants import ENABLE_MCP
+from application_sdk.constants import APPLICATION_MODE, ENABLE_MCP, ApplicationMode
 from application_sdk.handlers.base import BaseHandler
 from application_sdk.interceptors.models import EventRegistration
+from application_sdk.observability.decorators.observability_decorator import (
+    observability,
+)
 from application_sdk.observability.logger_adaptor import get_logger
+from application_sdk.observability.metrics_adaptor import get_metrics
+from application_sdk.observability.traces_adaptor import get_traces
 from application_sdk.server import ServerInterface
 from application_sdk.server.fastapi import APIServer, HttpWorkflowTrigger
 from application_sdk.server.fastapi.models import EventWorkflowTrigger
@@ -15,6 +22,8 @@ from application_sdk.worker import Worker
 from application_sdk.workflows import WorkflowInterface
 
 logger = get_logger(__name__)
+metrics = get_metrics()
+traces = get_traces()
 
 
 class BaseApplication:
@@ -110,6 +119,40 @@ class BaseApplication:
 
         self.event_subscriptions[event_id].workflow_class = workflow_class
 
+    async def start(
+        self,
+        workflow_class: Type[WorkflowInterface],
+        ui_enabled: bool = True,
+        has_configmap: bool = False,
+    ):
+        """Start the application based on the configured APPLICATION_MODE.
+
+        Args:
+            workflow_class: The workflow class to register with the server.
+            ui_enabled: Whether to enable the UI. Defaults to True.
+            has_configmap: Whether the application has a configmap. Defaults to False.
+
+        Behavior based on APPLICATION_MODE:
+            - LOCAL: Starts worker in daemon mode and server (for local development)
+            - WORKER: Starts only the worker in non-daemon mode (for production worker pods)
+            - SERVER: Starts only the server (for production API server pods)
+
+        Raises:
+            ValueError: If APPLICATION_MODE is not a valid ApplicationMode value.
+        """
+        if APPLICATION_MODE in (ApplicationMode.LOCAL, ApplicationMode.WORKER):
+            await self._start_worker(
+                daemon=APPLICATION_MODE == ApplicationMode.LOCAL,
+            )
+
+        if APPLICATION_MODE in (ApplicationMode.LOCAL, ApplicationMode.SERVER):
+            await self._setup_server(
+                workflow_class=workflow_class,
+                ui_enabled=ui_enabled,
+                has_configmap=has_configmap,
+            )
+            await self._start_server()
+
     async def setup_workflow(
         self,
         workflow_and_activities_classes: List[
@@ -167,7 +210,12 @@ class BaseApplication:
             raise ValueError("Workflow client not initialized")
         return await self.workflow_client.start_workflow(workflow_args, workflow_class)  # type: ignore
 
+    @deprecated("Use application.start(). Deprecated since v2.3.0.")
+    @observability(logger=logger, metrics=metrics, traces=traces)
     async def start_worker(self, daemon: bool = True):
+        return await self._start_worker(daemon=daemon)
+
+    async def _start_worker(self, daemon: bool = True):
         """
         Start the worker for the application.
 
@@ -178,11 +226,25 @@ class BaseApplication:
             raise ValueError("Worker not initialized")
         await self.worker.start(daemon=daemon)
 
+    @deprecated("Use application.start(). Deprecated since v2.3.0.")
+    @observability(logger=logger, metrics=metrics, traces=traces)
     async def setup_server(
         self,
         workflow_class: Type[WorkflowInterface],
         ui_enabled: bool = True,
         has_configmap: bool = False,
+    ):
+        return await self._setup_server(
+            workflow_class=workflow_class,
+            ui_enabled=ui_enabled,
+            has_configmap=has_configmap,
+        )
+
+    async def _setup_server(
+        self,
+        workflow_class: Type[WorkflowInterface],
+        ui_enabled: bool = True,
+        has_configmap: bool = False,
     ):
         """
         Set up FastAPI server and automatically mount MCP if enabled.
@@ -239,7 +301,12 @@ class BaseApplication:
             triggers=[HttpWorkflowTrigger()],
         )
 
+    @deprecated("Use application.start(). Deprecated since v2.3.0.")
+    @observability(logger=logger, metrics=metrics, traces=traces)
     async def start_server(self):
+        return await self._start_server()
+
+    async def _start_server(self):
         """
         Start the FastAPI server for the application.
 

application_sdk/application/metadata_extraction/sql.py

@@ -1,6 +1,8 @@
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Dict, List, Optional, Tuple, Type
 
+from typing_extensions import deprecated
+
 from application_sdk.application import BaseApplication
 from application_sdk.clients.sql import BaseSQLClient
 from application_sdk.clients.utils import get_workflow_client
@@ -146,8 +148,33 @@ class BaseSQLMetadataExtractionApplication(BaseApplication):
         )
         return workflow_response
 
+    @observability(logger=logger, metrics=metrics, traces=traces)
+    async def start(
+        self,
+        workflow_class: Type = BaseSQLMetadataExtractionWorkflow,
+        ui_enabled: bool = True,
+        has_configmap: bool = False,
+    ):
+        """Start the SQL metadata extraction application.
+
+        Args:
+            workflow_class: The workflow class to register. Defaults to BaseSQLMetadataExtractionWorkflow.
+            ui_enabled: Whether to enable the UI. Defaults to True.
+            has_configmap: Whether the application has a configmap. Defaults to False.
+        """
+        await super().start(
+            workflow_class=workflow_class,
+            ui_enabled=ui_enabled,
+            has_configmap=has_configmap,
+        )
+
+    @deprecated("Use application.start(). Deprecated since v2.3.0.")
     @observability(logger=logger, metrics=metrics, traces=traces)
     async def start_worker(self, daemon: bool = True):
+        return await self._start_worker(daemon=daemon)
+
+    @observability(logger=logger, metrics=metrics, traces=traces)
+    async def _start_worker(self, daemon: bool = True):
         """
         Start the worker for the SQL metadata extraction application.
         """
@@ -155,12 +182,27 @@ class BaseSQLMetadataExtractionApplication(BaseApplication):
             raise ValueError("Worker not initialized")
         await self.worker.start(daemon=daemon)
 
+    @deprecated("Use application.start(). Deprecated since v2.3.0.")
     @observability(logger=logger, metrics=metrics, traces=traces)
     async def setup_server(
+        self,
+        workflow_class: Type = BaseSQLMetadataExtractionWorkflow,
+        ui_enabled: bool = True,
+        has_configmap: bool = False,
+    ):
+        return await self._setup_server(
+            workflow_class=workflow_class,
+            ui_enabled=ui_enabled,
+            has_configmap=has_configmap,
+        )
+
+    @observability(logger=logger, metrics=metrics, traces=traces)
+    async def _setup_server(
         self,
         workflow_class: Type[
             BaseSQLMetadataExtractionWorkflow
         ] = BaseSQLMetadataExtractionWorkflow,
+        ui_enabled: bool = True,
         has_configmap: bool = False,
     ) -> Any:
         """
@@ -168,6 +210,7 @@ class BaseSQLMetadataExtractionApplication(BaseApplication):
 
         Args:
             workflow_class (Type): Workflow class to register with the server. Defaults to BaseSQLMetadataExtractionWorkflow.
+            ui_enabled (bool): Whether to enable the UI. Defaults to True.
             has_configmap (bool): Whether the application has a configmap. Defaults to False.
 
         Returns:
@@ -180,6 +223,7 @@ class BaseSQLMetadataExtractionApplication(BaseApplication):
         self.server = APIServer(
             handler=self.handler_class(sql_client=self.client_class()),
             workflow_client=self.workflow_client,
+            ui_enabled=ui_enabled,
             has_configmap=has_configmap,
         )
 

application_sdk/clients/azure/__init__.py

@@ -0,0 +1,6 @@
+"""Azure client module for the application-sdk framework."""
+
+# Azure Management API endpoint for token acquisition
+AZURE_MANAGEMENT_API_ENDPOINT = "https://management.azure.com/.default"
+
+__all__ = ["AZURE_MANAGEMENT_API_ENDPOINT"]

application_sdk/clients/azure/auth.py

@@ -0,0 +1,288 @@
+"""
+Azure authentication provider for the application-sdk framework.
+
+This module provides the AzureAuthProvider class that handles Azure
+Service Principal authentication for the application-sdk framework.
+
+Example:
+    >>> from application_sdk.clients.azure.auth import AzureAuthProvider
+    >>> import asyncio
+    >>>
+    >>> # Create authentication provider
+    >>> auth_provider = AzureAuthProvider()
+    >>>
+    >>> # Authenticate with Service Principal credentials
+    >>> credentials = {
+    ...     "tenant_id": "your-tenant-id",
+    ...     "client_id": "your-client-id",
+    ...     "client_secret": "your-client-secret"
+    ... }
+    >>>
+    >>> # Create credential
+    >>> credential = await auth_provider.create_credential(
+    ...     auth_type="service_principal",
+    ...     credentials=credentials
+    ... )
+    >>>
+    >>> # Alternative credential key formats are also supported
+    >>> alt_credentials = {
+    ...     "tenantId": "your-tenant-id",      # camelCase
+    ...     "clientId": "your-client-id",      # camelCase
+    ...     "clientSecret": "your-client-secret"  # camelCase
+    ... }
+    >>>
+    >>> credential = await auth_provider.create_credential(
+    ...     auth_type="service_principal",
+    ...     credentials=alt_credentials
+    ... )
+    >>>
+    >>> # Error handling for missing credentials
+    >>> try:
+    ...     await auth_provider.create_credential(
+    ...         auth_type="service_principal",
+    ...         credentials={"tenant_id": "only-tenant"}  # Missing client_id and client_secret
+    ...     )
+    ... except CommonError as e:
+    ...     print(f"Authentication failed: {e}")
+    ...     # Output: Authentication failed: Missing required credential keys: client_id, client_secret
+    >>>
+    >>> # Unsupported authentication type
+    >>> try:
+    ...     await auth_provider.create_credential(
+    ...         auth_type="unsupported_type",
+    ...         credentials=credentials
+    ...     )
+    ... except CommonError as e:
+    ...     print(f"Authentication failed: {e}")
+    ...     # Output: Authentication failed: Only 'service_principal' authentication is supported. Received: unsupported_type
+"""
+
+from typing import Any, Dict, Optional
+
+from azure.core.credentials import TokenCredential
+from azure.core.exceptions import ClientAuthenticationError
+from azure.identity import ClientSecretCredential
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
+
+from application_sdk.clients.azure import AZURE_MANAGEMENT_API_ENDPOINT
+from application_sdk.common.error_codes import CommonError
+from application_sdk.common.utils import run_sync
+from application_sdk.observability.logger_adaptor import get_logger
+
+logger = get_logger(__name__)
+
+
+class ServicePrincipalCredentials(BaseModel):
+    """
+    Pydantic model for Azure Service Principal credentials.
+
+    Supports both snake_case and camelCase field names through field aliases.
+    All fields are required for service principal authentication.
+
+    Attributes:
+        tenant_id: Azure tenant ID (also accepts 'tenantId').
+        client_id: Azure client ID (also accepts 'clientId').
+        client_secret: Azure client secret (also accepts 'clientSecret').
+    """
+
+    tenant_id: str = Field(
+        ...,
+        alias="tenantId",
+        description="Azure tenant ID for service principal authentication",
+    )
+    client_id: str = Field(
+        ...,
+        alias="clientId",
+        description="Azure client ID for service principal authentication",
+    )
+    client_secret: str = Field(
+        ...,
+        alias="clientSecret",
+        description="Azure client secret for service principal authentication",
+    )
+
+    model_config = ConfigDict(
+        populate_by_name=True,  # Allow both field name and alias
+        extra="ignore",  # Ignore additional fields (Azure client may need extra fields like storage_account_name, network_config, etc.)
+        validate_assignment=True,  # Validate on assignment
+    )
+
+
+class AzureAuthProvider:
+    """
+    Azure authentication provider for handling Service Principal authentication.
+
+    This class provides a unified interface for creating Azure credentials
+    using Service Principal authentication with Azure SDK.
+
+    Supported authentication method:
+    - service_principal: Using client ID, client secret, and tenant ID
+    """
+
+    def __init__(self):
+        """Initialize the Azure authentication provider."""
+        pass
+
+    async def create_credential(
+        self,
+        auth_type: str = "service_principal",
+        credentials: Optional[Dict[str, Any]] = None,
+    ) -> TokenCredential:
+        """
+        Create Azure credential using Service Principal authentication.
+
+        Args:
+            auth_type (str): Type of authentication to use.
+                Currently only supports 'service_principal'.
+            credentials (Optional[Dict[str, Any]]): Service Principal credentials.
+                Required fields: tenant_id, client_id, client_secret.
+
+        Returns:
+            TokenCredential: Azure credential instance.
+
+        Raises:
+            CommonError: If authentication type is not supported or credentials are invalid.
+            ClientAuthenticationError: If credential creation fails.
+        """
+        try:
+            logger.debug(f"Creating Azure credential with auth type: {auth_type}")
+
+            if auth_type.lower() != "service_principal":
+                raise CommonError(
+                    f"{CommonError.CREDENTIALS_PARSE_ERROR}: "
+                    f"Only 'service_principal' authentication is supported. "
+                    f"Received: {auth_type}"
+                )
+
+            if not credentials:
+                raise CommonError(
+                    f"{CommonError.CREDENTIALS_PARSE_ERROR}: "
+                    "Credentials required for service principal authentication"
+                )
+
+            return await self._create_service_principal_credential(credentials)
+
+        except ClientAuthenticationError as e:
+            logger.error(f"Azure authentication failed: {str(e)}")
+            raise CommonError(f"{CommonError.AZURE_CREDENTIAL_ERROR}: {str(e)}")
+        except ValueError as e:
+            logger.error(f"Invalid Azure credential parameters: {str(e)}")
+            raise CommonError(
+                f"{CommonError.CREDENTIALS_PARSE_ERROR}: Invalid credential parameters - {str(e)}"
+            )
+        except TypeError as e:
+            logger.error(f"Wrong Azure credential parameter types: {str(e)}")
+            raise CommonError(
+                f"{CommonError.CREDENTIALS_PARSE_ERROR}: Invalid credential parameter types - {str(e)}"
+            )
+        except Exception as e:
+            logger.error(f"Unexpected error creating Azure credential: {str(e)}")
+            raise CommonError(
+                f"{CommonError.CREDENTIALS_PARSE_ERROR}: Unexpected error - {str(e)}"
+            )
+
+    async def _create_service_principal_credential(
+        self, credentials: Dict[str, Any]
+    ) -> ClientSecretCredential:
+        """
+        Create service principal credential.
+
+        Args:
+            credentials (Dict[str, Any]): Service principal credentials.
+                Must include tenant_id, client_id, and client_secret.
+
+        Returns:
+            ClientSecretCredential: Service principal credential.
+
+        Raises:
+            CommonError: If required credentials are missing or invalid.
+        """
+        if not credentials:
+            raise CommonError(
+                f"{CommonError.CREDENTIALS_PARSE_ERROR}: "
+                "Credentials required for service principal authentication"
+            )
+
+        try:
+            # Validate credentials using Pydantic model
+            validated_credentials = ServicePrincipalCredentials(**credentials)
+        except ValidationError as e:
+            # Pydantic provides detailed error messages for all validation errors
+            # Format errors into a user-friendly message
+            error_details = "; ".join(
+                [
+                    f"{'.'.join(str(loc) for loc in err['loc'])}: {err['msg']}"
+                    for err in e.errors()
+                ]
+            )
+            error_message = f"Invalid credential parameters: {error_details}"
+            logger.error(f"Azure credential validation failed: {error_message}")
+            raise CommonError(f"{CommonError.CREDENTIALS_PARSE_ERROR}: {error_message}")
+
+        logger.debug(
+            f"Creating service principal credential for tenant: {validated_credentials.tenant_id}"
+        )
+
+        try:
+            return await run_sync(ClientSecretCredential)(
+                validated_credentials.tenant_id,
+                validated_credentials.client_id,
+                validated_credentials.client_secret,
+            )
+        except ValueError as e:
+            logger.error(f"Invalid Azure credential parameters: {str(e)}")
+            raise CommonError(
+                f"{CommonError.CREDENTIALS_PARSE_ERROR}: Invalid credential parameters - {str(e)}"
+            )
+        except TypeError as e:
+            logger.error(f"Wrong Azure credential parameter types: {str(e)}")
+            raise CommonError(
+                f"{CommonError.CREDENTIALS_PARSE_ERROR}: Invalid credential parameter types - {str(e)}"
+            )
+        except ClientAuthenticationError as e:
+            logger.error(f"Azure authentication failed: {str(e)}")
+            raise CommonError(f"{CommonError.AZURE_CREDENTIAL_ERROR}: {str(e)}")
+        except Exception as e:
+            logger.error(f"Unexpected error creating Azure credential: {str(e)}")
+            raise CommonError(
+                f"{CommonError.CREDENTIALS_PARSE_ERROR}: Unexpected error - {str(e)}"
+            )
+
+    async def validate_credential(self, credential: TokenCredential) -> bool:
+        """
+        Validate Azure credential by attempting to get a token.
+
+        Args:
+            credential (TokenCredential): Azure credential to validate.
+
+        Returns:
+            bool: True if credential is valid, False otherwise.
+        """
+        try:
+            logger.debug("Validating Azure credential")
+
+            # Try to get a token for Azure Management API
+            token = await run_sync(credential.get_token)(AZURE_MANAGEMENT_API_ENDPOINT)
+
+            if token and hasattr(token, "token"):
+                logger.debug("Azure credential validation successful")
+                return True
+            else:
+                logger.warning("Azure credential validation failed: No token received")
+                return False
+
+        except ClientAuthenticationError as e:
+            logger.error(
+                f"Azure credential validation failed - authentication error: {str(e)}"
+            )
+            return False
+        except ValueError as e:
+            logger.error(
+                f"Azure credential validation failed - invalid parameters: {str(e)}"
+            )
+            return False
+        except Exception as e:
+            logger.error(
+                f"Azure credential validation failed - unexpected error: {str(e)}"
+            )
+            return False