Add Azure Virtual Machines operators, sensor, and trigger (#49796)

docs/spelling_wordlist.txt

@@ -264,6 +264,7 @@ comparator
 compat
 Compute
 compute
+ComputeManagementClient
 Computenodes
 ComputeNodeState
 concat
@@ -406,6 +407,9 @@ dbt
 dbutils
 ddl
 de
+deallocate
+deallocated
+Debounce
 debuggability
 declaratively
 decommissioning

providers/microsoft/azure/docs/index.rst

@@ -128,6 +128,7 @@ PIP package                                 Version required
 ``azure-kusto-data``                        ``>=4.1.0,!=4.6.0,!=5.0.0``
 ``azure-mgmt-datafactory``                  ``>=2.0.0``
 ``azure-mgmt-containerregistry``            ``>=8.0.0``
+``azure-mgmt-compute``                      ``>=33.0.0``
 ``azure-mgmt-containerinstance``            ``>=10.1.0``
 ``msgraph-core``                            ``>=1.3.3``
 ``msgraphfs``                               ``>=0.3.0``

providers/microsoft/azure/docs/operators/compute.rst

@@ -0,0 +1,108 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+ ..   http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+
+Azure Virtual Machine Operators
+================================
+
+Waiting Strategy
+----------------
+The VM action operators support two patterns:
+
+* ``wait_for_completion=True`` (default): operator blocks until Azure operation finishes.
+* ``wait_for_completion=False``: operator submits the operation and returns quickly.
+
+When you want to reduce worker slot usage for long VM state transitions, use
+``wait_for_completion=False`` together with
+:class:`~airflow.providers.microsoft.azure.sensors.compute.AzureVirtualMachineStateSensor`
+in ``deferrable=True`` mode to move the waiting to the triggerer.
+
+.. _howto/operator:AzureVirtualMachineStartOperator:
+
+AzureVirtualMachineStartOperator
+---------------------------------
+Use the
+:class:`~airflow.providers.microsoft.azure.operators.compute.AzureVirtualMachineStartOperator`
+to start an Azure Virtual Machine.
+
+Below is an example of using this operator to start a VM:
+
+.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_compute.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_operator_azure_vm_start]
+    :end-before: [END howto_operator_azure_vm_start]
+
+
+.. _howto/operator:AzureVirtualMachineStopOperator:
+
+AzureVirtualMachineStopOperator
+--------------------------------
+Use the
+:class:`~airflow.providers.microsoft.azure.operators.compute.AzureVirtualMachineStopOperator`
+to stop (deallocate) an Azure Virtual Machine. This releases compute resources and stops billing.
+
+Below is an example of using this operator to stop a VM:
+
+.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_compute.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_operator_azure_vm_stop]
+    :end-before: [END howto_operator_azure_vm_stop]
+
+
+.. _howto/operator:AzureVirtualMachineRestartOperator:
+
+AzureVirtualMachineRestartOperator
+-----------------------------------
+Use the
+:class:`~airflow.providers.microsoft.azure.operators.compute.AzureVirtualMachineRestartOperator`
+to restart an Azure Virtual Machine.
+
+Below is an example of using this operator to restart a VM:
+
+.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_compute.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_operator_azure_vm_restart]
+    :end-before: [END howto_operator_azure_vm_restart]
+
+
+.. _howto/sensor:AzureVirtualMachineStateSensor:
+
+AzureVirtualMachineStateSensor
+-------------------------------
+Use the
+:class:`~airflow.providers.microsoft.azure.sensors.compute.AzureVirtualMachineStateSensor`
+to poll a VM until it reaches a target power state (e.g., ``running``, ``deallocated``).
+This sensor supports deferrable mode.
+
+Below is an example of using this sensor:
+
+.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_compute.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_sensor_azure_vm_state]
+    :end-before: [END howto_sensor_azure_vm_state]
+
+
+Reference
+---------
+
+For further information, look at:
+
+* `Azure Virtual Machines Documentation <https://azure.microsoft.com/en-us/products/virtual-machines/>`__

providers/microsoft/azure/provider.yaml

