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

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

application_sdk/clients/azure/client.py

@@ -0,0 +1,336 @@
+"""
+Azure client implementation for the application-sdk framework.
+
+This module provides the main AzureClient class that serves as a unified interface
+for connecting to and interacting with Azure Storage services. It supports Service Principal
+authentication and provides service-specific subclients.
+
+Example:
+    >>> from application_sdk.clients.azure.client import AzureClient
+    >>> from application_sdk.clients.azure.auth import AzureAuthProvider
+    >>>
+    >>> # Create Azure client with Service Principal credentials
+    >>> credentials = {
+    ...     "tenant_id": "your-tenant-id",
+    ...     "client_id": "your-client-id",
+    ...     "client_secret": "your-client-secret"
+    ... }
+    >>>
+    >>> client = AzureClient(credentials)
+    >>> await client.load()
+    >>>
+    >>> # Check client health
+    >>> health_status = await client.health_check()
+    >>> print(f"Overall health: {health_status.overall_health}")
+    >>> print(f"Connection health: {health_status.connection_health}")
+    >>>
+    >>> # Access health status details
+    >>> for service_name, service_health in health_status.services.items():
+    ...     print(f"{service_name}: {service_health.status}")
+    ...     if service_health.error:
+    ...         print(f"  Error: {service_health.error}")
+"""
+
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Dict, Optional
+
+from azure.core.credentials import TokenCredential
+from azure.core.exceptions import AzureError, ClientAuthenticationError
+from pydantic import BaseModel
+
+from application_sdk.clients import ClientInterface
+from application_sdk.clients.azure import AZURE_MANAGEMENT_API_ENDPOINT
+from application_sdk.clients.azure.auth import AzureAuthProvider
+from application_sdk.common.error_codes import ClientError
+from application_sdk.common.utils import run_sync
+from application_sdk.observability.logger_adaptor import get_logger
+
+logger = get_logger(__name__)
+
+
+class ServiceHealth(BaseModel):
+    """Model for individual service health status.
+
+    Attributes:
+        status: The health status of the service (e.g., "healthy", "error", "unknown")
+        error: Optional error message if the service is unhealthy
+    """
+
+    status: str
+    error: Optional[str] = None
+
+
+class HealthStatus(BaseModel):
+    """Model for overall Azure client health status.
+
+    Attributes:
+        connection_health: Whether the Azure connection is healthy
+        services: Dictionary mapping service names to their health status
+        overall_health: Overall health status considering connection and services
+    """
+
+    connection_health: bool
+    services: Dict[str, ServiceHealth]
+    overall_health: bool
+
+
+class AzureClient(ClientInterface):
+    """
+    Main Azure client for the application-sdk framework.
+
+    This client provides a unified interface for connecting to and interacting
+    with Azure services. It supports Service Principal authentication
+    and provides service-specific subclients for different Azure services.
+
+    Attributes:
+        credentials (Dict[str, Any]): Azure connection credentials
+        resolved_credentials (Dict[str, Any]): Resolved credentials after processing
+        credential (DefaultAzureCredential): Azure credential instance
+        auth_provider (AzureAuthProvider): Authentication provider instance
+        _services (Dict[str, Any]): Cache of service clients
+        _executor (ThreadPoolExecutor): Thread pool for async operations
+        _connection_health (bool): Connection health status
+    """
+
+    def __init__(
+        self,
+        credentials: Optional[Dict[str, Any]] = None,
+        max_workers: int = 10,
+        **kwargs: Any,
+    ):
+        """
+        Initialize the Azure client.
+
+        Args:
+            credentials (Optional[Dict[str, Any]]): Azure Service Principal credentials.
+                Must include tenant_id, client_id, and client_secret.
+            max_workers (int): Maximum number of worker threads for async operations.
+            **kwargs: Additional keyword arguments passed to service clients.
+        """
+        self.credentials = credentials or {}
+        self.resolved_credentials: Dict[str, Any] = {}
+        self.credential: Optional[TokenCredential] = None
+        self.auth_provider = AzureAuthProvider()
+        self._services: Dict[str, Any] = {}
+        self._executor = ThreadPoolExecutor(max_workers=max_workers)
+        self._connection_health = False
+        self._kwargs = kwargs
+
+    async def load(self, credentials: Optional[Dict[str, Any]] = None) -> None:
+        """
+        Load and establish Azure connection using Service Principal authentication.
+
+        Args:
+            credentials (Optional[Dict[str, Any]]): Azure Service Principal credentials.
+                If provided, will override the credentials passed to __init__.
+                Must include tenant_id, client_id, and client_secret.
+
+        Raises:
+            ClientError: If connection fails due to authentication or connection issues
+        """
+        if credentials:
+            self.credentials = credentials
+
+        try:
+            logger.info("Loading Azure client...")
+
+            # Handle credential resolution
+            if "credential_guid" in self.credentials:
+                # If we have a credential_guid, use the async get_credentials function
+                from application_sdk.services.secretstore import SecretStore
+
+                self.resolved_credentials = await SecretStore.get_credentials(
+                    self.credentials["credential_guid"]
+                )
+            else:
+                # If credentials are already resolved (direct format), use them as-is
+                # Check if credentials appear to need resolution but no credential_guid provided
+                if (
+                    "secret-path" in self.credentials
+                    or "credentialSource" in self.credentials
+                ):
+                    logger.warning(
+                        "Credentials appear to need resolution but no credential_guid provided. Using as-is."
+                    )
+                # Credentials are already in the correct format
+                self.resolved_credentials = self.credentials
+
+            # Create Azure credential using Service Principal authentication
+            self.credential = await self.auth_provider.create_credential(
+                auth_type="service_principal", credentials=self.resolved_credentials
+            )
+
+            # Test the connection
+            await self._test_connection()
+
+            self._connection_health = True
+            logger.info("Azure client loaded successfully")
+
+        except ClientAuthenticationError as e:
+            logger.error(f"Azure authentication failed: {str(e)}")
+            raise ClientError(f"{ClientError.CLIENT_AUTH_ERROR}: {str(e)}")
+        except AzureError as e:
+            logger.error(f"Azure connection error: {str(e)}")
+            raise ClientError(f"{ClientError.CLIENT_AUTH_ERROR}: {str(e)}")
+        except ValueError as e:
+            logger.error(f"Invalid Azure client parameters: {str(e)}")
+            raise ClientError(
+                f"{ClientError.INPUT_VALIDATION_ERROR}: Invalid parameters - {str(e)}"
+            )
+        except TypeError as e:
+            logger.error(f"Wrong Azure client parameter types: {str(e)}")
+            raise ClientError(
+                f"{ClientError.INPUT_VALIDATION_ERROR}: Invalid parameter types - {str(e)}"
+            )
+        except Exception as e:
+            logger.error(f"Unexpected error loading Azure client: {str(e)}")
+            raise ClientError(
+                f"{ClientError.CLIENT_AUTH_ERROR}: Unexpected error - {str(e)}"
+            )
+
+    async def close(self) -> None:
+        """Close Azure connections and clean up resources."""
+        try:
+            logger.info("Closing Azure client...")
+
+            # Close all service clients
+            for service_name, service_client in self._services.items():
+                try:
+                    if hasattr(service_client, "close"):
+                        await service_client.close()
+                    elif hasattr(service_client, "disconnect"):
+                        await service_client.disconnect()
+                except Exception as e:
+                    logger.warning(f"Error closing {service_name} client: {str(e)}")
+
+            # Clear service cache
+            self._services.clear()
+
+            # Shutdown executor
+            self._executor.shutdown(wait=True)
+
+            # Reset connection health
+            self._connection_health = False
+
+            logger.info("Azure client closed successfully")
+
+        except Exception as e:
+            logger.error(f"Error closing Azure client: {str(e)}")
+
+    async def health_check(self) -> HealthStatus:
+        """
+        Perform health check on Azure connection and services.
+
+        Returns:
+            HealthStatus: Health status information.
+        """
+        health_status = HealthStatus(
+            connection_health=self._connection_health,
+            services={},
+            overall_health=False,
+        )
+
+        if not self._connection_health:
+            return health_status
+
+        # Check each service
+        for service_name, service_client in self._services.items():
+            try:
+                if hasattr(service_client, "health_check"):
+                    service_health = await service_client.health_check()
+                    # Handle different return types from service health checks
+                    if isinstance(service_health, dict):
+                        status = service_health.get("status", "unknown")
+                        error = service_health.get("error")
+                    elif hasattr(service_health, "status"):
+                        # Handle Pydantic models or objects with status attribute
+                        status = getattr(service_health, "status", "unknown")
+                        error = getattr(service_health, "error", None)
+                    else:
+                        # Fallback for unexpected return types
+                        status = "unknown"
+                        error = f"Unexpected health check return type: {type(service_health)}"
+                else:
+                    status = "unknown"
+                    error = None
+
+                health_status.services[service_name] = ServiceHealth(
+                    status=status,
+                    error=error,
+                )
+            except Exception as e:
+                health_status.services[service_name] = ServiceHealth(
+                    status="error", error=str(e)
+                )
+
+        # Overall health is True if connection is healthy and at least one service is available
+        health_status.overall_health = (
+            self._connection_health and len(health_status.services) > 0
+        )
+
+        return health_status
+
+    async def _test_connection(self) -> None:
+        """
+        Test the Azure connection by attempting to get a token.
+
+        Raises:
+            ClientAuthenticationError: If connection test fails.
+        """
+        if not self.credential:
+            raise ClientError(
+                f"{ClientError.AUTH_CREDENTIALS_ERROR}: No credential available for connection test"
+            )
+
+        try:
+            # Test the credential by getting a token
+            await run_sync(self.credential.get_token)(AZURE_MANAGEMENT_API_ENDPOINT)
+        except ClientAuthenticationError as e:
+            logger.error(
+                f"Azure connection test failed - authentication error: {str(e)}"
+            )
+            raise ClientError(f"{ClientError.CLIENT_AUTH_ERROR}: {str(e)}")
+        except AzureError as e:
+            logger.error(f"Azure connection test failed - service error: {str(e)}")
+            raise ClientError(f"{ClientError.CLIENT_AUTH_ERROR}: {str(e)}")
+        except ValueError as e:
+            logger.error(f"Azure connection test failed - invalid parameters: {str(e)}")
+            raise ClientError(
+                f"{ClientError.INPUT_VALIDATION_ERROR}: Invalid parameters - {str(e)}"
+            )
+        except Exception as e:
+            logger.error(f"Azure connection test failed - unexpected error: {str(e)}")
+            raise ClientError(
+                f"{ClientError.CLIENT_AUTH_ERROR}: Unexpected error - {str(e)}"
+            )
+
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        # Note: This is a synchronous context manager.
+        # For proper async cleanup, use the async context manager instead.
+        # This method is kept for backward compatibility but doesn't guarantee cleanup.
+        logger.warning(
+            "Using synchronous context manager. For proper async cleanup, "
+            "use 'async with AzureClient() as client:' instead."
+        )
+        # Schedule cleanup but don't wait for it
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                loop.create_task(self.close())
+        except RuntimeError:
+            # No event loop running, can't schedule async cleanup
+            logger.warning("No event loop running, async cleanup not possible")
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.close()

