apache-airflow-providers-standard 1.2.0__py3-none-any.whl → 1.3.0rc1__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 (30) hide show
  1. airflow/providers/standard/__init__.py +1 -1
  2. airflow/providers/standard/example_dags/__init__.py +16 -0
  3. airflow/providers/standard/example_dags/example_bash_decorator.py +114 -0
  4. airflow/providers/standard/example_dags/example_bash_operator.py +74 -0
  5. airflow/providers/standard/example_dags/example_branch_datetime_operator.py +105 -0
  6. airflow/providers/standard/example_dags/example_branch_day_of_week_operator.py +61 -0
  7. airflow/providers/standard/example_dags/example_branch_operator.py +166 -0
  8. airflow/providers/standard/example_dags/example_branch_operator_decorator.py +142 -0
  9. airflow/providers/standard/example_dags/example_external_task_child_deferrable.py +34 -0
  10. airflow/providers/standard/example_dags/example_external_task_marker_dag.py +98 -0
  11. airflow/providers/standard/example_dags/example_external_task_parent_deferrable.py +64 -0
  12. airflow/providers/standard/example_dags/example_latest_only.py +40 -0
  13. airflow/providers/standard/example_dags/example_python_decorator.py +132 -0
  14. airflow/providers/standard/example_dags/example_python_operator.py +147 -0
  15. airflow/providers/standard/example_dags/example_sensor_decorator.py +66 -0
  16. airflow/providers/standard/example_dags/example_sensors.py +135 -0
  17. airflow/providers/standard/example_dags/example_short_circuit_decorator.py +60 -0
  18. airflow/providers/standard/example_dags/example_short_circuit_operator.py +66 -0
  19. airflow/providers/standard/example_dags/example_trigger_controller_dag.py +46 -0
  20. airflow/providers/standard/example_dags/sql/__init__.py +16 -0
  21. airflow/providers/standard/example_dags/sql/sample.sql +24 -0
  22. airflow/providers/standard/operators/python.py +15 -9
  23. airflow/providers/standard/sensors/date_time.py +10 -4
  24. airflow/providers/standard/sensors/external_task.py +7 -6
  25. airflow/providers/standard/sensors/time.py +52 -40
  26. airflow/providers/standard/sensors/time_delta.py +47 -20
  27. {apache_airflow_providers_standard-1.2.0.dist-info → apache_airflow_providers_standard-1.3.0rc1.dist-info}/METADATA +7 -7
  28. {apache_airflow_providers_standard-1.2.0.dist-info → apache_airflow_providers_standard-1.3.0rc1.dist-info}/RECORD +30 -10
  29. {apache_airflow_providers_standard-1.2.0.dist-info → apache_airflow_providers_standard-1.3.0rc1.dist-info}/WHEEL +0 -0
  30. {apache_airflow_providers_standard-1.2.0.dist-info → apache_airflow_providers_standard-1.3.0rc1.dist-info}/entry_points.txt +0 -0
