apache-airflow-providers-microsoft-azure 12.0.0rc1__py3-none-any.whl → 12.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (37) hide show
  1. airflow/providers/microsoft/azure/LICENSE +0 -52
  2. airflow/providers/microsoft/azure/__init__.py +1 -1
  3. airflow/providers/microsoft/azure/fs/adls.py +1 -2
  4. airflow/providers/microsoft/azure/get_provider_info.py +52 -46
  5. airflow/providers/microsoft/azure/hooks/adx.py +6 -7
  6. airflow/providers/microsoft/azure/hooks/asb.py +237 -8
  7. airflow/providers/microsoft/azure/hooks/base_azure.py +2 -3
  8. airflow/providers/microsoft/azure/hooks/batch.py +1 -2
  9. airflow/providers/microsoft/azure/hooks/container_instance.py +3 -4
  10. airflow/providers/microsoft/azure/hooks/container_registry.py +2 -3
  11. airflow/providers/microsoft/azure/hooks/container_volume.py +2 -3
  12. airflow/providers/microsoft/azure/hooks/cosmos.py +4 -5
  13. airflow/providers/microsoft/azure/hooks/data_factory.py +7 -7
  14. airflow/providers/microsoft/azure/hooks/data_lake.py +8 -9
  15. airflow/providers/microsoft/azure/hooks/fileshare.py +1 -2
  16. airflow/providers/microsoft/azure/hooks/msgraph.py +102 -35
  17. airflow/providers/microsoft/azure/hooks/synapse.py +4 -5
  18. airflow/providers/microsoft/azure/hooks/wasb.py +9 -9
  19. airflow/providers/microsoft/azure/log/wasb_task_handler.py +1 -2
  20. airflow/providers/microsoft/azure/operators/adx.py +1 -2
  21. airflow/providers/microsoft/azure/operators/asb.py +50 -62
  22. airflow/providers/microsoft/azure/operators/batch.py +1 -2
  23. airflow/providers/microsoft/azure/operators/container_instances.py +7 -7
  24. airflow/providers/microsoft/azure/operators/msgraph.py +44 -12
  25. airflow/providers/microsoft/azure/operators/powerbi.py +34 -5
  26. airflow/providers/microsoft/azure/operators/synapse.py +1 -2
  27. airflow/providers/microsoft/azure/secrets/key_vault.py +3 -4
  28. airflow/providers/microsoft/azure/sensors/msgraph.py +21 -5
  29. airflow/providers/microsoft/azure/triggers/data_factory.py +1 -2
  30. airflow/providers/microsoft/azure/triggers/msgraph.py +4 -0
  31. airflow/providers/microsoft/azure/triggers/powerbi.py +55 -11
  32. airflow/providers/microsoft/azure/utils.py +2 -1
  33. {apache_airflow_providers_microsoft_azure-12.0.0rc1.dist-info → apache_airflow_providers_microsoft_azure-12.1.0.dist-info}/METADATA +21 -38
  34. apache_airflow_providers_microsoft_azure-12.1.0.dist-info/RECORD +58 -0
  35. apache_airflow_providers_microsoft_azure-12.0.0rc1.dist-info/RECORD +0 -58
  36. {apache_airflow_providers_microsoft_azure-12.0.0rc1.dist-info → apache_airflow_providers_microsoft_azure-12.1.0.dist-info}/WHEEL +0 -0
  37. {apache_airflow_providers_microsoft_azure-12.0.0rc1.dist-info → apache_airflow_providers_microsoft_azure-12.1.0.dist-info}/entry_points.txt +0 -0
@@ -23,6 +23,13 @@ from collections import namedtuple
23
23
  from collections.abc import Sequence
24
24
  from typing import TYPE_CHECKING, Any
25
25
 
