You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by as...@apache.org on 2022/01/07 22:11:45 UTC
[airflow] branch main updated: Speed up creation of DagRun for large DAGs (5k+ tasks) by 25-130% (#20722)
This is an automated email from the ASF dual-hosted git repository.
ash pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f2039b4 Speed up creation of DagRun for large DAGs (5k+ tasks) by 25-130% (#20722)
f2039b4 is described below
commit f2039b4c9e15b514661d4facbd710791fe0a2ef4
Author: Ash Berlin-Taylor <as...@firemirror.com>
AuthorDate: Fri Jan 7 22:11:03 2022 +0000
Speed up creation of DagRun for large DAGs (5k+ tasks) by 25-130% (#20722)
* Speed up creation of DagRun for large DAGs (5k+ tasks) by 15-40%
This uses the "bulk" operation API of SQLAlchemy to get a big speed
up. Due to the `task_instance_mutation_hook` we still need to keep
actual TaskInstance objects around.
For postgresql we have enabled to "batch operation helpers"[1] which
makes it even faster. The default page sizes are chosen somewhat
randomly based on the SQLA docs.
To make these options configurable I have added (and used here and in
KubeConfig) a new `getjson` option to AirflowConfigParser class.
Postgresql is over 77% faster with bulk_save_objects:
Before:
```
number_of_tis=1 mean=0.004397215199423954 per=0.004397215199423954 times=[0.009390181003254838, 0.002814065999700688, 0.00284132499655243, 0.0036120269942330196, 0.0033284770033787936]
number_of_tis=10 mean=0.008078816600027494 per=0.0008078816600027494 times=[0.011014281000825576, 0.008476420000079088, 0.00741832799394615, 0.006857775995740667, 0.006627278009545989]
number_of_tis=50 mean=0.01927847799670417 per=0.00038556955993408336 times=[0.02556803499464877, 0.01935569499619305, 0.01662322599440813, 0.01840184700267855, 0.01644358699559234]
number_of_tis=100 mean=0.03301511880126782 per=0.00033015118801267817 times=[0.04117956099798903, 0.030890661000739783, 0.03007458901265636, 0.03125198099587578, 0.03167880199907813]
number_of_tis=500 mean=0.15320950179593637 per=0.0003064190035918727 times=[0.20054609200451523, 0.14052859699586406, 0.14509809199080337, 0.1365471329918364, 0.1433275949966628]
number_of_tis=1000 mean=0.2929377429973101 per=0.0002929377429973101 times=[0.3517978919990128, 0.2807794280088274, 0.2806490379880415, 0.27710555399244186, 0.27435680299822707]
number_of_tis=3000 mean=0.9935687056015012 per=0.00033118956853383374 times=[1.2047388390055858, 0.8248025969951414, 0.8685875020019012, 0.9017027500085533, 1.1680118399963249]
number_of_tis=5000 mean=1.5349355740036117 per=0.00030698711480072236 times=[1.8663743910001358, 1.5182018500054255, 1.5446484510030132, 1.3932801040064078, 1.3521730740030762]
number_of_tis=10000 mean=3.7448632712010292 per=0.0003744863271201029 times=[4.135914924001554, 3.4411147559876554, 3.526543836007477, 3.7195197630062466, 3.9012230770022143]
number_of_tis=15000 mean=6.3099766838044165 per=0.00042066511225362775 times=[6.552250057997298, 6.1369703890086384, 6.8749958210100885, 6.067943914007628, 5.917723236998427]
number_of_tis=20000 mean=8.317583500797628 per=0.00041587917503988143 times=[8.720249108009739, 8.0188543760014, 8.328030352990027, 8.398350054994808, 8.122433611992165]
```
When using bulk_save_objects:
```
number_of_tis=20000 mean=4.678154367001843 per=0.00023390771835009216 times=[4.465847548010061, 4.571855771995615, 4.749505186002352, 4.724330568002188, 4.8792327609990025]
```
MySQL is only 10-15% faster (and a lot noisier)
Before:
```
number_of_tis=1 mean=0.006164804595755413 per=0.006164804595755413 times=[0.013516580002033152, 0.00427598599344492, 0.004508020996581763, 0.004067091998877004, 0.004456343987840228]
number_of_tis=10 mean=0.007822793803643435 per=0.0007822793803643434 times=[0.0081135170039488, 0.00719467100861948, 0.009007985994685441, 0.00758794900320936, 0.007209846007754095]
number_of_tis=50 mean=0.020377356800599954 per=0.00040754713601199905 times=[0.02612382399092894, 0.018950315003166907, 0.019109474000288174, 0.018008680999628268, 0.019694490008987486]
number_of_tis=100 mean=0.040682651600218375 per=0.00040682651600218374 times=[0.05449078499805182, 0.037430580996442586, 0.039291110006161034, 0.03625023599306587, 0.035950546007370576]
number_of_tis=500 mean=0.18646696420037187 per=0.00037293392840074375 times=[0.24278165798750706, 0.17090376401029062, 0.1837275660072919, 0.16893767600413412, 0.1659841569926357]
number_of_tis=1000 mean=0.5903461098030676 per=0.0005903461098030675 times=[0.6001852740009781, 0.5642872750031529, 0.686630773008801, 0.5578094649972627, 0.5428177620051429]
number_of_tis=3000 mean=1.9076304554007948 per=0.0006358768184669316 times=[2.042052763994434, 2.1137778090051142, 1.7461599689995637, 1.7260139089921722, 1.9101478260126896]
number_of_tis=5000 mean=2.9185905692051164 per=0.0005837181138410233 times=[2.9221124830073677, 3.2889883980096783, 2.7569778940087417, 2.973596281008213, 2.651277789991582]
number_of_tis=10000 mean=8.880191986600403 per=0.0008880191986600403 times=[7.3548113360011484, 9.13715232499817, 9.568511486999341, 8.80206210000324, 9.538422685000114]
number_of_tis=15000 mean=15.426499317999696 per=0.0010284332878666464 times=[14.944712879005237, 15.38737604500784, 15.409629273999599, 15.852925243991194, 15.53785314799461]
number_of_tis=20000 mean=20.579332908798825 per=0.0010289666454399414 times=[20.362008597003296, 19.878823954990366, 20.73281196100288, 20.837948996995692, 21.085071034001885]
```
After:
```
number_of_tis=20000 mean=18.36637533060275 per=0.0009183187665301375 times=[17.728908119010157, 18.62269214099797, 18.936747477011522, 17.74613195299753, 18.797396962996572]
```
[1]: https://docs.sqlalchemy.org/en/13/dialects/postgresql.html#psycopg2-batch-mode
* Use bulk_insert_mappings for even more speed where possible.
It gives us an extra speed up over bulk_save_objects, but we can't
use it when the task_instance_mutation_hook does anything, as that hook
needs an actual object.
So _when_ we know that hook won't do anything we switch in to
insert_mappings mode.
New speeds (vs baseline, not vs bulk_save_objects) when using
bulk_insert_mappings
PostgreSQL now 130% faster:
```
number_of_tis=1 mean=0.028053103599813767 per=0.028053103599813767 times=[0.03762496300623752, 0.02637488600157667, 0.025065611000172794, 0.024561002996051684, 0.026639054995030165]
number_of_tis=10 mean=0.02647183560184203 per=0.002647183560184203 times=[0.02698062499985099, 0.026417658998980187, 0.027347976007149555, 0.025797458001761697, 0.025815460001467727]
number_of_tis=50 mean=0.03149963079486042 per=0.0006299926158972085 times=[0.03810671299288515, 0.03055680700344965, 0.029733988994848914, 0.03016914198815357, 0.02893150299496483]
number_of_tis=100 mean=0.033998635396710594 per=0.0003399863539671059 times=[0.0351028829900315, 0.03299884400621522, 0.03358584298985079, 0.03295094799250364, 0.03535465900495183]
number_of_tis=500 mean=0.07903424859978259 per=0.00015806849719956516 times=[0.08279920800123364, 0.08588568199775182, 0.07312070899934042, 0.07360191999759991, 0.07976372400298715]
number_of_tis=1000 mean=0.12571056479937398 per=0.00012571056479937398 times=[0.12573593499837443, 0.12141938100103289, 0.12616568499652203, 0.12907471299695317, 0.12615711000398733]
number_of_tis=3000 mean=0.36025245799683037 per=0.00012008415266561012 times=[0.36071603700111154, 0.3470657339930767, 0.3373015969991684, 0.3337128989951452, 0.42246602299564984]
number_of_tis=5000 mean=0.6916533229988999 per=0.00013833066459977998 times=[0.9647149289958179, 0.6451378140045563, 0.5970188640058041, 0.5849326960014878, 0.6664623119868338]
number_of_tis=10000 mean=2.071472014003666 per=0.00020714720140036663 times=[2.957865878008306, 1.9388906149979448, 1.766649461002089, 1.8647991580073722, 1.8291549580026185]
number_of_tis=15000 mean=2.866650845797267 per=0.00019111005638648446 times=[3.3783503199956613, 2.657773957995232, 2.707275656008278, 2.7875704979960574, 2.802283796991105]
number_of_tis=20000 mean=3.5886989389982773 per=0.00017943494694991387 times=[3.969436354993377, 3.436962780993781, 3.9078941010084236, 3.6387251569976797, 2.9904763009981252]
```
MySQL is (only) 27% faster:
```
number_of_tis=1 mean=0.035956257799989545 per=0.035956257799989545 times=[0.03932315899874084, 0.03545605999534018, 0.03535486999317072, 0.034727805003058165, 0.03491939500963781]
number_of_tis=10 mean=0.036957260797498746 per=0.0036957260797498745 times=[0.040442515004542656, 0.0379129799985094, 0.03494819799379911, 0.03562593398964964, 0.03585667700099293]
number_of_tis=50 mean=0.04745422120031435 per=0.0009490844240062871 times=[0.06965546800347511, 0.04221734800375998, 0.04038520700123627, 0.040363031992455944, 0.04465005100064445]
number_of_tis=100 mean=0.0528092162014218 per=0.000528092162014218 times=[0.06113427500531543, 0.04883724599494599, 0.05276876600692049, 0.047688748003565706, 0.05361704599636141]
number_of_tis=500 mean=0.16223246100416872 per=0.0003244649220083374 times=[0.24469116200634744, 0.1407806619972689, 0.14792052800476085, 0.14703868801007047, 0.13073126500239596]
number_of_tis=1000 mean=0.285728433605982 per=0.00028572843360598197 times=[0.3230128890136257, 0.27035739900020417, 0.3003890450054314, 0.2638379510026425, 0.2710448840080062]
number_of_tis=3000 mean=1.1824120475997915 per=0.0003941373491999305 times=[1.3103130240051541, 1.286688863998279, 1.1455156929878285, 1.1072918410063721, 1.062250816001324]
number_of_tis=5000 mean=1.9416745471942705 per=0.0003883349094388541 times=[2.3746965279860888, 1.9103765429899795, 2.0542518720030785, 1.7706374429981224, 1.598410349994083]
number_of_tis=10000 mean=5.059874459402636 per=0.0005059874459402636 times=[5.431018351999228, 5.262124675995437, 5.174487816999317, 4.423381198008428, 5.008360254010768]
number_of_tis=15000 mean=9.717965700797503 per=0.0006478643800531668 times=[7.884617075993447, 9.466949063993525, 10.005758297003922, 10.105231182998978, 11.127272883997648]
number_of_tis=20000 mean=16.2008618004038 per=0.00081004309002019 times=[14.645835625007749, 16.304637463006657, 16.255490412993822, 16.830263861003914, 16.968081640006858]
```
---
airflow/config_templates/config.yml | 8 ++++
airflow/config_templates/default_airflow.cfg | 4 ++
airflow/configuration.py | 26 +++++++++++
airflow/kubernetes/kube_config.py | 27 +++++++-----
airflow/models/dagrun.py | 38 +++++++++++++---
airflow/models/taskinstance.py | 20 +++++++++
airflow/settings.py | 35 ++++++++++++---
airflow/utils/log/log_reader.py | 2 +-
airflow/utils/platform.py | 3 ++
tests/api_connexion/endpoints/test_log_endpoint.py | 7 +--
tests/core/test_configuration.py | 49 +++++++++++++++++++++
tests/core/test_sqlalchemy_config.py | 3 +-
tests/models/test_dagrun.py | 50 +++++++++++++++-------
.../microsoft/azure/log/test_wasb_task_handler.py | 3 +-
tests/utils/log/test_log_reader.py | 17 ++++----
tests/www/test_app.py | 18 --------
tests/www/views/test_views_log.py | 2 +
17 files changed, 239 insertions(+), 73 deletions(-)
diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml
index 3c9fc5d..4407c71 100644
--- a/airflow/config_templates/config.yml
+++ b/airflow/config_templates/config.yml
@@ -71,6 +71,14 @@
sensitive: true
example: ~
default: "sqlite:///{AIRFLOW_HOME}/airflow.db"
+ - name: sql_alchemy_engine_args
+ description: |
+ Extra engine specific keyword args passed to SQLAlchemy's create_engine, as a JSON-encoded value
+ version_added: 2.3.0
+ type: string
+ sensitive: true
+ example: '{"arg1": True}'
+ default: ~
- name: sql_engine_encoding
description: |
The encoding for the databases
diff --git a/airflow/config_templates/default_airflow.cfg b/airflow/config_templates/default_airflow.cfg
index 5dacfd8..76df6f7 100644
--- a/airflow/config_templates/default_airflow.cfg
+++ b/airflow/config_templates/default_airflow.cfg
@@ -59,6 +59,10 @@ executor = SequentialExecutor
# http://airflow.apache.org/docs/apache-airflow/stable/howto/set-up-database.html#database-uri
sql_alchemy_conn = sqlite:///{AIRFLOW_HOME}/airflow.db
+# Extra engine specific keyword args passed to SQLAlchemy's create_engine, as a JSON-encoded value
+# Example: sql_alchemy_engine_args = {{"arg1": True}}
+# sql_alchemy_engine_args =
+
# The encoding for the databases
sql_engine_encoding = utf-8
diff --git a/airflow/configuration.py b/airflow/configuration.py
index 0ec252e..0a113b4 100644
--- a/airflow/configuration.py
+++ b/airflow/configuration.py
@@ -493,6 +493,32 @@ class AirflowConfigParser(ConfigParser):
f'Current value: "{full_qualified_path}".'
)
+ def getjson(self, section, key, fallback=_UNSET, **kwargs) -> Union[dict, list, str, int, float, None]:
+ """
+ Return a config value parsed from a JSON string.
+
+ ``fallback`` is *not* JSON parsed but used verbatim when no config value is given.
+ """
+ # get always returns the fallback value as a string, so for this if
+ # someone gives us an object we want to keep that
+ default = _UNSET
+ if fallback is not _UNSET:
+ default = fallback
+ fallback = _UNSET
+
+ try:
+ data = self.get(section=section, key=key, fallback=fallback, **kwargs)
+ except (NoSectionError, NoOptionError):
+ return default
+
+ if len(data) == 0:
+ return default if default is not _UNSET else None
+
+ try:
+ return json.loads(data)
+ except JSONDecodeError as e:
+ raise AirflowConfigException(f'Unable to parse [{section}] {key!r} as valid json') from e
+
def read(self, filenames, encoding=None):
super().read(filenames=filenames, encoding=encoding)
diff --git a/airflow/kubernetes/kube_config.py b/airflow/kubernetes/kube_config.py
index e5aca2b..2b55021 100644
--- a/airflow/kubernetes/kube_config.py
+++ b/airflow/kubernetes/kube_config.py
@@ -14,9 +14,9 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-import json
from airflow.configuration import conf
+from airflow.exceptions import AirflowConfigException
from airflow.settings import AIRFLOW_HOME
@@ -73,19 +73,24 @@ class KubeConfig:
self.kubernetes_section, 'worker_pods_queued_check_interval'
)
- kube_client_request_args = conf.get(self.kubernetes_section, 'kube_client_request_args')
- if kube_client_request_args:
- self.kube_client_request_args = json.loads(kube_client_request_args)
+ self.kube_client_request_args = conf.getjson(
+ self.kubernetes_section, 'kube_client_request_args', fallback={}
+ )
+ if not isinstance(self.kube_client_request_args, dict):
+ raise AirflowConfigException(
+ f"[{self.kubernetes_section}] 'kube_client_request_args' expected a JSON dict, got "
+ + type(self.kube_client_request_args).__name__
+ )
+ if self.kube_client_request_args:
if '_request_timeout' in self.kube_client_request_args and isinstance(
self.kube_client_request_args['_request_timeout'], list
):
self.kube_client_request_args['_request_timeout'] = tuple(
self.kube_client_request_args['_request_timeout']
)
- else:
- self.kube_client_request_args = {}
- delete_option_kwargs = conf.get(self.kubernetes_section, 'delete_option_kwargs')
- if delete_option_kwargs:
- self.delete_option_kwargs = json.loads(delete_option_kwargs)
- else:
- self.delete_option_kwargs = {}
+ self.delete_option_kwargs = conf.getjson(self.kubernetes_section, 'delete_option_kwargs', fallback={})
+ if not isinstance(self.delete_option_kwargs, dict):
+ raise AirflowConfigException(
+ f"[{self.kubernetes_section}] 'delete_option_kwargs' expected a JSON dict, got "
+ + type(self.delete_option_kwargs).__name__
+ )
diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py
index 4b8cb46..bd94ad1 100644
--- a/airflow/models/dagrun.py
+++ b/airflow/models/dagrun.py
@@ -17,6 +17,7 @@
# under the License.
import os
import warnings
+from collections import defaultdict
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, NamedTuple, Optional, Tuple, Union
@@ -57,6 +58,7 @@ from airflow.utils.state import DagRunState, State, TaskInstanceState
from airflow.utils.types import NOTSET, ArgNotSet, DagRunType
if TYPE_CHECKING:
+ from airflow.models.baseoperator import BaseOperator
from airflow.models.dag import DAG
@@ -804,18 +806,40 @@ class DagRun(Base, LoggingMixin):
ti.state = State.NONE
session.merge(ti)
- # check for missing tasks
- for task in dag.task_dict.values():
- if task.start_date > self.execution_date and not self.is_backfill:
- continue
+ def task_filter(task: "BaseOperator"):
+ return task.task_id not in task_ids and (
+ self.is_backfill or task.start_date <= self.execution_date
+ )
+
+ created_counts: Dict[str, int] = defaultdict(int)
+
+ # Set for the empty default in airflow.settings -- if it's not set this means it has been changed
+ hook_is_noop = getattr(task_instance_mutation_hook, 'is_noop', False)
+
+ if hook_is_noop:
+
+ def create_ti_mapping(task: "BaseOperator"):
+ created_counts[task.task_type] += 1
+ return TI.insert_mapping(self.run_id, task)
- if task.task_id not in task_ids:
- Stats.incr(f"task_instance_created-{task.task_type}", 1, 1)
+ else:
+
+ def create_ti(task: "BaseOperator") -> TI:
ti = TI(task, run_id=self.run_id)
task_instance_mutation_hook(ti)
- session.add(ti)
+ created_counts[ti.operator] += 1
+ return ti
+ # Create missing tasks
+ tasks = list(filter(task_filter, dag.task_dict.values()))
try:
+ if hook_is_noop:
+ session.bulk_insert_mappings(TI, map(create_ti_mapping, tasks))
+ else:
+ session.bulk_save_objects(map(create_ti, tasks))
+
+ for task_type, count in created_counts.items():
+ Stats.incr(f"task_instance_created-{task_type}", count)
session.flush()
except IntegrityError as err:
self.log.info(str(err))
diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index ceead7f..5015f46 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -477,6 +477,26 @@ class TaskInstance(Base, LoggingMixin):
# can be changed when calling 'run'
self.test_mode = False
+ @staticmethod
+ def insert_mapping(run_id: str, task: "BaseOperator") -> dict:
+ """:meta private:"""
+ return {
+ 'dag_id': task.dag_id,
+ 'task_id': task.task_id,
+ 'run_id': run_id,
+ '_try_number': 0,
+ 'hostname': '',
+ 'unixname': getuser(),
+ 'queue': task.queue,
+ 'pool': task.pool,
+ 'pool_slots': task.pool_slots,
+ 'priority_weight': task.priority_weight_total,
+ 'run_as_user': task.run_as_user,
+ 'max_tries': task.retries,
+ 'executor_config': task.executor_config,
+ 'operator': task.task_type,
+ }
+
@reconstructor
def init_on_load(self):
"""Initialize the attributes that aren't stored in the DB"""
diff --git a/airflow/settings.py b/airflow/settings.py
index ef23f65..eea4d5b 100644
--- a/airflow/settings.py
+++ b/airflow/settings.py
@@ -181,6 +181,9 @@ def task_instance_mutation_hook(task_instance):
"""
+task_instance_mutation_hook.is_noop = True # type: ignore
+
+
def pod_mutation_hook(pod):
"""
This setting allows altering ``kubernetes.client.models.V1Pod`` object
@@ -238,10 +241,6 @@ def configure_orm(disable_connection_pool=False):
global Session
engine_args = prepare_engine_args(disable_connection_pool)
- # Allow the user to specify an encoding for their DB otherwise default
- # to utf-8 so jobs & users with non-latin1 characters can still use us.
- engine_args['encoding'] = conf.get('core', 'SQL_ENGINE_ENCODING', fallback='utf-8')
-
if conf.has_option('core', 'sql_alchemy_connect_args'):
connect_args = conf.getimport('core', 'sql_alchemy_connect_args')
else:
@@ -284,11 +283,26 @@ def configure_orm(disable_connection_pool=False):
session.close()
+DEFAULT_ENGINE_ARGS = {
+ 'postgresql': {
+ 'executemany_mode': 'values',
+ 'executemany_values_page_size': 10000,
+ 'executemany_batch_page_size': 2000,
+ },
+}
+
+
def prepare_engine_args(disable_connection_pool=False):
"""Prepare SQLAlchemy engine args"""
- engine_args = {}
- pool_connections = conf.getboolean('core', 'SQL_ALCHEMY_POOL_ENABLED')
- if disable_connection_pool or not pool_connections:
+ default_args = {}
+ for dialect, default in DEFAULT_ENGINE_ARGS.items():
+ if SQL_ALCHEMY_CONN.startswith(dialect):
+ default_args = default.copy()
+ break
+
+ engine_args: dict = conf.getjson('core', 'sql_alchemy_engine_args', fallback=default_args) # type: ignore
+
+ if disable_connection_pool or not conf.getboolean('core', 'SQL_ALCHEMY_POOL_ENABLED'):
engine_args['poolclass'] = NullPool
log.debug("settings.prepare_engine_args(): Using NullPool")
elif not SQL_ALCHEMY_CONN.startswith('sqlite'):
@@ -349,6 +363,10 @@ def prepare_engine_args(disable_connection_pool=False):
if SQL_ALCHEMY_CONN.startswith(('mysql', 'mssql')):
engine_args['isolation_level'] = 'READ COMMITTED'
+ # Allow the user to specify an encoding for their DB otherwise default
+ # to utf-8 so jobs & users with non-latin1 characters can still use us.
+ engine_args['encoding'] = conf.get('core', 'SQL_ENGINE_ENCODING', fallback='utf-8')
+
return engine_args
@@ -487,6 +505,9 @@ def import_local_settings():
globals()["task_policy"] = globals()["policy"]
del globals()["policy"]
+ if not hasattr(task_instance_mutation_hook, 'is_noop'):
+ task_instance_mutation_hook.is_noop = False
+
log.info("Loaded airflow_local_settings from %s .", airflow_local_settings.__file__)
except ModuleNotFoundError as e:
if e.name == "airflow_local_settings":
diff --git a/airflow/utils/log/log_reader.py b/airflow/utils/log/log_reader.py
index 03ed353..623af90 100644
--- a/airflow/utils/log/log_reader.py
+++ b/airflow/utils/log/log_reader.py
@@ -85,7 +85,7 @@ class TaskLogReader:
while 'end_of_log' not in metadata or not metadata['end_of_log']:
logs, metadata = self.read_log_chunks(ti, current_try_number, metadata)
for host, log in logs[0]:
- yield "\n".join([host, log]) + "\n"
+ yield "\n".join([host or '', log]) + "\n"
@cached_property
def log_handler(self):
diff --git a/airflow/utils/platform.py b/airflow/utils/platform.py
index 0b36946..3697f87 100644
--- a/airflow/utils/platform.py
+++ b/airflow/utils/platform.py
@@ -23,6 +23,8 @@ import pkgutil
import platform
import sys
+from airflow.compat.functools import cache
+
IS_WINDOWS = platform.system() == 'Windows'
log = logging.getLogger(__name__)
@@ -63,6 +65,7 @@ def get_airflow_git_version():
return git_version
+@cache
def getuser() -> str:
"""
Gets the username associated with the current user, or error with a nice
diff --git a/tests/api_connexion/endpoints/test_log_endpoint.py b/tests/api_connexion/endpoints/test_log_endpoint.py
index cd71963..f5d2e5d 100644
--- a/tests/api_connexion/endpoints/test_log_endpoint.py
+++ b/tests/api_connexion/endpoints/test_log_endpoint.py
@@ -96,6 +96,7 @@ class TestGetLog:
self.ti = dr.task_instances[0]
self.ti.try_number = 1
+ self.ti.hostname = 'localhost'
@pytest.fixture
def configure_loggers(self, tmp_path):
@@ -139,7 +140,7 @@ class TestGetLog:
)
assert (
response.json['content']
- == f"[('', '*** Reading local file: {expected_filename}\\nLog for testing.')]"
+ == f"[('localhost', '*** Reading local file: {expected_filename}\\nLog for testing.')]"
)
info = serializer.loads(response.json['continuation_token'])
assert info == {'end_of_log': True}
@@ -162,7 +163,7 @@ class TestGetLog:
assert 200 == response.status_code
assert (
response.data.decode('utf-8')
- == f"\n*** Reading local file: {expected_filename}\nLog for testing.\n"
+ == f"localhost\n*** Reading local file: {expected_filename}\nLog for testing.\n"
)
def test_get_logs_of_removed_task(self):
@@ -188,7 +189,7 @@ class TestGetLog:
assert 200 == response.status_code
assert (
response.data.decode('utf-8')
- == f"\n*** Reading local file: {expected_filename}\nLog for testing.\n"
+ == f"localhost\n*** Reading local file: {expected_filename}\nLog for testing.\n"
)
def test_get_logs_response_with_ti_equal_to_none(self):
diff --git a/tests/core/test_configuration.py b/tests/core/test_configuration.py
index da1d736..d35267d 100644
--- a/tests/core/test_configuration.py
+++ b/tests/core/test_configuration.py
@@ -367,6 +367,55 @@ key2 = 1.23
assert isinstance(test_conf.getfloat('valid', 'key2'), float)
assert 1.23 == test_conf.getfloat('valid', 'key2')
+ @pytest.mark.parametrize(
+ ("config_str", "expected"),
+ [
+ pytest.param('{"a": 123}', {'a': 123}, id='dict'),
+ pytest.param('[1,2,3]', [1, 2, 3], id='list'),
+ pytest.param('"abc"', 'abc', id='str'),
+ pytest.param('2.1', 2.1, id='num'),
+ pytest.param('', None, id='empty'),
+ ],
+ )
+ def test_getjson(self, config_str, expected):
+ config = textwrap.dedent(
+ f"""
+ [test]
+ json = {config_str}
+ """
+ )
+ test_conf = AirflowConfigParser()
+ test_conf.read_string(config)
+
+ assert test_conf.getjson('test', 'json') == expected
+
+ def test_getjson_empty_with_fallback(self):
+ config = textwrap.dedent(
+ """
+ [test]
+ json =
+ """
+ )
+ test_conf = AirflowConfigParser()
+ test_conf.read_string(config)
+
+ assert test_conf.getjson('test', 'json', fallback={}) == {}
+ assert test_conf.getjson('test', 'json') is None
+
+ @pytest.mark.parametrize(
+ ("fallback"),
+ [
+ pytest.param({"a": "b"}, id='dict'),
+ # fallback is _NOT_ json parsed, but used verbatim
+ pytest.param('{"a": "b"}', id='str'),
+ pytest.param(None, id='None'),
+ ],
+ )
+ def test_getjson_fallback(self, fallback):
+ test_conf = AirflowConfigParser()
+
+ assert test_conf.getjson('test', 'json', fallback=fallback) == fallback
+
def test_has_option(self):
test_config = '''[test]
key1 = value1
diff --git a/tests/core/test_sqlalchemy_config.py b/tests/core/test_sqlalchemy_config.py
index 39b45ec..37d4cfe 100644
--- a/tests/core/test_sqlalchemy_config.py
+++ b/tests/core/test_sqlalchemy_config.py
@@ -72,11 +72,12 @@ class TestSqlAlchemySettings(unittest.TestCase):
'core',
'sql_alchemy_connect_args',
): 'tests.core.test_sqlalchemy_config.SQL_ALCHEMY_CONNECT_ARGS',
+ ('core', 'sql_alchemy_engine_args'): '{"arg": 1}',
('core', 'sql_alchemy_pool_enabled'): 'False',
}
with conf_vars(config):
settings.configure_orm()
- engine_args = {}
+ engine_args = {'arg': 1}
if settings.SQL_ALCHEMY_CONN.startswith(('mysql', 'mssql')):
engine_args['isolation_level'] = 'READ COMMITTED'
mock_create_engine.assert_called_once_with(
diff --git a/tests/models/test_dagrun.py b/tests/models/test_dagrun.py
index 00799be..90c47da 100644
--- a/tests/models/test_dagrun.py
+++ b/tests/models/test_dagrun.py
@@ -21,6 +21,7 @@ import unittest
from unittest import mock
from unittest.mock import call
+import pytest
from parameterized import parameterized
from airflow import models, settings
@@ -581,21 +582,6 @@ class TestDagRun(unittest.TestCase):
if dagrun.dag_id == 'test_latest_runs_1':
assert dagrun.execution_date == timezone.datetime(2015, 1, 2)
- def test_is_backfill(self):
- dag = DAG(dag_id='test_is_backfill', start_date=DEFAULT_DATE)
-
- dagrun = self.create_dag_run(dag, execution_date=DEFAULT_DATE)
- dagrun.run_type = DagRunType.BACKFILL_JOB
-
- dagrun2 = self.create_dag_run(dag, execution_date=DEFAULT_DATE + datetime.timedelta(days=1))
-
- dagrun3 = self.create_dag_run(dag, execution_date=DEFAULT_DATE + datetime.timedelta(days=2))
- dagrun3.run_id = None
-
- assert dagrun.is_backfill
- assert not dagrun2.is_backfill
- assert not dagrun3.is_backfill
-
def test_removed_task_instances_can_be_restored(self):
def with_all_tasks_removed(dag):
return DAG(dag_id=dag.dag_id, start_date=dag.start_date)
@@ -640,7 +626,7 @@ class TestDagRun(unittest.TestCase):
assert State.NONE == first_ti.state
@parameterized.expand([(state,) for state in State.task_states])
- @mock.patch('airflow.settings.task_instance_mutation_hook')
+ @mock.patch.object(settings, 'task_instance_mutation_hook', autospec=True)
def test_task_instance_mutation_hook(self, state, mock_hook):
def mutate_task_instance(task_instance):
if task_instance.queue == 'queue1':
@@ -852,3 +838,35 @@ class TestDagRun(unittest.TestCase):
ti_failed = dag_run.get_task_instance(dag_task_failed.task_id)
assert ti_success.state in State.success_states
assert ti_failed.state in State.failed_states
+
+
+@pytest.mark.parametrize(
+ ('run_type', 'expected_tis'),
+ [
+ pytest.param(DagRunType.MANUAL, 1, id='manual'),
+ pytest.param(DagRunType.BACKFILL_JOB, 2, id='backfill'),
+ ],
+)
+@mock.patch.object(Stats, 'incr')
+def test_verify_integrity_task_start_date(Stats_incr, session, run_type, expected_tis):
+ """Test that tasks with specific start dates are only created for backfill runs"""
+ with DAG('test', start_date=DEFAULT_DATE) as dag:
+ DummyOperator(task_id='without')
+ DummyOperator(task_id='with_startdate', start_date=DEFAULT_DATE + datetime.timedelta(1))
+
+ dag_run = DagRun(
+ dag_id=dag.dag_id,
+ run_type=run_type,
+ execution_date=DEFAULT_DATE,
+ run_id=DagRun.generate_run_id(run_type, DEFAULT_DATE),
+ )
+ dag_run.dag = dag
+
+ session.add(dag_run)
+ session.flush()
+ dag_run.verify_integrity(session)
+
+ tis = dag_run.task_instances
+ assert len(tis) == expected_tis
+
+ Stats_incr.assert_called_with('task_instance_created-DummyOperator', expected_tis)
diff --git a/tests/providers/microsoft/azure/log/test_wasb_task_handler.py b/tests/providers/microsoft/azure/log/test_wasb_task_handler.py
index 595d2e5..4fe9676 100644
--- a/tests/providers/microsoft/azure/log/test_wasb_task_handler.py
+++ b/tests/providers/microsoft/azure/log/test_wasb_task_handler.py
@@ -41,6 +41,7 @@ class TestWasbTaskHandler:
state=State.RUNNING,
)
ti.try_number = 1
+ ti.hostname = 'localhost'
ti.raw = False
yield ti
clear_db_runs()
@@ -110,7 +111,7 @@ class TestWasbTaskHandler:
[
[
(
- '',
+ 'localhost',
'*** Reading remote log from wasb://container/remote/log/location/1.log.\n'
'Log line\n',
)
diff --git a/tests/utils/log/test_log_reader.py b/tests/utils/log/test_log_reader.py
index a3bb9b7..4e6e942 100644
--- a/tests/utils/log/test_log_reader.py
+++ b/tests/utils/log/test_log_reader.py
@@ -102,6 +102,7 @@ class TestLogView:
state=TaskInstanceState.RUNNING,
)
ti.try_number = 3
+ ti.hostname = 'localhost'
self.ti = ti
yield
clear_db_runs()
@@ -113,7 +114,7 @@ class TestLogView:
assert [
(
- '',
+ 'localhost',
f"*** Reading local file: "
f"{self.log_dir}/dag_log_reader/task_log_reader/2017-09-01T00.00.00+00.00/1.log\n"
f"try_number=1.\n",
@@ -128,7 +129,7 @@ class TestLogView:
assert [
[
(
- '',
+ 'localhost',
"*** Reading local file: "
f"{self.log_dir}/dag_log_reader/task_log_reader/2017-09-01T00.00.00+00.00/1.log\n"
"try_number=1.\n",
@@ -136,7 +137,7 @@ class TestLogView:
],
[
(
- '',
+ 'localhost',
f"*** Reading local file: "
f"{self.log_dir}/dag_log_reader/task_log_reader/2017-09-01T00.00.00+00.00/2.log\n"
f"try_number=2.\n",
@@ -144,7 +145,7 @@ class TestLogView:
],
[
(
- '',
+ 'localhost',
f"*** Reading local file: "
f"{self.log_dir}/dag_log_reader/task_log_reader/2017-09-01T00.00.00+00.00/3.log\n"
f"try_number=3.\n",
@@ -158,7 +159,7 @@ class TestLogView:
stream = task_log_reader.read_log_stream(ti=self.ti, try_number=1, metadata={})
assert [
- "\n*** Reading local file: "
+ "localhost\n*** Reading local file: "
f"{self.log_dir}/dag_log_reader/task_log_reader/2017-09-01T00.00.00+00.00/1.log\n"
"try_number=1.\n"
"\n"
@@ -168,15 +169,15 @@ class TestLogView:
task_log_reader = TaskLogReader()
stream = task_log_reader.read_log_stream(ti=self.ti, try_number=None, metadata={})
assert [
- "\n*** Reading local file: "
+ "localhost\n*** Reading local file: "
f"{self.log_dir}/dag_log_reader/task_log_reader/2017-09-01T00.00.00+00.00/1.log\n"
"try_number=1.\n"
"\n",
- "\n*** Reading local file: "
+ "localhost\n*** Reading local file: "
f"{self.log_dir}/dag_log_reader/task_log_reader/2017-09-01T00.00.00+00.00/2.log\n"
"try_number=2.\n"
"\n",
- "\n*** Reading local file: "
+ "localhost\n*** Reading local file: "
f"{self.log_dir}/dag_log_reader/task_log_reader/2017-09-01T00.00.00+00.00/3.log\n"
"try_number=3.\n"
"\n",
diff --git a/tests/www/test_app.py b/tests/www/test_app.py
index 7a06010..0e65c17 100644
--- a/tests/www/test_app.py
+++ b/tests/www/test_app.py
@@ -211,24 +211,6 @@ class TestApp(unittest.TestCase):
@conf_vars(
{
- ('core', 'sql_alchemy_pool_enabled'): 'True',
- ('core', 'sql_alchemy_pool_size'): '3',
- ('core', 'sql_alchemy_max_overflow'): '5',
- ('core', 'sql_alchemy_pool_recycle'): '120',
- ('core', 'sql_alchemy_pool_pre_ping'): 'True',
- }
- )
- @dont_initialize_flask_app_submodules
- @pytest.mark.backend("mysql", "postgres")
- def test_should_set_sqlalchemy_engine_options(self):
- app = application.cached_app(testing=True)
- engine_params = {'pool_size': 3, 'pool_recycle': 120, 'pool_pre_ping': True, 'max_overflow': 5}
- if app.config['SQLALCHEMY_DATABASE_URI'].startswith('mysql'):
- engine_params['isolation_level'] = 'READ COMMITTED'
- assert app.config['SQLALCHEMY_ENGINE_OPTIONS'] == engine_params
-
- @conf_vars(
- {
('webserver', 'session_lifetime_minutes'): '3600',
}
)
diff --git a/tests/www/views/test_views_log.py b/tests/www/views/test_views_log.py
index 6dc4580..20b0690 100644
--- a/tests/www/views/test_views_log.py
+++ b/tests/www/views/test_views_log.py
@@ -144,6 +144,7 @@ def tis(dags, session):
)
(ti,) = dagrun.task_instances
ti.try_number = 1
+ ti.hostname = 'localhost'
dagrun_removed = dag_removed.create_dagrun(
run_type=DagRunType.SCHEDULED,
execution_date=DEFAULT_DATE,
@@ -229,6 +230,7 @@ def test_get_logs_with_metadata_as_download_file(log_admin_client):
assert f'{DAG_ID}/{TASK_ID}/{DEFAULT_DATE.isoformat()}/{try_number}.log' in content_disposition
assert 200 == response.status_code
assert 'Log for testing.' in response.data.decode('utf-8')
+ assert 'localhost\n' in response.data.decode('utf-8')
DIFFERENT_LOG_FILENAME = "{{ ti.dag_id }}/{{ ti.run_id }}/{{ ti.task_id }}/{{ try_number }}.log"