@@ -0,0 +1,135 @@
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
+ from __future__ import annotations
19
+
20
+ import datetime
21
+
22
+ import pendulum
23
+
24
+ from airflow.providers.standard.operators.bash import BashOperator
25
+ from airflow.providers.standard.sensors.bash import BashSensor
26
+ from airflow.providers.standard.sensors.filesystem import FileSensor
27
+ from airflow.providers.standard.sensors.python import PythonSensor
28
+ from airflow.providers.standard.sensors.time import TimeSensor
29
+ from airflow.providers.standard.sensors.time_delta import TimeDeltaSensor
30
+ from airflow.providers.standard.sensors.weekday import DayOfWeekSensor
31
+ from airflow.providers.standard.utils.weekday import WeekDay
32
+ from airflow.sdk import DAG
33
+ from airflow.utils.trigger_rule import TriggerRule
34
+
35
+
36
+ # [START example_callables]
37
+ def success_callable():
38
+ return True
39
+
40
+
41
+ def failure_callable():
42
+ return False
43
+
44
+
45
+ # [END example_callables]
46
+
47
+
48
+ with DAG(
49
+ dag_id="example_sensors",
50
+ schedule=None,
51
+ start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
52
+ catchup=False,
53
+ tags=["example"],
54
+ ) as dag:
55
+ # [START example_time_delta_sensor]
56
+ t0 = TimeDeltaSensor(task_id="wait_some_seconds", delta=datetime.timedelta(seconds=2))
57
+ # [END example_time_delta_sensor]
58
+
59
+ # [START example_time_delta_sensor_async]
60
+ t0a = TimeDeltaSensor(
61
+ task_id="wait_some_seconds_async", delta=datetime.timedelta(seconds=2), deferrable=True
62
+ )
63
+ # [END example_time_delta_sensor_async]
64
+
65
+ # [START example_time_sensors]
66
+ t1 = TimeSensor(
67
+ task_id="fire_immediately", target_time=datetime.datetime.now(tz=datetime.timezone.utc).time()
68
+ )
69
+
70
+ t2 = TimeSensor(
71
+ task_id="timeout_after_second_date_in_the_future",
72
+ timeout=1,
73
+ soft_fail=True,
74
+ target_time=(datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(hours=1)).time(),
75
+ )
76
+
77
+ t1a = TimeSensor(
78
+ task_id="fire_immediately_async",
79
+ target_time=datetime.datetime.now(tz=datetime.timezone.utc).time(),
80
+ deferrable=True,
81
+ )
82
+
83
+ t2a = TimeSensor(
84
+ task_id="timeout_after_second_date_in_the_future_async",
85
+ timeout=1,
86
+ soft_fail=True,
87
+ target_time=(datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(hours=1)).time(),
88
+ deferrable=True,
89
+ )
90
+ # [END example_time_sensors]
91
+
92
+ # [START example_bash_sensors]
93
+ t3 = BashSensor(task_id="Sensor_succeeds", bash_command="exit 0")
94
+
95
+ t4 = BashSensor(task_id="Sensor_fails_after_3_seconds", timeout=3, soft_fail=True, bash_command="exit 1")
96
+ # [END example_bash_sensors]
97
+
98
+ t5 = BashOperator(task_id="remove_file", bash_command="rm -rf /tmp/temporary_file_for_testing")
99
+
100
+ # [START example_file_sensor]
101
+ t6 = FileSensor(task_id="wait_for_file", filepath="/tmp/temporary_file_for_testing")
102
+ # [END example_file_sensor]
103
+
104
+ # [START example_file_sensor_async]
105
+ t7 = FileSensor(
106
+ task_id="wait_for_file_async", filepath="/tmp/temporary_file_for_testing", deferrable=True
107
+ )
108
+ # [END example_file_sensor_async]
109
+
110
+ t8 = BashOperator(
111
+ task_id="create_file_after_3_seconds", bash_command="sleep 3; touch /tmp/temporary_file_for_testing"
112
+ )
113
+
114
+ # [START example_python_sensors]
115
+ t9 = PythonSensor(task_id="success_sensor_python", python_callable=success_callable)
116
+
117
+ t10 = PythonSensor(
118
+ task_id="failure_timeout_sensor_python", timeout=3, soft_fail=True, python_callable=failure_callable
119
+ )
120
+ # [END example_python_sensors]
121
+
122
+ # [START example_day_of_week_sensor]
123
+ t11 = DayOfWeekSensor(
124
+ task_id="week_day_sensor_failing_on_timeout", timeout=3, soft_fail=True, week_day=WeekDay.MONDAY
125
+ )
126
+ # [END example_day_of_week_sensor]
127
+
128
+ tx = BashOperator(task_id="print_date_in_bash", bash_command="date")
129
+
130
+ tx.trigger_rule = TriggerRule.NONE_FAILED
131
+ [t0, t0a, t1, t1a, t2, t2a, t3, t4] >> tx
132
+ t5 >> t6 >> t7 >> tx
133
+ t8 >> tx
134
+ [t9, t10] >> tx
135
+ t11 >> tx
@@ -0,0 +1,60 @@
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
+ """Example DAG demonstrating the usage of the `@task.short_circuit()` TaskFlow decorator."""
18
+
19
+ from __future__ import annotations
20
+
21
+ import pendulum
22
+
23
+ from airflow.providers.standard.operators.empty import EmptyOperator
24
+ from airflow.sdk import chain, dag, task
25
+ from airflow.utils.trigger_rule import TriggerRule
26
+
27
+
28
+ @dag(schedule=None, start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), catchup=False, tags=["example"])
29
+ def example_short_circuit_decorator():
30
+ # [START howto_operator_short_circuit]
31
+ @task.short_circuit()
32
+ def check_condition(condition):
33
+ return condition
34
+
35
+ ds_true = [EmptyOperator(task_id=f"true_{i}") for i in [1, 2]]
36
+ ds_false = [EmptyOperator(task_id=f"false_{i}") for i in [1, 2]]
37
+
38
+ condition_is_true = check_condition.override(task_id="condition_is_true")(condition=True)
39
+ condition_is_false = check_condition.override(task_id="condition_is_false")(condition=False)
40
+
41
+ chain(condition_is_true, *ds_true)
42
+ chain(condition_is_false, *ds_false)
43
+ # [END howto_operator_short_circuit]
44
+
45
+ # [START howto_operator_short_circuit_trigger_rules]
46
+ [task_1, task_2, task_3, task_4, task_5, task_6] = [
47
+ EmptyOperator(task_id=f"task_{i}") for i in range(1, 7)
48
+ ]
49
+
50
+ task_7 = EmptyOperator(task_id="task_7", trigger_rule=TriggerRule.ALL_DONE)
51
+
52
+ short_circuit = check_condition.override(task_id="short_circuit", ignore_downstream_trigger_rules=False)(
53
+ condition=False
54
+ )
55
+
56
+ chain(task_1, [task_2, short_circuit], [task_3, task_4], [task_5, task_6], task_7)
57
+ # [END howto_operator_short_circuit_trigger_rules]
58
+
59
+
60
+ example_dag = example_short_circuit_decorator()
@@ -0,0 +1,66 @@
1
+ #
2
+ # Licensed to the Apache Software Foundation (ASF) under one
3
+ # or more contributor license agreements. See the NOTICE file
4
+ # distributed with this work for additional information
5
+ # regarding copyright ownership. The ASF licenses this file
6
+ # to you under the Apache License, Version 2.0 (the
7
+ # "License"); you may not use this file except in compliance
8
+ # with the License. You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing,
13
+ # software distributed under the License is distributed on an
14
+ # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15
+ # KIND, either express or implied. See the License for the
16
+ # specific language governing permissions and limitations
17
+ # under the License.
18
+ """Example DAG demonstrating the usage of the ShortCircuitOperator."""
19
+
20
+ from __future__ import annotations
21
+
22
+ import pendulum
23
+
24
+ from airflow.providers.standard.operators.empty import EmptyOperator
25
+ from airflow.providers.standard.operators.python import ShortCircuitOperator
26
+ from airflow.sdk import DAG, chain
27
+ from airflow.utils.trigger_rule import TriggerRule
28
+
29
+ with DAG(
30
+ dag_id="example_short_circuit_operator",
31
+ schedule=None,
32
+ start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
33
+ catchup=False,
34
+ tags=["example"],
35
+ ) as dag:
36
+ # [START howto_operator_short_circuit]
37
+ cond_true = ShortCircuitOperator(
38
+ task_id="condition_is_True",
39
+ python_callable=lambda: True,
40
+ )
41
+
42
+ cond_false = ShortCircuitOperator(
43
+ task_id="condition_is_False",
44
+ python_callable=lambda: False,
45
+ )
46
+
47
+ ds_true = [EmptyOperator(task_id=f"true_{i}") for i in [1, 2]]
48
+ ds_false = [EmptyOperator(task_id=f"false_{i}") for i in [1, 2]]
49
+
50
+ chain(cond_true, *ds_true)
51
+ chain(cond_false, *ds_false)
52
+ # [END howto_operator_short_circuit]
53
+
54
+ # [START howto_operator_short_circuit_trigger_rules]
55
+ [task_1, task_2, task_3, task_4, task_5, task_6] = [
56
+ EmptyOperator(task_id=f"task_{i}") for i in range(1, 7)
57
+ ]
58
+
59
+ task_7 = EmptyOperator(task_id="task_7", trigger_rule=TriggerRule.ALL_DONE)
60
+
61
+ short_circuit = ShortCircuitOperator(
62
+ task_id="short_circuit", ignore_downstream_trigger_rules=False, python_callable=lambda: False
63
+ )
64
+
65
+ chain(task_1, [task_2, short_circuit], [task_3, task_4], [task_5, task_6], task_7)
66
+ # [END howto_operator_short_circuit_trigger_rules]
@@ -0,0 +1,46 @@
1
+ #
2
+ # Licensed to the Apache Software Foundation (ASF) under one
3
+ # or more contributor license agreements. See the NOTICE file
4
+ # distributed with this work for additional information
5
+ # regarding copyright ownership. The ASF licenses this file
6
+ # to you under the Apache License, Version 2.0 (the
7
+ # "License"); you may not use this file except in compliance
8
+ # with the License. You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing,
13
+ # software distributed under the License is distributed on an
14
+ # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15
+ # KIND, either express or implied. See the License for the
16
+ # specific language governing permissions and limitations
17
+ # under the License.
18
+ """
19
+ Example usage of the TriggerDagRunOperator. This example holds 2 DAGs:
20
+ 1. 1st DAG (example_trigger_controller_dag) holds a TriggerDagRunOperator, which will trigger the 2nd DAG
21
+ 2. 2nd DAG (example_trigger_target_dag) which will be triggered by the TriggerDagRunOperator in the 1st DAG
22
+ """
23
+
24
+ from __future__ import annotations
25
+
26
+ import pendulum
27
+
28
+ from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator
29
+ from airflow.sdk import DAG
30
+
31
+ with DAG(
32
+ dag_id="example_trigger_controller_dag",
33
+ start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
34
+ catchup=False,
35
+ schedule="@once",
36
+ tags=["example"],
37
+ ) as dag:
38
+ # [START howto_operator_trigger_dagrun]
39
+
40
+ trigger = TriggerDagRunOperator(
41
+ task_id="test_trigger_dagrun",
42
+ trigger_dag_id="example_trigger_target_dag", # Ensure this equals the dag_id of the DAG to trigger
43
+ conf={"message": "Hello World"},
44
+ )
45
+
46
+ # [END howto_operator_trigger_dagrun]
@@ -0,0 +1,16 @@
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.
@@ -0,0 +1,24 @@
1
+ /*
2
+ Licensed to the Apache Software Foundation (ASF) under one
3
+ or more contributor license agreements. See the NOTICE file
4
+ distributed with this work for additional information
5
+ regarding copyright ownership. The ASF licenses this file
6
+ to you under the Apache License, Version 2.0 (the
7
+ "License"); you may not use this file except in compliance
8
+ with the License. You may obtain a copy of the License at
9
+
10
+ http://www.apache.org/licenses/LICENSE-2.0
11
+
12
+ Unless required by applicable law or agreed to in writing,
13
+ software distributed under the License is distributed on an
14
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15
+ KIND, either express or implied. See the License for the
16
+ specific language governing permissions and limitations
17
+ under the License.
18
+ */
19
+
20
+ CREATE TABLE Orders (
21
+ order_id INT PRIMARY KEY,
22
+ name TEXT,
23
+ description TEXT
24
+ )
@@ -31,6 +31,7 @@ import warnings
31
31
  from abc import ABCMeta, abstractmethod