26
+ from msrestazure.azure_exceptions import CloudError
27
+
28
+ from airflow.exceptions import AirflowException, AirflowTaskTimeout
29
+ from airflow.models import BaseOperator
30
+ from airflow.providers.microsoft.azure.hooks.container_instance import AzureContainerInstanceHook
31
+ from airflow.providers.microsoft.azure.hooks.container_registry import AzureContainerRegistryHook
32
+ from airflow.providers.microsoft.azure.hooks.container_volume import AzureContainerVolumeHook
26
33
  from azure.mgmt.containerinstance.models import (
27
34
  Container,
28
35
  ContainerGroup,
@@ -37,13 +44,6 @@ from azure.mgmt.containerinstance.models import (
37
44
  Volume as _AzureVolume,
38
45
  VolumeMount,
39
46
  )
40
- from msrestazure.azure_exceptions import CloudError
41
-
42
- from airflow.exceptions import AirflowException, AirflowTaskTimeout
43
- from airflow.models import BaseOperator
44
- from airflow.providers.microsoft.azure.hooks.container_instance import AzureContainerInstanceHook
45
- from airflow.providers.microsoft.azure.hooks.container_registry import AzureContainerRegistryHook
46
- from airflow.providers.microsoft.azure.hooks.container_volume import AzureContainerVolumeHook
47
47
 
48
48
  if TYPE_CHECKING:
49
49
  from airflow.utils.context import Context
@@ -17,6 +17,7 @@
17
17
  # under the License.
18
18
  from __future__ import annotations
19
19
 
20
+ import warnings
20
21
  from collections.abc import Sequence
21
22
  from copy import deepcopy
22
23
  from typing import (
@@ -25,7 +26,7 @@ from typing import (
25
26
  Callable,
26
27
  )
27
28
 
28
- from airflow.exceptions import AirflowException, TaskDeferred
29
+ from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning, TaskDeferred
29
30
  from airflow.models import BaseOperator
30
31
  from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook
31
32
  from airflow.providers.microsoft.azure.triggers.msgraph import (
@@ -44,7 +45,7 @@ if TYPE_CHECKING:
44
45
  from airflow.utils.context import Context
45
46
 
46
47
 
47
- def default_event_handler(context: Context, event: dict[Any, Any] | None = None) -> Any:
48
+ def default_event_handler(event: dict[Any, Any] | None = None, **context) -> Any:
48
49
  if event:
49
50
  if event.get("status") == "failure":
50
51
  raise AirflowException(event.get("message"))
@@ -52,6 +53,23 @@ def default_event_handler(context: Context, event: dict[Any, Any] | None = None)
52
53
  return event.get("response")
53
54
 
54
55
 
56
+ def execute_callable(
57
+ func: Callable[[dict[Any, Any] | None, Context], Any] | Callable[[dict[Any, Any] | None, Any], Any],
58
+ value: Any,
59
+ context: Context,
60
+ message: str,
61
+ ) -> Any:
62
+ try:
63
+ return func(value, **context) # type: ignore
64
+ except TypeError:
65
+ warnings.warn(
66
+ message,
67
+ AirflowProviderDeprecationWarning,
68
+ stacklevel=2,
69
+ )
70
+ return func(context, value) # type: ignore
71
+
72
+
55
73
  class MSGraphAsyncOperator(BaseOperator):
56
74
  """
57
75
  A Microsoft Graph API operator which allows you to execute REST call to the Microsoft Graph API.
@@ -71,11 +89,12 @@ class MSGraphAsyncOperator(BaseOperator):
71
89
  :param timeout: The HTTP timeout being used by the `KiotaRequestAdapter` (default is None).
72
90
  When no timeout is specified or set to None then there is no HTTP timeout on each request.
73
91
  :param proxies: A dict defining the HTTP proxies to be used (default is None).
92
+ :param scopes: The scopes to be used (default is ["https://graph.microsoft.com/.default"]).
74
93
  :param api_version: The API version of the Microsoft Graph API to be used (default is v1).
75
94
  You can pass an enum named APIVersion which has 2 possible members v1 and beta,
76
95
  or you can pass a string as `v1.0` or `beta`.
77
96
  :param result_processor: Function to further process the response from MS Graph API
78
- (default is lambda: context, response: response). When the response returned by the
97
+ (default is lambda: response, context: response). When the response returned by the
79
98
  `KiotaRequestAdapterHook` are bytes, then those will be base64 encoded into a string.
80
99
  :param event_handler: Function to process the event returned from `MSGraphTrigger`. By default, when the
81
100
  event returned by the `MSGraphTrigger` has a failed status, an AirflowException is being raised with
@@ -110,10 +129,11 @@ class MSGraphAsyncOperator(BaseOperator):
110
129
  key: str = XCOM_RETURN_KEY,
111
130
  timeout: float | None = None,
112
131
  proxies: dict | None = None,
132
+ scopes: str | list[str] | None = None,
113
133
  api_version: APIVersion | str | None = None,
114
134
  pagination_function: Callable[[MSGraphAsyncOperator, dict, Context], tuple[str, dict]] | None = None,
115
- result_processor: Callable[[Context, Any], Any] = lambda context, result: result,
116
- event_handler: Callable[[Context, dict[Any, Any] | None], Any] | None = None,
135
+ result_processor: Callable[[Any, Context], Any] = lambda result, **context: result,
136
+ event_handler: Callable[[dict[Any, Any] | None, Context], Any] | None = None,
117
137
  serializer: type[ResponseSerializer] = ResponseSerializer,
118
138
  **kwargs: Any,
119
139
  ):
@@ -130,6 +150,7 @@ class MSGraphAsyncOperator(BaseOperator):
130
150
  self.key = key
131
151
  self.timeout = timeout
132
152
  self.proxies = proxies
153
+ self.scopes = scopes
133
154
  self.api_version = api_version
134
155
  self.pagination_function = pagination_function or self.paginate
135
156
  self.result_processor = result_processor
@@ -150,6 +171,7 @@ class MSGraphAsyncOperator(BaseOperator):
150
171
  conn_id=self.conn_id,
151
172
  timeout=self.timeout,
152
173
  proxies=self.proxies,
174
+ scopes=self.scopes,
153
175
  api_version=self.api_version,
154
176
  serializer=type(self.serializer),
155
177
  ),
@@ -171,7 +193,12 @@ class MSGraphAsyncOperator(BaseOperator):
171
193
  if event:
172
194
  self.log.debug("%s completed with %s: %s", self.task_id, event.get("status"), event)
173
195
 
174
- response = self.event_handler(context, event)
196
+ response = execute_callable(
197
+ self.event_handler, # type: ignore
198
+ event,
199
+ context,
200
+ "event_handler signature has changed, event parameter should be defined before context!",
201
+ )
175
202
 
176
203
  self.log.debug("response: %s", response)
177
204
 
@@ -182,7 +209,12 @@ class MSGraphAsyncOperator(BaseOperator):
182
209
 
183
210
  self.log.debug("deserialize response: %s", response)
184
211
 
185
- result = self.result_processor(context, response)
212
+ result = execute_callable(
213
+ self.result_processor,
214
+ response,
215
+ context,
216
+ "result_processor signature has changed, result parameter should be defined before context!",
217
+ )
186
218
 
187
219
  self.log.debug("processed response: %s", result)
188
220
 
@@ -230,7 +262,7 @@ class MSGraphAsyncOperator(BaseOperator):
230
262
  return result
231
263
  return results
232
264
 
233
- def pull_xcom(self, context: Context) -> list:
265
+ def pull_xcom(self, context: Context | dict[str, Any]) -> list:
234
266
  map_index = context["ti"].map_index
235
267
  value = list(
236
268
  context["ti"].xcom_pull(
@@ -262,7 +294,7 @@ class MSGraphAsyncOperator(BaseOperator):
262
294
 
263
295
  return value
264
296
 
265
- def push_xcom(self, context: Context, value) -> None:
297
+ def push_xcom(self, context: Any, value) -> None:
266
298
  self.log.debug("do_xcom_push: %s", self.do_xcom_push)
267
299
  if self.do_xcom_push:
268
300
  self.log.info("Pushing XCom with key '%s': %s", self.key, value)
@@ -270,7 +302,7 @@ class MSGraphAsyncOperator(BaseOperator):
270
302
 
271
303
  @staticmethod
272
304
  def paginate(
273
- operator: MSGraphAsyncOperator, response: dict, context: Context
305
+ operator: MSGraphAsyncOperator, response: dict, **context
274
306
  ) -> tuple[Any, dict[str, Any] | None]:
275
307
  odata_count = response.get("@odata.count")
276
308
  if odata_count and operator.query_parameters:
@@ -279,7 +311,7 @@ class MSGraphAsyncOperator(BaseOperator):
279
311
 
280
312
  if top and odata_count:
281
313
  if len(response.get("value", [])) == top and context:
282
- results = operator.pull_xcom(context=context)
314
+ results = operator.pull_xcom(context)
283
315
  skip = sum([len(result["value"]) for result in results]) + top if results else top # type: ignore
284
316
  query_parameters["$skip"] = skip
285
317
  return operator.url, query_parameters
@@ -287,7 +319,7 @@ class MSGraphAsyncOperator(BaseOperator):
287
319
 
288
320
  def trigger_next_link(self, response, method_name: str, context: Context) -> None:
289
321
  if isinstance(response, dict):
290
- url, query_parameters = self.pagination_function(self, response, context)
322
+ url, query_parameters = self.pagination_function(self, response, **dict(context.items())) # type: ignore
291
323
 
292
324
  self.log.debug("url: %s", url)
293
325
  self.log.debug("query_parameters: %s", query_parameters)
@@ -114,6 +114,35 @@ class PowerBIDatasetRefreshOperator(BaseOperator):
114
114
  check_interval=self.check_interval,
115
115
  wait_for_termination=self.wait_for_termination,
116
116
  ),
117
+ method_name=self.get_refresh_status.__name__,
118
+ )
119
+
120
+ def get_refresh_status(self, context: Context, event: dict[str, str] | None = None):
121
+ """Push the refresh Id to XCom then runs the Trigger to wait for refresh completion."""
122
+ if event:
123
+ if event["status"] == "error":
124
+ raise AirflowException(event["message"])
125
+
126
+ dataset_refresh_id = event["dataset_refresh_id"]
127
+
128
+ if dataset_refresh_id:
129
+ self.xcom_push(
130
+ context=context,
131
+ key=f"{self.task_id}.powerbi_dataset_refresh_Id",
132
+ value=dataset_refresh_id,
133
+ )
134
+ self.defer(
135
+ trigger=PowerBITrigger(
136
+ conn_id=self.conn_id,
137
+ group_id=self.group_id,
138
+ dataset_id=self.dataset_id,
139
+ dataset_refresh_id=dataset_refresh_id,
140
+ timeout=self.timeout,
141
+ proxies=self.proxies,
142
+ api_version=self.api_version,
143
+ check_interval=self.check_interval,
144
+ wait_for_termination=self.wait_for_termination,
145
+ ),
117
146
  method_name=self.execute_complete.__name__,
118
147
  )
119
148
 
@@ -124,10 +153,10 @@ class PowerBIDatasetRefreshOperator(BaseOperator):
124
153
  Relies on trigger to throw an exception, otherwise it assumes execution was successful.
125
154
  """
126
155
  if event:
127
- if event["status"] == "error":
128
- raise AirflowException(event["message"])
129
-
130
156
  self.xcom_push(
131
- context=context, key="powerbi_dataset_refresh_Id", value=event["dataset_refresh_id"]
157
+ context=context,
158
+ key=f"{self.task_id}.powerbi_dataset_refresh_status",
159
+ value=event["dataset_refresh_status"],
132
160
  )
133
- self.xcom_push(context=context, key="powerbi_dataset_refresh_status", value=event["status"])
161
+ if event["status"] == "error":
162
+ raise AirflowException(event["message"])
@@ -33,10 +33,9 @@ from airflow.providers.microsoft.azure.hooks.synapse import (
33
33
  )
34
34
 
35
35
  if TYPE_CHECKING:
36
- from azure.synapse.spark.models import SparkBatchJobOptions
37
-
38
36
  from airflow.models.taskinstancekey import TaskInstanceKey
39
37
  from airflow.utils.context import Context
38
+ from azure.synapse.spark.models import SparkBatchJobOptions
40
39
 
41
40
 
42
41
  class AzureSynapseRunSparkBatchOperator(BaseOperator):
@@ -28,13 +28,12 @@ import logging
28
28
  import os
29
29
  from functools import cached_property
30
30
 
31
- from azure.core.exceptions import ResourceNotFoundError
32
- from azure.identity import ClientSecretCredential, DefaultAzureCredential
33
- from azure.keyvault.secrets import SecretClient
34
-
35
31
  from airflow.providers.microsoft.azure.utils import get_sync_default_azure_credential
36
32
  from airflow.secrets import BaseSecretsBackend
37
33
  from airflow.utils.log.logging_mixin import LoggingMixin
34
+ from azure.core.exceptions import ResourceNotFoundError
35
+ from azure.identity import ClientSecretCredential, DefaultAzureCredential
36
+ from azure.keyvault.secrets import SecretClient
38
37
 
39
38
 
40
39
  class AzureKeyVaultBackend(BaseSecretsBackend, LoggingMixin):
@@ -23,6 +23,7 @@ from typing import TYPE_CHECKING, Any, Callable
23
23
  from airflow.exceptions import AirflowException
24
24
  from airflow.providers.common.compat.standard.triggers import TimeDeltaTrigger
25
25
  from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook
26
+ from airflow.providers.microsoft.azure.operators.msgraph import execute_callable
26
27
  from airflow.providers.microsoft.azure.triggers.msgraph import MSGraphTrigger, ResponseSerializer
27
28
  from airflow.sensors.base import BaseSensorOperator
28
29
 
@@ -47,6 +48,7 @@ class MSGraphSensor(BaseSensorOperator):
47
48
  :param method: The HTTP method being used to do the REST call (default is GET).
48
49
  :param conn_id: The HTTP Connection ID to run the operator against (templated).
49
50
  :param proxies: A dict defining the HTTP proxies to be used (default is None).
51
+ :param scopes: The scopes to be used (default is ["https://graph.microsoft.com/.default"]).
50
52
  :param api_version: The API version of the Microsoft Graph API to be used (default is v1).
51
53
  You can pass an enum named APIVersion which has 2 possible members v1 and beta,
52
54
  or you can pass a string as `v1.0` or `beta`.
@@ -54,7 +56,7 @@ class MSGraphSensor(BaseSensorOperator):
54
56
  `default_event_processor` method) and returns a boolean. When the result is True, the sensor
55
57
  will stop poking, otherwise it will continue until it's True or times out.
56
58
  :param result_processor: Function to further process the response from MS Graph API
57
- (default is lambda: context, response: response). When the response returned by the
59
+ (default is lambda: response, context: response). When the response returned by the
58
60
  `KiotaRequestAdapterHook` are bytes, then those will be base64 encoded into a string.
59
61
  :param serializer: Class which handles response serialization (default is ResponseSerializer).
60
62
  Bytes will be base64 encoded into a string, so it can be stored as an XCom.
@@ -83,9 +85,10 @@ class MSGraphSensor(BaseSensorOperator):
83
85
  data: dict[str, Any] | str | BytesIO | None = None,
84
86
  conn_id: str = KiotaRequestAdapterHook.default_conn_name,
85
87
  proxies: dict | None = None,
88
+ scopes: str | list[str] | None = None,
86
89
  api_version: APIVersion | str | None = None,
87
- event_processor: Callable[[Context, Any], bool] = lambda context, e: e.get("status") == "Succeeded",
88
- result_processor: Callable[[Context, Any], Any] = lambda context, result: result,
90
+ event_processor: Callable[[Any, Context], bool] = lambda e, **context: e.get("status") == "Succeeded",
91
+ result_processor: Callable[[Any, Context], Any] = lambda result, **context: result,
89
92
  serializer: type[ResponseSerializer] = ResponseSerializer,
90
93
  retry_delay: timedelta | float = 60,
91
94
  **kwargs,
@@ -101,6 +104,7 @@ class MSGraphSensor(BaseSensorOperator):
101
104
  self.data = data
102
105
  self.conn_id = conn_id
103
106
  self.proxies = proxies
107
+ self.scopes = scopes
104
108
  self.api_version = api_version
105
109
  self.event_processor = event_processor
106
110
  self.result_processor = result_processor
@@ -120,6 +124,7 @@ class MSGraphSensor(BaseSensorOperator):
120
124
  conn_id=self.conn_id,
121
125
  timeout=self.timeout,
122
126
  proxies=self.proxies,
127
+ scopes=self.scopes,
123
128
  api_version=self.api_version,
124
129
  serializer=type(self.serializer),
125
130
  ),
@@ -129,6 +134,7 @@ class MSGraphSensor(BaseSensorOperator):
129
134
  def retry_execute(
130
135
  self,
131
136
  context: Context,
137
+ **kwargs,
132
138
  ) -> Any:
133
139
  self.execute(context=context)
134
140
 
@@ -159,12 +165,22 @@ class MSGraphSensor(BaseSensorOperator):
159
165
 
160
166
  self.log.debug("deserialize response: %s", response)
161
167
 
162
- is_done = self.event_processor(context, response)
168
+ is_done = execute_callable(
169
+ self.event_processor,
170
+ response,
171
+ context,
172
+ "event_processor signature has changed, event parameter should be defined before context!",
173
+ )
163
174
 
164
175
  self.log.debug("is_done: %s", is_done)
165
176
 
166
177
  if is_done:
167
- result = self.result_processor(context, response)
178
+ result = execute_callable(
179
+ self.result_processor,
180
+ response,
181
+ context,
182
+ "result_processor signature has changed, result parameter should be defined before context!",
183
+ )
168
184
 
169
185
  self.log.debug("processed response: %s", result)
170
186
 
@@ -21,13 +21,12 @@ import time
21
21
  from collections.abc import AsyncIterator
22
22
  from typing import Any
23
23
 
24
- from azure.core.exceptions import ServiceRequestError
25
-
26
24
  from airflow.providers.microsoft.azure.hooks.data_factory import (
27
25
  AzureDataFactoryAsyncHook,
28
26
  AzureDataFactoryPipelineRunStatus,
29
27
  )
30
28
  from airflow.triggers.base import BaseTrigger, TriggerEvent
29
+ from azure.core.exceptions import ServiceRequestError
31
30
 
32
31
 
33
32
  class ADFPipelineRunStatusSensorTrigger(BaseTrigger):
@@ -90,6 +90,7 @@ class MSGraphTrigger(BaseTrigger):
90
90
  :param timeout: The HTTP timeout being used by the `KiotaRequestAdapter` (default is None).
91
91
  When no timeout is specified or set to None then there is no HTTP timeout on each request.
92
92
  :param proxies: A dict defining the HTTP proxies to be used (default is None).
93
+ :param scopes: The scopes to be used (default is ["https://graph.microsoft.com/.default"]).
93
94
  :param api_version: The API version of the Microsoft Graph API to be used (default is v1).
94
95
  You can pass an enum named APIVersion which has 2 possible members v1 and beta,
95
96
  or you can pass a string as `v1.0` or `beta`.
@@ -121,6 +122,7 @@ class MSGraphTrigger(BaseTrigger):
121
122
  conn_id: str = KiotaRequestAdapterHook.default_conn_name,
122
123
  timeout: float | None = None,
123
124
  proxies: dict | None = None,
125
+ scopes: str | list[str] | None = None,
124
126
  api_version: APIVersion | str | None = None,
125
127
  serializer: type[ResponseSerializer] = ResponseSerializer,
126
128
  ):