@@ -118,6 +118,12 @@ integrations:
       - /docs/apache-airflow-providers-microsoft-azure/operators/batch.rst
     logo: /docs/integration-logos/Microsoft-Azure-Batch.png
     tags: [azure]
+  - integration-name: Microsoft Azure Compute
+    external-doc-url: https://azure.microsoft.com/en-us/products/virtual-machines/
+    how-to-guide:
+      - /docs/apache-airflow-providers-microsoft-azure/operators/compute.rst
+    logo: /docs/integration-logos/Microsoft-Azure.png
+    tags: [azure]
   - integration-name: Microsoft Azure Blob Storage
     external-doc-url: https://azure.microsoft.com/en-us/services/storage/blobs/
     how-to-guide:
@@ -191,6 +197,9 @@ integrations:
     tags: [azure]
 
 operators:
+  - integration-name: Microsoft Azure Compute
+    python-modules:
+      - airflow.providers.microsoft.azure.operators.compute
   - integration-name: Microsoft Azure Data Lake Storage
     python-modules:
       - airflow.providers.microsoft.azure.operators.adls
@@ -226,6 +235,9 @@ operators:
       - airflow.providers.microsoft.azure.operators.powerbi
 
 sensors:
+  - integration-name: Microsoft Azure Compute
+    python-modules:
+      - airflow.providers.microsoft.azure.sensors.compute
   - integration-name: Microsoft Azure Cosmos DB
     python-modules:
       - airflow.providers.microsoft.azure.sensors.cosmos
@@ -244,6 +256,9 @@ filesystems:
   - airflow.providers.microsoft.azure.fs.msgraph
 
 hooks:
+  - integration-name: Microsoft Azure Compute
+    python-modules:
+      - airflow.providers.microsoft.azure.hooks.compute
   - integration-name: Microsoft Azure Container Instances
     python-modules:
       - airflow.providers.microsoft.azure.hooks.container_volume
@@ -290,6 +305,9 @@ hooks:
       - airflow.providers.microsoft.azure.hooks.powerbi
 
 triggers:
+  - integration-name: Microsoft Azure Compute
+    python-modules:
+      - airflow.providers.microsoft.azure.triggers.compute
   - integration-name: Microsoft Azure Data Factory
     python-modules:
       - airflow.providers.microsoft.azure.triggers.data_factory
@@ -362,6 +380,36 @@ connection-types:
         label: Workload Identity Tenant ID
         schema:
           type: ["string", "null"]
+  - hook-class-name: airflow.providers.microsoft.azure.hooks.compute.AzureComputeHook
+    connection-type: azure_compute
+    ui-field-behaviour:
+      hidden-fields: ["schema", "port", "host"]
+      relabeling:
+        login: Client ID
+        password: Client Secret
+      placeholders:
+        extra: '{"key_path": "path to json file for auth", "key_json": "specifies json dict for auth"}'
+        login: client_id (token credentials auth)
+        password: secret (token credentials auth)
+        tenantId: tenantId (token credentials auth)
+        subscriptionId: subscriptionId (token credentials auth)
+    conn-fields:
+      tenantId:
+        label: Azure Tenant ID
+        schema:
+          type: ["string", "null"]
+      subscriptionId:
+        label: Azure Subscription ID
+        schema:
+          type: ["string", "null"]
+      managed_identity_client_id:
+        label: Managed Identity Client ID
+        schema:
+          type: ["string", "null"]
+      workload_identity_tenant_id:
+        label: Workload Identity Tenant ID
+        schema:
+          type: ["string", "null"]
   - hook-class-name: airflow.providers.microsoft.azure.hooks.adx.AzureDataExplorerHook
     connection-type: azure_data_explorer
     ui-field-behaviour:

providers/microsoft/azure/pyproject.toml

@@ -81,6 +81,7 @@ dependencies = [
     "azure-kusto-data>=4.1.0,!=4.6.0,!=5.0.0",
     "azure-mgmt-datafactory>=2.0.0",
     "azure-mgmt-containerregistry>=8.0.0",
+    "azure-mgmt-compute>=33.0.0",
     "azure-mgmt-containerinstance>=10.1.0",
     "msgraph-core>=1.3.3",
     "msgraphfs>=0.3.0",

providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py

@@ -34,6 +34,13 @@ def get_provider_info():
                 "logo": "/docs/integration-logos/Microsoft-Azure-Batch.png",
                 "tags": ["azure"],
             },