32
32
  from collections.abc import Collection, Container, Iterable, Mapping, Sequence
33
33
  from functools import cache
34
+ from itertools import chain
34
35
  from pathlib import Path
35
36
  from tempfile import TemporaryDirectory
36
37
  from typing import TYPE_CHECKING, Any, Callable, NamedTuple, cast
@@ -494,9 +495,21 @@ class _BasePythonVirtualenvOperator(PythonOperator, metaclass=ABCMeta):
494
495
  return textwrap.dedent(inspect.getsource(self.python_callable))
495
496
 
496
497
  def _write_args(self, file: Path):
498
+ def resolve_proxies(obj):
499
+ """Recursively replaces lazy_object_proxy.Proxy instances with their resolved values."""
500
+ if isinstance(obj, lazy_object_proxy.Proxy):
501
+ return obj.__wrapped__ # force evaluation
502
+ if isinstance(obj, dict):
503
+ return {k: resolve_proxies(v) for k, v in obj.items()}
504
+ if isinstance(obj, list):
505
+ return [resolve_proxies(v) for v in obj]
506
+ return obj
507
+
497
508
  if self.op_args or self.op_kwargs:
498
509
  self.log.info("Use %r as serializer.", self.serializer)
499
- file.write_bytes(self.pickling_library.dumps({"args": self.op_args, "kwargs": self.op_kwargs}))
510
+ file.write_bytes(
511
+ self.pickling_library.dumps({"args": self.op_args, "kwargs": resolve_proxies(self.op_kwargs)})
512
+ )
500
513
 