@@ -129,6 +131,7 @@ class MSGraphTrigger(BaseTrigger):
129
131
  conn_id=conn_id,
130
132
  timeout=timeout,
131
133
  proxies=proxies,
134
+ scopes=scopes,
132
135
  api_version=api_version,
133
136
  )
134
137
  self.url = url
@@ -157,6 +160,7 @@ class MSGraphTrigger(BaseTrigger):
157
160
  "conn_id": self.conn_id,
158
161
  "timeout": self.timeout,
159
162
  "proxies": self.proxies,
163
+ "scopes": self.hook.scopes,
160
164
  "api_version": self.api_version,
161
165
  "serializer": f"{self.serializer.__class__.__module__}.{self.serializer.__class__.__name__}",
162
166
  "url": self.url,
@@ -22,7 +22,13 @@ import time
22
22
  from collections.abc import AsyncIterator
23
23
  from typing import TYPE_CHECKING
24
24
 
25
- from airflow.providers.microsoft.azure.hooks.powerbi import PowerBIDatasetRefreshStatus, PowerBIHook
25
+ import tenacity
26
+
27
+ from airflow.providers.microsoft.azure.hooks.powerbi import (
28
+ PowerBIDatasetRefreshException,
29
+ PowerBIDatasetRefreshStatus,
30
+ PowerBIHook,
31
+ )
26
32
  from airflow.triggers.base import BaseTrigger, TriggerEvent
