Skip to content

Commit 0b7ebb7

Browse files
nailo2cnailo2c
authored
Add Azure Virtual Machines operators, sensor, and trigger (#49796) (#62391)
fix CI MyPy checks error Replace AirflowException with standard Python exceptions per contribution guidelines Add Azure Compute connection fields and UI behavior for authentication Refactor Azure Virtual Machine operators to use a base class and improve connection handling Add ComputeManagementClient to spelling wordlist Re-trigger CI
1 parent a6ff00f commit 0b7ebb7

File tree

15 files changed

+1552
-0
lines changed

15 files changed

+1552
-0
lines changed

‎docs/spelling_wordlist.txt‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -265,6 +265,7 @@ comparator
265265
compat
266266
Compute
267267
compute
268+
ComputeManagementClient
268269
Computenodes
269270
ComputeNodeState
270271
concat
@@ -407,6 +408,9 @@ dbt
407408
dbutils
408409
ddl
409410
de
411+
deallocate
412+
deallocated
413+
Debounce
410414
debuggability
411415
declaratively
412416
decommissioning

‎providers/microsoft/azure/docs/index.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,7 @@ PIP package Version required
128128
``azure-kusto-data`` ``>=4.1.0,!=4.6.0,!=5.0.0``
129129
``azure-mgmt-datafactory`` ``>=2.0.0``
130130
``azure-mgmt-containerregistry`` ``>=8.0.0``
131+
``azure-mgmt-compute`` ``>=33.0.0``
131132
``azure-mgmt-containerinstance`` ``>=10.1.0``
132133
``msgraph-core`` ``>=1.3.3``
133134
``msgraphfs`` ``>=0.3.0``

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

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
.. Licensed to the Apache Software Foundation (ASF) under one
2+
or more contributor license agreements. See the NOTICE file
3+
distributed with this work for additional information
4+
regarding copyright ownership. The ASF licenses this file
5+
to you under the Apache License, Version 2.0 (the
6+
"License"); you may not use this file except in compliance
7+
with the License. You may obtain a copy of the License at
8+
9+
.. http://www.apache.org/licenses/LICENSE-2.0
10+
11+
.. Unless required by applicable law or agreed to in writing,
12+
software distributed under the License is distributed on an
13+
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
KIND, either express or implied. See the License for the
15+
specific language governing permissions and limitations
16+
under the License.
17+
18+
19+
Azure Virtual Machine Operators
20+
================================
21+
22+
Waiting Strategy
23+
----------------
24+
The VM action operators support two patterns:
25+
26+
* ``wait_for_completion=True`` (default): operator blocks until Azure operation finishes.
27+
* ``wait_for_completion=False``: operator submits the operation and returns quickly.
28+
29+
When you want to reduce worker slot usage for long VM state transitions, use
30+
``wait_for_completion=False`` together with
31+
:class:`~airflow.providers.microsoft.azure.sensors.compute.AzureVirtualMachineStateSensor`
32+
in ``deferrable=True`` mode to move the waiting to the triggerer.
33+
34+
.. _howto/operator:AzureVirtualMachineStartOperator:
35+
36+
AzureVirtualMachineStartOperator
37+
---------------------------------
38+
Use the
39+
:class:`~airflow.providers.microsoft.azure.operators.compute.AzureVirtualMachineStartOperator`
40+
to start an Azure Virtual Machine.
41+
42+
Below is an example of using this operator to start a VM:
43+
44+
.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_compute.py
45+
:language: python
46+
:dedent: 4
47+
:start-after: [START howto_operator_azure_vm_start]
48+
:end-before: [END howto_operator_azure_vm_start]
49+
50+
51+
.. _howto/operator:AzureVirtualMachineStopOperator:
52+
53+
AzureVirtualMachineStopOperator
54+
--------------------------------
55+
Use the
56+
:class:`~airflow.providers.microsoft.azure.operators.compute.AzureVirtualMachineStopOperator`
57+
to stop (deallocate) an Azure Virtual Machine. This releases compute resources and stops billing.
58+
59+
Below is an example of using this operator to stop a VM:
60+
61+
.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_compute.py
62+
:language: python
63+
:dedent: 4
64+
:start-after: [START howto_operator_azure_vm_stop]
65+
:end-before: [END howto_operator_azure_vm_stop]
66+
67+
68+
.. _howto/operator:AzureVirtualMachineRestartOperator:
69+
70+
AzureVirtualMachineRestartOperator
71+
-----------------------------------
72+
Use the
73+
:class:`~airflow.providers.microsoft.azure.operators.compute.AzureVirtualMachineRestartOperator`
74+
to restart an Azure Virtual Machine.
75+
76+
Below is an example of using this operator to restart a VM:
77+
78+
.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_compute.py
79+
:language: python
80+
:dedent: 4
81+
:start-after: [START howto_operator_azure_vm_restart]
82+
:end-before: [END howto_operator_azure_vm_restart]
83+
84+
85+
.. _howto/sensor:AzureVirtualMachineStateSensor:
86+
87+
AzureVirtualMachineStateSensor
88+
-------------------------------
89+
Use the
90+
:class:`~airflow.providers.microsoft.azure.sensors.compute.AzureVirtualMachineStateSensor`
91+
to poll a VM until it reaches a target power state (e.g., ``running``, ``deallocated``).
92+
This sensor supports deferrable mode.
93+
94+
Below is an example of using this sensor:
95+
96+
.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_compute.py
97+
:language: python
98+
:dedent: 4
99+
:start-after: [START howto_sensor_azure_vm_state]
100+
:end-before: [END howto_sensor_azure_vm_state]
101+
102+
103+
Reference
104+
---------
105+
106+
For further information, look at:
107+
108+
* `Azure Virtual Machines Documentation <https://azure.microsoft.com/en-us/products/virtual-machines/>`__

‎providers/microsoft/azure/provider.yaml‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,12 @@ integrations:
118118
- /docs/apache-airflow-providers-microsoft-azure/operators/batch.rst
119119
logo: /docs/integration-logos/Microsoft-Azure-Batch.png
120120
tags: [azure]
121+
- integration-name: Microsoft Azure Compute
122+
external-doc-url: https://azure.microsoft.com/en-us/products/virtual-machines/
123+
how-to-guide:
124+
- /docs/apache-airflow-providers-microsoft-azure/operators/compute.rst
125+
logo: /docs/integration-logos/Microsoft-Azure.png
126+
tags: [azure]
121127
- integration-name: Microsoft Azure Blob Storage
122128
external-doc-url: https://azure.microsoft.com/en-us/services/storage/blobs/
123129
how-to-guide:
@@ -191,6 +197,9 @@ integrations:
191197
tags: [azure]
192198

193199
operators:
200+
- integration-name: Microsoft Azure Compute
201+
python-modules:
202+
- airflow.providers.microsoft.azure.operators.compute
194203
- integration-name: Microsoft Azure Data Lake Storage
195204
python-modules:
196205
- airflow.providers.microsoft.azure.operators.adls
@@ -226,6 +235,9 @@ operators:
226235
- airflow.providers.microsoft.azure.operators.powerbi
227236

228237
sensors:
238+
- integration-name: Microsoft Azure Compute
239+
python-modules:
240+
- airflow.providers.microsoft.azure.sensors.compute
229241
- integration-name: Microsoft Azure Cosmos DB
230242
python-modules:
231243
- airflow.providers.microsoft.azure.sensors.cosmos
@@ -244,6 +256,9 @@ filesystems:
244256
- airflow.providers.microsoft.azure.fs.msgraph
245257

246258
hooks:
259+
- integration-name: Microsoft Azure Compute
260+
python-modules:
261+
- airflow.providers.microsoft.azure.hooks.compute
247262
- integration-name: Microsoft Azure Container Instances
248263
python-modules:
249264
- airflow.providers.microsoft.azure.hooks.container_volume
@@ -290,6 +305,9 @@ hooks:
290305
- airflow.providers.microsoft.azure.hooks.powerbi
291306

292307
triggers:
308+
- integration-name: Microsoft Azure Compute
309+
python-modules:
310+
- airflow.providers.microsoft.azure.triggers.compute
293311
- integration-name: Microsoft Azure Data Factory
294312
python-modules:
295313
- airflow.providers.microsoft.azure.triggers.data_factory
@@ -362,6 +380,36 @@ connection-types:
362380
label: Workload Identity Tenant ID
363381
schema:
364382
type: ["string", "null"]
383+
- hook-class-name: airflow.providers.microsoft.azure.hooks.compute.AzureComputeHook
384+
connection-type: azure_compute
385+
ui-field-behaviour:
386+
hidden-fields: ["schema", "port", "host"]
387+
relabeling:
388+
login: Client ID
389+
password: Client Secret
390+
placeholders:
391+
extra: '{"key_path": "path to json file for auth", "key_json": "specifies json dict for auth"}'
392+
login: client_id (token credentials auth)
393+
password: secret (token credentials auth)
394+
tenantId: tenantId (token credentials auth)
395+
subscriptionId: subscriptionId (token credentials auth)
396+
conn-fields:
397+
tenantId:
398+
label: Azure Tenant ID
399+
schema:
400+
type: ["string", "null"]
401+
subscriptionId:
402+
label: Azure Subscription ID
403+
schema:
404+
type: ["string", "null"]
405+
managed_identity_client_id:
406+
label: Managed Identity Client ID
407+
schema:
408+
type: ["string", "null"]
409+
workload_identity_tenant_id:
410+
label: Workload Identity Tenant ID
411+
schema:
412+
type: ["string", "null"]
365413
- hook-class-name: airflow.providers.microsoft.azure.hooks.adx.AzureDataExplorerHook
366414
connection-type: azure_data_explorer
367415
ui-field-behaviour:

‎providers/microsoft/azure/pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ dependencies = [
8181
"azure-kusto-data>=4.1.0,!=4.6.0,!=5.0.0",
8282
"azure-mgmt-datafactory>=2.0.0",
8383
"azure-mgmt-containerregistry>=8.0.0",
84+
"azure-mgmt-compute>=33.0.0",
8485
"azure-mgmt-containerinstance>=10.1.0",
8586
"msgraph-core>=1.3.3",
8687
"msgraphfs>=0.3.0",

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

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,13 @@ def get_provider_info():
3434
"logo": "/docs/integration-logos/Microsoft-Azure-Batch.png",
3535
"tags": ["azure"],
3636
},
37+
{
38+
"integration-name": "Microsoft Azure Compute",
39+
"external-doc-url": "https://azure.microsoft.com/en-us/products/virtual-machines/",
40+
"how-to-guide": ["/docs/apache-airflow-providers-microsoft-azure/operators/compute.rst"],
41+
"logo": "/docs/integration-logos/Microsoft-Azure.png",
42+
"tags": ["azure"],
43+
},
3744
{
3845
"integration-name": "Microsoft Azure Blob Storage",
3946
"external-doc-url": "https://azure.microsoft.com/en-us/services/storage/blobs/",
@@ -135,6 +142,10 @@ def get_provider_info():
135142
},
136143
],
137144
"operators": [
145+
{
146+
"integration-name": "Microsoft Azure Compute",
147+
"python-modules": ["airflow.providers.microsoft.azure.operators.compute"],
148+
},
138149
{
139150
"integration-name": "Microsoft Azure Data Lake Storage",
140151
"python-modules": ["airflow.providers.microsoft.azure.operators.adls"],
@@ -181,6 +192,10 @@ def get_provider_info():
181192
},
182193
],
183194
"sensors": [
195+
{
196+
"integration-name": "Microsoft Azure Compute",
197+
"python-modules": ["airflow.providers.microsoft.azure.sensors.compute"],
198+
},
184199
{
185200
"integration-name": "Microsoft Azure Cosmos DB",
186201
"python-modules": ["airflow.providers.microsoft.azure.sensors.cosmos"],
@@ -203,6 +218,10 @@ def get_provider_info():
203218
"airflow.providers.microsoft.azure.fs.msgraph",
204219
],
205220
"hooks": [
221+
{
222+
"integration-name": "Microsoft Azure Compute",
223+
"python-modules": ["airflow.providers.microsoft.azure.hooks.compute"],
224+
},
206225
{
207226
"integration-name": "Microsoft Azure Container Instances",
208227
"python-modules": [
@@ -265,6 +284,10 @@ def get_provider_info():
265284
},
266285
],
267286
"triggers": [
287+
{
288+
"integration-name": "Microsoft Azure Compute",
289+
"python-modules": ["airflow.providers.microsoft.azure.triggers.compute"],
290+
},
268291
{
269292
"integration-name": "Microsoft Azure Data Factory",
270293
"python-modules": ["airflow.providers.microsoft.azure.triggers.data_factory"],
@@ -349,6 +372,36 @@ def get_provider_info():
349372
},
350373
},
351374
},
375+
{
376+
"hook-class-name": "airflow.providers.microsoft.azure.hooks.compute.AzureComputeHook",
377+
"connection-type": "azure_compute",
378+
"ui-field-behaviour": {
379+
"hidden-fields": ["schema", "port", "host"],
380+
"relabeling": {"login": "Client ID", "password": "Client Secret"},
381+
"placeholders": {
382+
"extra": '{"key_path": "path to json file for auth", "key_json": "specifies json dict for auth"}',
383+
"login": "client_id (token credentials auth)",
384+
"password": "secret (token credentials auth)",
385+
"tenantId": "tenantId (token credentials auth)",
386+
"subscriptionId": "subscriptionId (token credentials auth)",
387+
},
388+
},
389+
"conn-fields": {
390+
"tenantId": {"label": "Azure Tenant ID", "schema": {"type": ["string", "null"]}},
391+
"subscriptionId": {
392+
"label": "Azure Subscription ID",
393+
"schema": {"type": ["string", "null"]},
394+
},
395+
"managed_identity_client_id": {
396+
"label": "Managed Identity Client ID",
397+
"schema": {"type": ["string", "null"]},
398+
},
399+
"workload_identity_tenant_id": {
400+
"label": "Workload Identity Tenant ID",
401+
"schema": {"type": ["string", "null"]},
402+
},
403+
},
404+
},
352405
{
353406
"hook-class-name": "airflow.providers.microsoft.azure.hooks.adx.AzureDataExplorerHook",
354407
"connection-type": "azure_data_explorer",

0 commit comments

Comments
 (0)