application_sdk/common/error_codes.py

@@ -288,6 +288,15 @@ class CommonError(AtlanError):
     CREDENTIALS_RESOLUTION_ERROR = ErrorCode(
         ErrorComponent.COMMON, "401", "03", "Credentials resolution error"
     )
+    AZURE_CREDENTIAL_ERROR = ErrorCode(
+        ErrorComponent.COMMON, "401", "04", "Azure credential error"
+    )
+    AZURE_CONNECTION_ERROR = ErrorCode(
+        ErrorComponent.COMMON, "503", "01", "Azure connection error"
+    )
+    AZURE_SERVICE_ERROR = ErrorCode(
+        ErrorComponent.COMMON, "500", "01", "Azure service error"
+    )
 
 
 class DocGenError(AtlanError):

application_sdk/server/fastapi/__init__.py

@@ -42,6 +42,7 @@ from application_sdk.server.fastapi.models import (
     HttpWorkflowTrigger,
     PreflightCheckRequest,
     PreflightCheckResponse,
+    Subscription,
     TestAuthRequest,
     TestAuthResponse,
     WorkflowConfigRequest,
@@ -90,12 +91,15 @@ class APIServer(ServerInterface):
     workflow_router: APIRouter
     dapr_router: APIRouter
     events_router: APIRouter
+    subscription_router: APIRouter
     handler: Optional[HandlerInterface]
     templates: Jinja2Templates
     duckdb_ui: DuckDBUI
 
     docs_directory_path: str = "docs"
     docs_export_path: str = "dist"
+    # List of subscriptions to be registered
+    subscriptions: List[Subscription] = []
 
     frontend_assets_path: str = "frontend/static"
 
@@ -112,6 +116,7 @@ class APIServer(ServerInterface):
         frontend_templates_path: str = "frontend/templates",
         ui_enabled: bool = True,
         has_configmap: bool = False,
+        subscriptions: List[Subscription] = [],
     ):
         """Initialize the FastAPI application.
 
@@ -138,7 +143,7 @@ class APIServer(ServerInterface):
         self.workflow_router = APIRouter()
         self.dapr_router = APIRouter()
         self.events_router = APIRouter()
-
+        self.subscriptions = subscriptions
         # Set up the application
         error_handler = internal_server_error_handler  # Store as local variable
         self.app.add_exception_handler(
@@ -205,6 +210,7 @@ class APIServer(ServerInterface):
         - Workflow router (/workflows/v1)
         - Pubsub router (/dapr)
         - Events router (/events/v1)
+        - Subscription router (/subscriptions/v1)
         """
         # Register all routes first
         self.register_routes()