27
33
 
28
34
  if TYPE_CHECKING:
@@ -43,6 +49,7 @@ class PowerBITrigger(BaseTrigger):
43
49
  You can pass an enum named APIVersion which has 2 possible members v1 and beta,
44
50
  or you can pass a string as `v1.0` or `beta`.
45
51
  :param dataset_id: The dataset Id to refresh.
52
+ :param dataset_refresh_id: The dataset refresh Id to poll for the status, if not provided a new refresh will be triggered.
46
53
  :param group_id: The workspace Id where dataset is located.
47
54
  :param end_time: Time in seconds when trigger should stop polling.
48
55
  :param check_interval: Time in seconds to wait between each poll.
@@ -55,6 +62,7 @@ class PowerBITrigger(BaseTrigger):
55
62
  dataset_id: str,
56
63
  group_id: str,
57
64
  timeout: float = 60 * 60 * 24 * 7,
65
+ dataset_refresh_id: str | None = None,
58
66
  proxies: dict | None = None,
59
67
  api_version: APIVersion | str | None = None,
60
68
  check_interval: int = 60,
@@ -63,6 +71,7 @@ class PowerBITrigger(BaseTrigger):
63
71
  super().__init__()
64
72
  self.hook = PowerBIHook(conn_id=conn_id, proxies=proxies, api_version=api_version, timeout=timeout)