+            {
+                "integration-name": "Microsoft Azure Compute",
+                "external-doc-url": "https://azure.microsoft.com/en-us/products/virtual-machines/",
+                "how-to-guide": ["/docs/apache-airflow-providers-microsoft-azure/operators/compute.rst"],
+                "logo": "/docs/integration-logos/Microsoft-Azure.png",
+                "tags": ["azure"],
+            },
             {
                 "integration-name": "Microsoft Azure Blob Storage",
                 "external-doc-url": "https://azure.microsoft.com/en-us/services/storage/blobs/",
@@ -135,6 +142,10 @@ def get_provider_info():
             },
         ],
         "operators": [
+            {
+                "integration-name": "Microsoft Azure Compute",
+                "python-modules": ["airflow.providers.microsoft.azure.operators.compute"],
+            },
             {
                 "integration-name": "Microsoft Azure Data Lake Storage",
                 "python-modules": ["airflow.providers.microsoft.azure.operators.adls"],
@@ -181,6 +192,10 @@ def get_provider_info():
             },
         ],
         "sensors": [
+            {
+                "integration-name": "Microsoft Azure Compute",
+                "python-modules": ["airflow.providers.microsoft.azure.sensors.compute"],
+            },
             {
                 "integration-name": "Microsoft Azure Cosmos DB",
                 "python-modules": ["airflow.providers.microsoft.azure.sensors.cosmos"],
@@ -203,6 +218,10 @@ def get_provider_info():
             "airflow.providers.microsoft.azure.fs.msgraph",
         ],
         "hooks": [
+            {
+                "integration-name": "Microsoft Azure Compute",
+                "python-modules": ["airflow.providers.microsoft.azure.hooks.compute"],
+            },
             {
                 "integration-name": "Microsoft Azure Container Instances",
                 "python-modules": [
@@ -265,6 +284,10 @@ def get_provider_info():
             },
         ],
         "triggers": [
+            {
+                "integration-name": "Microsoft Azure Compute",
+                "python-modules": ["airflow.providers.microsoft.azure.triggers.compute"],
+            },
             {
                 "integration-name": "Microsoft Azure Data Factory",
                 "python-modules": ["airflow.providers.microsoft.azure.triggers.data_factory"],
@@ -349,6 +372,36 @@ def get_provider_info():
                     },
                 },
             },
+            {
+                "hook-class-name": "airflow.providers.microsoft.azure.hooks.compute.AzureComputeHook",
+                "connection-type": "azure_compute",
+                "ui-field-behaviour": {
+                    "hidden-fields": ["schema", "port", "host"],
+                    "relabeling": {"login": "Client ID", "password": "Client Secret"},
+                    "placeholders": {
+                        "extra": '{"key_path": "path to json file for auth", "key_json": "specifies json dict for auth"}',
+                        "login": "client_id (token credentials auth)",
+                        "password": "secret (token credentials auth)",
+                        "tenantId": "tenantId (token credentials auth)",
+                        "subscriptionId": "subscriptionId (token credentials auth)",
+                    },
+                },
+                "conn-fields": {
+                    "tenantId": {"label": "Azure Tenant ID", "schema": {"type": ["string", "null"]}},
+                    "subscriptionId": {
+                        "label": "Azure Subscription ID",
+                        "schema": {"type": ["string", "null"]},
+                    },
+                    "managed_identity_client_id": {
+                        "label": "Managed Identity Client ID",
+                        "schema": {"type": ["string", "null"]},
+                    },
+                    "workload_identity_tenant_id": {
+                        "label": "Workload Identity Tenant ID",
+                        "schema": {"type": ["string", "null"]},
+                    },
+                },
+            },
             {
                 "hook-class-name": "airflow.providers.microsoft.azure.hooks.adx.AzureDataExplorerHook",
                 "connection-type": "azure_data_explorer",