501
514
  def _write_string_args(self, file: Path):
502
515
  file.write_text("\n".join(map(str, self.string_args)))
@@ -859,14 +872,7 @@ class PythonVirtualenvOperator(_BasePythonVirtualenvOperator):
859
872
  # If we're using system packages, assume both are present
860
873
  found_airflow = found_pendulum = True
861
874
  else:
862
- requirements_iterable = []
863
- if isinstance(self.requirements, str):
864
- requirements_iterable = self.requirements.splitlines()
865
- else:
866
- for item in self.requirements:
867
- requirements_iterable.extend(item.splitlines())
868
-
869
- for raw_str in requirements_iterable:
875
+ for raw_str in chain.from_iterable(req.splitlines() for req in self.requirements):
870
876
  line = raw_str.strip()
871
877
  # Skip blank lines and full‐line comments
872
878
  if not line or line.startswith("#"):
@@ -25,6 +25,7 @@ from typing import TYPE_CHECKING, Any, NoReturn
25
25
  from airflow.providers.standard.triggers.temporal import DateTimeTrigger
26
26
  from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
27
27
  from airflow.sensors.base import BaseSensorOperator
28
+ from airflow.utils import timezone
28
29
 
29
30
  try:
30
31
  from airflow.triggers.base import StartTriggerArgs
