Apache Airflow Microsoft Azure Provider

from airflow.providers.microsoft.azure.hooks.base_azure import AzureBaseHook

# Azure Blob Storage
from airflow.providers.microsoft.azure.hooks.wasb import WasbHook
from airflow.providers.microsoft.azure.operators.wasb_delete_blob import WasbDeleteBlobOperator
from airflow.providers.microsoft.azure.sensors.wasb import WasbBlobSensor

# Azure Data Factory
from airflow.providers.microsoft.azure.hooks.data_factory import AzureDataFactoryHook
from airflow.providers.microsoft.azure.operators.data_factory import AzureDataFactoryRunPipelineOperator

# Azure Cosmos DB
from airflow.providers.microsoft.azure.hooks.cosmos import AzureCosmosDBHook
from airflow.providers.microsoft.azure.operators.cosmos import AzureCosmosInsertDocumentOperator

from airflow import DAG
from airflow.providers.microsoft.azure.operators.wasb_delete_blob import WasbDeleteBlobOperator
from airflow.providers.microsoft.azure.sensors.wasb import WasbBlobSensor
from datetime import datetime, timedelta

# Define DAG
dag = DAG(
    'azure_workflow_example',
    default_args={
        'owner': 'data-team',
        'retries': 1,
        'retry_delay': timedelta(minutes=5)
    },
    description='Example Azure workflow',
    schedule_interval=timedelta(days=1),
    start_date=datetime(2024, 1, 1),
    catchup=False
)

# Wait for a blob to exist
wait_for_file = WasbBlobSensor(
    task_id='wait_for_input_file',
    container_name='input-data',
    blob_name='daily_data.csv',
    azure_conn_id='azure_default',
    dag=dag
)

# Delete processed blob
cleanup_blob = WasbDeleteBlobOperator(
    task_id='cleanup_processed_file',
    container_name='processed-data',
    blob_name='processed_data.csv',
    azure_conn_id='azure_default',
    dag=dag
)

wait_for_file >> cleanup_blob

class WasbHook(AzureBaseHook):
    """Hook for Azure Blob Storage operations."""
    
    def get_conn(self) -> BlobServiceClient: ...
    def check_for_blob(self, container_name: str, blob_name: str) -> bool: ...
    def load_file(self, file_path: str, container_name: str, blob_name: str) -> None: ...
    def load_string(self, string_data: str, container_name: str, blob_name: str) -> None: ...
    def read_file(self, container_name: str, blob_name: str) -> bytes: ...
    def delete_file(self, container_name: str, blob_name: str) -> None: ...

class AzureDataFactoryHook(BaseHook):
    """Hook for Azure Data Factory operations."""
    
    def get_conn(self) -> DataFactoryManagementClient: ...
    def run_pipeline(self, pipeline_name: str, resource_group_name: str, factory_name: str, **config: Any) -> CreateRunResponse: ...
    def get_pipeline_run(self, run_id: str, resource_group_name: str, factory_name: str) -> PipelineRun: ...
    def cancel_pipeline_run(self, run_id: str, resource_group_name: str, factory_name: str) -> None: ...

class AzureCosmosDBHook(BaseHook):
    """Hook for Azure Cosmos DB operations."""
    
    def get_conn(self) -> CosmosClient: ...
    def create_database(self, database_name: str) -> None: ...
    def create_collection(self, collection_name: str, database_name: str) -> None: ...
    def upsert_document(self, document: dict, database_name: str, collection_name: str) -> dict: ...
    def get_document(self, document_id: str, database_name: str, collection_name: str) -> dict: ...

class AzureDataLakeHook(BaseHook):
    """Hook for Azure Data Lake Storage Gen1."""
    
    def get_conn(self) -> core.AzureDLFileSystem: ...
    def upload_file(self, local_path: str, remote_path: str, overwrite: bool = True) -> None: ...
    def download_file(self, local_path: str, remote_path: str, overwrite: bool = True) -> None: ...

class AzureDataLakeStorageV2Hook(BaseHook):
    """Hook for Azure Data Lake Storage Gen2."""
    
    def get_conn(self) -> DataLakeServiceClient: ...
    def create_file_system(self, file_system_name: str) -> None: ...
    def upload_file(self, file_system_name: str, file_name: str, file_path: str) -> DataLakeFileClient: ...

class AdminClientHook(BaseAzureServiceBusHook):
    """Hook for Azure Service Bus administrative operations."""
    
    def create_queue(self, queue_name: str, **kwargs: Any) -> None: ...
    def create_topic(self, topic_name: str, **kwargs: Any) -> None: ...
    def create_subscription(self, topic_name: str, subscription_name: str, **kwargs: Any) -> None: ...

class MessageHook(BaseAzureServiceBusHook):
    """Hook for Azure Service Bus message operations."""
    
    def send_message(self, queue_name: str, message: str | ServiceBusMessage, **kwargs: Any) -> None: ...
    def receive_message(self, queue_name: str, **kwargs: Any) -> list[ServiceBusReceivedMessage]: ...