65
73
  self.dataset_id = dataset_id
74
+ self.dataset_refresh_id = dataset_refresh_id
66
75
  self.timeout = timeout
67
76
  self.group_id = group_id
68
77
  self.check_interval = check_interval
@@ -77,6 +86,7 @@ class PowerBITrigger(BaseTrigger):
77
86
  "proxies": self.proxies,
78
87
  "api_version": self.api_version,
79
88
  "dataset_id": self.dataset_id,
89
+ "dataset_refresh_id": self.dataset_refresh_id,
80
90
  "group_id": self.group_id,
81
91
  "timeout": self.timeout,
82
92
  "check_interval": self.check_interval,
@@ -98,19 +108,53 @@ class PowerBITrigger(BaseTrigger):
98
108
 
99
109
  async def run(self) -> AsyncIterator[TriggerEvent]:
100
110
  """Make async connection to the PowerBI and polls for the dataset refresh status."""
101
- self.dataset_refresh_id = await self.hook.trigger_dataset_refresh(
102
- dataset_id=self.dataset_id,
103
- group_id=self.group_id,
104
- )
105
-
106
- async def fetch_refresh_status_and_error() -> tuple[str, str]:
107
- """Fetch the current status and error of the dataset refresh."""
108
- refresh_details = await self.hook.get_refresh_details_by_refresh_id(
111
+ if not self.dataset_refresh_id:
112
+ # Trigger the dataset refresh
113
+ dataset_refresh_id = await self.hook.trigger_dataset_refresh(
109
114
  dataset_id=self.dataset_id,
110
115
  group_id=self.group_id,
111
- refresh_id=self.dataset_refresh_id,
112
116
  )
113
- return refresh_details["status"], refresh_details["error"]
117
+
118
+ if dataset_refresh_id:
119
+ self.log.info("Triggered dataset refresh %s", dataset_refresh_id)
120
+ yield TriggerEvent(
121
+ {
122
+ "status": "success",
123
+ "dataset_refresh_status": None,
124
+ "message": f"The dataset refresh {dataset_refresh_id} has been triggered.",
125
+ "dataset_refresh_id": dataset_refresh_id,
126
+ }
127
+ )
128
+ return
129
+
130
+ yield TriggerEvent(
131
+ {
132
+ "status": "error",
133
+ "dataset_refresh_status": None,
134
+ "message": "Failed to trigger the dataset refresh.",
135
+ "dataset_refresh_id": None,
136
+ }
137
+ )
138
+ return
139
+
140
+ # The dataset refresh is already triggered. Poll for the dataset refresh status.
141
+ @tenacity.retry(
142
+ stop=tenacity.stop_after_attempt(3),
143
+ wait=tenacity.wait_exponential(min=5, multiplier=2),
144
+ reraise=True,
145
+ retry=tenacity.retry_if_exception_type(PowerBIDatasetRefreshException),
146
+ )
147
+ async def fetch_refresh_status_and_error() -> tuple[str, str]:
148
+ """Fetch the current status and error of the dataset refresh."""
149
+ if self.dataset_refresh_id:
150
+ refresh_details = await self.hook.get_refresh_details_by_refresh_id(
151
+ dataset_id=self.dataset_id,
152
+ group_id=self.group_id,
153
+ refresh_id=self.dataset_refresh_id,
154
+ )
155
+ return refresh_details["status"], refresh_details["error"]
156
+
157
+ raise PowerBIDatasetRefreshException("Dataset refresh Id is missing.")
114
158
 
115
159
  try:
116
160
  dataset_refresh_status, dataset_refresh_error = await fetch_refresh_status_and_error()
@@ -21,12 +21,13 @@ import warnings
21
21
  from functools import partial, wraps
22
22
  from urllib.parse import urlparse, urlunparse
23
23
 
24
+ from msrest.authentication import BasicTokenAuthentication
25
+
24
26
  from azure.core.pipeline import PipelineContext, PipelineRequest
25
27
  from azure.core.pipeline.policies import BearerTokenCredentialPolicy
26
28
  from azure.core.pipeline.transport import HttpRequest
27
29
  from azure.identity import DefaultAzureCredential
28
30
  from azure.identity.aio import DefaultAzureCredential as AsyncDefaultAzureCredential
29
- from msrest.authentication import BasicTokenAuthentication
30
31
 
31
32
 
32
33
  def get_field(*, conn_id: str, conn_type: str, extras: dict, field_name: str):
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: apache-airflow-providers-microsoft-azure
3
- Version: 12.0.0rc1
3
+ Version: 12.1.0
4
4
  Summary: Provider package apache-airflow-providers-microsoft-azure for Apache Airflow
5
5
  Keywords: airflow-provider,microsoft.azure,airflow,integration
6
6
  Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -20,40 +20,40 @@ Classifier: Programming Language :: Python :: 3.10
20
20
  Classifier: Programming Language :: Python :: 3.11
21
21
  Classifier: Programming Language :: Python :: 3.12
22
22
  Classifier: Topic :: System :: Monitoring
23
- Requires-Dist: adal>=1.2.7
23
+ Requires-Dist: apache-airflow>=2.9.0
24
24
  Requires-Dist: adlfs>=2023.10.0
25
- Requires-Dist: apache-airflow>=2.9.0rc0
26
25
  Requires-Dist: azure-batch>=8.0.0
27
26
  Requires-Dist: azure-cosmos>=4.6.0
27
+ Requires-Dist: azure-mgmt-cosmosdb>=3.0.0
28
28
  Requires-Dist: azure-datalake-store>=0.0.45
29
29
  Requires-Dist: azure-identity>=1.3.1
30
30
  Requires-Dist: azure-keyvault-secrets>=4.1.0
31
- Requires-Dist: azure-kusto-data>=4.1.0,!=4.6.0
32
- Requires-Dist: azure-mgmt-containerinstance>=10.1.0
33
- Requires-Dist: azure-mgmt-containerregistry>=8.0.0
34
- Requires-Dist: azure-mgmt-cosmosdb>=3.0.0
35
- Requires-Dist: azure-mgmt-datafactory>=2.0.0
36
31
  Requires-Dist: azure-mgmt-datalake-store>=0.5.0
37
32
  Requires-Dist: azure-mgmt-resource>=2.2.0
38
- Requires-Dist: azure-mgmt-storage>=16.0.0
39
- Requires-Dist: azure-servicebus>=7.12.1
40
33
  Requires-Dist: azure-storage-blob>=12.14.0
41
- Requires-Dist: azure-storage-file-datalake>=12.9.1
34
+ Requires-Dist: azure-mgmt-storage>=16.0.0
42
35
  Requires-Dist: azure-storage-file-share>=12.7.0
43
- Requires-Dist: azure-synapse-artifacts>=0.17.0
36
+ Requires-Dist: azure-servicebus>=7.12.1
44
37
  Requires-Dist: azure-synapse-spark>=0.2.0
45
- Requires-Dist: microsoft-kiota-abstractions<1.4.0
38
+ Requires-Dist: azure-synapse-artifacts>=0.17.0
39
+ Requires-Dist: adal>=1.2.7
40
+ Requires-Dist: azure-storage-file-datalake>=12.9.1
41
+ Requires-Dist: azure-kusto-data>=4.1.0,!=4.6.0
42
+ Requires-Dist: azure-mgmt-datafactory>=2.0.0
43
+ Requires-Dist: azure-mgmt-containerregistry>=8.0.0
44
+ Requires-Dist: azure-mgmt-containerinstance>=10.1.0
45
+ Requires-Dist: msgraph-core>=1.0.0,!=1.1.8
46
46
  Requires-Dist: microsoft-kiota-http>=1.3.0,!=1.3.4
47
47
  Requires-Dist: microsoft-kiota-serialization-json==1.0.0
48
48
  Requires-Dist: microsoft-kiota-serialization-text==1.0.0
49
- Requires-Dist: msgraph-core>=1.0.0,!=1.1.8
49
+ Requires-Dist: microsoft-kiota-abstractions<1.4.0
50
50
  Requires-Dist: apache-airflow-providers-amazon ; extra == "amazon"
51
51
  Requires-Dist: apache-airflow-providers-common-compat ; extra == "common-compat"
52
52
  Requires-Dist: apache-airflow-providers-oracle ; extra == "oracle"
53
53
  Requires-Dist: apache-airflow-providers-sftp ; extra == "sftp"
54
54
  Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
55
- Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/12.0.0/changelog.html
56
- Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/12.0.0
55
+ Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/12.1.0/changelog.html
56
+ Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/12.1.0
57
57
  Project-URL: Slack Chat, https://s.apache.org/airflow-slack
58
58
  Project-URL: Source Code, https://github.com/apache/airflow
59
59
  Project-URL: Twitter, https://x.com/ApacheAirflow
@@ -64,23 +64,6 @@ Provides-Extra: oracle
64
64
  Provides-Extra: sftp
65
65
 
66
66
 
67
- .. Licensed to the Apache Software Foundation (ASF) under one
68
- or more contributor license agreements. See the NOTICE file
69
- distributed with this work for additional information
70
- regarding copyright ownership. The ASF licenses this file
71
- to you under the Apache License, Version 2.0 (the
72
- "License"); you may not use this file except in compliance
73
- with the License. You may obtain a copy of the License at
74
-
75
- .. http://www.apache.org/licenses/LICENSE-2.0
76
-
77
- .. Unless required by applicable law or agreed to in writing,
78
- software distributed under the License is distributed on an
79
- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
80
- KIND, either express or implied. See the License for the
81
- specific language governing permissions and limitations
82
- under the License.
83
-
84
67
  .. Licensed to the Apache Software Foundation (ASF) under one
85
68
  or more contributor license agreements. See the NOTICE file
86
69
  distributed with this work for additional information
@@ -98,8 +81,7 @@ Provides-Extra: sftp
98
81
  specific language governing permissions and limitations
99
82
  under the License.
100
83
 
101
- .. NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE
102
- OVERWRITTEN WHEN PREPARING PACKAGES.
84
+ .. NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE OVERWRITTEN!
103
85
 
104
86
  .. IF YOU WANT TO MODIFY TEMPLATE FOR THIS FILE, YOU SHOULD MODIFY THE TEMPLATE
105
87
  `PROVIDER_README_TEMPLATE.rst.jinja2` IN the `dev/breeze/src/airflow_breeze/templates` DIRECTORY
@@ -107,7 +89,7 @@ Provides-Extra: sftp
107
89
 
108
90
  Package ``apache-airflow-providers-microsoft-azure``
109
91
 
110
- Release: ``12.0.0.rc1``
92
+ Release: ``12.1.0``
111
93
 
112
94
 
113
95
  `Microsoft Azure <https://azure.microsoft.com/>`__
@@ -120,7 +102,7 @@ This is a provider package for ``microsoft.azure`` provider. All classes for thi
120
102
  are in ``airflow.providers.microsoft.azure`` python package.
121
103
 
122
104
  You can find package information and changelog for the provider
123
- in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/12.0.0/>`_.
105
+ in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/12.1.0/>`_.
124
106
 
125
107
  Installation
126
108
  ------------
@@ -189,4 +171,5 @@ Dependent package
189
171
  ================================================================================================================== =================
190
172
 
191
173
  The changelog for the provider package can be found in the
192
- `changelog <https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/12.0.0/changelog.html>`_.
174
+ `changelog <https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure/12.1.0/changelog.html>`_.
175
+