@@ -41,8 +42,6 @@ except ImportError:
41
42
  timeout: datetime.timedelta | None = None
42
43
 
43
44
 
44
- from airflow.utils import timezone
45
-
46
45
  if TYPE_CHECKING:
47
46
  try:
48
47
  from airflow.sdk.definitions.context import Context
@@ -99,6 +98,13 @@ class DateTimeSensor(BaseSensorOperator):
99
98
  self.log.info("Checking if the time (%s) has come", self.target_time)
100
99
  return timezone.utcnow() > timezone.parse(self.target_time)
101
100
 
101
+ @property
102
+ def _moment(self) -> datetime.datetime:
103
+ if isinstance(self.target_time, datetime.datetime):
104
+ return self.target_time
105
+
106
+ return timezone.parse(self.target_time)
107
+
102
108
 
103
109
  class DateTimeSensorAsync(DateTimeSensor):
104
110
  """
@@ -145,11 +151,11 @@ class DateTimeSensorAsync(DateTimeSensor):
145
151
  self.defer(
146
152
  method_name="execute_complete",
147
153
  trigger=DateTimeTrigger(
148
- moment=timezone.parse(self.target_time),
154
+ moment=self._moment,
149
155
  end_from_trigger=self.end_from_trigger,
150
156
  )
151
157
  if AIRFLOW_V_3_0_PLUS
152
- else DateTimeTrigger(moment=timezone.parse(self.target_time)),
158
+ else DateTimeTrigger(moment=self._moment),
153
159
  )
154
160
 
155
161
  def execute_complete(self, context: Context, event: Any = None) -> None:
@@ -260,12 +260,13 @@ class ExternalTaskSensor(BaseSensorOperator):
260
260
 
261
261
  def _get_dttm_filter(self, context):
262
262
  logical_date = context.get("logical_date")
263
- if logical_date is None:
264
- dag_run = context.get("dag_run")
265
- if TYPE_CHECKING:
266
- assert dag_run
263
+ if AIRFLOW_V_3_0_PLUS:
264
+ if logical_date is None:
265
+ dag_run = context.get("dag_run")
266
+ if TYPE_CHECKING:
267
+ assert dag_run
267
268
 
268
- logical_date = dag_run.run_after
269
+ logical_date = dag_run.run_after
269
270
  if self.execution_delta:
270
271
  dttm = logical_date - self.execution_delta
271
272
  elif self.execution_date_fn:
@@ -428,7 +429,7 @@ class ExternalTaskSensor(BaseSensorOperator):
428
429
  else:
429
430
  dttm_filter = self._get_dttm_filter(context)
430
431
  logical_or_execution_dates = (
431
- {"logical_dates": dttm_filter} if AIRFLOW_V_3_0_PLUS else {"execution_date": dttm_filter}
432
+ {"logical_dates": dttm_filter} if AIRFLOW_V_3_0_PLUS else {"execution_dates": dttm_filter}
432
433
  )
433
434
  self.defer(
434
435
  timeout=self.execution_timeout,
@@ -18,9 +18,12 @@
18
18
  from __future__ import annotations
19
19
 
20
20
  import datetime
21
+ import warnings
21
22
  from dataclasses import dataclass
22
- from typing import TYPE_CHECKING, Any, NoReturn
23
+ from typing import TYPE_CHECKING, Any
23
24
 
25
+ from airflow.configuration import conf
26
+ from airflow.exceptions import AirflowProviderDeprecationWarning
24
27
  from airflow.providers.standard.triggers.temporal import DateTimeTrigger
25
28
  from airflow.sensors.base import BaseSensorOperator
26
29
 
@@ -54,6 +57,7 @@ class TimeSensor(BaseSensorOperator):
54
57
  Waits until the specified time of the day.
55
58
 
56
59
  :param target_time: time after which the job succeeds
60
+ :param deferrable: whether to defer execution
57
61
 
58
62
  .. seealso::
59
63
  For more information on how to use this sensor, take a look at the guide:
@@ -61,32 +65,6 @@ class TimeSensor(BaseSensorOperator):
61
65
 
62
66
  """
