You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by je...@apache.org on 2021/10/20 21:26:45 UTC

[airflow] branch v2-2-test updated (fe30eb7 -> e9c4926)

This is an automated email from the ASF dual-hosted git repository.

jedcunningham pushed a change to branch v2-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git.


    from fe30eb7  Upgrade old DAG/task param format when deserializing from the DB (#18986)
     new 4f26959  Allow Param to support a default value of ``None`` (#19034)
     new c28affb  Rename trigger page label to Logical Date (#19061)
     new bb7ebb9  Relax packaging requirement (#19087)
     new bbe9012  Ensure task state doesn't change when marked as failed/success/skipped (#19095)
     new 44c37e4  Change `ds`, `ts`, etc. back to use logical date (#19088)
     new f3f0eb2  Workaround libstdcpp TLS error (#19010)
     new 478f94b  Don't install SQLAlchemy/Pendulum adapters for other DBs (#18745)
     new 8223e89  Prevent scheduler crash when serialized dag is missing (#19113)
     new 69048f7  Fix catchup by limiting queued dagrun creation using max_active_runs (#18897)
     new e9c4926  Update taskinstance REST API schema to include dag_run_id field (#19105)

The 10 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails.  The revisions
listed as "add" were already present in the repository and have only
been added to this reference.


Summary of changes:
 Dockerfile                                         |   9 ++
 Dockerfile.ci                                      |  10 ++
 UPDATING.md                                        |  19 +++
 airflow/api_connexion/openapi/v1.yaml              |   6 +
 .../api_connexion/schemas/task_instance_schema.py  |  47 ++++---
 airflow/config_templates/config.yml                |   8 --
 airflow/config_templates/default_airflow.cfg       |   4 -
 airflow/jobs/scheduler_job.py                      |  75 ++++++-----
 airflow/models/dag.py                              |  25 ++--
 airflow/models/dagrun.py                           |  18 ++-
 airflow/models/param.py                            |  39 ++++--
 airflow/models/taskinstance.py                     |  50 +++++---
 airflow/serialization/serialized_objects.py        |  18 +--
 airflow/settings.py                                |  26 ++--
 airflow/www/templates/airflow/trigger.html         |   2 +-
 docs/apache-airflow/templates-ref.rst              |  20 +--
 setup.cfg                                          |   2 +-
 tests/api_connexion/endpoints/test_dag_endpoint.py |   6 +-
 .../api_connexion/endpoints/test_task_endpoint.py  |   6 +-
 .../endpoints/test_task_instance_endpoint.py       |   3 +
 tests/api_connexion/schemas/test_dag_schema.py     |   2 +-
 .../schemas/test_task_instance_schema.py           |   2 +
 tests/api_connexion/schemas/test_task_schema.py    |   2 +-
 tests/jobs/test_scheduler_job.py                   | 137 ++++++++++++++++++---
 tests/models/test_dag.py                           |  40 ++++--
 tests/models/test_param.py                         |  31 ++++-
 tests/models/test_taskinstance.py                  |  25 +++-
 tests/serialization/test_dag_serialization.py      |  17 +++
 tests/utils/log/test_log_reader.py                 |  54 +++++++-
 29 files changed, 525 insertions(+), 178 deletions(-)

[airflow] 03/10: Relax packaging requirement (#19087)

Posted by je...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

jedcunningham pushed a commit to branch v2-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit bb7ebb9bc7a115ef5adcd222f2f4fb235cdf30c3
Author: Tzu-ping Chung <tp...@astronomer.io>
AuthorDate: Wed Oct 20 04:19:25 2021 +0800

    Relax packaging requirement (#19087)
    
    (cherry picked from commit 454819905523fcb4be64374282c0e8acca1313fd)
---
 setup.cfg | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.cfg b/setup.cfg
index 56a18f1..2978518 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -126,7 +126,7 @@ install_requires =
     marshmallow-oneofschema>=2.0.1
     # Required by vendored-in connexion
     openapi-spec-validator>=0.2.4
-    packaging~=21.0
+    packaging>=14.0
     pendulum~=2.0
     pep562~=1.0;python_version<"3.7"
     psutil>=4.2.0, <6.0.0

[airflow] 05/10: Change `ds`, `ts`, etc. back to use logical date (#19088)

Posted by je...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

jedcunningham pushed a commit to branch v2-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit 44c37e4eeefcc62adb9a19ab101f7630f36f5de6
Author: Tzu-ping Chung <tp...@astronomer.io>
AuthorDate: Thu Oct 21 01:26:08 2021 +0800

    Change `ds`, `ts`, etc. back to use logical date (#19088)
    
    (cherry picked from commit 25a50bb1fbf0e228706c7927cb36570921881adb)
---
 airflow/models/taskinstance.py        | 43 +++++++++++++++++-----------
 docs/apache-airflow/templates-ref.rst | 20 +++++++------
 tests/utils/log/test_log_reader.py    | 54 ++++++++++++++++++++++++++++++++++-
 3 files changed, 91 insertions(+), 26 deletions(-)

diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index 46945c7..71ab107 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -27,7 +27,7 @@ from collections import defaultdict
 from datetime import datetime, timedelta
 from functools import partial
 from tempfile import NamedTemporaryFile
-from typing import IO, TYPE_CHECKING, Any, Dict, Iterable, List, NamedTuple, Optional, Tuple, Union
+from typing import IO, TYPE_CHECKING, Any, Callable, Dict, Iterable, List, NamedTuple, Optional, Tuple, Union
 from urllib.parse import quote
 
 import dill
@@ -1788,17 +1788,16 @@ class TaskInstance(Base, LoggingMixin):
         if conf.getboolean('core', 'dag_run_conf_overrides_params'):
             self.overwrite_params_with_dag_run_conf(params=params, dag_run=dag_run)
 
-        interval_start = dag.get_run_data_interval(dag_run).start
-        ds = interval_start.strftime('%Y-%m-%d')
+        logical_date = timezone.coerce_datetime(self.execution_date)
+        ds = logical_date.strftime('%Y-%m-%d')
+        ds_nodash = ds.replace('-', '')
+        ts = logical_date.isoformat()
+        ts_nodash = logical_date.strftime('%Y%m%dT%H%M%S')
+        ts_nodash_with_tz = ts.replace('-', '').replace(':', '')
 
         # Now validates Params and convert them into a simple dict
         task.params = params.validate()
 
-        ds_nodash = ds.replace('-', '')
-        ts = interval_start.isoformat()
-        ts_nodash = interval_start.strftime('%Y%m%dT%H%M%S')
-        ts_nodash_with_tz = ts.replace('-', '').replace(':', '')
-
         @cache  # Prevent multiple database access.
         def _get_previous_dagrun_success() -> Optional["DagRun"]:
             return self.get_previous_dagrun(state=State.SUCCESS, session=session)
@@ -1910,14 +1909,23 @@ class TaskInstance(Base, LoggingMixin):
 
         # Create lazy proxies for deprecated stuff.
 
-        def deprecated_proxy(func, *, key, replacement=None) -> lazy_object_proxy.Proxy:
+        def deprecated_proxy(
+            func: Callable[[], Any],
+            *,
+            key: str,
+            replacements: Optional[List[str]] = None,
+        ) -> lazy_object_proxy.Proxy:
             def deprecated_func():
                 message = (
                     f"Accessing {key!r} from the template is deprecated and "
                     f"will be removed in a future version."
                 )
-                if replacement:
-                    message += f" Please use {replacement!r} instead."
+                if replacements:
+                    display_except_last = ", ".join(repr(r) for r in replacements[:-1])
+                    if display_except_last:
+                        message += f" Please use {display_except_last} or {replacements[-1]!r} instead."
+                    else:
+                        message += f" Please use {replacements[-1]!r} instead."
                 warnings.warn(message, DeprecationWarning)
                 return func()
 
@@ -1995,22 +2003,23 @@ class TaskInstance(Base, LoggingMixin):
             'ds': ds,
             'ds_nodash': ds_nodash,
             'execution_date': deprecated_proxy(
-                lambda: timezone.coerce_datetime(self.execution_date),
+                lambda: logical_date,
                 key='execution_date',
-                replacement='data_interval_start',
+                replacements=['logical_date', 'data_interval_start'],
             ),
             'inlets': task.inlets,
+            'logical_date': logical_date,
             'macros': macros,
-            'next_ds': deprecated_proxy(get_next_ds, key="next_ds", replacement="data_interval_end | ds"),
+            'next_ds': deprecated_proxy(get_next_ds, key="next_ds", replacements=["data_interval_end | ds"]),
             'next_ds_nodash': deprecated_proxy(
                 get_next_ds_nodash,
                 key="next_ds_nodash",
-                replacement="data_interval_end | ds_nodash",
+                replacements=["data_interval_end | ds_nodash"],
             ),
             'next_execution_date': deprecated_proxy(
                 get_next_execution_date,
                 key='next_execution_date',
-                replacement='data_interval_end',
+                replacements=['data_interval_end'],
             ),
             'outlets': task.outlets,
             'params': task.params,
@@ -2022,7 +2031,7 @@ class TaskInstance(Base, LoggingMixin):
             'prev_execution_date_success': deprecated_proxy(
                 lambda: self.get_previous_execution_date(state=State.SUCCESS, session=session),
                 key='prev_execution_date_success',
-                replacement='prev_data_interval_start_success',
+                replacements=['prev_data_interval_start_success'],
             ),
             'prev_start_date_success': lazy_object_proxy.Proxy(get_prev_start_date_success),
             'run_id': self.run_id,
diff --git a/docs/apache-airflow/templates-ref.rst b/docs/apache-airflow/templates-ref.rst
index bc97315..a74c0aa 100644
--- a/docs/apache-airflow/templates-ref.rst
+++ b/docs/apache-airflow/templates-ref.rst
@@ -38,15 +38,14 @@ Variable                                    Description
 ==========================================  ====================================
 ``{{ data_interval_start }}``               Start of the data interval (`pendulum.DateTime`_).
 ``{{ data_interval_end }}``                 End of the data interval (`pendulum.DateTime`_).
-``{{ ds }}``                                Start of the data interval as ``YYYY-MM-DD``.
-                                            Same as ``{{ data_interval_start | ds }}``.
-``{{ ds_nodash }}``                         Start of the data interval as ``YYYYMMDD``.
-                                            Same as ``{{ data_interval_start | ds_nodash }}``.
-``{{ ts }}``                                Same as ``{{ data_interval_start | ts }}``.
+``{{ ds }}``                                The DAG run's logical date as ``YYYY-MM-DD``.
+                                            Same as ``{{ dag_run.logical_date | ds }}``.
+``{{ ds_nodash }}``                         Same as ``{{ dag_run.logical_date | ds_nodash }}``.
+``{{ ts }}``                                Same as ``{{ dag_run.logical_date | ts }}``.
                                             Example: ``2018-01-01T00:00:00+00:00``.
-``{{ ts_nodash_with_tz }}``                 Same as ``{{ data_interval_start | ts_nodash_with_tz }}``.
+``{{ ts_nodash_with_tz }}``                 Same as ``{{ dag_run.logical_date | ts_nodash_with_tz }}``.
                                             Example: ``20180101T000000+0000``.
-``{{ ts_nodash }}``                         Same as ``{{ data_interval_start | ts_nodash }}``.
+``{{ ts_nodash }}``                         Same as ``{{ dag_run.logical_date | ts_nodash }}``.
                                             Example: ``20180101T000000``.
 ``{{ prev_data_interval_start_success }}``  Start of the data interval from prior successful DAG run
                                             (`pendulum.DateTime`_ or ``None``).
@@ -78,13 +77,18 @@ Variable                                    Description
                                             subcommand.
 ==========================================  ====================================
 
+.. note::
+
+    The DAG run's logical date, and values derived from it, such as ``ds`` and
+    ``ts``, **should not** be considered unique in a DAG. Use ``run_id`` instead.
+
 The following variables are deprecated. They are kept for backward compatibility, but you should convert
 existing code to use other variables instead.
 
 =====================================   ====================================
 Deprecated Variable                     Description
 =====================================   ====================================
-``{{ execution_date }}``                the execution date (logical date), same as ``logical_date``
+``{{ execution_date }}``                the execution date (logical date), same as ``dag_run.logical_date``
 ``{{ next_execution_date }}``           the next execution date (if available) (`pendulum.DateTime`_)
                                         if ``{{ execution_date }}`` is ``2018-01-01 00:00:00`` and
                                         ``schedule_interval`` is ``@weekly``, ``{{ next_execution_date }}``
diff --git a/tests/utils/log/test_log_reader.py b/tests/utils/log/test_log_reader.py
index f891b2f..a3bb9b7 100644
--- a/tests/utils/log/test_log_reader.py
+++ b/tests/utils/log/test_log_reader.py
@@ -16,16 +16,21 @@
 # under the License.
 
 import copy
+import datetime
 import logging
 import os
 import sys
 import tempfile
 from unittest import mock
 
+import pendulum
 import pytest
 
 from airflow import settings
 from airflow.config_templates.airflow_local_settings import DEFAULT_LOGGING_CONFIG
+from airflow.models import DagRun
+from airflow.operators.python import PythonOperator
+from airflow.timetables.base import DataInterval
 from airflow.utils import timezone
 from airflow.utils.log.log_reader import TaskLogReader
 from airflow.utils.log.logging_mixin import ExternalLoggingMixin
@@ -87,7 +92,7 @@ class TestLogView:
                 f.flush()
 
     @pytest.fixture(autouse=True)
-    def prepare_db(self, session, create_task_instance):
+    def prepare_db(self, create_task_instance):
         ti = create_task_instance(
             dag_id=self.DAG_ID,
             task_id=self.TASK_ID,
@@ -238,3 +243,50 @@ class TestLogView:
 
         mock_prop.return_value = True
         assert task_log_reader.supports_external_link
+
+    def test_task_log_filename_unique(self, dag_maker):
+        """Ensure the default log_filename_template produces a unique filename.
+
+        See discussion in apache/airflow#19058 [1]_ for how uniqueness may
+        change in a future Airflow release. For now, the logical date is used
+        to distinguish DAG runs. This test should be modified when the logical
+        date is no longer used to ensure uniqueness.
+
+        [1]: https://github.com/apache/airflow/issues/19058
+        """
+        dag_id = "test_task_log_filename_ts_corresponds_to_logical_date"
+        task_id = "echo_run_type"
+
+        def echo_run_type(dag_run: DagRun, **kwargs):
+            print(dag_run.run_type)
+
+        with dag_maker(dag_id, start_date=self.DEFAULT_DATE, schedule_interval="@daily") as dag:
+            PythonOperator(task_id=task_id, python_callable=echo_run_type)
+
+        start = pendulum.datetime(2021, 1, 1)
+        end = start + datetime.timedelta(days=1)
+        trigger_time = end + datetime.timedelta(hours=4, minutes=29)  # Arbitrary.
+
+        # Create two DAG runs that have the same data interval, but not the same
+        # execution date, to check if they correctly use different log files.
+        scheduled_dagrun: DagRun = dag_maker.create_dagrun(
+            run_type=DagRunType.SCHEDULED,
+            execution_date=start,
+            data_interval=DataInterval(start, end),
+        )
+        manual_dagrun: DagRun = dag_maker.create_dagrun(
+            run_type=DagRunType.MANUAL,
+            execution_date=trigger_time,
+            data_interval=DataInterval(start, end),
+        )
+
+        scheduled_ti = scheduled_dagrun.get_task_instance(task_id)
+        manual_ti = manual_dagrun.get_task_instance(task_id)
+        assert scheduled_ti is not None
+        assert manual_ti is not None
+
+        scheduled_ti.refresh_from_task(dag.get_task(task_id))
+        manual_ti.refresh_from_task(dag.get_task(task_id))
+
+        reader = TaskLogReader()
+        assert reader.render_log_filename(scheduled_ti, 1) != reader.render_log_filename(manual_ti, 1)

[airflow] 01/10: Allow Param to support a default value of ``None`` (#19034)

Posted by je...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

jedcunningham pushed a commit to branch v2-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit 4f269594366f501974c6c37bde5d19f615761859
Author: Jed Cunningham <66...@users.noreply.github.com>
AuthorDate: Tue Oct 19 14:30:15 2021 -0600

    Allow Param to support a default value of ``None`` (#19034)
    
    (cherry picked from commit ee4fb073b6d3c64437fc27b673ed9b6568953f46)
---
 UPDATING.md                                        | 14 ++++++++
 airflow/models/dag.py                              |  4 +--
 airflow/models/param.py                            | 39 +++++++++++++++++-----
 airflow/serialization/serialized_objects.py        | 18 +++-------
 tests/api_connexion/endpoints/test_dag_endpoint.py |  6 ++--
 .../api_connexion/endpoints/test_task_endpoint.py  |  6 ++--
 tests/api_connexion/schemas/test_dag_schema.py     |  2 +-
 tests/api_connexion/schemas/test_task_schema.py    |  2 +-
 tests/models/test_dag.py                           | 16 +++++++++
 tests/models/test_param.py                         | 31 +++++++++++++++--
 tests/serialization/test_dag_serialization.py      | 17 ++++++++++
 11 files changed, 119 insertions(+), 36 deletions(-)

diff --git a/UPDATING.md b/UPDATING.md
index 8ffd87d..9c0c4fd 100644
--- a/UPDATING.md
+++ b/UPDATING.md
@@ -77,6 +77,20 @@ https://developers.google.com/style/inclusive-documentation
 
 -->
 
+### ``Param``'s default value for ``default`` removed
+
+``Param``, introduced in Airflow 2.2.0, accidentally set the default value to ``None``. This default has been removed. If you want ``None`` as your default, explicitly set it as such. For example:
+
+```python
+Param(None, type=["null", "string"])
+```
+
+Now if you resolve a ``Param`` without a default and don't pass a value, you will get an ``TypeError``. For Example:
+
+```python
+Param().resolve()  # raises TypeError
+```
+
 ## Airflow 2.2.0
 
 ### `worker_log_server_port` configuration has been moved to the ``logging`` section.
diff --git a/airflow/models/dag.py b/airflow/models/dag.py
index 129a30a..b372cfe 100644
--- a/airflow/models/dag.py
+++ b/airflow/models/dag.py
@@ -2631,8 +2631,8 @@ class DAG(LoggingMixin):
             return
 
         for k, v in self.params.items():
-            # As type can be an array, we would check if `null` is a allowed type or not
-            if v.default is None and ("type" not in v.schema or "null" not in v.schema["type"]):
+            # As type can be an array, we would check if `null` is an allowed type or not
+            if not v.has_value and ("type" not in v.schema or "null" not in v.schema["type"]):
                 raise AirflowException(
                     "DAG Schedule must be None, if there are any required params without default values"
                 )
diff --git a/airflow/models/param.py b/airflow/models/param.py
index 2869481..1ae01dc 100644
--- a/airflow/models/param.py
+++ b/airflow/models/param.py
@@ -24,6 +24,16 @@ from jsonschema.exceptions import ValidationError
 from airflow.exceptions import AirflowException
 
 
+class NoValueSentinel:
+    """Sentinel class used to distinguish between None and no passed value"""
+
+    def __str__(self):
+        return "NoValueSentinel"
+
+    def __repr__(self):
+        return "NoValueSentinel"
+
+
 class Param:
     """
     Class to hold the default value of a Param and rule set to do the validations. Without the rule set
@@ -38,22 +48,25 @@ class Param:
     :type schema: dict
     """
 
-    def __init__(self, default: Any = None, description: str = None, **kwargs):
-        self.default = default
+    __NO_VALUE_SENTINEL = NoValueSentinel()
+
+    def __init__(self, default: Any = __NO_VALUE_SENTINEL, description: str = None, **kwargs):
+        self.value = default
         self.description = description
         self.schema = kwargs.pop('schema') if 'schema' in kwargs else kwargs
 
-        # If default is not None, then validate it once, may raise ValueError
-        if default:
+        # If we have a value, validate it once. May raise ValueError.
+        if self.has_value:
             try:
-                jsonschema.validate(self.default, self.schema, format_checker=FormatChecker())
+                jsonschema.validate(self.value, self.schema, format_checker=FormatChecker())
             except ValidationError as err:
                 raise ValueError(err)
 
-    def resolve(self, value: Optional[Any] = None, suppress_exception: bool = False) -> Any:
+    def resolve(self, value: Optional[Any] = __NO_VALUE_SENTINEL, suppress_exception: bool = False) -> Any:
         """
         Runs the validations and returns the Param's final value.
-        May raise ValueError on failed validations.
+        May raise ValueError on failed validations, or TypeError
+        if no value is passed and no value already exists.
 
         :param value: The value to be updated for the Param
         :type value: Optional[Any]
@@ -61,14 +74,18 @@ class Param:
             If true and validations fails, the return value would be None.
         :type suppress_exception: bool
         """
+        final_val = value if value != self.__NO_VALUE_SENTINEL else self.value
+        if isinstance(final_val, NoValueSentinel):
+            if suppress_exception:
+                return None
+            raise TypeError("No value passed and Param has no default value")
         try:
-            final_val = value or self.default
             jsonschema.validate(final_val, self.schema, format_checker=FormatChecker())
-            self.default = final_val
         except ValidationError as err:
             if suppress_exception:
                 return None
             raise ValueError(err) from None
+        self.value = final_val
         return final_val
 
     def dump(self) -> dict:
@@ -77,6 +94,10 @@ class Param:
         out_dict.update(self.__dict__)
         return out_dict
 
+    @property
+    def has_value(self) -> bool:
+        return not isinstance(self.value, NoValueSentinel)
+
 
 class ParamsDict(dict):
     """
diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py
index 3d30ca4..6ec9770 100644
--- a/airflow/serialization/serialized_objects.py
+++ b/airflow/serialization/serialized_objects.py
@@ -416,7 +416,9 @@ class BaseSerialization:
         for k, v in params.items():
             # TODO: As of now, we would allow serialization of params which are of type Param only
             if f'{v.__module__}.{v.__class__.__name__}' == 'airflow.models.param.Param':
-                serialized_params[k] = v.dump()
+                kwargs = v.dump()
+                kwargs['default'] = kwargs.pop('value')
+                serialized_params[k] = kwargs
             else:
                 raise ValueError('Params to a DAG or a Task can be only of type airflow.models.param.Param')
         return serialized_params
@@ -543,7 +545,7 @@ class SerializedBaseOperator(BaseOperator, BaseSerialization):
                     serialize_op[template_field] = serialize_template_field(value)
 
         if op.params:
-            serialize_op['params'] = cls._serialize_operator_params(op.params)
+            serialize_op['params'] = cls._serialize_params_dict(op.params)
 
         return serialize_op
 
@@ -747,18 +749,6 @@ class SerializedBaseOperator(BaseOperator, BaseSerialization):
 
         return serialize_operator_extra_links
 
-    @classmethod
-    def _serialize_operator_params(cls, op_params: ParamsDict):
-        """Serialize Params dict of a operator"""
-        serialized_params = {}
-        for k, v in op_params.items():
-            # TODO: As of now, we would allow serialization of params which are of type Param only
-            if f'{v.__module__}.{v.__class__.__name__}' == 'airflow.models.param.Param':
-                serialized_params[k] = v.dump()
-            else:
-                raise ValueError('Params to a Task can be only of type airflow.models.param.Param')
-        return serialized_params
-
 
 class SerializedDAG(DAG, BaseSerialization):
     """
diff --git a/tests/api_connexion/endpoints/test_dag_endpoint.py b/tests/api_connexion/endpoints/test_dag_endpoint.py
index e9602fa..73dac72 100644
--- a/tests/api_connexion/endpoints/test_dag_endpoint.py
+++ b/tests/api_connexion/endpoints/test_dag_endpoint.py
@@ -240,7 +240,7 @@ class TestGetDagDetails(TestDagEndpoint):
             "params": {
                 "foo": {
                     '__class': 'airflow.models.param.Param',
-                    'default': 1,
+                    'value': 1,
                     'description': None,
                     'schema': {},
                 }
@@ -353,7 +353,7 @@ class TestGetDagDetails(TestDagEndpoint):
             "params": {
                 "foo": {
                     '__class': 'airflow.models.param.Param',
-                    'default': 1,
+                    'value': 1,
                     'description': None,
                     'schema': {},
                 }
@@ -400,7 +400,7 @@ class TestGetDagDetails(TestDagEndpoint):
             "params": {
                 "foo": {
                     '__class': 'airflow.models.param.Param',
-                    'default': 1,
+                    'value': 1,
                     'description': None,
                     'schema': {},
                 }
diff --git a/tests/api_connexion/endpoints/test_task_endpoint.py b/tests/api_connexion/endpoints/test_task_endpoint.py
index 242df7f..2630d49 100644
--- a/tests/api_connexion/endpoints/test_task_endpoint.py
+++ b/tests/api_connexion/endpoints/test_task_endpoint.py
@@ -100,7 +100,7 @@ class TestGetTask(TestTaskEndpoint):
             'params': {
                 'foo': {
                     '__class': 'airflow.models.param.Param',
-                    'default': 'bar',
+                    'value': 'bar',
                     'description': None,
                     'schema': {},
                 }
@@ -150,7 +150,7 @@ class TestGetTask(TestTaskEndpoint):
             'params': {
                 'foo': {
                     '__class': 'airflow.models.param.Param',
-                    'default': 'bar',
+                    'value': 'bar',
                     'description': None,
                     'schema': {},
                 }
@@ -215,7 +215,7 @@ class TestGetTasks(TestTaskEndpoint):
                     'params': {
                         'foo': {
                             '__class': 'airflow.models.param.Param',
-                            'default': 'bar',
+                            'value': 'bar',
                             'description': None,
                             'schema': {},
                         }
diff --git a/tests/api_connexion/schemas/test_dag_schema.py b/tests/api_connexion/schemas/test_dag_schema.py
index 3f6d3b0..f12c06d 100644
--- a/tests/api_connexion/schemas/test_dag_schema.py
+++ b/tests/api_connexion/schemas/test_dag_schema.py
@@ -133,7 +133,7 @@ class TestDAGDetailSchema:
             'params': {
                 'foo': {
                     '__class': 'airflow.models.param.Param',
-                    'default': 1,
+                    'value': 1,
                     'description': None,
                     'schema': {},
                 }
diff --git a/tests/api_connexion/schemas/test_task_schema.py b/tests/api_connexion/schemas/test_task_schema.py
index f4a4867..65006ed 100644
--- a/tests/api_connexion/schemas/test_task_schema.py
+++ b/tests/api_connexion/schemas/test_task_schema.py
@@ -81,7 +81,7 @@ class TestTaskCollectionSchema:
                     'params': {
                         'foo': {
                             '__class': 'airflow.models.param.Param',
-                            'default': 'bar',
+                            'value': 'bar',
                             'description': None,
                             'schema': {},
                         }
diff --git a/tests/models/test_dag.py b/tests/models/test_dag.py
index e77f271..dfa1a0a 100644
--- a/tests/models/test_dag.py
+++ b/tests/models/test_dag.py
@@ -1799,14 +1799,30 @@ class TestDag(unittest.TestCase):
 
     def test_validate_params_on_trigger_dag(self):
         dag = models.DAG('dummy-dag', schedule_interval=None, params={'param1': Param(type="string")})
+        with pytest.raises(TypeError, match="No value passed and Param has no default value"):
+            dag.create_dagrun(
+                run_id="test_dagrun_missing_param",
+                state=State.RUNNING,
+                execution_date=TEST_DATE,
+            )
 
+        dag = models.DAG('dummy-dag', schedule_interval=None, params={'param1': Param(type="string")})
         with pytest.raises(ValueError, match="Invalid input for param param1: None is not of type 'string'"):
             dag.create_dagrun(
                 run_id="test_dagrun_missing_param",
                 state=State.RUNNING,
                 execution_date=TEST_DATE,
+                conf={"param1": None},
             )
 
+        dag = models.DAG('dummy-dag', schedule_interval=None, params={'param1': Param(type="string")})
+        dag.create_dagrun(
+            run_id="test_dagrun_missing_param",
+            state=State.RUNNING,
+            execution_date=TEST_DATE,
+            conf={"param1": "hello"},
+        )
+
 
 class TestDagModel:
     def test_dags_needing_dagruns_not_too_early(self):
diff --git a/tests/models/test_param.py b/tests/models/test_param.py
index ea25f21..a9f260b 100644
--- a/tests/models/test_param.py
+++ b/tests/models/test_param.py
@@ -31,15 +31,25 @@ class TestParam(unittest.TestCase):
         p = Param('test')
         assert p.resolve() == 'test'
 
-        p.default = 10
+        p.value = 10
         assert p.resolve() == 10
 
     def test_null_param(self):
         p = Param()
+        with pytest.raises(TypeError, match='No value passed and Param has no default value'):
+            p.resolve()
+        assert p.resolve(None) is None
+
+        p = Param(None)
         assert p.resolve() is None
+        assert p.resolve(None) is None
 
         p = Param(type="null")
+        p = Param(None, type='null')
         assert p.resolve() is None
+        assert p.resolve(None) is None
+        with pytest.raises(ValueError):
+            p.resolve('test')
 
     def test_string_param(self):
         p = Param('test', type='string')
@@ -51,8 +61,10 @@ class TestParam(unittest.TestCase):
         p = Param('10.0.0.0', type='string', format='ipv4')
         assert p.resolve() == '10.0.0.0'
 
+        p = Param(type='string')
         with pytest.raises(ValueError):
-            p = Param(type='string')
+            p.resolve(None)
+        with pytest.raises(TypeError, match='No value passed and Param has no default value'):
             p.resolve()
 
     def test_int_param(self):
@@ -96,7 +108,7 @@ class TestParam(unittest.TestCase):
         p = Param('abc', type='string', minLength=2, maxLength=4)
         assert p.resolve() == 'abc'
 
-        p.default = 'long_string'
+        p.value = 'long_string'
         assert p.resolve(suppress_exception=True) is None
 
     def test_explicit_schema(self):
@@ -115,6 +127,19 @@ class TestParam(unittest.TestCase):
         with pytest.raises(ValueError):
             p = S3Param("file://not_valid/s3_path")
 
+    def test_value_saved(self):
+        p = Param("hello", type="string")
+        assert p.resolve("world") == "world"
+        assert p.resolve() == "world"
+
+    def test_dump(self):
+        p = Param('hello', description='world', type='string', minLength=2)
+        dump = p.dump()
+        assert dump['__class'] == 'airflow.models.param.Param'
+        assert dump['value'] == 'hello'
+        assert dump['description'] == 'world'
+        assert dump['schema'] == {'type': 'string', 'minLength': 2}
+
 
 class TestParamsDict(unittest.TestCase):
     def test_params_dict(self):
diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py
index 22eea07..afba96a 100644
--- a/tests/serialization/test_dag_serialization.py
+++ b/tests/serialization/test_dag_serialization.py
@@ -1450,6 +1450,23 @@ class TestStringifiedDAGs:
         assert isinstance(dict.__getitem__(dag.params, "none"), Param)
         assert dag.params["str"] == "str"
 
+    def test_params_serialize_default(self):
+        serialized = {
+            "__version": 1,
+            "dag": {
+                "_dag_id": "simple_dag",
+                "fileloc": __file__,
+                "tasks": [],
+                "timezone": "UTC",
+                "params": {"str": {"__class": "airflow.models.param.Param", "default": "str"}},
+            },
+        }
+        SerializedDAG.validate_schema(serialized)
+        dag = SerializedDAG.from_dict(serialized)
+
+        assert isinstance(dict.__getitem__(dag.params, "str"), Param)
+        assert dag.params["str"] == "str"
+
 
 def test_kubernetes_optional():
     """Serialisation / deserialisation continues to work without kubernetes installed"""

[airflow] 08/10: Prevent scheduler crash when serialized dag is missing (#19113)

Posted by je...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

jedcunningham pushed a commit to branch v2-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit 8223e8975197e6d96068127b4f5ff5b29ceafee5
Author: Ephraim Anierobi <sp...@gmail.com>
AuthorDate: Wed Oct 20 21:25:46 2021 +0100

    Prevent scheduler crash when serialized dag is missing (#19113)
    
    Scheduler._send_dag_callbacks_to_processor calls dag_run.get_dag which
    raises exception. This PR changes to calling dagbag.get_dag and changing
    Scheduler._send_dag_callbacks_to_processor args to accept dag instead of dag_run.
    
    (cherry picked from commit 5dc375aa7744f37c7a09f322cd9f4a221aa4ccbe)
---
 airflow/jobs/scheduler_job.py    | 14 ++++++++------
 tests/jobs/test_scheduler_job.py |  4 +---
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/airflow/jobs/scheduler_job.py b/airflow/jobs/scheduler_job.py
index 7684032..17fc55c 100644
--- a/airflow/jobs/scheduler_job.py
+++ b/airflow/jobs/scheduler_job.py
@@ -766,7 +766,12 @@ class SchedulerJob(BaseJob):
 
             # Send the callbacks after we commit to ensure the context is up to date when it gets run
             for dag_run, callback_to_run in callback_tuples:
-                self._send_dag_callbacks_to_processor(dag_run, callback_to_run)
+                dag = self.dagbag.get_dag(dag_run.dag_id, session=session)
+                if not dag:
+                    self.log.error("DAG '%s' not found in serialized_dag table", dag_run.dag_id)
+                    continue
+
+                self._send_dag_callbacks_to_processor(dag, callback_to_run)
 
             # Without this, the session has an invalid view of the DB
             session.expunge_all()
@@ -990,7 +995,7 @@ class SchedulerJob(BaseJob):
             )
 
             # Send SLA & DAG Success/Failure Callbacks to be executed
-            self._send_dag_callbacks_to_processor(dag_run, callback_to_execute)
+            self._send_dag_callbacks_to_processor(dag, callback_to_execute)
 
             return 0
 
@@ -1026,13 +1031,10 @@ class SchedulerJob(BaseJob):
         # Verify integrity also takes care of session.flush
         dag_run.verify_integrity(session=session)
 
-    def _send_dag_callbacks_to_processor(
-        self, dag_run: DagRun, callback: Optional[DagCallbackRequest] = None
-    ):
+    def _send_dag_callbacks_to_processor(self, dag: DAG, callback: Optional[DagCallbackRequest] = None):
         if not self.processor_agent:
             raise ValueError("Processor agent is not started.")
 
-        dag = dag_run.get_dag()
         self._send_sla_callbacks_to_processor(dag)
         if callback:
             self.processor_agent.send_callback_to_execute(callback)
diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py
index 3813265..11a2fca 100644
--- a/tests/jobs/test_scheduler_job.py
+++ b/tests/jobs/test_scheduler_job.py
@@ -1363,7 +1363,6 @@ class TestSchedulerJob:
         self.scheduler_job._send_dag_callbacks_to_processor.assert_called_once()
         call_args = self.scheduler_job._send_dag_callbacks_to_processor.call_args[0]
         assert call_args[0].dag_id == dr.dag_id
-        assert call_args[0].execution_date == dr.execution_date
         assert call_args[1] is None
 
         session.rollback()
@@ -1394,11 +1393,10 @@ class TestSchedulerJob:
         with mock.patch.object(settings, "USE_JOB_SCHEDULE", False):
             self.scheduler_job._do_scheduling(session)
 
-        # Verify Callback is not set (i.e is None) when no callbacks are set on DAG
+        # Verify Callback is set (i.e is None) when no callbacks are set on DAG
         self.scheduler_job._send_dag_callbacks_to_processor.assert_called_once()
         call_args = self.scheduler_job._send_dag_callbacks_to_processor.call_args[0]
         assert call_args[0].dag_id == dr.dag_id
-        assert call_args[0].execution_date == dr.execution_date
         assert call_args[1] is not None
         assert call_args[1].msg == msg
         session.rollback()

[airflow] 10/10: Update taskinstance REST API schema to include dag_run_id field (#19105)

Posted by je...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

jedcunningham pushed a commit to branch v2-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit e9c49265ef9b3c3e4fd2fa7caf29799c81a8dda5
Author: Ephraim Anierobi <sp...@gmail.com>
AuthorDate: Wed Oct 20 19:56:44 2021 +0100

    Update taskinstance REST API schema to include dag_run_id field (#19105)
    
    This PR adds dag_run_id field to taskinstance schema
    
    (cherry picked from commit cc3b062a2bdca16a7b239e73c4dc9e2a3a43c4f0)
---
 airflow/api_connexion/openapi/v1.yaml              |  6 +++
 .../api_connexion/schemas/task_instance_schema.py  | 47 +++++++++++++---------
 .../endpoints/test_task_instance_endpoint.py       |  3 ++
 .../schemas/test_task_instance_schema.py           |  2 +
 4 files changed, 38 insertions(+), 20 deletions(-)

diff --git a/airflow/api_connexion/openapi/v1.yaml b/airflow/api_connexion/openapi/v1.yaml
index a663e82..f34597e 100644
--- a/airflow/api_connexion/openapi/v1.yaml
+++ b/airflow/api_connexion/openapi/v1.yaml
@@ -2291,6 +2291,12 @@ components:
           type: string
         dag_id:
           type: string
+        dag_run_id:
+          type: string
+          description: |
+            The DagRun ID for this task instance
+
+            *New in version 2.2.1*
         execution_date:
           type: string
           format: datetime
diff --git a/airflow/api_connexion/schemas/task_instance_schema.py b/airflow/api_connexion/schemas/task_instance_schema.py
index 89ae9a6..956246b 100644
--- a/airflow/api_connexion/schemas/task_instance_schema.py
+++ b/airflow/api_connexion/schemas/task_instance_schema.py
@@ -19,6 +19,7 @@ from typing import List, NamedTuple, Optional, Tuple
 
 from marshmallow import Schema, ValidationError, fields, validate, validates_schema
 from marshmallow.utils import get_value
+from marshmallow_sqlalchemy import SQLAlchemySchema, auto_field
 
 from airflow.api_connexion.parameters import validate_istimezone
 from airflow.api_connexion.schemas.enum_schemas import TaskInstanceStateField
@@ -27,28 +28,34 @@ from airflow.models import SlaMiss, TaskInstance
 from airflow.utils.state import State
 
 
-class TaskInstanceSchema(Schema):
+class TaskInstanceSchema(SQLAlchemySchema):
     """Task instance schema"""
 
-    task_id = fields.Str()
-    dag_id = fields.Str()
-    execution_date = fields.DateTime(validate=validate_istimezone)
-    start_date = fields.DateTime(validate=validate_istimezone)
-    end_date = fields.DateTime(validate=validate_istimezone)
-    duration = fields.Float()
-    state = TaskInstanceStateField()
-    _try_number = fields.Int(data_key="try_number")
-    max_tries = fields.Int()
-    hostname = fields.Str()
-    unixname = fields.Str()
-    pool = fields.Str()
-    pool_slots = fields.Int()
-    queue = fields.Str()
-    priority_weight = fields.Int()
-    operator = fields.Str()
-    queued_dttm = fields.DateTime(data_key="queued_when")
-    pid = fields.Int()
-    executor_config = fields.Str()
+    class Meta:
+        """Meta"""
+
+        model = TaskInstance
+
+    task_id = auto_field()
+    dag_id = auto_field()
+    run_id = auto_field(data_key="dag_run_id")
+    execution_date = auto_field()
+    start_date = auto_field()
+    end_date = auto_field()
+    duration = auto_field()
+    state = auto_field()
+    _try_number = auto_field(data_key="try_number")
+    max_tries = auto_field()
+    hostname = auto_field()
+    unixname = auto_field()
+    pool = auto_field()
+    pool_slots = auto_field()
+    queue = auto_field()
+    priority_weight = auto_field()
+    operator = auto_field()
+    queued_dttm = auto_field(data_key="queued_when")
+    pid = auto_field()
+    executor_config = auto_field()
     sla_miss = fields.Nested(SlaMissSchema, dump_default=None)
 
     def get_attribute(self, obj, attr, default):
diff --git a/tests/api_connexion/endpoints/test_task_instance_endpoint.py b/tests/api_connexion/endpoints/test_task_instance_endpoint.py
index 719ae35..18b0bf1 100644
--- a/tests/api_connexion/endpoints/test_task_instance_endpoint.py
+++ b/tests/api_connexion/endpoints/test_task_instance_endpoint.py
@@ -197,6 +197,7 @@ class TestGetTaskInstance(TestTaskInstanceEndpoint):
             "task_id": "print_the_context",
             "try_number": 0,
             "unixname": getuser(),
+            "dag_run_id": "TEST_DAG_RUN_ID",
         }
 
     def test_should_respond_200_with_task_state_in_removed(self, session):
@@ -227,6 +228,7 @@ class TestGetTaskInstance(TestTaskInstanceEndpoint):
             "task_id": "print_the_context",
             "try_number": 0,
             "unixname": getuser(),
+            "dag_run_id": "TEST_DAG_RUN_ID",
         }
 
     def test_should_respond_200_task_instance_with_sla(self, session):
@@ -275,6 +277,7 @@ class TestGetTaskInstance(TestTaskInstanceEndpoint):
             "task_id": "print_the_context",
             "try_number": 0,
             "unixname": getuser(),
+            "dag_run_id": "TEST_DAG_RUN_ID",
         }
 
     def test_should_raises_401_unauthenticated(self):
diff --git a/tests/api_connexion/schemas/test_task_instance_schema.py b/tests/api_connexion/schemas/test_task_instance_schema.py
index 883d936..e13defb 100644
--- a/tests/api_connexion/schemas/test_task_instance_schema.py
+++ b/tests/api_connexion/schemas/test_task_instance_schema.py
@@ -88,6 +88,7 @@ class TestTaskInstanceSchema:
             "task_id": "TEST_TASK_ID",
             "try_number": 0,
             "unixname": getuser(),
+            "dag_run_id": None,
         }
         assert serialized_ti == expected_json
 
@@ -132,6 +133,7 @@ class TestTaskInstanceSchema:
             "task_id": "TEST_TASK_ID",
             "try_number": 0,
             "unixname": getuser(),
+            "dag_run_id": None,
         }
         assert serialized_ti == expected_json
 

[airflow] 09/10: Fix catchup by limiting queued dagrun creation using max_active_runs (#18897)

Posted by je...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

jedcunningham pushed a commit to branch v2-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit 69048f788983ae5fef5a1eacb8d1118ba92ae47a
Author: Ephraim Anierobi <sp...@gmail.com>
AuthorDate: Wed Oct 20 21:33:45 2021 +0100

    Fix catchup by limiting queued dagrun creation using max_active_runs (#18897)
    
    Currently, when catchup is True, we create a lot of dagruns limited by
    max_queued_runs_per_dag setting. This is not efficient as some dagruns takes
    longer to run.
    
    This PR brings back the old behavior of not creating dagruns once max_active_runs
    is reached thereby solving the catchup issue.
    
    Now, the dagruns appears as though they were created in running state
    
    (cherry picked from commit 05eea002fbd6dfc97ffc207f2a7f73ccf49aacde)
---
 UPDATING.md                                  |   5 ++
 airflow/config_templates/config.yml          |   8 --
 airflow/config_templates/default_airflow.cfg |   4 -
 airflow/jobs/scheduler_job.py                |  58 +++++++-------
 airflow/models/dag.py                        |  21 +++--
 airflow/models/dagrun.py                     |  18 ++++-
 tests/jobs/test_scheduler_job.py             | 110 +++++++++++++++++++++++----
 tests/models/test_dag.py                     |  24 +++---
 8 files changed, 179 insertions(+), 69 deletions(-)

diff --git a/UPDATING.md b/UPDATING.md
index 9c0c4fd..443afe9 100644
--- a/UPDATING.md
+++ b/UPDATING.md
@@ -91,6 +91,11 @@ Now if you resolve a ``Param`` without a default and don't pass a value, you wil
 Param().resolve()  # raises TypeError
 ```
 
+### `max_queued_runs_per_dag` configuration has been removed
+
+The `max_queued_runs_per_dag` configuration option in `[core]` section has been removed. Previously, this controlled the number of queued dagrun
+the scheduler can create in a dag. Now, the maximum number is controlled internally by the DAG's `max_active_runs`
+
 ## Airflow 2.2.0
 
 ### `worker_log_server_port` configuration has been moved to the ``logging`` section.
diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml
index 0e58f26..7aefbe0 100644
--- a/airflow/config_templates/config.yml
+++ b/airflow/config_templates/config.yml
@@ -195,14 +195,6 @@
       type: string
       example: ~
       default: "16"
-    - name: max_queued_runs_per_dag
-      description: |
-        The maximum number of queued dagruns for a single DAG. The scheduler will not create more DAG runs
-        if it reaches the limit. This is not configurable at the DAG level.
-      version_added: 2.1.4
-      type: string
-      example: ~
-      default: "16"
     - name: load_examples
       description: |
         Whether to load the DAG examples that ship with Airflow. It's good to
diff --git a/airflow/config_templates/default_airflow.cfg b/airflow/config_templates/default_airflow.cfg
index 41998ea..d03a06b 100644
--- a/airflow/config_templates/default_airflow.cfg
+++ b/airflow/config_templates/default_airflow.cfg
@@ -131,10 +131,6 @@ dags_are_paused_at_creation = True
 # which is defaulted as ``max_active_runs_per_dag``.
 max_active_runs_per_dag = 16
 
-# The maximum number of queued dagruns for a single DAG. The scheduler will not create more DAG runs
-# if it reaches the limit. This is not configurable at the DAG level.
-max_queued_runs_per_dag = 16
-
 # Whether to load the DAG examples that ship with Airflow. It's good to
 # get started, but you probably want to set this to ``False`` in a production
 # environment
diff --git a/airflow/jobs/scheduler_job.py b/airflow/jobs/scheduler_job.py
index 17fc55c..bf2410a 100644
--- a/airflow/jobs/scheduler_job.py
+++ b/airflow/jobs/scheduler_job.py
@@ -848,30 +848,19 @@ class SchedulerJob(BaseJob):
         existing_dagruns = (
             session.query(DagRun.dag_id, DagRun.execution_date).filter(existing_dagruns_filter).all()
         )
-        max_queued_dagruns = conf.getint('core', 'max_queued_runs_per_dag')
 
-        queued_runs_of_dags = defaultdict(
+        active_runs_of_dags = defaultdict(
             int,
-            session.query(DagRun.dag_id, func.count('*'))
-            .filter(  # We use `list` here because SQLA doesn't accept a set
-                # We use set to avoid duplicate dag_ids
-                DagRun.dag_id.in_(list({dm.dag_id for dm in dag_models})),
-                DagRun.state == State.QUEUED,
-            )
-            .group_by(DagRun.dag_id)
-            .all(),
+            DagRun.active_runs_of_dags(dag_ids=(dm.dag_id for dm in dag_models), session=session),
         )
 
         for dag_model in dag_models:
-            # Lets quickly check if we have exceeded the number of queued dagruns per dags
-            total_queued = queued_runs_of_dags[dag_model.dag_id]
-            if total_queued >= max_queued_dagruns:
-                continue
 
             dag = self.dagbag.get_dag(dag_model.dag_id, session=session)
             if not dag:
                 self.log.error("DAG '%s' not found in serialized_dag table", dag_model.dag_id)
                 continue
+
             dag_hash = self.dagbag.dags_hash.get(dag.dag_id)
 
             data_interval = dag.get_next_data_interval(dag_model)
@@ -894,12 +883,28 @@ class SchedulerJob(BaseJob):
                     dag_hash=dag_hash,
                     creating_job_id=self.id,
                 )
-                queued_runs_of_dags[dag_model.dag_id] += 1
-            dag_model.calculate_dagrun_date_fields(dag, data_interval)
-
+                active_runs_of_dags[dag.dag_id] += 1
+            self._update_dag_next_dagruns(dag, dag_model, active_runs_of_dags[dag.dag_id])
         # TODO[HA]: Should we do a session.flush() so we don't have to keep lots of state/object in
         # memory for larger dags? or expunge_all()
 
+    def _update_dag_next_dagruns(self, dag, dag_model: DagModel, total_active_runs) -> None:
+        """
+        Update the next_dagrun, next_dagrun_data_interval_start/end
+        and next_dagrun_create_after for this dag.
+        """
+        if total_active_runs >= dag_model.max_active_runs:
+            self.log.info(
+                "DAG %s is at (or above) max_active_runs (%d of %d), not creating any more runs",
+                dag_model.dag_id,
+                total_active_runs,
+                dag_model.max_active_runs,
+            )
+            dag_model.next_dagrun_create_after = None
+        else:
+            data_interval = dag.get_next_data_interval(dag_model)
+            dag_model.calculate_dagrun_date_fields(dag, data_interval)
+
     def _start_queued_dagruns(
         self,
         session: Session,
@@ -908,15 +913,8 @@ class SchedulerJob(BaseJob):
         dag_runs = self._get_next_dagruns_to_examine(State.QUEUED, session)
 
         active_runs_of_dags = defaultdict(
-            lambda: 0,
-            session.query(DagRun.dag_id, func.count('*'))
-            .filter(  # We use `list` here because SQLA doesn't accept a set
-                # We use set to avoid duplicate dag_ids
-                DagRun.dag_id.in_(list({dr.dag_id for dr in dag_runs})),
-                DagRun.state == State.RUNNING,
-            )
-            .group_by(DagRun.dag_id)
-            .all(),
+            int,
+            DagRun.active_runs_of_dags((dr.dag_id for dr in dag_runs), only_running=True, session=session),
         )
 
         def _update_state(dag: DAG, dag_run: DagRun):
@@ -967,6 +965,7 @@ class SchedulerJob(BaseJob):
         if not dag:
             self.log.error("Couldn't find dag %s in DagBag/DB!", dag_run.dag_id)
             return 0
+        dag_model = DM.get_dagmodel(dag.dag_id, session)
 
         if (
             dag_run.start_date
@@ -985,6 +984,9 @@ class SchedulerJob(BaseJob):
                 session.merge(task_instance)
             session.flush()
             self.log.info("Run %s of %s has timed-out", dag_run.run_id, dag_run.dag_id)
+            active_runs = dag.get_num_active_runs(only_running=False, session=session)
+            # Work out if we should allow creating a new DagRun now?
+            self._update_dag_next_dagruns(dag, dag_model, active_runs)
 
             callback_to_execute = DagCallbackRequest(
                 full_filepath=dag.fileloc,
@@ -1006,6 +1008,10 @@ class SchedulerJob(BaseJob):
         self._verify_integrity_if_dag_changed(dag_run=dag_run, session=session)
         # TODO[HA]: Rename update_state -> schedule_dag_run, ?? something else?
         schedulable_tis, callback_to_run = dag_run.update_state(session=session, execute_callbacks=False)
+        if dag_run.state in State.finished:
+            active_runs = dag.get_num_active_runs(only_running=False, session=session)
+            # Work out if we should allow creating a new DagRun now?
+            self._update_dag_next_dagruns(dag, dag_model, active_runs)
 
         # This will do one query per dag run. We "could" build up a complex
         # query to update all the TIs across all the execution dates and dag
diff --git a/airflow/models/dag.py b/airflow/models/dag.py
index b372cfe..55be0ec 100644
--- a/airflow/models/dag.py
+++ b/airflow/models/dag.py
@@ -1138,7 +1138,7 @@ class DAG(LoggingMixin):
         return active_dates
 
     @provide_session
-    def get_num_active_runs(self, external_trigger=None, session=None):
+    def get_num_active_runs(self, external_trigger=None, only_running=True, session=None):
         """
         Returns the number of active "running" dag runs
 
@@ -1148,11 +1148,11 @@ class DAG(LoggingMixin):
         :return: number greater than 0 for active dag runs
         """
         # .count() is inefficient
-        query = (
-            session.query(func.count())
-            .filter(DagRun.dag_id == self.dag_id)
-            .filter(DagRun.state == State.RUNNING)
-        )
+        query = session.query(func.count()).filter(DagRun.dag_id == self.dag_id)
+        if only_running:
+            query = query.filter(DagRun.state == State.RUNNING)
+        else:
+            query = query.filter(DagRun.state.in_({State.RUNNING, State.QUEUED}))
 
         if external_trigger is not None:
             query = query.filter(
@@ -2425,6 +2425,10 @@ class DAG(LoggingMixin):
         )
         most_recent_runs = {run.dag_id: run for run in most_recent_runs_iter}
 
+        # Get number of active dagruns for all dags we are processing as a single query.
+
+        num_active_runs = DagRun.active_runs_of_dags(dag_ids=existing_dag_ids, session=session)
+
         filelocs = []
 
         for orm_dag in sorted(orm_dags, key=lambda d: d.dag_id):
@@ -2453,7 +2457,10 @@ class DAG(LoggingMixin):
                 data_interval = None
             else:
                 data_interval = dag.get_run_data_interval(run)
-            orm_dag.calculate_dagrun_date_fields(dag, data_interval)
+            if num_active_runs.get(dag.dag_id, 0) >= orm_dag.max_active_runs:
+                orm_dag.next_dagrun_create_after = None
+            else:
+                orm_dag.calculate_dagrun_date_fields(dag, data_interval)
 
             for orm_tag in list(orm_dag.tags):
                 if orm_tag.name not in set(dag.tags):
diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py
index 2b651c5..1d53265 100644
--- a/airflow/models/dagrun.py
+++ b/airflow/models/dagrun.py
@@ -17,7 +17,7 @@
 # under the License.
 import warnings
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Iterable, List, NamedTuple, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, NamedTuple, Optional, Tuple, Union
 
 from sqlalchemy import (
     Boolean,
@@ -208,6 +208,22 @@ class DagRun(Base, LoggingMixin):
         self.state = dr.state
 
     @classmethod
+    @provide_session
+    def active_runs_of_dags(cls, dag_ids=None, only_running=False, session=None) -> Dict[str, int]:
+        """Get the number of active dag runs for each dag."""
+        query = session.query(cls.dag_id, func.count('*'))
+        if dag_ids is not None:
+            # 'set' called to avoid duplicate dag_ids, but converted back to 'list'
+            # because SQLAlchemy doesn't accept a set here.
+            query = query.filter(cls.dag_id.in_(list(set(dag_ids))))
+        if only_running:
+            query = query.filter(cls.state == State.RUNNING)
+        else:
+            query = query.filter(cls.state.in_([State.RUNNING, State.QUEUED]))
+        query = query.group_by(cls.dag_id)
+        return {dag_id: count for dag_id, count in query.all()}
+
+    @classmethod
     def next_dagruns_to_examine(
         cls,
         state: DagRunState,
diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py
index 11a2fca..16b2921 100644
--- a/tests/jobs/test_scheduler_job.py
+++ b/tests/jobs/test_scheduler_job.py
@@ -49,7 +49,7 @@ from airflow.utils import timezone
 from airflow.utils.callback_requests import DagCallbackRequest
 from airflow.utils.file import list_py_file_paths
 from airflow.utils.session import create_session, provide_session
-from airflow.utils.state import State
+from airflow.utils.state import DagRunState, State, TaskInstanceState
 from airflow.utils.types import DagRunType
 from tests.test_utils.asserts import assert_queries_count
 from tests.test_utils.config import conf_vars, env_vars
@@ -1105,9 +1105,9 @@ class TestSchedulerJob:
         self.scheduler_job.executor.end.assert_called_once()
         mock_processor_agent.return_value.end.reset_mock(side_effect=True)
 
-    def test_theres_limit_to_queued_dagruns_in_a_dag(self, dag_maker):
-        """This tests that there's limit to the number of queued dagrun scheduler can create in a dag"""
-        with dag_maker() as dag:
+    def test_queued_dagruns_stops_creating_when_max_active_is_reached(self, dag_maker):
+        """This tests that queued dagruns stops creating once max_active_runs is reached"""
+        with dag_maker(max_active_runs=10) as dag:
             DummyOperator(task_id='mytask')
 
         session = settings.Session()
@@ -1122,13 +1122,20 @@ class TestSchedulerJob:
         assert orm_dag is not None
         for _ in range(20):
             self.scheduler_job._create_dag_runs([orm_dag], session)
-        assert session.query(DagRun).count() == 16
+        drs = session.query(DagRun).all()
+        assert len(drs) == 10
 
-        with conf_vars({('core', 'max_queued_runs_per_dag'): '5'}):
-            clear_db_runs()
-            for i in range(20):
-                self.scheduler_job._create_dag_runs([orm_dag], session)
-        assert session.query(DagRun).count() == 5
+        for dr in drs:
+            dr.state = State.RUNNING
+            session.merge(dr)
+        session.commit()
+        assert session.query(DagRun.state).filter(DagRun.state == State.RUNNING).count() == 10
+        for _ in range(20):
+            self.scheduler_job._create_dag_runs([orm_dag], session)
+        assert session.query(DagRun).count() == 10
+        assert session.query(DagRun.state).filter(DagRun.state == State.RUNNING).count() == 10
+        assert session.query(DagRun.state).filter(DagRun.state == State.QUEUED).count() == 0
+        assert orm_dag.next_dagrun_create_after is None
 
     def test_dagrun_timeout_verify_max_active_runs(self, dag_maker):
         """
@@ -1160,9 +1167,7 @@ class TestSchedulerJob:
         assert len(drs) == 1
         dr = drs[0]
 
-        # This should have a value since we control max_active_runs
-        # by DagRun State.
-        assert orm_dag.next_dagrun_create_after
+        assert orm_dag.next_dagrun_create_after is None
         # But we should record the date of _what run_ it would be
         assert isinstance(orm_dag.next_dagrun, datetime.datetime)
         assert isinstance(orm_dag.next_dagrun_data_interval_start, datetime.datetime)
@@ -1242,6 +1247,40 @@ class TestSchedulerJob:
         session.rollback()
         session.close()
 
+    def test_dagrun_timeout_fails_run_and_update_next_dagrun(self, dag_maker):
+        """
+        Test that dagrun timeout fails run and update the next dagrun
+        """
+        session = settings.Session()
+        with dag_maker(
+            max_active_runs=1,
+            dag_id='test_scheduler_fail_dagrun_timeout',
+            dagrun_timeout=datetime.timedelta(seconds=60),
+        ):
+            DummyOperator(task_id='dummy')
+
+        dr = dag_maker.create_dagrun(start_date=timezone.utcnow() - datetime.timedelta(days=1))
+        # check that next_dagrun is dr.execution_date
+        dag_maker.dag_model.next_dagrun == dr.execution_date
+        self.scheduler_job = SchedulerJob(subdir=os.devnull)
+        self.scheduler_job.dagbag = dag_maker.dagbag
+
+        # Mock that processor_agent is started
+        self.scheduler_job.processor_agent = mock.Mock()
+        self.scheduler_job.processor_agent.send_callback_to_execute = mock.Mock()
+
+        self.scheduler_job._schedule_dag_run(dr, session)
+        session.flush()
+        session.refresh(dr)
+        assert dr.state == State.FAILED
+        # check that next_dagrun has been updated by Schedulerjob._update_dag_next_dagruns
+        assert dag_maker.dag_model.next_dagrun == dr.execution_date + timedelta(days=1)
+        # check that no running/queued runs yet
+        assert (
+            session.query(DagRun).filter(DagRun.state.in_([DagRunState.RUNNING, DagRunState.QUEUED])).count()
+            == 0
+        )
+
     @pytest.mark.parametrize(
         "state, expected_callback_msg", [(State.SUCCESS, "success"), (State.FAILED, "task_failure")]
     )
@@ -2756,6 +2795,51 @@ class TestSchedulerJob:
         ti.refresh_from_db(session=session)
         assert ti.state == State.QUEUED
 
+    def test_more_runs_are_not_created_when_max_active_runs_is_reached(self, dag_maker, caplog):
+        """
+        This tests that when max_active_runs is reached, _create_dag_runs doesn't create
+        more dagruns
+        """
+        with dag_maker(max_active_runs=1):
+            DummyOperator(task_id='task')
+        self.scheduler_job = SchedulerJob(subdir=os.devnull)
+        self.scheduler_job.executor = MockExecutor(do_update=False)
+        self.scheduler_job.processor_agent = mock.MagicMock(spec=DagFileProcessorAgent)
+        session = settings.Session()
+        assert session.query(DagRun).count() == 0
+        dag_models = DagModel.dags_needing_dagruns(session).all()
+        self.scheduler_job._create_dag_runs(dag_models, session)
+        dr = session.query(DagRun).one()
+        dr.state == DagRunState.QUEUED
+        assert session.query(DagRun).count() == 1
+        assert dag_maker.dag_model.next_dagrun_create_after is None
+        session.flush()
+        # dags_needing_dagruns query should not return any value
+        assert len(DagModel.dags_needing_dagruns(session).all()) == 0
+        self.scheduler_job._create_dag_runs(dag_models, session)
+        assert session.query(DagRun).count() == 1
+        assert dag_maker.dag_model.next_dagrun_create_after is None
+        assert dag_maker.dag_model.next_dagrun == DEFAULT_DATE
+        # set dagrun to success
+        dr = session.query(DagRun).one()
+        dr.state = DagRunState.SUCCESS
+        ti = dr.get_task_instance('task', session)
+        ti.state = TaskInstanceState.SUCCESS
+        session.merge(ti)
+        session.merge(dr)
+        session.flush()
+        # check that next_dagrun is set properly by Schedulerjob._update_dag_next_dagruns
+        self.scheduler_job._schedule_dag_run(dr, session)
+        session.flush()
+        assert len(DagModel.dags_needing_dagruns(session).all()) == 1
+        # assert next_dagrun has been updated correctly
+        assert dag_maker.dag_model.next_dagrun == DEFAULT_DATE + timedelta(days=1)
+        # assert no dagruns is created yet
+        assert (
+            session.query(DagRun).filter(DagRun.state.in_([DagRunState.RUNNING, DagRunState.QUEUED])).count()
+            == 0
+        )
+
     def test_do_schedule_max_active_runs_and_manual_trigger(self, dag_maker):
         """
         Make sure that when a DAG is already at max_active_runs, that manually triggered
diff --git a/tests/models/test_dag.py b/tests/models/test_dag.py
index dfa1a0a..5bea2b7 100644
--- a/tests/models/test_dag.py
+++ b/tests/models/test_dag.py
@@ -726,7 +726,7 @@ class TestDag(unittest.TestCase):
         clear_db_dags()
         dags = [DAG(f'dag-bulk-sync-{i}', start_date=DEFAULT_DATE, tags=["test-dag"]) for i in range(0, 4)]
 
-        with assert_queries_count(4):
+        with assert_queries_count(5):
             DAG.bulk_write_to_db(dags)
         with create_session() as session:
             assert {'dag-bulk-sync-0', 'dag-bulk-sync-1', 'dag-bulk-sync-2', 'dag-bulk-sync-3'} == {
@@ -743,14 +743,14 @@ class TestDag(unittest.TestCase):
                 assert row[0] is not None
 
         # Re-sync should do fewer queries
-        with assert_queries_count(3):
+        with assert_queries_count(4):
             DAG.bulk_write_to_db(dags)
-        with assert_queries_count(3):
+        with assert_queries_count(4):
             DAG.bulk_write_to_db(dags)
         # Adding tags
         for dag in dags:
             dag.tags.append("test-dag2")
-        with assert_queries_count(4):
+        with assert_queries_count(5):
             DAG.bulk_write_to_db(dags)
         with create_session() as session:
             assert {'dag-bulk-sync-0', 'dag-bulk-sync-1', 'dag-bulk-sync-2', 'dag-bulk-sync-3'} == {
@@ -769,7 +769,7 @@ class TestDag(unittest.TestCase):
         # Removing tags
         for dag in dags:
             dag.tags.remove("test-dag")
-        with assert_queries_count(4):
+        with assert_queries_count(5):
             DAG.bulk_write_to_db(dags)
         with create_session() as session:
             assert {'dag-bulk-sync-0', 'dag-bulk-sync-1', 'dag-bulk-sync-2', 'dag-bulk-sync-3'} == {
@@ -785,7 +785,8 @@ class TestDag(unittest.TestCase):
             for row in session.query(DagModel.last_parsed_time).all():
                 assert row[0] is not None
 
-    def test_bulk_write_to_db_max_active_runs(self):
+    @parameterized.expand([State.RUNNING, State.QUEUED])
+    def test_bulk_write_to_db_max_active_runs(self, state):
         """
         Test that DagModel.next_dagrun_create_after is set to NULL when the dag cannot be created due to max
         active runs being hit.
@@ -805,7 +806,7 @@ class TestDag(unittest.TestCase):
         assert model.next_dagrun_create_after == DEFAULT_DATE + timedelta(days=1)
 
         dr = dag.create_dagrun(
-            state=State.RUNNING,
+            state=state,
             execution_date=model.next_dagrun,
             run_type=DagRunType.SCHEDULED,
             session=session,
@@ -814,9 +815,12 @@ class TestDag(unittest.TestCase):
         DAG.bulk_write_to_db([dag])
 
         model = session.query(DagModel).get((dag.dag_id,))
-        assert model.next_dagrun == DEFAULT_DATE + timedelta(days=1)
-        # Next dagrun after is not None because the dagrun would be in queued state
-        assert model.next_dagrun_create_after is not None
+        # We signal "at max active runs" by saying this run is never eligible to be created
+        assert model.next_dagrun_create_after is None
+        # test that bulk_write_to_db again doesn't update next_dagrun_create_after
+        DAG.bulk_write_to_db([dag])
+        model = session.query(DagModel).get((dag.dag_id,))
+        assert model.next_dagrun_create_after is None
 
     def test_sync_to_db(self):
         dag = DAG(

[airflow] 07/10: Don't install SQLAlchemy/Pendulum adapters for other DBs (#18745)

Posted by je...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

jedcunningham pushed a commit to branch v2-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit 478f94b843a2963fde3150c9eec8846d797abd9a
Author: Ash Berlin-Taylor <as...@firemirror.com>
AuthorDate: Wed Oct 20 19:56:10 2021 +0100

    Don't install SQLAlchemy/Pendulum adapters for other DBs (#18745)
    
    This stops the MySQL libs being imported "unnecessarily" when Postgres
    is in use -- and there have been a few confusing reports of the mysql
    client libs causing problems in rare cases, so lets avoid the import if
    we can.
    
    (cherry picked from commit d75cf4d60ddbff5b88bfe348cb83f9d173187744)
---
 airflow/settings.py | 26 +++++++++++++-------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/airflow/settings.py b/airflow/settings.py
index 9a456ba..7f5a255 100644
--- a/airflow/settings.py
+++ b/airflow/settings.py
@@ -344,24 +344,24 @@ def configure_adapters():
     """Register Adapters and DB Converters"""
     from pendulum import DateTime as Pendulum
 
-    try:
+    if SQL_ALCHEMY_CONN.startswith('sqlite'):
         from sqlite3 import register_adapter
 
         register_adapter(Pendulum, lambda val: val.isoformat(' '))
-    except ImportError:
-        pass
-    try:
-        import MySQLdb.converters
 
-        MySQLdb.converters.conversions[Pendulum] = MySQLdb.converters.DateTime2literal
-    except ImportError:
-        pass
-    try:
-        import pymysql.converters
+    if SQL_ALCHEMY_CONN.startswith('mysql'):
+        try:
+            import MySQLdb.converters
 
-        pymysql.converters.conversions[Pendulum] = pymysql.converters.escape_datetime
-    except ImportError:
-        pass
+            MySQLdb.converters.conversions[Pendulum] = MySQLdb.converters.DateTime2literal
+        except ImportError:
+            pass
+        try:
+            import pymysql.converters
+
+            pymysql.converters.conversions[Pendulum] = pymysql.converters.escape_datetime
+        except ImportError:
+            pass
 
 
 def validate_session():

[airflow] 02/10: Rename trigger page label to Logical Date (#19061)

Posted by je...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

jedcunningham pushed a commit to branch v2-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit c28affb7cfc7ce522fb93e4a749a38a9181cbddb
Author: Tzu-ping Chung <tp...@astronomer.io>
AuthorDate: Wed Oct 20 05:01:30 2021 +0800

    Rename trigger page label to Logical Date (#19061)
    
    (cherry picked from commit eb5cae0757a52774906f1bd9aa4326208126fc5f)
---
 airflow/www/templates/airflow/trigger.html | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/www/templates/airflow/trigger.html b/airflow/www/templates/airflow/trigger.html
index fb1fd88..efc1650 100644
--- a/airflow/www/templates/airflow/trigger.html
+++ b/airflow/www/templates/airflow/trigger.html
@@ -36,7 +36,7 @@
     <input type="hidden" name="dag_id" value="{{ dag_id }}">
     <input type="hidden" name="origin" value="{{ origin }}">
   <div class="form-group">
-      <label class="sr-only" for="execution_date">Execution date</label>
+      <label class="sr-only" for="execution_date">Logical date</label>
       <div class="input-group">
         {{ form.execution_date(class_="form-control", disabled=False) }}
       </div>

[airflow] 04/10: Ensure task state doesn't change when marked as failed/success/skipped (#19095)

Posted by je...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

jedcunningham pushed a commit to branch v2-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit bbe9012b9fb8929a2a971529b78f508a4817e85a
Author: Ephraim Anierobi <sp...@gmail.com>
AuthorDate: Wed Oct 20 15:10:39 2021 +0100

    Ensure task state doesn't change when marked as failed/success/skipped (#19095)
    
    Currently, when a dagrun is marked as success in the UI, the expected
    sigterm is sent but the task continues running, changing the dagrun to running
    before eventually failing. Same as well when marked as failed.
    
    Also, queued tasks continue running even when the dagrun was failed
    
    The same thing happens when a dagrun times out. The tasks that are marked
    skipped, starts again, and then fails.
    
    This PR fixes these issues
    
    (cherry picked from commit 8d703ae7db3c2a08b94c824a6f4287c3dd29cebf)
---
 airflow/jobs/scheduler_job.py     |  3 +++
 airflow/models/taskinstance.py    |  7 +++++--
 tests/jobs/test_scheduler_job.py  | 23 +++++++++++++++++++++++
 tests/models/test_taskinstance.py | 25 ++++++++++++++++++++++++-
 4 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/airflow/jobs/scheduler_job.py b/airflow/jobs/scheduler_job.py
index 1e74e99..7684032 100644
--- a/airflow/jobs/scheduler_job.py
+++ b/airflow/jobs/scheduler_job.py
@@ -432,6 +432,9 @@ class SchedulerJob(BaseJob):
         """
         # actually enqueue them
         for ti in task_instances:
+            if ti.dag_run.state in State.finished:
+                ti.set_state(State.NONE)
+                continue
             command = ti.command_as_list(
                 local=True,
                 pickle_id=ti.dag_model.pickle_id,
diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index 2563109..46945c7 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -1363,9 +1363,12 @@ class TaskInstance(Base, LoggingMixin):
             session.commit()
             raise
         except AirflowException as e:
+            if not test_mode:
+                self.refresh_from_db(lock_for_update=True, session=session)
             # for case when task is marked as success/failed externally
-            # current behavior doesn't hit the success callback
-            if self.state in {State.SUCCESS, State.FAILED}:
+            # or dagrun timed out and task is marked as skipped
+            # current behavior doesn't hit the callbacks
+            if self.state in State.finished:
                 return
             else:
                 self.handle_failure(e, test_mode, error_file=error_file, session=session)
diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py
index 87ebd51..3813265 100644
--- a/tests/jobs/test_scheduler_job.py
+++ b/tests/jobs/test_scheduler_job.py
@@ -874,6 +874,29 @@ class TestSchedulerJob:
         assert mock_queue_command.called
         session.rollback()
 
+    @pytest.mark.parametrize("state", [State.FAILED, State.SUCCESS])
+    def test_enqueue_task_instances_sets_ti_state_to_None_if_dagrun_in_finish_state(self, state, dag_maker):
+        """This tests that task instances whose dagrun is in finished state are not queued"""
+        dag_id = 'SchedulerJobTest.test_enqueue_task_instances_with_queued_state'
+        task_id_1 = 'dummy'
+        session = settings.Session()
+        with dag_maker(dag_id=dag_id, start_date=DEFAULT_DATE, session=session):
+            task1 = DummyOperator(task_id=task_id_1)
+
+        self.scheduler_job = SchedulerJob(subdir=os.devnull)
+
+        dr1 = dag_maker.create_dagrun(state=state)
+        ti = dr1.get_task_instance(task1.task_id, session)
+        ti.state = State.SCHEDULED
+        session.merge(ti)
+        session.commit()
+
+        with patch.object(BaseExecutor, 'queue_command') as mock_queue_command:
+            self.scheduler_job._enqueue_task_instances_with_queued_state([ti])
+        ti.refresh_from_db()
+        assert ti.state == State.NONE
+        mock_queue_command.assert_not_called()
+
     def test_critical_section_execute_task_instances(self, dag_maker):
         dag_id = 'SchedulerJobTest.test_execute_task_instances'
         task_id_1 = 'dummy_task'
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 4f7d3f9..bf58aac 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -439,7 +439,6 @@ class TestTaskInstance:
         """
 
         def task_function(ti):
-            # pylint: disable=unused-argument
             os.kill(ti.pid, signal.SIGTERM)
 
         with dag_maker('test_mark_failure_2'):
@@ -458,6 +457,30 @@ class TestTaskInstance:
         ti.refresh_from_db()
         assert ti.state == State.UP_FOR_RETRY
 
+    @pytest.mark.parametrize("state", [State.SUCCESS, State.FAILED, State.SKIPPED])
+    def test_task_sigterm_doesnt_change_state_of_finished_tasks(self, state, dag_maker):
+        session = settings.Session()
+
+        def task_function(ti):
+            ti.state = state
+            session.merge(ti)
+            session.commit()
+            raise AirflowException()
+
+        with dag_maker('test_mark_failure_2'):
+            task = PythonOperator(
+                task_id='test_on_failure',
+                python_callable=task_function,
+            )
+
+        dr = dag_maker.create_dagrun()
+        ti = dr.task_instances[0]
+        ti.task = task
+
+        ti.run()
+        ti.refresh_from_db()
+        ti.state == state
+
     def test_task_retry_wipes_next_fields(self, session, dag_maker):
         """
         Test that ensures that tasks wipe their next_method and next_kwargs

[airflow] 06/10: Workaround libstdcpp TLS error (#19010)

Posted by je...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

jedcunningham pushed a commit to branch v2-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit f3f0eb2fd219655a4aa37d579a8480313af0d1e9
Author: Jarek Potiuk <ja...@potiuk.com>
AuthorDate: Tue Oct 19 00:03:58 2021 -0100

    Workaround libstdcpp TLS error (#19010)
    
    Workaround https://github.com/apache/airflow/issues/17546 issue with
    /usr/lib/x86_64-linux-gnu/libstdc++.so.6: cannot allocate memory in
    static TLS block.
    
    We do not yet a more "correct" solution to the problem but in order to
    avoid raising new issues by users of the prod image, we implement the
    workaround now.
    
    The side effect of this is slightly (in the range of 100s of
    milliseconds) slower load for any binary started and a little memory
    used for Heap allocated by initialization of libstdc++.
    
    This overhead is not happening for binaries that already link
    dynamically libstdc++.
    
    (cherry picked from commit 88583095c408ef9ea60f793e7072e3fd4b88e329)
---
 Dockerfile    |  9 +++++++++
 Dockerfile.ci | 10 ++++++++++
 2 files changed, 19 insertions(+)

diff --git a/Dockerfile b/Dockerfile
index 4e4bbe8..088d561 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -487,5 +487,14 @@ LABEL org.apache.airflow.distro="debian" \
 # to learn more about the way how signals are handled by the image
 ENV DUMB_INIT_SETSID="1"
 
+# This one is to workaround https://github.com/apache/airflow/issues/17546
+# issue with /usr/lib/x86_64-linux-gnu/libstdc++.so.6: cannot allocate memory in static TLS block
+# We do not yet a more "correct" solution to the problem but in order to avoid raising new issues
+# by users of the prod image, we implement the workaround now.
+# The side effect of this is slightly (in the range of 100s of milliseconds) slower load for any
+# binary started and a little memory used for Heap allocated by initialization of libstdc++
+# This overhead is not happening for binaries that already link dynamically libstdc++
+ENV LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libstdc++.so.6"
+
 ENTRYPOINT ["/usr/bin/dumb-init", "--", "/entrypoint"]
 CMD []
diff --git a/Dockerfile.ci b/Dockerfile.ci
index 0e665f3..b4eb465 100644
--- a/Dockerfile.ci
+++ b/Dockerfile.ci
@@ -400,6 +400,16 @@ LABEL org.apache.airflow.distro="debian" \
   org.opencontainers.image.title="Continuous Integration Airflow Image" \
   org.opencontainers.image.description="Installed Apache Airflow with Continuous Integration dependencies"
 
+# This one is to workaround https://github.com/apache/airflow/issues/17546
+# issue with /usr/lib/x86_64-linux-gnu/libstdc++.so.6: cannot allocate memory in static TLS block
+# We do not yet a more "correct" solution to the problem but in order to avoid raising new issues
+# by users of the prod image, we implement the workaround now.
+# The side effect of this is slightly (in the range of 100s of milliseconds) slower load for any
+# binary started and a little memory used for Heap allocated by initialization of libstdc++
+# This overhead is not happening for binaries that already link dynamically libstdc++
+ENV LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libstdc++.so.6"
+
+
 EXPOSE 8080
 
 ENTRYPOINT ["/usr/bin/dumb-init", "--", "/entrypoint"]