class AzureContainerInstanceHook(BaseHook):
    """Hook for Azure Container Instances management."""
    
    def get_conn(self) -> ContainerInstanceManagementClient: ...
    def create_or_update(self, resource_group_name: str, container_group_name: str, container_group: ContainerGroup) -> ContainerGroup: ...
    def get_logs(self, resource_group_name: str, container_group_name: str, container_name: str) -> str: ...

class AzureSynapseHook(BaseAzureSynapseHook):
    """Hook for Azure Synapse Spark operations."""
    
    def get_conn(self) -> SparkClient: ...
    def run_spark_job(self, payload: dict) -> dict: ...
    def get_job_run_status(self, job_id: int) -> str: ...

class AzureSynapsePipelineHook(BaseAzureSynapseHook):
    """Hook for Azure Synapse Pipeline operations."""
    
    def run_pipeline(self, pipeline_name: str, **config: Any) -> CreateRunResponse: ...
    def get_pipeline_run_status(self, run_id: str) -> str: ...

class KiotaRequestAdapterHook(BaseHook):
    """Hook for Microsoft Graph API using Kiota request adapter."""
    
    def get_conn(self) -> RequestAdapter: ...
    def test_connection(self) -> tuple[bool, str]: ...

class LocalFilesystemToWasbOperator(BaseOperator):
    """Transfer files from local filesystem to Azure Blob Storage."""
    
    def __init__(self, file_path: str, container_name: str, blob_name: str, **kwargs): ...

class S3ToAzureBlobStorageOperator(BaseOperator):
    """Transfer objects from AWS S3 to Azure Blob Storage."""
    
    def __init__(self, s3_source_key: str, container_name: str, blob_name: str, **kwargs): ...

class AzureDataExplorerHook(BaseHook):
    """Hook for Azure Data Explorer (Kusto) operations."""
    
    def get_conn(self) -> KustoClient: ...
    def run_query(self, query: str, database: str, options: dict | None = None) -> KustoResponseDataSet: ...

class AzureDataExplorerQueryOperator(BaseOperator):
    """Operator for querying Azure Data Explorer (Kusto)."""
    
    def __init__(self, *, query: str, database: str, **kwargs): ...

class PowerBIHook(KiotaRequestAdapterHook):
    """Hook for Power BI operations via Microsoft Graph API."""
    
    async def trigger_dataset_refresh(self, dataset_id: str, group_id: str, **kwargs) -> str: ...
    async def get_refresh_details_by_refresh_id(self, dataset_id: str, group_id: str, dataset_refresh_id: str) -> dict: ...

class PowerBIDatasetRefreshOperator(BaseOperator):
    """Refreshes a Power BI dataset."""
    
    def __init__(self, *, dataset_id: str, group_id: str, **kwargs): ...

class AzureBatchHook(BaseHook):
    """Hook for Azure Batch APIs."""
    
    def get_conn(self) -> BatchServiceClient: ...
    def create_pool(self, pool: PoolAddParameter) -> None: ...
    def create_job(self, job: JobAddParameter) -> None: ...
    def add_single_task_to_job(self, job_id: str, task: TaskAddParameter) -> None: ...

class AzureBatchOperator(BaseOperator):
    """Executes a job on Azure Batch Service."""
    
    def __init__(self, *, batch_pool_id: str, batch_job_id: str, batch_task_command_line: str, **kwargs): ...

class AzureFileShareHook(BaseHook):
    """Hook for Azure File Share operations."""
    
    def create_share(self, share_name: str, **kwargs) -> bool: ...
    def create_directory(self, **kwargs) -> Any: ...
    def load_file(self, file_path: str, **kwargs) -> None: ...
    def get_file(self, file_path: str, **kwargs) -> None: ...

# Base Azure connection information
class AzureBaseHook(BaseHook):
    conn_name_attr: str = "azure_conn_id"
    default_conn_name: str = "azure_default"
    conn_type: str = "azure"

# Common authentication credentials
AzureCredentials = Union[
    ServicePrincipal,
    ManagedIdentity, 
    WorkloadIdentity,
    DefaultAzureCredential
]

# Pipeline run statuses
class AzureDataFactoryPipelineRunStatus:
    QUEUED: str = "Queued"
    IN_PROGRESS: str = "InProgress" 
    SUCCEEDED: str = "Succeeded"
    FAILED: str = "Failed"
    CANCELLED: str = "Cancelled"

# Synapse job statuses  
class AzureSynapseSparkBatchRunStatus:
    NOT_STARTED: str = "not_started"
    STARTING: str = "starting"
    RUNNING: str = "running"
    IDLE: str = "idle"
    BUSY: str = "busy"
    SHUTTING_DOWN: str = "shutting_down"
    ERROR: str = "error"
    DEAD: str = "dead"
    KILLED: str = "killed"
    SUCCESS: str = "success"