63
67
 
64
- def __init__(self, *, target_time: datetime.time, **kwargs) -> None:
65
- super().__init__(**kwargs)
66
- self.target_time = target_time
67
-
68
- def poke(self, context: Context) -> bool:
69
- self.log.info("Checking if the time (%s) has come", self.target_time)
70
- return timezone.make_naive(timezone.utcnow(), self.dag.timezone).time() > self.target_time
71
-
72
-
73
- class TimeSensorAsync(BaseSensorOperator):
74
- """
75
- Waits until the specified time of the day.
76
-
77
- This frees up a worker slot while it is waiting.
78
-
79
- :param target_time: time after which the job succeeds
80
- :param start_from_trigger: Start the task directly from the triggerer without going into the worker.
81
- :param end_from_trigger: End the task directly from the triggerer without going into the worker.
82
- :param trigger_kwargs: The keyword arguments passed to the trigger when start_from_trigger is set to True
83
- during dynamic task mapping. This argument is not used in standard usage.
84
-
85
- .. seealso::
86
- For more information on how to use this sensor, take a look at the guide:
87
- :ref:`howto/operator:TimeSensorAsync`
88
- """
89
-
90
68
  start_trigger_args = StartTriggerArgs(
91
69
  trigger_cls="airflow.providers.standard.triggers.temporal.DateTimeTrigger",
92
70
  trigger_kwargs={"moment": "", "end_from_trigger": False},
@@ -100,32 +78,66 @@ class TimeSensorAsync(BaseSensorOperator):
100
78
  self,
101
79
  *,
102
80
  target_time: datetime.time,
81
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
103
82
  start_from_trigger: bool = False,
104
- trigger_kwargs: dict[str, Any] | None = None,
105
83
  end_from_trigger: bool = False,
84
+ trigger_kwargs: dict[str, Any] | None = None,
106
85
  **kwargs,
107
86
  ) -> None:
108
87
  super().__init__(**kwargs)
109
- self.start_from_trigger = start_from_trigger
110
- self.end_from_trigger = end_from_trigger
111
- self.target_time = target_time
112
88
 
89
+ # Create a "date-aware" timestamp that will be used as the "target_datetime". This is a requirement
90
+ # of the DateTimeTrigger
91
+
92
+ # Get date considering dag.timezone
113
93
  aware_time = timezone.coerce_datetime(
114
- datetime.datetime.combine(datetime.datetime.today(), self.target_time, self.dag.timezone)
94
+ datetime.datetime.combine(
95
+ datetime.datetime.now(self.dag.timezone), target_time, self.dag.timezone
96
+ )
115
97
  )
116
98
 
99
+ # Now that the dag's timezone has made the datetime timezone aware, we need to convert to UTC
117
100
  self.target_datetime = timezone.convert_to_utc(aware_time)
101
+ self.deferrable = deferrable
102
+ self.start_from_trigger = start_from_trigger
103
+ self.end_from_trigger = end_from_trigger
104
+
118
105
  if self.start_from_trigger:
119
106
  self.start_trigger_args.trigger_kwargs = dict(
120
107
  moment=self.target_datetime, end_from_trigger=self.end_from_trigger
121
108
  )
122
109
 
123
- def execute(self, context: Context) -> NoReturn:
124
- self.defer(
125
- trigger=DateTimeTrigger(moment=self.target_datetime, end_from_trigger=self.end_from_trigger),
126
- method_name="execute_complete",
127
- )
110
+ def execute(self, context: Context) -> None:
111
+ if self.deferrable:
112
+ self.defer(
113
+ trigger=DateTimeTrigger(
114
+ moment=self.target_datetime, # This needs to be an aware timestamp
115
+ end_from_trigger=self.end_from_trigger,
116
+ ),
117
+ method_name="execute_complete",
118
+ )
119
+
120
+ def execute_complete(self, context: Context) -> None:
121
+ return
128
122
 
129
- def execute_complete(self, context: Context, event: Any = None) -> None:
130
- """Handle the event when the trigger fires and return immediately."""
131
- return None
123
+ def poke(self, context: Context) -> bool:
124
+ self.log.info("Checking if the time (%s) has come", self.target_datetime)
125
+
126
+ # self.target_date has been converted to UTC, so we do not need to convert timezone
127
+ return timezone.utcnow() > self.target_datetime
128
+
129
+
130
+ class TimeSensorAsync(TimeSensor):
131
+ """
132
+ Deprecated. Use TimeSensor with deferrable=True instead.
133
+
134
+ :sphinx-autoapi-skip:
135
+ """
136
+
137
+ def __init__(self, **kwargs) -> None:
138
+ warnings.warn(
139
+ "TimeSensorAsync is deprecated and will be removed in a future version. Use `TimeSensor` with deferrable=True instead.",
140
+ AirflowProviderDeprecationWarning,
141
+ stacklevel=2,
142
+ )
143
+ super().__init__(deferrable=True, **kwargs)