@@ -215,6 +221,16 @@ class APIServer(ServerInterface):
         self.app.include_router(self.dapr_router, prefix="/dapr")
         self.app.include_router(self.events_router, prefix="/events/v1")
 
+        # Register subscription routes from subscriptions with handler callbacks
+        subscription_router = APIRouter()
+        for subscription in self.subscriptions:
+            subscription_router.add_api_route(
+                f"/{subscription.route}",
+                subscription.handler,
+                methods=["POST"],
+            )
+        self.app.include_router(subscription_router, prefix="/subscriptions/v1")
+
     def fallback_home(self, request: Request) -> HTMLResponse:
         return self.templates.TemplateResponse(
             "index.html",
@@ -432,6 +448,19 @@ class APIServer(ServerInterface):
         """
 
         subscriptions: List[dict[str, Any]] = []
+        for subscription in self.subscriptions:
+            subscription_dict: dict[str, Any] = {
+                "pubsubname": subscription.component_name,
+                "topic": subscription.topic,
+                "route": f"/subscriptions/v1/{subscription.route}",
+            }
+            if subscription.bulk_config:
+                subscription_dict["bulkSubscribe"] = (
+                    subscription.bulk_config.model_dump(by_alias=True)
+                )
+            if subscription.dead_letter_topic:
+                subscription_dict["deadLetterTopic"] = subscription.dead_letter_topic
+            subscriptions.append(subscription_dict)
         for event_trigger in self.event_triggers:
             filters = [
                 f"({event_filter.path} {event_filter.operator} '{event_filter.value}')"

application_sdk/server/fastapi/models.py

@@ -1,7 +1,7 @@
 # Request/Response DTOs for workflows
 
 from enum import Enum
-from typing import Any, Dict, List, Optional, Type
+from typing import Any, Callable, Coroutine, Dict, List, Optional, Type, Union
 
 from pydantic import BaseModel, Field, RootModel
 
@@ -240,3 +240,64 @@ class EventWorkflowTrigger(WorkflowTrigger):
 
     def should_trigger_workflow(self, event: Event) -> bool:
         return True
+
+
+class Subscription(BaseModel):
+    """Subscription configuration for Dapr messaging.
+
+    Attributes:
+        component_name: Name of the Dapr pubsub component
+        topic: Topic to subscribe to
+        route: Route path for the message handler endpoint
+        handler: Required callback function to handle incoming messages
+        bulk_config: Optional bulk subscribe configuration
+        dead_letter_topic: Optional dead letter topic for failed messages
+
+    Nested Classes:
+        BulkConfig: Configuration for bulk message processing
+        MessageStatus: Status codes for handler responses (SUCCESS, RETRY, DROP)
+    """
+
+    class BulkConfig(BaseModel):
+        """Bulk configuration for Dapr messaging.
+
+        Attributes:
+            enabled: Whether bulk subscribe is enabled
+            max_messages_count: Maximum number of messages to receive in a batch
+            max_await_duration_ms: Maximum time to wait for messages in milliseconds
+        """
+
+        enabled: bool = False
+        max_messages_count: int = Field(
+            default=100, serialization_alias="maxMessagesCount"
+        )
+        max_await_duration_ms: int = Field(
+            default=40, serialization_alias="maxAwaitDurationMs"
+        )
+
+    class MessageStatus(str, Enum):
+        """Status codes for Dapr pub/sub subscription message handler responses.
+
+        Used in subscription handler responses to indicate how Dapr should handle the message.
+        Based on Dapr docs: https://docs.dapr.io/reference/api/pubsub_api/#expected-http-response
+
+        Attributes:
+            SUCCESS: Message was processed successfully.
+            RETRY: Message processing failed, should be retried.
+            DROP: Message should be dropped (sent to dead letter topic if configured).
+        """
+
+        SUCCESS = "SUCCESS"
+        RETRY = "RETRY"
+        DROP = "DROP"
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    component_name: str
+    topic: str
+    route: str
+    handler: Union[
+        Callable[[Any], Any], Callable[[Any], Coroutine[Any, Any, Any]]
+    ]  # Required callback function (sync or async)
+    bulk_config: Optional[BulkConfig] = None
+    dead_letter_topic: Optional[str] = None

application_sdk/transformers/query/__init__.py

@@ -205,12 +205,25 @@ class QueryBasedTransformer(TransformerInterface):
         Returns:
             Optional[daft.Expression]: The constructed struct expression or None if all fields are null
         """
+
+        # Check if level is None
+        if level is None:
+            logger.error("ERROR: level is None in _build_struct!")
+            raise ValueError("level cannot be None in _build_struct")
+
+        # Check if prefix is None
+        if prefix is None:
+            logger.error("ERROR: prefix is None in _build_struct!")
+            raise ValueError("prefix cannot be None in _build_struct")
+
         struct_fields = []
         non_null_fields = []
 
         # Handle columns at this level
         if "columns" in level:
+            logger.debug(f"Processing columns at level: {level['columns']}")
             for full_col, suffix in level["columns"]:
+                logger.debug(f"Processing column: {full_col} -> {suffix}")
                 field = daft.col(full_col).alias(suffix)
                 struct_fields.append(field)
                 # Add to non_null check by negating is_null()
@@ -219,6 +232,7 @@ class QueryBasedTransformer(TransformerInterface):
         # Handle nested levels
         for component, sub_level in level.items():
             if component != "columns":  # Skip the columns key
+                logger.debug(f"Processing nested component: {component}")
                 nested_struct = self._build_struct(sub_level, component)
                 if nested_struct is not None:
                     struct_fields.append(nested_struct)
@@ -227,6 +241,7 @@ class QueryBasedTransformer(TransformerInterface):
 
         # Only create a struct if we have fields
         if struct_fields:
+            logger.debug(f"Creating struct with {len(struct_fields)} fields")
             # Create the struct first
             struct = to_struct(*struct_fields)
 
@@ -242,6 +257,7 @@ class QueryBasedTransformer(TransformerInterface):
 
             return struct.alias(prefix)
 
+        logger.warning(f"No fields found for level {level}")
         return None
 
     def get_grouped_dataframe_by_prefix(
@@ -282,12 +298,18 @@ class QueryBasedTransformer(TransformerInterface):
         try:
             # Get all column names
             columns = dataframe.column_names
+            logger.debug("=== DEBUG: get_grouped_dataframe_by_prefix ===")
+            logger.debug(f"Input DataFrame columns: {columns}")
 
             # Group columns by their path components
             path_groups = {}
             standalone_columns = []
 
             for col in columns:
+                if col is None:
+                    logger.error(f"Found None column in DataFrame columns: {columns}")
+                    continue
+
                 if "." in col:
                     # Split the full path into components
                     path_components = col.split(".")
@@ -313,8 +335,12 @@ class QueryBasedTransformer(TransformerInterface):
             for col in standalone_columns:
                 new_columns.append(daft.col(col))
 
+            logger.debug(f"path_groups: {path_groups}")
+            logger.debug(f"standalone_columns: {standalone_columns}")
+
             # Build nested structs starting from the root level
             for prefix, level in path_groups.items():
+                logger.debug(f"Building struct for prefix: {prefix}, level: {level}")
                 struct_expr = self._build_struct(level, prefix)
                 new_columns.append(struct_expr)
 

application_sdk/version.py

@@ -2,4 +2,4 @@
 Version information for the application_sdk package.
 """
 
-__version__ = "2.1.1"
+__version__ = "2.2.0"

{atlan_application_sdk-2.1.1.dist-info → atlan_application_sdk-2.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atlan-application-sdk
-Version: 2.1.1
+Version: 2.2.0
 Summary: Atlan Application SDK is a Python library for developing applications on the Atlan Platform
 Project-URL: Repository, https://github.com/atlanhq/application-sdk
 Project-URL: Documentation, https://github.com/atlanhq/application-sdk/README.md
@@ -30,6 +30,10 @@ Requires-Dist: pyatlan<8.5.0,>=8.0.2
 Requires-Dist: pydantic<2.13.0,>=2.10.6
 Requires-Dist: python-dotenv<1.3.0,>=1.1.0
 Requires-Dist: uvloop<0.23.0,>=0.21.0; sys_platform != 'win32'
+Provides-Extra: azure
+Requires-Dist: azure-identity>=1.15.0; extra == 'azure'
+Requires-Dist: azure-storage-blob>=12.19.0; extra == 'azure'
+Requires-Dist: azure-storage-file-datalake>=12.19.0; extra == 'azure'
 Provides-Extra: daft
 Requires-Dist: daft<0.8.0,>=0.7.1; extra == 'daft'
 Provides-Extra: distributed-lock

{atlan_application_sdk-2.1.1.dist-info → atlan_application_sdk-2.2.0.dist-info}/RECORD

@@ -1,6 +1,6 @@
 application_sdk/__init__.py,sha256=2e2mvmLJ5dxmJGPELtb33xwP-j6JMdoIuqKycEn7hjg,151
 application_sdk/constants.py,sha256=TvdmKQShVWBNQZdVF2y-fxuE31FmeraTnqQ9jT_n5XY,11567
-application_sdk/version.py,sha256=sNbvXviG7NgxM58lOHKhbZfERat5qAJNr3UZy_toVQs,84
+application_sdk/version.py,sha256=q2hQsJUqlYmQIe9BcRqShBDGfZcI_1e30HROiiS-LdU,84
 application_sdk/worker.py,sha256=DLMocpHvvwpdAopyXhxwM7ftaNlKvZMQfkgy1MFyiik,7561
 application_sdk/activities/__init__.py,sha256=i7iY6aL1VFg185n2rLLvD_sI2BA9zJ33jL5rD_sY__U,12350
 application_sdk/activities/lock_management.py,sha256=6Wdf3jMKitoarHQP91PIJOoGFz4aaOLS_40c7n1yAOA,3902
@@ -28,9 +28,12 @@ application_sdk/clients/temporal.py,sha256=7ZkQSwSSZTFkBkhwFlqmLBFxlD2-jLS8QWpSi
 application_sdk/clients/utils.py,sha256=zLFOJbTr_6TOqnjfVFGY85OtIXZ4FQy_rquzjaydkbY,779
 application_sdk/clients/workflow.py,sha256=6bSqmA3sNCk9oY68dOjBUDZ9DhNKQxPD75qqE0cfldc,6104
 application_sdk/clients/.cursor/BUGBOT.md,sha256=7nEDUqWBEMI_uU6eK1jCSZGeXoQtLQcKwOrDn8AIDWo,10595
+application_sdk/clients/azure/__init__.py,sha256=koHM4dpsLDf6NQfBWh1axIwF8DGFePcc_nu8_4a-D1g,233
+application_sdk/clients/azure/auth.py,sha256=lvdyNt0V099XlSuTapprNx9xmtpr1El5I_2pKtEnDX0,11129
+application_sdk/clients/azure/client.py,sha256=7f2kanMWgF61xM3WXeFnsgrZaG3Pdl9a1Vt8818e8HQ,13538
 application_sdk/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 application_sdk/common/aws_utils.py,sha256=xlSMIQyjvQ-CydEXaxXrnPUygv7AAbCLsxhZ2wtKnzg,11219
-application_sdk/common/error_codes.py,sha256=bxgvugN_0H5b8VXfJw-44mybgX5I9lRJbRdYjtPjqDI,14561
+application_sdk/common/error_codes.py,sha256=BGdOSfhKe4pd9oM9T2N8GASZeIJU2Be1NcThL8LkKBQ,14900
 application_sdk/common/file_converter.py,sha256=ta0PVh7uIEGJg0BTPUJnSjj55ve2iVAOkqwAeg96_-g,3079
 application_sdk/common/types.py,sha256=qkVr3SAR1zn5_0w9hFt18vGtcnaPLKlMJLSBhTSKePU,134
 application_sdk/common/utils.py,sha256=czcWvqoe2PawDvHOahk_AI88Zqth-CM3KzdGmehHQJ4,19286
@@ -78,8 +81,8 @@ application_sdk/observability/utils.py,sha256=-02GAFom8Bg4SNyCTNYySmen2dzvLfTu43
 application_sdk/observability/decorators/observability_decorator.py,sha256=yd6qfrg1MmH5KcZ5Ydzb0RaBzmxx5FrmiI9qwvZx3EU,8963
 application_sdk/server/__init__.py,sha256=KTqE1YPw_3WDVMWatJUuf9OOiobLM2K5SMaBrI62sCo,1568
 application_sdk/server/.cursor/BUGBOT.md,sha256=p_MMoWUW5G1894WfOKYReZKWCuyJT_OJz3rL5g21NbI,16566
-application_sdk/server/fastapi/__init__.py,sha256=BVqf63z1hxEdpJqLU4LXpFTbk5q8dVkjEbWbu_vbW_Y,29578
-application_sdk/server/fastapi/models.py,sha256=h0hMtMg_p5G0Ug2MBkmBcT94W025VKCLMFyh0FciNoQ,7559
+application_sdk/server/fastapi/__init__.py,sha256=Ce1eXW9Y_8icuEL3RawYfoqcJbjpyfEK2nHq9y-o04k,30941
+application_sdk/server/fastapi/models.py,sha256=ZKz0RWISrJKQ50kepPiPsVqZ6li44I_hFIy5GXochLc,9796
 application_sdk/server/fastapi/utils.py,sha256=WoDGDmq6E1kwS2FN5pjIuzygKNTHpA-tg2SQJZgJOOI,1415
 application_sdk/server/fastapi/middleware/logmiddleware.py,sha256=sKKi-ysI9XQDT0uKW3sfw2r3XQq2iJT4_XrNGXZWdjI,2907
 application_sdk/server/fastapi/middleware/metrics.py,sha256=F_EhayzNgEn9KJvITv8VFxwXhBZVaOQkDptyfa9nQK4,1820
@@ -138,7 +141,7 @@ application_sdk/transformers/atlas/__init__.py,sha256=fw3D8bBtt61SseAfYut3JZddpX
 application_sdk/transformers/atlas/sql.py,sha256=rkQXNZ7oebts5oF5E_Bw8NpcHHKScU0TmKciH_1l_k4,50419
 application_sdk/transformers/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 application_sdk/transformers/common/utils.py,sha256=4ISMIQ0Gzghmi31p51FOFm5KLF7XF-fmH9PVT7i0DFE,4899
-application_sdk/transformers/query/__init__.py,sha256=4uVCU-NfDe08PlffjWQ5p4smQa7c518IL2rDgIk6694,17446
+application_sdk/transformers/query/__init__.py,sha256=asp3IC0zLjrwM8wyUxI8FSFeXmBT5IoFZ23IHn67Muw,18711
 application_sdk/transformers/query/templates/column.yaml,sha256=EXLYwGXN7LKT-v51n2EZnY99o6vHucyFaVSpM-sUSXw,7679
 application_sdk/transformers/query/templates/database.yaml,sha256=SD1hJg5LI7gsBHQL5mW341sa51EkhcsIDDFlIOi9zdk,1374
 application_sdk/transformers/query/templates/extras-procedure.yaml,sha256=XhAfVY4zm99K8fcgkYA1XPLv4ks-SA6SzMO3SMtQ60s,2298
@@ -152,8 +155,8 @@ application_sdk/workflows/metadata_extraction/__init__.py,sha256=jHUe_ZBQ66jx8bg
 application_sdk/workflows/metadata_extraction/sql.py,sha256=6ZaVt84n-8U2ZvR9GR7uIJKv5v8CuyQjhlnoRJvDszc,12435
 application_sdk/workflows/query_extraction/__init__.py,sha256=n066_CX5RpJz6DIxGMkKS3eGSRg03ilaCtsqfJWQb7Q,117
 application_sdk/workflows/query_extraction/sql.py,sha256=kT_JQkLCRZ44ZpaC4QvPL6DxnRIIVh8gYHLqRbMI-hA,4826
-atlan_application_sdk-2.1.1.dist-info/METADATA,sha256=Vc2uG2FMhuNXyZFXmGMmvc_LRpCBaNTcQEHpSV8NpOE,5805
-atlan_application_sdk-2.1.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-atlan_application_sdk-2.1.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-atlan_application_sdk-2.1.1.dist-info/licenses/NOTICE,sha256=A-XVVGt3KOYuuMmvSMIFkg534F1vHiCggEBp4Ez3wGk,1041
-atlan_application_sdk-2.1.1.dist-info/RECORD,,
+atlan_application_sdk-2.2.0.dist-info/METADATA,sha256=PiV9n04V1F27b28MfIUEsL62baglFRHbi8gxQ4yeG_k,6014
+atlan_application_sdk-2.2.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+atlan_application_sdk-2.2.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+atlan_application_sdk-2.2.0.dist-info/licenses/NOTICE,sha256=A-XVVGt3KOYuuMmvSMIFkg534F1vHiCggEBp4Ez3wGk,1041
+atlan_application_sdk-2.2.0.dist-info/RECORD,,