diff --git a/.dockerignore b/.dockerignore index d10cfbcaae575..f6113e2bd6d33 100644 --- a/.dockerignore +++ b/.dockerignore @@ -40,9 +40,6 @@ !scripts/in_container !scripts/docker -# Add provider packages to the context -!provider_packages - # Add tests and kubernetes_tests to context. !tests !kubernetes_tests @@ -129,3 +126,4 @@ airflow/www/static/docs # Exclude docs generated files docs/_build/ docs/_api/ +docs/_doctrees/ diff --git a/.github/workflows/build-images.yml b/.github/workflows/build-images.yml index f29e199e4fb84..ec8f4354d0908 100644 --- a/.github/workflows/build-images.yml +++ b/.github/workflows/build-images.yml @@ -148,7 +148,6 @@ jobs: BACKEND: postgres PYTHON_MAJOR_MINOR_VERSION: ${{ matrix.python-version }} UPGRADE_TO_NEWER_DEPENDENCIES: ${{ needs.build-info.outputs.upgradeToNewerDependencies }} - CONTINUE_ON_PIP_CHECK_FAILURE: "true" DOCKER_CACHE: ${{ needs.build-info.outputs.cacheDirective }} CHECK_IF_BASE_PYTHON_IMAGE_UPDATED: > ${{ github.event_name == 'pull_request_target' && 'false' || 'true' }} diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index bd8c2d23b735f..36c8fc9d40292 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -563,7 +563,7 @@ ${{ hashFiles('.pre-commit-config.yaml') }}" PACKAGE_FORMAT: "sdist" tests-helm: - timeout-minutes: 20 + timeout-minutes: 40 name: "Python unit tests for helm chart" runs-on: ${{ fromJson(needs.build-info.outputs.runsOn) }} needs: [build-info, ci-images] diff --git a/BREEZE.rst b/BREEZE.rst index 90663657a85ad..683d2fb320f98 100644 --- a/BREEZE.rst +++ b/BREEZE.rst @@ -1280,9 +1280,6 @@ This is the current syntax for `./breeze <./breeze>`_: --upgrade-to-newer-dependencies Upgrades PIP packages to latest versions available without looking at the constraints. - --continue-on-pip-check-failure - Continue even if 'pip check' fails. - -I, --production-image Use production image for entering the environment and builds (not for tests). @@ -2393,9 +2390,9 @@ This is the current syntax for `./breeze <./breeze>`_: Helm version - only used in case one of kind-cluster commands is used. One of: - v3.2.4 + v3.6.3 - Default: v3.2.4 + Default: v3.6.3 --executor EXECUTOR Executor to use in a kubernetes cluster. @@ -2446,9 +2443,6 @@ This is the current syntax for `./breeze <./breeze>`_: --upgrade-to-newer-dependencies Upgrades PIP packages to latest versions available without looking at the constraints. - --continue-on-pip-check-failure - Continue even if 'pip check' fails. - **************************************************************************************************** Use different Airflow version at runtime in CI image diff --git a/CONTRIBUTORS_QUICK_START.rst b/CONTRIBUTORS_QUICK_START.rst index af91b9bac9c40..0313852dab842 100644 --- a/CONTRIBUTORS_QUICK_START.rst +++ b/CONTRIBUTORS_QUICK_START.rst @@ -276,14 +276,14 @@ Setting up Breeze $ ./breeze --python 3.8 --backend mysql -4. Creating airflow tables and users. ``airflow db reset`` is required to execute at least once for Airflow Breeze to get - the database/tables created. +4. Once the breeze environment is initialized, create airflow tables and users from the breeze CLI. ``airflow db reset`` + is required to execute at least once for Airflow Breeze to get the database/tables created. .. code-block:: bash - $ airflow db reset - $ airflow users create --role Admin --username admin --password admin --email admin@example.com --firstname\ - foo --lastname bar + root@b76fcb399bb6:/opt/airflow# airflow db reset + root@b76fcb399bb6:/opt/airflow# airflow users create --role Admin --username admin --password admin \ + --email admin@example.com --firstname foo --lastname bar 5. Closing Breeze environment. After successfully finishing above command will leave you in container, diff --git a/Dockerfile b/Dockerfile index 9a7f8ecdc1fec..782e5b4bbd9a3 100644 --- a/Dockerfile +++ b/Dockerfile @@ -44,7 +44,8 @@ ARG AIRFLOW_GID="50000" ARG PYTHON_BASE_IMAGE="python:3.6-slim-buster" -ARG AIRFLOW_PIP_VERSION=21.1.2 +ARG AIRFLOW_PIP_VERSION=21.2.2 +ARG AIRFLOW_IMAGE_REPOSITORY="https://github.com/apache/airflow" # By default PIP has progress bar but you can disable it. ARG PIP_PROGRESS_BAR="on" @@ -108,12 +109,13 @@ ARG DEV_APT_COMMAND="\ && curl https://dl.yarnpkg.com/debian/pubkey.gpg | apt-key add - > /dev/null \ && echo 'deb https://dl.yarnpkg.com/debian/ stable main' > /etc/apt/sources.list.d/yarn.list" ARG ADDITIONAL_DEV_APT_COMMAND="echo" +ARG ADDITIONAL_DEV_APT_ENV="" ENV DEV_APT_DEPS=${DEV_APT_DEPS} \ ADDITIONAL_DEV_APT_DEPS=${ADDITIONAL_DEV_APT_DEPS} \ DEV_APT_COMMAND=${DEV_APT_COMMAND} \ ADDITIONAL_DEV_APT_COMMAND=${ADDITIONAL_DEV_APT_COMMAND} \ - ADDITIONAL_DEV_APT_ENV="" + ADDITIONAL_DEV_APT_ENV=${ADDITIONAL_DEV_APT_ENV} # Note missing man directories on debian-buster # https://bugs.debian.org/cgi-bin/bugreport.cgi?bug=863199 @@ -216,7 +218,7 @@ ENV AIRFLOW_PRE_CACHED_PIP_PACKAGES=${AIRFLOW_PRE_CACHED_PIP_PACKAGES} \ RUN bash /scripts/docker/install_pip_version.sh; \ if [[ ${AIRFLOW_PRE_CACHED_PIP_PACKAGES} == "true" && \ ${UPGRADE_TO_NEWER_DEPENDENCIES} == "false" ]]; then \ - bash /scripts/docker/install_airflow_from_branch_tip.sh; \ + bash /scripts/docker/install_airflow_dependencies_from_branch_tip.sh; \ fi COPY ${AIRFLOW_SOURCES_FROM} ${AIRFLOW_SOURCES_TO} @@ -236,14 +238,11 @@ ARG INSTALL_FROM_PYPI="true" # * pyjwt<2.0.0: flask-jwt-extended requires it # * dill<0.3.3 required by apache-beam ARG EAGER_UPGRADE_ADDITIONAL_REQUIREMENTS="pyjwt<2.0.0 dill<0.3.3 certifi<2021.0.0" -ARG CONTINUE_ON_PIP_CHECK_FAILURE="false" - ENV ADDITIONAL_PYTHON_DEPS=${ADDITIONAL_PYTHON_DEPS} \ INSTALL_FROM_DOCKER_CONTEXT_FILES=${INSTALL_FROM_DOCKER_CONTEXT_FILES} \ INSTALL_FROM_PYPI=${INSTALL_FROM_PYPI} \ - EAGER_UPGRADE_ADDITIONAL_REQUIREMENTS=${EAGER_UPGRADE_ADDITIONAL_REQUIREMENTS} \ - CONTINUE_ON_PIP_CHECK_FAILURE=${CONTINUE_ON_PIP_CHECK_FAILURE} + EAGER_UPGRADE_ADDITIONAL_REQUIREMENTS=${EAGER_UPGRADE_ADDITIONAL_REQUIREMENTS} WORKDIR /opt/airflow @@ -276,7 +275,7 @@ RUN if [[ -f /docker-context-files/requirements.txt ]]; then \ ARG BUILD_ID ARG COMMIT_SHA -ARG AIRFLOW_IMAGE_REPOSITORY="https://github.com/apache/airflow" +ARG AIRFLOW_IMAGE_REPOSITORY ARG AIRFLOW_IMAGE_DATE_CREATED ENV BUILD_ID=${BUILD_ID} COMMIT_SHA=${COMMIT_SHA} @@ -293,15 +292,14 @@ LABEL org.apache.airflow.distro="debian" \ org.opencontainers.image.created=${AIRFLOW_IMAGE_DATE_CREATED} \ org.opencontainers.image.authors="dev@airflow.apache.org" \ org.opencontainers.image.url="https://airflow.apache.org" \ - org.opencontainers.image.documentation="https://airflow.apache.org/docs/apache-airflow/stable/production-deployment.html" \ - org.opencontainers.image.source="https://github.com/apache/airflow" \ + org.opencontainers.image.documentation="https://airflow.apache.org/docs/docker-stack/index.html" \ org.opencontainers.image.version="${AIRFLOW_VERSION}" \ org.opencontainers.image.revision="${COMMIT_SHA}" \ org.opencontainers.image.vendor="Apache Software Foundation" \ org.opencontainers.image.licenses="Apache-2.0" \ org.opencontainers.image.ref.name="airflow-build-image" \ org.opencontainers.image.title="Build Image Segment for Production Airflow Image" \ - org.opencontainers.image.description="Installed Apache Airflow with build-time dependencies" + org.opencontainers.image.description="Reference build-time dependencies image for production-ready Apache Airflow image" ############################################################################################## # This is the actual Airflow image - much smaller than the build one. We copy @@ -379,7 +377,7 @@ ARG AIRFLOW_HOME ARG AIRFLOW_INSTALLATION_METHOD="apache-airflow" ARG BUILD_ID ARG COMMIT_SHA -ARG AIRFLOW_IMAGE_REPOSITORY="https://github.com/apache/airflow" +ARG AIRFLOW_IMAGE_REPOSITORY ARG AIRFLOW_IMAGE_DATE_CREATED # By default PIP will install everything in ~/.local ARG PIP_USER="true" @@ -468,15 +466,14 @@ LABEL org.apache.airflow.distro="debian" \ org.opencontainers.image.created=${AIRFLOW_IMAGE_DATE_CREATED} \ org.opencontainers.image.authors="dev@airflow.apache.org" \ org.opencontainers.image.url="https://airflow.apache.org" \ - org.opencontainers.image.documentation="https://airflow.apache.org/docs/apache-airflow/stable/production-deployment.html" \ - org.opencontainers.image.source="https://github.com/apache/airflow" \ + org.opencontainers.image.documentation="https://airflow.apache.org/docs/docker-stack/index.html" \ org.opencontainers.image.version="${AIRFLOW_VERSION}" \ org.opencontainers.image.revision="${COMMIT_SHA}" \ org.opencontainers.image.vendor="Apache Software Foundation" \ org.opencontainers.image.licenses="Apache-2.0" \ org.opencontainers.image.ref.name="airflow" \ org.opencontainers.image.title="Production Airflow Image" \ - org.opencontainers.image.description="Installed Apache Airflow" + org.opencontainers.image.description="Reference, production-ready Apache Airflow image" ENTRYPOINT ["/usr/bin/dumb-init", "--", "/entrypoint"] diff --git a/Dockerfile.ci b/Dockerfile.ci index 552afd8f4d418..13e9107137576 100644 --- a/Dockerfile.ci +++ b/Dockerfile.ci @@ -22,6 +22,8 @@ SHELL ["/bin/bash", "-o", "pipefail", "-e", "-u", "-x", "-c"] ARG PYTHON_BASE_IMAGE="python:3.6-slim-buster" ARG AIRFLOW_VERSION="2.2.0.dev0" +ARG AIRFLOW_IMAGE_REPOSITORY="https://github.com/apache/airflow" + # By increasing this number we can do force build of all dependencies ARG DEPENDENCIES_EPOCH_NUMBER="6" @@ -126,6 +128,15 @@ ARG RUNTIME_APT_DEPS="\ unzip \ vim \ xxd" + +# Install Helm +ARG HELM_VERSION="v3.6.3" + +RUN SYSTEM=$(uname -s | tr '[:upper:]' '[:lower:]') \ + && HELM_URL="https://get.helm.sh/helm-${HELM_VERSION}-${SYSTEM}-amd64.tar.gz" \ + && curl --location "${HELM_URL}" | tar -xvz -O "${SYSTEM}"-amd64/helm > /usr/local/bin/helm \ + && chmod +x /usr/local/bin/helm + ARG ADDITIONAL_RUNTIME_APT_DEPS="" ARG RUNTIME_APT_COMMAND="" ARG ADDITIONAL_RUNTIME_APT_COMMAND="" @@ -219,7 +230,7 @@ ARG AIRFLOW_PRE_CACHED_PIP_PACKAGES="true" # By default in the image, we are installing all providers when installing from sources ARG INSTALL_PROVIDERS_FROM_SOURCES="true" ARG INSTALL_FROM_PYPI="true" -ARG AIRFLOW_PIP_VERSION=21.1.2 +ARG AIRFLOW_PIP_VERSION=21.2.2 # Setup PIP # By default PIP install run without cache to make image smaller ARG PIP_NO_CACHE_DIR="true" @@ -282,7 +293,7 @@ ENV EAGER_UPGRADE_ADDITIONAL_REQUIREMENTS=${EAGER_UPGRADE_ADDITIONAL_REQUIREMENT RUN bash /scripts/docker/install_pip_version.sh; \ if [[ ${AIRFLOW_PRE_CACHED_PIP_PACKAGES} == "true" && \ ${UPGRADE_TO_NEWER_DEPENDENCIES} == "false" ]]; then \ - bash /scripts/docker/install_airflow_from_branch_tip.sh; \ + bash /scripts/docker/install_airflow_dependencies_from_branch_tip.sh; \ fi # Generate random hex dump file so that we can determine whether it's faster to rebuild the image @@ -311,8 +322,6 @@ COPY setup.cfg ${AIRFLOW_SOURCES}/setup.cfg COPY airflow/__init__.py ${AIRFLOW_SOURCES}/airflow/__init__.py -ARG CONTINUE_ON_PIP_CHECK_FAILURE="false" - # The goal of this line is to install the dependencies from the most current setup.py from sources # This will be usually incremental small set of packages in CI optimized build, so it will be very fast # In non-CI optimized build this will install all dependencies before installing sources. @@ -337,11 +346,13 @@ RUN chmod a+x /entrypoint COPY scripts/docker/load.bash /opt/bats/lib/ -# We can copy everything here. The Context is filtered by dockerignore. This makes sure we are not -# copying over stuff that is accidentally generated or that we do not need (such as egg-info) -# if you want to add something that is missing and you expect to see it in the image you can -# add it with ! in .dockerignore next to the airflow, test etc. directories there -COPY . ${AIRFLOW_SOURCES}/ +# Additional python deps to install +ARG ADDITIONAL_PYTHON_DEPS="" + +RUN bash /scripts/docker/install_pip_version.sh; \ + if [[ -n "${ADDITIONAL_PYTHON_DEPS}" ]]; then \ + bash /scripts/docker/install_additional_dependencies.sh; \ + fi # Install autocomplete for airflow RUN if command -v airflow; then \ @@ -351,27 +362,16 @@ RUN if command -v airflow; then \ # Install autocomplete for Kubectl RUN echo "source /etc/bash_completion" >> ~/.bashrc -WORKDIR ${AIRFLOW_SOURCES} - -# Install Helm -ARG HELM_VERSION="v3.2.4" - -RUN SYSTEM=$(uname -s | tr '[:upper:]' '[:lower:]') \ - && HELM_URL="https://get.helm.sh/helm-${HELM_VERSION}-${SYSTEM}-amd64.tar.gz" \ - && curl --location "${HELM_URL}" | tar -xvz -O "${SYSTEM}"-amd64/helm > /usr/local/bin/helm \ - && chmod +x /usr/local/bin/helm - -# Additional python deps to install -ARG ADDITIONAL_PYTHON_DEPS="" +# We can copy everything here. The Context is filtered by dockerignore. This makes sure we are not +# copying over stuff that is accidentally generated or that we do not need (such as egg-info) +# if you want to add something that is missing and you expect to see it in the image you can +# add it with ! in .dockerignore next to the airflow, test etc. directories there +COPY . ${AIRFLOW_SOURCES}/ -RUN bash /scripts/docker/install_pip_version.sh; \ - if [[ -n "${ADDITIONAL_PYTHON_DEPS}" ]]; then \ - bash /scripts/docker/install_additional_dependencies.sh; \ - fi +WORKDIR ${AIRFLOW_SOURCES} ARG BUILD_ID ARG COMMIT_SHA -ARG AIRFLOW_IMAGE_REPOSITORY="https://github.com/apache/airflow" ARG AIRFLOW_IMAGE_DATE_CREATED ENV PATH="/files/bin/:/opt/airflow/scripts/in_container/bin/:${HOME}:${PATH}" \ diff --git a/IMAGES.rst b/IMAGES.rst index 82e69895f3842..bc34e6c6abe37 100644 --- a/IMAGES.rst +++ b/IMAGES.rst @@ -445,12 +445,6 @@ The following build arguments (``--build-arg`` in docker build command) can be u | | | upgraded to newer versions matching | | | | setup.py before installation. | +------------------------------------------+------------------------------------------+------------------------------------------+ -| ``CONTINUE_ON_PIP_CHECK_FAILURE`` | ``false`` | By default the image will fail if pip | -| | | check fails for it. This is good for | -| | | interactive building but on CI the | -| | | image should be built regardless - we | -| | | have a separate step to verify image. | -+------------------------------------------+------------------------------------------+------------------------------------------+ | ``AIRFLOW_PRE_CACHED_PIP_PACKAGES`` | ``true`` | Allows to pre-cache airflow PIP packages | | | | from the GitHub of Apache Airflow | | | | This allows to optimize iterations for | diff --git a/INSTALL b/INSTALL index a195478d03446..554af5c25338e 100644 --- a/INSTALL +++ b/INSTALL @@ -104,3 +104,8 @@ ssh, statsd, tableau, telegram, trino, vertica, virtualenv, webhdfs, winrm, yand # END EXTRAS HERE # For installing Airflow in development environments - see CONTRIBUTING.rst + +# COMPILING FRONT-END ASSETS (in case you see "Please make sure to build the frontend in static/ directory and then restart the server") +# Optional : Installing yarn - https://classic.yarnpkg.com/en/docs/install + +python setup.py compile_assets diff --git a/INTHEWILD.md b/INTHEWILD.md index 3b63811d55db4..6f491d7f947b3 100644 --- a/INTHEWILD.md +++ b/INTHEWILD.md @@ -24,11 +24,13 @@ the platform. Please send a PR with your company name and @githubhandle. Currently, **officially** using Airflow: +1. [2RP Net](https://www.2rpnet.com.br/en) 1. [4G Capital](http://www.4g-capital.com/) [[@posei](https://github.com/posei)] 1. [6play](https://www.6play.fr) [[@lemourA](https://github.com/lemoura), [@achaussende](https://github.com/achaussende), [@d-nguyen](https://github.com/d-nguyen), [@julien-gm](https://github.com/julien-gm)] 1. [8fit](https://8fit.com/) [[@nicor88](https://github.com/nicor88), [@frnzska](https://github.com/frnzska)] 1. [90 Seconds](https://90seconds.tv/) [[@aaronmak](https://github.com/aaronmak)] 1. [99](https://99taxis.com) [[@fbenevides](https://github.com/fbenevides), [@gustavoamigo](https://github.com/gustavoamigo) & [@mmmaia](https://github.com/mmmaia)] +1. [Accenture](https://www.accenture.com/au-en) [[@nijanthanvijayakumar](https://github.com/nijanthanvijayakumar)] 1. [AdBOOST](https://www.adboost.sk) [[AdBOOST](https://github.com/AdBOOST)] 1. [Adobe](https://www.adobe.com/) [[@mishikaSingh](https://github.com/mishikaSingh), [@ramandumcs](https://github.com/ramandumcs), [@vardancse](https://github.com/vardancse)] 1. [Agari](https://github.com/agaridata) [[@r39132](https://github.com/r39132)] @@ -53,7 +55,7 @@ Currently, **officially** using Airflow: 1. [Arrive](https://www.arrive.com/) 1. [Artelys](https://www.artelys.com/) [[@fortierq](https://github.com/fortierq)] 1. [Asana](https://asana.com/) [[@chang](https://github.com/chang), [@dima-asana](https://github.com/dima-asana), [@jdavidheiser](https://github.com/jdavidheiser), [@ricardoandresrojas](https://github.com/ricardoandresrojas)] -1. [Astronomer](https://www.astronomer.io) [[@schnie](https://github.com/schnie), [@ashb](https://github.com/ashb), [@kaxil](https://github.com/kaxil), [@dimberman](https://github.com/dimberman), [@andriisoldatenko](https://github.com/andriisoldatenko), [@ryw](https://github.com/ryw), [@ryanahamilton](https://github.com/ryanahamilton), [@jhtimmins](https://github.com/jhtimmins), [@vikramkoka](https://github.com/vikramkoka), [@jedcunningham](https://github.com/jedcunningham)] +1. [Astronomer](https://www.astronomer.io) [[@schnie](https://github.com/schnie), [@ashb](https://github.com/ashb), [@kaxil](https://github.com/kaxil), [@dimberman](https://github.com/dimberman), [@andriisoldatenko](https://github.com/andriisoldatenko), [@ryw](https://github.com/ryw), [@ryanahamilton](https://github.com/ryanahamilton), [@jhtimmins](https://github.com/jhtimmins), [@vikramkoka](https://github.com/vikramkoka), [@jedcunningham](https://github.com/jedcunningham), [@BasPH](https://github.com/basph)] 1. [Auth0](https://auth0.com) [[@scottypate](https://github.com/scottypate)], [[@dm03514](https://github.com/dm03514)], [[@karangale](https://github.com/karangale)] 1. [Automattic](https://automattic.com/) [[@anandnalya](https://github.com/anandnalya), [@bperson](https://github.com/bperson), [@khrol](https://github.com/Khrol), [@xyu](https://github.com/xyu)] 1. [Avesta Technologies](https://avestatechnologies.com) [[@TheRum](https://github.com/TheRum)] @@ -186,7 +188,7 @@ Currently, **officially** using Airflow: 1. [GitLab](https://about.gitlab.com/) [[@tayloramurphy](https://gitlab.com/tayloramurphy) & [@m_walker](https://gitlab.com/m_walker)] 1. [Glassdoor](https://github.com/Glassdoor) [[@syvineckruyk](https://github.com/syvineckruyk) & [@sid88in](https://github.com/sid88in)] 1. [Global Fashion Group](http://global-fashion-group.com) [[@GFG](https://github.com/GFG)] -1. [GoDataDriven](https://godatadriven.com/) [[@BasPH](https://github.com/basph), [@danielvdende](https://github.com/danielvdende), [@ffinfo](https://github.com/ffinfo), [@Fokko](https://github.com/Fokko), [@gglanzani](https://github.com/gglanzani), [@hgrif](https://github.com/hgrif), [@jrderuiter](https://github.com/jrderuiter), [@NielsZeilemaker](https://github.com/NielsZeilemaker)] +1. [GoDataDriven](https://godatadriven.com/) [[@danielvdende](https://github.com/danielvdende), [@ffinfo](https://github.com/ffinfo), [@Fokko](https://github.com/Fokko), [@gglanzani](https://github.com/gglanzani), [@hgrif](https://github.com/hgrif), [@jrderuiter](https://github.com/jrderuiter), [@NielsZeilemaker](https://github.com/NielsZeilemaker)] 1. [Gojek](https://gojek.com/) [[@gojek](https://github.com/gojek), [@rootcss](https://github.com/rootcss)] 1. [GovTech GDS](https://gds-gov.tech) [[@chrissng](https://github.com/chrissng) & [@datagovsg](https://github.com/datagovsg)] 1. [Grab](https://www.grab.com/sg/) [[@calvintran](https://github.com/canhtran)] @@ -305,6 +307,7 @@ Currently, **officially** using Airflow: 1. [PayPal](https://www.paypal.com/) [[@r39132](https://github.com/r39132) & [@jhsenjaliya](https://github.com/jhsenjaliya)] 1. [Pecan](https://www.pecan.ai) [[@ohadmata](https://github.com/ohadmata)] 1. [Pernod-Ricard](https://www.pernod-ricard.com/) [[@romain-nio](https://github.com/romain-nio)] +1. [PEXA](https://www.pexa.com.au/) [[@andriyfedorov](https://github.com/andriyfedorov)] 1. [Plaid](https://www.plaid.com/) [[@plaid](https://github.com/plaid), [@AustinBGibbons](https://github.com/AustinBGibbons) & [@jeeyoungk](https://github.com/jeeyoungk)] 1. [Playbuzz](https://www.playbuzz.com/) [[@clintonboys](https://github.com/clintonboys) & [@dbn](https://github.com/dbn)] 1. [Playsimple Games](https://playsimple.in/) [[@joshi95](https://github.com/joshi95)] @@ -411,6 +414,7 @@ Currently, **officially** using Airflow: 1. [Vidora](https://www.vidora.com/) 1. [Ville de Montréal](http://ville.montreal.qc.ca/) [[@VilledeMontreal](https://github.com/VilledeMontreal/)] 1. [Vnomics](https://github.com/vnomics) [[@lpalum](https://github.com/lpalum)] +1. [Vodafone](https://www.vodafone.com) [[@nijanthanvijayakumar](https://github.com/nijanthanvijayakumar)] 1. [Walmart Labs](https://www.walmartlabs.com) [[@bharathpalaksha](https://github.com/bharathpalaksha), [@vipul007ravi](https://github.com/vipul007ravi), [@pateash](https://github.com/pateash)] 1. [Waze](https://www.waze.com) [[@waze](https://github.com/wazeHQ)] 1. [WePay](http://www.wepay.com) [[@criccomini](https://github.com/criccomini) & [@mtagle](https://github.com/mtagle)] diff --git a/UPDATING.md b/UPDATING.md index 69cd409fef3bf..b18dff578e76c 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -73,6 +73,12 @@ https://developers.google.com/style/inclusive-documentation --> +### Dummy trigger rule has been deprecated + +`TriggerRule.DUMMY` is replaced by `TriggerRule.ALWAYS`. +This is only name change, no functionality changes made. +This change is backward compatible however `TriggerRule.DUMMY` will be removed in next major release. + ### DAG concurrency settings have been renamed `[core] dag_concurrency` setting in `airflow.cfg` has been renamed to `[core] max_active_tasks_per_dag` @@ -127,6 +133,10 @@ If you are using DAGs Details API endpoint, use `max_active_tasks` instead of `c When marking a task success/failed in Graph View, its downstream tasks that are in failed/upstream_failed state are automatically cleared. +### Clearing a running task sets its state to `RESTARTING` + +Previously, clearing a running task sets its state to `SHUTDOWN`. The task gets killed and goes into `FAILED` state. After [#16681](https://github.com/apache/airflow/pull/16681), clearing a running task sets its state to `RESTARTING`. The task is eligible for retry without going into `FAILED` state. + ### Remove `TaskInstance.log_filepath` attribute This method returned incorrect values for a long time, because it did not take into account the different @@ -151,6 +161,12 @@ not have any effect in an existing deployment where the ``default_pool`` already Previously this was controlled by `non_pooled_task_slot_count` in `[core]` section, which was not documented. +### Webserver DAG refresh buttons removed + +Now that the DAG parser syncs DAG permissions there is no longer a need for manually refreshing DAGs. As such, the buttons to refresh a DAG have been removed from the UI. + +In addition, the `/refresh` and `/refresh_all` webserver endpoints have also been removed. + ## Airflow 2.1.1 ### `activate_dag_runs` argument of the function `clear_task_instances` is replaced with `dag_run_state` diff --git a/airflow/__init__.py b/airflow/__init__.py index a448491057467..9f9073eb796d4 100644 --- a/airflow/__init__.py +++ b/airflow/__init__.py @@ -74,11 +74,13 @@ def __getattr__(name): if not settings.LAZY_LOAD_PROVIDERS: from airflow import providers_manager - providers_manager.ProvidersManager().initialize_providers_manager() + manager = providers_manager.ProvidersManager() + manager.initialize_providers_list() + manager.initialize_providers_hooks() + manager.initialize_providers_extra_links() # This is never executed, but tricks static analyzers (PyDev, PyCharm,) -# into knowing the types of these symbols, and what # they contain. STATICA_HACK = True globals()['kcah_acitats'[::-1].upper()] = False diff --git a/airflow/cli/cli_parser.py b/airflow/cli/cli_parser.py index b6528384d1fa0..1cca4c2a0dd3d 100644 --- a/airflow/cli/cli_parser.py +++ b/airflow/cli/cli_parser.py @@ -1383,7 +1383,7 @@ class GroupCommand(NamedTuple): name='stop', help="Stop the Celery worker gracefully", func=lazy_load_command('airflow.cli.commands.celery_command.stop_worker'), - args=(), + args=(ARG_PID,), ), ) diff --git a/airflow/cli/commands/celery_command.py b/airflow/cli/commands/celery_command.py index ba3c45e58a111..e9c3e38fd5a15 100644 --- a/airflow/cli/commands/celery_command.py +++ b/airflow/cli/commands/celery_command.py @@ -183,7 +183,10 @@ def worker(args): def stop_worker(args): """Sends SIGTERM to Celery worker""" # Read PID from file - pid_file_path, _, _, _ = setup_locations(process=WORKER_PROCESS_NAME) + if args.pid: + pid_file_path = args.pid + else: + pid_file_path, _, _, _ = setup_locations(process=WORKER_PROCESS_NAME) pid = read_pid_from_pidfile(pid_file_path) # Send SIGTERM diff --git a/airflow/cli/commands/kubernetes_command.py b/airflow/cli/commands/kubernetes_command.py index 3c3c8e68cfca3..2660daeb38bfb 100644 --- a/airflow/cli/commands/kubernetes_command.py +++ b/airflow/cli/commands/kubernetes_command.py @@ -96,16 +96,8 @@ def cleanup_pods(args): 'try_number', 'airflow_version', ] - list_kwargs = { - "namespace": namespace, - "limit": 500, - "label_selector": client.V1LabelSelector( - match_expressions=[ - client.V1LabelSelectorRequirement(key=label, operator="Exists") - for label in airflow_pod_labels - ] - ), - } + list_kwargs = {"namespace": namespace, "limit": 500, "label_selector": ','.join(airflow_pod_labels)} + while True: pod_list = kube_client.list_namespaced_pod(**list_kwargs) for pod in pod_list.items: diff --git a/airflow/cli/commands/scheduler_command.py b/airflow/cli/commands/scheduler_command.py index 368db6f756e37..44674f04c6af4 100644 --- a/airflow/cli/commands/scheduler_command.py +++ b/airflow/cli/commands/scheduler_command.py @@ -29,17 +29,21 @@ from airflow.utils.cli import process_subdir, setup_locations, setup_logging, sigint_handler, sigquit_handler +def _create_scheduler_job(args): + job = SchedulerJob( + subdir=process_subdir(args.subdir), + num_runs=args.num_runs, + do_pickle=args.do_pickle, + ) + return job + + @cli_utils.action_logging def scheduler(args): """Starts Airflow Scheduler""" skip_serve_logs = args.skip_serve_logs print(settings.HEADER) - job = SchedulerJob( - subdir=process_subdir(args.subdir), - num_runs=args.num_runs, - do_pickle=args.do_pickle, - ) if args.daemon: pid, stdout, stderr, log_file = setup_locations( @@ -54,9 +58,11 @@ def scheduler(args): stderr=stderr_handle, ) with ctx: + job = _create_scheduler_job(args) sub_proc = _serve_logs(skip_serve_logs) job.run() else: + job = _create_scheduler_job(args) signal.signal(signal.SIGINT, sigint_handler) signal.signal(signal.SIGTERM, sigint_handler) signal.signal(signal.SIGQUIT, sigquit_handler) diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml index fd60a7a29dc16..c83c1c48dd1f0 100644 --- a/airflow/config_templates/config.yml +++ b/airflow/config_templates/config.yml @@ -1831,36 +1831,30 @@ default: "True" - name: max_dagruns_to_create_per_loop description: | - Max number of DAGs to create DagRuns for per scheduler loop - - Default: 10 + Max number of DAGs to create DagRuns for per scheduler loop. example: ~ version_added: 2.0.0 type: string - default: ~ + default: "10" see_also: ":ref:`scheduler:ha:tunables`" - name: max_dagruns_per_loop_to_schedule description: | How many DagRuns should a scheduler examine (and lock) when scheduling and queuing tasks. - - Default: 20 example: ~ version_added: 2.0.0 type: string - default: ~ + default: "20" see_also: ":ref:`scheduler:ha:tunables`" - name: schedule_after_task_execution description: | Should the Task supervisor process perform a "mini scheduler" to attempt to schedule more tasks of the same DAG. Leaving this on will mean tasks in the same DAG execute quicker, but might starve out other dags in some circumstances - - Default: True example: ~ version_added: 2.0.0 type: boolean - default: ~ + default: "True" - name: parsing_processes description: | The scheduler can run multiple processes in parallel to parse dags. diff --git a/airflow/config_templates/default_airflow.cfg b/airflow/config_templates/default_airflow.cfg index f3b4589da3216..d558422426f84 100644 --- a/airflow/config_templates/default_airflow.cfg +++ b/airflow/config_templates/default_airflow.cfg @@ -917,23 +917,17 @@ max_tis_per_query = 512 # scheduler at once use_row_level_locking = True -# Max number of DAGs to create DagRuns for per scheduler loop -# -# Default: 10 -# max_dagruns_to_create_per_loop = +# Max number of DAGs to create DagRuns for per scheduler loop. +max_dagruns_to_create_per_loop = 10 # How many DagRuns should a scheduler examine (and lock) when scheduling # and queuing tasks. -# -# Default: 20 -# max_dagruns_per_loop_to_schedule = +max_dagruns_per_loop_to_schedule = 20 # Should the Task supervisor process perform a "mini scheduler" to attempt to schedule more tasks of the # same DAG. Leaving this on will mean tasks in the same DAG execute quicker, but might starve out other # dags in some circumstances -# -# Default: True -# schedule_after_task_execution = +schedule_after_task_execution = True # The scheduler can run multiple processes in parallel to parse dags. # This defines how many processes will run. diff --git a/airflow/configuration.py b/airflow/configuration.py index a6f67cc73da5d..b74393b3c2458 100644 --- a/airflow/configuration.py +++ b/airflow/configuration.py @@ -242,7 +242,7 @@ def _validate_config_dependencies(self): if StrictVersion(sqlite3.sqlite_version) < StrictVersion(min_sqlite_version): raise AirflowConfigException( f"error: sqlite C library version too old (< {min_sqlite_version}). " - f"See {get_docs_url('howto/set-up-database.rst#setting-up-a-sqlite-database')}" + f"See {get_docs_url('howto/set-up-database.html#setting-up-a-sqlite-database')}" ) if self.has_option('core', 'mp_start_method'): diff --git a/airflow/jobs/base_job.py b/airflow/jobs/base_job.py index 18893f2110a4f..745f248fc4da0 100644 --- a/airflow/jobs/base_job.py +++ b/airflow/jobs/base_job.py @@ -202,7 +202,7 @@ def heartbeat(self, only_if_necessary: bool = False): session.merge(self) previous_heartbeat = self.latest_heartbeat - if self.state == State.SHUTDOWN: + if self.state in State.terminating_states: self.kill() # Figure out how long to sleep for diff --git a/airflow/jobs/local_task_job.py b/airflow/jobs/local_task_job.py index 6852576fc8a90..203a7a82b20b3 100644 --- a/airflow/jobs/local_task_job.py +++ b/airflow/jobs/local_task_job.py @@ -78,12 +78,9 @@ def _execute(self): def signal_handler(signum, frame): """Setting kill signal handler""" self.log.error("Received SIGTERM. Terminating subprocesses") - self.on_kill() - self.task_instance.refresh_from_db() - if self.task_instance.state not in State.finished: - self.task_instance.set_state(State.FAILED) - self.task_instance._run_finished_callback(error="task received sigterm") - raise AirflowException("LocalTaskJob received SIGTERM signal") + self.task_runner.terminate() + self.handle_task_exit(128 + signum) + return signal.signal(signal.SIGTERM, signal_handler) @@ -148,16 +145,19 @@ def signal_handler(signum, frame): self.on_kill() def handle_task_exit(self, return_code: int) -> None: - """Handle case where self.task_runner exits by itself""" + """Handle case where self.task_runner exits by itself or is externally killed""" + # Without setting this, heartbeat may get us + self.terminating = True self.log.info("Task exited with return code %s", return_code) self.task_instance.refresh_from_db() - # task exited by itself, so we need to check for error file + + if self.task_instance.state == State.RUNNING: + # This is for a case where the task received a SIGKILL + # while running or the task runner received a sigterm + self.task_instance.handle_failure(error=None) + # We need to check for error file # in case it failed due to runtime exception/error error = None - if self.task_instance.state == State.RUNNING: - # This is for a case where the task received a sigkill - # while running - self.task_instance.set_state(State.FAILED) if self.task_instance.state != State.SUCCESS: error = self.task_runner.deserialize_run_error() self.task_instance._run_finished_callback(error=error) @@ -192,11 +192,17 @@ def heartbeat_callback(self, session=None): ) raise AirflowException("Hostname of job runner does not match") current_pid = self.task_runner.process.pid - same_process = ti.pid == current_pid - if ti.run_as_user: - same_process = psutil.Process(ti.pid).ppid() == current_pid - if ti.pid is not None and not same_process: - self.log.warning("Recorded pid %s does not match " "the current pid %s", ti.pid, current_pid) + recorded_pid = ti.pid + same_process = recorded_pid == current_pid + + if ti.run_as_user or self.task_runner.run_as_user: + recorded_pid = psutil.Process(ti.pid).ppid() + same_process = recorded_pid == current_pid + + if recorded_pid is not None and not same_process: + self.log.warning( + "Recorded pid %s does not match the current pid %s", recorded_pid, current_pid + ) raise AirflowException("PID of job runner does not match") elif self.task_runner.return_code() is None and hasattr(self.task_runner, 'process'): self.log.warning( diff --git a/airflow/jobs/scheduler_job.py b/airflow/jobs/scheduler_job.py index 5c201e39a4163..d2a19cae93f3c 100644 --- a/airflow/jobs/scheduler_job.py +++ b/airflow/jobs/scheduler_job.py @@ -496,18 +496,8 @@ def _enqueue_task_instances_with_queued_state(self, task_instances: List[TI]) -> """ # actually enqueue them for ti in task_instances: - command = TI.generate_command( - ti.dag_id, - ti.task_id, - ti.execution_date, + command = ti.command_as_list( local=True, - mark_success=False, - ignore_all_deps=False, - ignore_depends_on_past=False, - ignore_task_deps=False, - ignore_ti_state=False, - pool=ti.pool, - file_path=ti.dag_model.fileloc, pickle_id=ti.dag_model.pickle_id, ) @@ -834,6 +824,7 @@ def _do_scheduling(self, session) -> int: # Bulk fetch the currently active dag runs for the dags we are # examining, rather than making one query per DagRun + callback_tuples = [] for dag_run in dag_runs: # Use try_except to not stop the Scheduler when a Serialized DAG is not found # This takes care of Dynamic DAGs especially @@ -842,13 +833,18 @@ def _do_scheduling(self, session) -> int: # But this would take care of the scenario when the Scheduler is restarted after DagRun is # created and the DAG is deleted / renamed try: - self._schedule_dag_run(dag_run, session) + callback_to_run = self._schedule_dag_run(dag_run, session) + callback_tuples.append((dag_run, callback_to_run)) except SerializedDagNotFound: self.log.exception("DAG '%s' not found in serialized_dag table", dag_run.dag_id) continue guard.commit() + # 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) + # Without this, the session has an invalid view of the DB session.expunge_all() # END: schedule TIs @@ -1010,12 +1006,12 @@ def _schedule_dag_run( self, dag_run: DagRun, session: Session, - ) -> int: + ) -> Optional[DagCallbackRequest]: """ Make scheduling decisions about an individual dag run :param dag_run: The DagRun to schedule - :return: Number of tasks scheduled + :return: Callback that needs to be executed """ dag = dag_run.dag = self.dagbag.get_dag(dag_run.dag_id, session=session) @@ -1062,13 +1058,13 @@ def _schedule_dag_run( # TODO[HA]: Rename update_state -> schedule_dag_run, ?? something else? schedulable_tis, callback_to_run = dag_run.update_state(session=session, execute_callbacks=False) - self._send_dag_callbacks_to_processor(dag_run, callback_to_run) - # 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 # IDs in a single query, but it turns out that can be _very very slow_ # see #11147/commit ee90807ac for more details - return dag_run.schedule_tis(schedulable_tis, session) + dag_run.schedule_tis(schedulable_tis, session) + + return callback_to_run @provide_session def _verify_integrity_if_dag_changed(self, dag_run: DagRun, session=None): diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py index 5018a761a0c38..f584154162f04 100644 --- a/airflow/models/baseoperator.py +++ b/airflow/models/baseoperator.py @@ -348,7 +348,7 @@ class derived from this one results in the creation of a task object, :param trigger_rule: defines the rule by which dependencies are applied for the task to get triggered. Options are: ``{ all_success | all_failed | all_done | one_success | - one_failed | none_failed | none_failed_or_skipped | none_skipped | dummy}`` + one_failed | none_failed | none_failed_or_skipped | none_skipped | always}`` default is ``all_success``. Options can be set as string or using the constants defined in the static class ``airflow.utils.TriggerRule`` @@ -542,6 +542,14 @@ def __init__( if end_date: self.end_date = timezone.convert_to_utc(end_date) + if trigger_rule == "dummy": + warnings.warn( + "dummy Trigger Rule is deprecated. Please use `TriggerRule.ALWAYS`.", + DeprecationWarning, + stacklevel=2, + ) + trigger_rule = TriggerRule.ALWAYS + if not TriggerRule.is_valid(trigger_rule): raise AirflowException( "The trigger_rule must be one of {all_triggers}," diff --git a/airflow/models/dag.py b/airflow/models/dag.py index 22e781c23f14e..82c4d587c7e7c 100644 --- a/airflow/models/dag.py +++ b/airflow/models/dag.py @@ -20,6 +20,7 @@ import functools import logging import os +import pathlib import pickle import re import sys @@ -236,13 +237,21 @@ class DAG(LoggingMixin): 'parent_dag', 'start_date', 'schedule_interval', - 'full_filepath', + 'fileloc', 'template_searchpath', 'last_loaded', } __serialized_fields: Optional[FrozenSet[str]] = None + fileloc: str + """ + File path that needs to be imported to load this DAG or subdag. + + This may not be an actual file on disk in the case when this DAG is loaded + from a ZIP file or other DAG distribution format. + """ + def __init__( self, dag_id: str, @@ -286,10 +295,16 @@ def __init__( self.params.update(self.default_args['params']) del self.default_args['params'] + if full_filepath: + warnings.warn( + "Passing full_filepath to DAG() is deprecated and has no effect", + DeprecationWarning, + stacklevel=2, + ) + validate_key(dag_id) self._dag_id = dag_id - self._full_filepath = full_filepath if full_filepath else '' if concurrency and not max_active_tasks: # TODO: Remove in Airflow 3.0 warnings.warn( @@ -655,11 +670,22 @@ def dag_id(self, value: str) -> None: @property def full_filepath(self) -> str: - return self._full_filepath + """:meta private:""" + warnings.warn( + "DAG.full_filepath is deprecated in favour of fileloc", + DeprecationWarning, + stacklevel=2, + ) + return self.fileloc @full_filepath.setter def full_filepath(self, value) -> None: - self._full_filepath = value + warnings.warn( + "DAG.full_filepath is deprecated in favour of fileloc", + DeprecationWarning, + stacklevel=2, + ) + self.fileloc = value @property def concurrency(self) -> int: @@ -735,15 +761,26 @@ def task_group(self) -> "TaskGroup": @property def filepath(self) -> str: - """File location of where the dag object is instantiated""" - fn = self.full_filepath.replace(settings.DAGS_FOLDER + '/', '') - fn = fn.replace(os.path.dirname(__file__) + '/', '') - return fn + """:meta private:""" + warnings.warn( + "filepath is deprecated, use relative_fileloc instead", DeprecationWarning, stacklevel=2 + ) + return str(self.relative_fileloc) + + @property + def relative_fileloc(self) -> pathlib.Path: + """File location of the importable dag 'file' relative to the configured DAGs folder.""" + path = pathlib.Path(self.fileloc) + try: + return path.relative_to(settings.DAGS_FOLDER) + except ValueError: + # Not relative to DAGS_FOLDER. + return path @property def folder(self) -> str: """Folder location of where the DAG object is instantiated.""" - return os.path.dirname(self.full_filepath) + return os.path.dirname(self.fileloc) @property def owner(self) -> str: @@ -2118,9 +2155,11 @@ def bulk_write_to_db(cls, dags: Collection["DAG"], session=None): .group_by(DagRun.dag_id) .all() ) + filelocs = [] for orm_dag in sorted(orm_dags, key=lambda d: d.dag_id): dag = dag_by_ids[orm_dag.dag_id] + filelocs.append(dag.fileloc) if dag.is_subdag: orm_dag.is_subdag = True orm_dag.fileloc = dag.parent_dag.fileloc # type: ignore @@ -2157,7 +2196,7 @@ def bulk_write_to_db(cls, dags: Collection["DAG"], session=None): session.add(dag_tag_orm) if settings.STORE_DAG_CODE: - DagCode.bulk_sync_to_db([dag.fileloc for dag in orm_dags]) + DagCode.bulk_sync_to_db(filelocs) # Issue SQL/finish "Unit of Work", but let @provide_session commit (or if passed a session, let caller # decide when to commit @@ -2274,7 +2313,6 @@ def get_serialized_fields(cls): '_old_context_manager_dags', 'safe_dag_id', 'last_loaded', - '_full_filepath', 'user_defined_filters', 'user_defined_macros', 'partial', @@ -2382,6 +2420,10 @@ class DagModel(Base): Index('idx_next_dagrun_create_after', next_dagrun_create_after, unique=False), ) + parent_dag = relationship( + "DagModel", remote_side=[dag_id], primaryjoin=root_dag_id == dag_id, foreign_keys=[root_dag_id] + ) + NUM_DAGS_PER_DAGRUN_QUERY = conf.getint('scheduler', 'max_dagruns_to_create_per_loop', fallback=10) def __init__(self, concurrency=None, **kwargs): @@ -2410,7 +2452,7 @@ def timezone(self): @staticmethod @provide_session def get_dagmodel(dag_id, session=None): - return session.query(DagModel).filter(DagModel.dag_id == dag_id).first() + return session.query(DagModel).options(joinedload(DagModel.parent_dag)).get(dag_id) @classmethod @provide_session @@ -2455,6 +2497,18 @@ def get_default_view(self) -> str: def safe_dag_id(self): return self.dag_id.replace('.', '__dot__') + @property + def relative_fileloc(self) -> Optional[pathlib.Path]: + """File location of the importable dag 'file' relative to the configured DAGs folder.""" + if self.fileloc is None: + return None + path = pathlib.Path(self.fileloc) + try: + return path.relative_to(settings.DAGS_FOLDER) + except ValueError: + # Not relative to DAGS_FOLDER. + return path + @provide_session def set_is_paused(self, is_paused: bool, including_subdags: bool = True, session=None) -> None: """ diff --git a/airflow/models/dagbag.py b/airflow/models/dagbag.py index 3e40678a912ba..b1249b518b9bd 100644 --- a/airflow/models/dagbag.py +++ b/airflow/models/dagbag.py @@ -381,16 +381,12 @@ def _load_modules_from_zip(self, filepath, safe_mode): def _process_modules(self, filepath, mods, file_last_changed_on_disk): from airflow.models.dag import DAG # Avoid circular import - is_zipfile = zipfile.is_zipfile(filepath) - top_level_dags = [o for m in mods for o in list(m.__dict__.values()) if isinstance(o, DAG)] + top_level_dags = ((o, m) for m in mods for o in m.__dict__.values() if isinstance(o, DAG)) found_dags = [] - for dag in top_level_dags: - if not dag.full_filepath: - dag.full_filepath = filepath - if dag.fileloc != filepath and not is_zipfile: - dag.fileloc = filepath + for (dag, mod) in top_level_dags: + dag.fileloc = mod.__file__ try: dag.is_subdag = False dag.timetable.validate() @@ -398,17 +394,17 @@ def _process_modules(self, filepath, mods, file_last_changed_on_disk): found_dags.append(dag) found_dags += dag.subdags except AirflowTimetableInvalid as exception: - self.log.exception("Failed to bag_dag: %s", dag.full_filepath) - self.import_errors[dag.full_filepath] = f"Invalid timetable expression: {exception}" - self.file_last_changed[dag.full_filepath] = file_last_changed_on_disk + self.log.exception("Failed to bag_dag: %s", dag.fileloc) + self.import_errors[dag.fileloc] = f"Invalid timetable expression: {exception}" + self.file_last_changed[dag.fileloc] = file_last_changed_on_disk except ( AirflowDagCycleException, AirflowDagDuplicatedIdException, AirflowClusterPolicyViolation, ) as exception: - self.log.exception("Failed to bag_dag: %s", dag.full_filepath) - self.import_errors[dag.full_filepath] = str(exception) - self.file_last_changed[dag.full_filepath] = file_last_changed_on_disk + self.log.exception("Failed to bag_dag: %s", dag.fileloc) + self.import_errors[dag.fileloc] = str(exception) + self.file_last_changed[dag.fileloc] = file_last_changed_on_disk return found_dags def bag_dag(self, dag, root_dag): @@ -444,17 +440,17 @@ def _bag_dag(self, *, dag, root_dag, recursive): # into further _bag_dag() calls. if recursive: for subdag in subdags: - subdag.full_filepath = dag.full_filepath + subdag.fileloc = dag.fileloc subdag.parent_dag = dag subdag.is_subdag = True self._bag_dag(dag=subdag, root_dag=root_dag, recursive=False) prev_dag = self.dags.get(dag.dag_id) - if prev_dag and prev_dag.full_filepath != dag.full_filepath: + if prev_dag and prev_dag.fileloc != dag.fileloc: raise AirflowDagDuplicatedIdException( dag_id=dag.dag_id, - incoming=dag.full_filepath, - existing=self.dags[dag.dag_id].full_filepath, + incoming=dag.fileloc, + existing=self.dags[dag.dag_id].fileloc, ) self.dags[dag.dag_id] = dag self.log.debug('Loaded DAG %s', dag) @@ -594,6 +590,7 @@ def _serialize_dag_capturing_errors(dag, session): except OperationalError: raise except Exception: + self.log.exception("Failed to write serialized DAG: %s", dag.full_filepath) return [(dag.fileloc, traceback.format_exc(limit=-self.dagbag_import_error_traceback_depth))] # Retry 'DAG.bulk_write_to_db' & 'SerializedDagModel.bulk_sync_to_db' in case diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py index e2057d5a34af0..1fdd26cd0a611 100644 --- a/airflow/models/dagrun.py +++ b/airflow/models/dagrun.py @@ -477,7 +477,7 @@ def task_instance_scheduling_decisions(self, session: Session = None) -> TISched schedulable_tis: List[TI] = [] changed_tis = False - tis = list(self.get_task_instances(session=session, state=State.task_states + (State.SHUTDOWN,))) + tis = list(self.get_task_instances(session=session, state=State.task_states)) self.log.debug("number of tis tasks for %s: %s task(s)", self, len(tis)) for ti in tis: try: diff --git a/airflow/models/serialized_dag.py b/airflow/models/serialized_dag.py index 1dabc0ed53d6b..98d933f29425f 100644 --- a/airflow/models/serialized_dag.py +++ b/airflow/models/serialized_dag.py @@ -90,7 +90,7 @@ class SerializedDagModel(Base): def __init__(self, dag: DAG): self.dag_id = dag.dag_id - self.fileloc = dag.full_filepath + self.fileloc = dag.fileloc self.fileloc_hash = DagCode.dag_fileloc_hash(self.fileloc) self.data = SerializedDAG.to_dict(dag) self.last_updated = timezone.utcnow() diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 6fb437b012b68..1b76145a469e2 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -26,7 +26,7 @@ from collections import defaultdict from datetime import datetime, timedelta from tempfile import NamedTemporaryFile -from typing import IO, Any, Dict, Iterable, List, NamedTuple, Optional, Tuple, Union +from typing import IO, TYPE_CHECKING, Any, Dict, Iterable, List, NamedTuple, Optional, Tuple, Union from urllib.parse import quote import dill @@ -91,6 +91,10 @@ log = logging.getLogger(__name__) +if TYPE_CHECKING: + from airflow.models.dag import DAG, DagModel + + @contextlib.contextmanager def set_current_context(context: Context): """ @@ -158,7 +162,9 @@ def clear_task_instances( for ti in tis: if ti.state == State.RUNNING: if ti.job_id: - ti.state = State.SHUTDOWN + # If a task is cleared when running, set its state to RESTARTING so that + # the task is terminated and becomes eligible for retry. + ti.state = State.RESTARTING job_ids.append(ti.job_id) else: task_id = ti.task_id @@ -211,7 +217,7 @@ def clear_task_instances( from airflow.jobs.base_job import BaseJob for job in session.query(BaseJob).filter(BaseJob.id.in_(job_ids)).all(): - job.state = State.SHUTDOWN + job.state = State.RESTARTING if activate_dag_runs is not None: warnings.warn( @@ -434,15 +440,25 @@ def command_as_list( installed. This command is part of the message sent to executors by the orchestrator. """ - dag = self.task.dag + dag: Union["DAG", "DagModel"] + # Use the dag if we have it, else fallback to the ORM dag_model, which might not be loaded + if hasattr(self, 'task') and hasattr(self.task, 'dag'): + dag = self.task.dag + else: + dag = self.dag_model should_pass_filepath = not pickle_id and dag - if should_pass_filepath and dag.full_filepath != dag.filepath: - path = f"DAGS_FOLDER/{dag.filepath}" - elif should_pass_filepath and dag.full_filepath: - path = dag.full_filepath - else: - path = None + path = None + if should_pass_filepath: + if dag.is_subdag: + path = dag.parent_dag.relative_fileloc + else: + path = dag.relative_fileloc + + if path: + if not path.is_absolute(): + path = 'DAGS_FOLDER' / path + path = str(path) return TaskInstance.generate_command( self.dag_id, @@ -1042,6 +1058,7 @@ def check_and_change_state_before_execution( self.refresh_from_db(session=session, lock_for_update=True) self.job_id = job_id self.hostname = get_hostname() + self.pid = None if not ignore_all_deps and not ignore_ti_state and self.state == State.SUCCESS: Stats.incr('previously_succeeded', 1, 1) @@ -1519,6 +1536,11 @@ def handle_failure_with_callback( def is_eligible_to_retry(self): """Is task instance is eligible for retry""" + if self.state == State.RESTARTING: + # If a task is cleared when running, it goes into RESTARTING state and is always + # eligible for retry + return True + return self.task.retries and self.try_number <= self.max_tries @provide_session diff --git a/airflow/operators/python.py b/airflow/operators/python.py index c13b002429598..511cb992d7611 100644 --- a/airflow/operators/python.py +++ b/airflow/operators/python.py @@ -18,6 +18,7 @@ import inspect import os import pickle +import shutil import sys import types import warnings @@ -325,6 +326,8 @@ def __init__( "Passing op_args or op_kwargs is not supported across different Python " "major versions for PythonVirtualenvOperator. Please use string_args." ) + if not shutil.which("virtualenv"): + raise AirflowException('PythonVirtualenvOperator requires virtualenv, please install it.') super().__init__( python_callable=python_callable, op_args=op_args, diff --git a/airflow/providers/airbyte/CHANGELOG.rst b/airflow/providers/airbyte/CHANGELOG.rst index 9a9371c0fb15f..1c42b91c6adfe 100644 --- a/airflow/providers/airbyte/CHANGELOG.rst +++ b/airflow/providers/airbyte/CHANGELOG.rst @@ -22,14 +22,6 @@ Changelog 2.1.0 ..... -Breaking changes -~~~~~~~~~~~~~~~~ - - -Features -~~~~~~~~ - - Bug Fixes ~~~~~~~~~ diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index eb8e48949f160..57010ed8097b8 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -37,6 +37,7 @@ import tenacity from botocore.config import Config from botocore.credentials import ReadOnlyCredentials +from slugify import slugify try: from functools import cached_property @@ -188,11 +189,14 @@ def _read_credentials_from_connection(self) -> Tuple[Optional[str], Optional[str self.log.info("No credentials retrieved from Connection") return aws_access_key_id, aws_secret_access_key + def _strip_invalid_session_name_characters(self, role_session_name: str) -> str: + return slugify(role_session_name, regex_pattern=r'[^\w+=,.@-]+') + def _assume_role(self, sts_client: boto3.client) -> Dict: assume_role_kwargs = self.extra_config.get("assume_role_kwargs", {}) if "external_id" in self.extra_config: # Backwards compatibility assume_role_kwargs["ExternalId"] = self.extra_config.get("external_id") - role_session_name = f"Airflow_{self.conn.conn_id}" + role_session_name = self._strip_invalid_session_name_characters(f"Airflow_{self.conn.conn_id}") self.log.info( "Doing sts_client.assume_role to role_arn=%s (role_session_name=%s)", self.role_arn, diff --git a/airflow/providers/amazon/aws/operators/ecs.py b/airflow/providers/amazon/aws/operators/ecs.py index 07f5702815ae9..43baa3c92662a 100644 --- a/airflow/providers/amazon/aws/operators/ecs.py +++ b/airflow/providers/amazon/aws/operators/ecs.py @@ -169,7 +169,7 @@ def __init__( group: Optional[str] = None, placement_constraints: Optional[list] = None, placement_strategy: Optional[list] = None, - platform_version: str = 'LATEST', + platform_version: Optional[str] = None, network_configuration: Optional[dict] = None, tags: Optional[dict] = None, awslogs_group: Optional[str] = None, @@ -254,11 +254,10 @@ def _start_task(self, context): if self.capacity_provider_strategy: run_opts['capacityProviderStrategy'] = self.capacity_provider_strategy - run_opts['platformVersion'] = self.platform_version elif self.launch_type: run_opts['launchType'] = self.launch_type - if self.launch_type == 'FARGATE': - run_opts['platformVersion'] = self.platform_version + if self.platform_version is not None: + run_opts['platformVersion'] = self.platform_version if self.group is not None: run_opts['group'] = self.group if self.placement_constraints is not None: diff --git a/airflow/providers/amazon/aws/sensors/sqs.py b/airflow/providers/amazon/aws/sensors/sqs.py index dc6217b5bd0a3..e2d04e1f1b0bd 100644 --- a/airflow/providers/amazon/aws/sensors/sqs.py +++ b/airflow/providers/amazon/aws/sensors/sqs.py @@ -16,7 +16,11 @@ # specific language governing permissions and limitations # under the License. """Reads and then deletes the message from SQS queue""" -from typing import Optional +import json +from typing import Any, Optional + +from jsonpath_ng import parse +from typing_extensions import Literal from airflow.exceptions import AirflowException from airflow.providers.amazon.aws.hooks.sqs import SQSHook @@ -37,9 +41,26 @@ class SQSSensor(BaseSensorOperator): :type max_messages: int :param wait_time_seconds: The time in seconds to wait for receiving messages (default: 1 second) :type wait_time_seconds: int + :param visibility_timeout: Visibility timeout, a period of time during which + Amazon SQS prevents other consumers from receiving and processing the message. + :type visibility_timeout: Optional[Int] + :param message_filtering: Specified how received messages should be filtered. Supported options are: + `None` (no filtering, default), `'literal'` (message Body literal match) or `'jsonpath'` + (message Body filtered using a JSONPath expression). + You may add further methods by overriding the relevant class methods. + :type message_filtering: Optional[Literal["literal", "jsonpath"]] + :param message_filtering_match_values: Optional value/s for the message filter to match on. + For example, with literal matching, if a message body matches any of the specified values + then it is included. For JSONPath matching, the result of the JSONPath expression is used + and may match any of the specified values. + :type message_filtering_match_values: Any + :param message_filtering_config: Additional configuration to pass to the message filter. + For example with JSONPath filtering you can pass a JSONPath expression string here, + such as `'foo[*].baz'`. Messages with a Body which does not match are ignored. + :type message_filtering_config: Any """ - template_fields = ('sqs_queue', 'max_messages') + template_fields = ('sqs_queue', 'max_messages', 'message_filtering_config') def __init__( self, @@ -48,6 +69,10 @@ def __init__( aws_conn_id: str = 'aws_default', max_messages: int = 5, wait_time_seconds: int = 1, + visibility_timeout: Optional[int] = None, + message_filtering: Optional[Literal["literal", "jsonpath"]] = None, + message_filtering_match_values: Any = None, + message_filtering_config: Any = None, **kwargs, ): super().__init__(**kwargs) @@ -55,6 +80,21 @@ def __init__( self.aws_conn_id = aws_conn_id self.max_messages = max_messages self.wait_time_seconds = wait_time_seconds + self.visibility_timeout = visibility_timeout + + self.message_filtering = message_filtering + + if message_filtering_match_values is not None: + if not isinstance(message_filtering_match_values, set): + message_filtering_match_values = set(message_filtering_match_values) + self.message_filtering_match_values = message_filtering_match_values + + if self.message_filtering == 'literal': + if self.message_filtering_match_values is None: + raise TypeError('message_filtering_match_values must be specified for literal matching') + + self.message_filtering_config = message_filtering_config + self.hook: Optional[SQSHook] = None def poke(self, context): @@ -69,31 +109,48 @@ def poke(self, context): self.log.info('SQSSensor checking for message on queue: %s', self.sqs_queue) - messages = sqs_conn.receive_message( - QueueUrl=self.sqs_queue, - MaxNumberOfMessages=self.max_messages, - WaitTimeSeconds=self.wait_time_seconds, - ) + receive_message_kwargs = { + 'QueueUrl': self.sqs_queue, + 'MaxNumberOfMessages': self.max_messages, + 'WaitTimeSeconds': self.wait_time_seconds, + } + if self.visibility_timeout is not None: + receive_message_kwargs['VisibilityTimeout'] = self.visibility_timeout - self.log.info("received message %s", str(messages)) + response = sqs_conn.receive_message(**receive_message_kwargs) - if 'Messages' in messages and messages['Messages']: - entries = [ - {'Id': message['MessageId'], 'ReceiptHandle': message['ReceiptHandle']} - for message in messages['Messages'] - ] + if "Messages" not in response: + return False - result = sqs_conn.delete_message_batch(QueueUrl=self.sqs_queue, Entries=entries) + messages = response['Messages'] + num_messages = len(messages) + self.log.info("Received %d messages", num_messages) - if 'Successful' in result: - context['ti'].xcom_push(key='messages', value=messages) - return True - else: - raise AirflowException( - 'Delete SQS Messages failed ' + str(result) + ' for messages ' + str(messages) - ) + if not num_messages: + return False - return False + if self.message_filtering: + messages = self.filter_messages(messages) + num_messages = len(messages) + self.log.info("There are %d messages left after filtering", num_messages) + + if not num_messages: + return False + + self.log.info("Deleting %d messages", num_messages) + + entries = [ + {'Id': message['MessageId'], 'ReceiptHandle': message['ReceiptHandle']} for message in messages + ] + response = sqs_conn.delete_message_batch(QueueUrl=self.sqs_queue, Entries=entries) + + if 'Successful' in response: + context['ti'].xcom_push(key='messages', value=messages) + return True + else: + raise AirflowException( + 'Delete SQS Messages failed ' + str(response) + ' for messages ' + str(messages) + ) def get_hook(self) -> SQSHook: """Create and return an SQSHook""" @@ -102,3 +159,37 @@ def get_hook(self) -> SQSHook: self.hook = SQSHook(aws_conn_id=self.aws_conn_id) return self.hook + + def filter_messages(self, messages): + if self.message_filtering == 'literal': + return self.filter_messages_literal(messages) + if self.message_filtering == 'jsonpath': + return self.filter_messages_jsonpath(messages) + else: + raise NotImplementedError('Override this method to define custom filters') + + def filter_messages_literal(self, messages): + filtered_messages = [] + for message in messages: + if message['Body'] in self.message_filtering_match_values: + filtered_messages.append(message) + return filtered_messages + + def filter_messages_jsonpath(self, messages): + jsonpath_expr = parse(self.message_filtering_config) + filtered_messages = [] + for message in messages: + body = message['Body'] + # Body is a string, deserialise to an object and then parse + body = json.loads(body) + results = jsonpath_expr.find(body) + if not results: + continue + if self.message_filtering_match_values is None: + filtered_messages.append(message) + continue + for result in results: + if result.value in self.message_filtering_match_values: + filtered_messages.append(message) + break + return filtered_messages diff --git a/airflow/providers/apache/drill/example_dags/__init__.py b/airflow/providers/apache/drill/example_dags/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow/providers/apache/drill/example_dags/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/airflow/providers/apache/drill/example_dags/example_drill_dag.py b/airflow/providers/apache/drill/example_dags/example_drill_dag.py index 60a35ee6835d9..5a9ee59dd690f 100644 --- a/airflow/providers/apache/drill/example_dags/example_drill_dag.py +++ b/airflow/providers/apache/drill/example_dags/example_drill_dag.py @@ -17,20 +17,14 @@ # under the License. """ -Example Airflow DAG to submit Apache Spark applications using -`SparkSubmitOperator`, `SparkJDBCOperator` and `SparkSqlOperator`. +Example Airflow DAG to execute SQL in an Apache Drill environment using the `DrillOperator`. """ from airflow.models import DAG from airflow.providers.apache.drill.operators.drill import DrillOperator from airflow.utils.dates import days_ago -args = { - 'owner': 'Airflow', -} - with DAG( dag_id='example_drill_dag', - default_args=args, schedule_interval=None, start_date=days_ago(2), tags=['example'], diff --git a/airflow/providers/apache/druid/CHANGELOG.rst b/airflow/providers/apache/druid/CHANGELOG.rst index 7700a2ef71849..817c46e438877 100644 --- a/airflow/providers/apache/druid/CHANGELOG.rst +++ b/airflow/providers/apache/druid/CHANGELOG.rst @@ -22,14 +22,6 @@ Changelog 2.0.1 ..... -Breaking changes -~~~~~~~~~~~~~~~~ - - -Features -~~~~~~~~ - - Bug Fixes ~~~~~~~~~ diff --git a/airflow/providers/cncf/kubernetes/CHANGELOG.rst b/airflow/providers/cncf/kubernetes/CHANGELOG.rst index 3135c57a1576f..33bf22797748e 100644 --- a/airflow/providers/cncf/kubernetes/CHANGELOG.rst +++ b/airflow/providers/cncf/kubernetes/CHANGELOG.rst @@ -22,13 +22,11 @@ Changelog 2.0.1 ..... -Breaking changes -~~~~~~~~~~~~~~~~ - Features ~~~~~~~~ +* ``Enable using custom pod launcher in Kubernetes Pod Operator (#16945)`` Bug Fixes ~~~~~~~~~ @@ -43,25 +41,6 @@ Bug Fixes * ``Prepare documentation for July release of providers. (#17015)`` * ``Updating task dependencies (#16624)`` * ``Removes pylint from our toolchain (#16682)`` - -2.1.0 -..... - - -Features -~~~~~~~~ - -* ``Enable using custom pod launcher in Kubernetes Pod Operator (#16945)`` - -Bug Fixes -~~~~~~~~~ - -* ``BugFix: Using 'json' string in template_field causes issue with K8s Operators (#16930)`` - -.. Below changes are excluded from the changelog. Move them to - appropriate section above if needed. Do not delete the lines(!): - * ``Updating task dependencies (#16624)`` - * ``Removes pylint from our toolchain (#16682)`` * ``Prepare documentation for July release of providers. (#17015)`` * ``Fixed wrongly escaped characters in amazon's changelog (#17020)`` diff --git a/airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py b/airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py index 36905fb3bf10b..7d15b944f89e4 100644 --- a/airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py +++ b/airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py @@ -80,6 +80,7 @@ class GcpTransferOperationStatus: NAME = 'name' OBJECT_CONDITIONS = 'object_conditions' OPERATIONS = 'operations' +PATH = 'path' PROJECT_ID = 'projectId' SCHEDULE = 'schedule' SCHEDULE_END_DATE = 'scheduleEndDate' diff --git a/airflow/providers/google/cloud/hooks/gcs.py b/airflow/providers/google/cloud/hooks/gcs.py index 5aba40f135d30..14e2bfaf4ff0f 100644 --- a/airflow/providers/google/cloud/hooks/gcs.py +++ b/airflow/providers/google/cloud/hooks/gcs.py @@ -1065,8 +1065,8 @@ def sync( source_bucket_obj = client.bucket(source_bucket) destination_bucket_obj = client.bucket(destination_bucket) # Normalize parameters when they are passed - source_object = self._normalize_directory_path(source_object) - destination_object = self._normalize_directory_path(destination_object) + source_object = _normalize_directory_path(source_object) + destination_object = _normalize_directory_path(destination_object) # Calculate the number of characters that remove from the name, because they contain information # about the parent's path source_object_prefix_len = len(source_object) if source_object else 0 @@ -1137,9 +1137,6 @@ def _calculate_sync_destination_path( else blob.name[source_object_prefix_len:] ) - def _normalize_directory_path(self, source_object: Optional[str]) -> Optional[str]: - return source_object + "/" if source_object and not source_object.endswith("/") else source_object - @staticmethod def _prepare_sync_plan( source_bucket: storage.Bucket, @@ -1207,3 +1204,7 @@ def _parse_gcs_url(gsurl: str) -> Tuple[str, str]: # Remove leading '/' but NOT trailing one blob = parsed_url.path.lstrip('/') return bucket, blob + + +def _normalize_directory_path(source_object: Optional[str]) -> Optional[str]: + return source_object + "/" if source_object and not source_object.endswith("/") else source_object diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py index 66980aee35cd2..4a0f6c007adc1 100644 --- a/airflow/providers/google/cloud/operators/bigquery.py +++ b/airflow/providers/google/cloud/operators/bigquery.py @@ -2196,7 +2196,7 @@ class BigQueryInsertJobOperator(BaseOperator): "impersonation_chain", ) template_ext = (".json",) - template_fields_renderers = {"configuration": "json"} + template_fields_renderers = {"configuration": "json", "configuration.query.query": "sql"} ui_color = BigQueryUIColors.QUERY.value def __init__( diff --git a/airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py b/airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py index db0eba368b4bd..712f1d5d13a39 100644 --- a/airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py +++ b/airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py @@ -39,6 +39,7 @@ MONTH, NAME, OBJECT_CONDITIONS, + PATH, PROJECT_ID, SCHEDULE, SCHEDULE_END_DATE, @@ -53,6 +54,7 @@ CloudDataTransferServiceHook, GcpTransferJobsStatus, ) +from airflow.providers.google.cloud.hooks.gcs import _normalize_directory_path class TransferJobPreprocessor: @@ -763,6 +765,8 @@ class CloudDataTransferServiceS3ToGCSOperator(BaseOperator): :param gcs_bucket: The destination Google Cloud Storage bucket where you want to store the files. (templated) :type gcs_bucket: str + :param gcs_path: Optional root path/prefix to transfer objects. (templated) + :type: gcs_path: str :param project_id: Optional ID of the Google Cloud Console project that owns the job :type project_id: str @@ -815,6 +819,7 @@ class CloudDataTransferServiceS3ToGCSOperator(BaseOperator): 'gcp_conn_id', 's3_bucket', 'gcs_bucket', + 'gcs_path', 'description', 'object_conditions', 'google_impersonation_chain', @@ -826,6 +831,7 @@ def __init__( *, s3_bucket: str, gcs_bucket: str, + gcs_path: Optional[str] = None, project_id: Optional[str] = None, aws_conn_id: str = 'aws_default', gcp_conn_id: str = 'google_cloud_default', @@ -844,6 +850,7 @@ def __init__( super().__init__(**kwargs) self.s3_bucket = s3_bucket self.gcs_bucket = gcs_bucket + self.gcs_path = _normalize_directory_path(gcs_path) self.project_id = project_id self.aws_conn_id = aws_conn_id self.gcp_conn_id = gcp_conn_id @@ -885,7 +892,7 @@ def _create_body(self) -> dict: STATUS: GcpTransferJobsStatus.ENABLED, TRANSFER_SPEC: { AWS_S3_DATA_SOURCE: {BUCKET_NAME: self.s3_bucket}, - GCS_DATA_SINK: {BUCKET_NAME: self.gcs_bucket}, + GCS_DATA_SINK: {BUCKET_NAME: self.gcs_bucket, PATH: self.gcs_path}, }, } @@ -935,6 +942,8 @@ class CloudDataTransferServiceGCSToGCSOperator(BaseOperator): :param destination_bucket: The destination Google Cloud Storage bucket where the object should be. (templated) :type destination_bucket: str + :param destination_path: Optional root path/prefix to transfer objects. (templated) + :type: destination_path: str :param project_id: The ID of the Google Cloud Console project that owns the job :type project_id: str @@ -985,6 +994,7 @@ class CloudDataTransferServiceGCSToGCSOperator(BaseOperator): 'gcp_conn_id', 'source_bucket', 'destination_bucket', + 'destination_path', 'description', 'object_conditions', 'google_impersonation_chain', @@ -996,6 +1006,7 @@ def __init__( *, source_bucket: str, destination_bucket: str, + destination_path: Optional[str] = None, project_id: Optional[str] = None, gcp_conn_id: str = 'google_cloud_default', delegate_to: Optional[str] = None, @@ -1013,6 +1024,7 @@ def __init__( super().__init__(**kwargs) self.source_bucket = source_bucket self.destination_bucket = destination_bucket + self.destination_path = _normalize_directory_path(destination_path) self.project_id = project_id self.gcp_conn_id = gcp_conn_id self.delegate_to = delegate_to @@ -1054,7 +1066,7 @@ def _create_body(self) -> dict: STATUS: GcpTransferJobsStatus.ENABLED, TRANSFER_SPEC: { GCS_DATA_SOURCE: {BUCKET_NAME: self.source_bucket}, - GCS_DATA_SINK: {BUCKET_NAME: self.destination_bucket}, + GCS_DATA_SINK: {BUCKET_NAME: self.destination_bucket, PATH: self.destination_path}, }, } diff --git a/airflow/providers/google/cloud/secrets/secret_manager.py b/airflow/providers/google/cloud/secrets/secret_manager.py index 919927fae8ecf..8c2cc9af3566b 100644 --- a/airflow/providers/google/cloud/secrets/secret_manager.py +++ b/airflow/providers/google/cloud/secrets/secret_manager.py @@ -16,6 +16,7 @@ # under the License. """Objects relating to sourcing connections from Google Cloud Secrets Manager""" +import logging from typing import Optional try: @@ -23,12 +24,16 @@ except ImportError: from cached_property import cached_property +from google.auth.exceptions import DefaultCredentialsError + from airflow.exceptions import AirflowException from airflow.providers.google.cloud._internal_client.secret_manager_client import _SecretManagerClient from airflow.providers.google.cloud.utils.credentials_provider import get_credentials_and_project_id from airflow.secrets import BaseSecretsBackend from airflow.utils.log.logging_mixin import LoggingMixin +log = logging.getLogger(__name__) + SECRET_ID_PATTERN = r"^[a-zA-Z0-9-_]*$" @@ -101,9 +106,17 @@ def __init__( "`connections_prefix`, `variables_prefix` and `sep` should " f"follows that pattern {SECRET_ID_PATTERN}" ) - self.credentials, self.project_id = get_credentials_and_project_id( - keyfile_dict=gcp_keyfile_dict, key_path=gcp_key_path, scopes=gcp_scopes - ) + try: + self.credentials, self.project_id = get_credentials_and_project_id( + keyfile_dict=gcp_keyfile_dict, key_path=gcp_key_path, scopes=gcp_scopes + ) + except (DefaultCredentialsError, FileNotFoundError): + log.exception( + 'Unable to load credentials for GCP Secret Manager. ' + 'Make sure that the keyfile path, dictionary, or GOOGLE_APPLICATION_CREDENTIALS ' + 'environment variable is correct and properly configured.' + ) + # In case project id provided if project_id: self.project_id = project_id diff --git a/airflow/providers/mysql/operators/mysql.py b/airflow/providers/mysql/operators/mysql.py index 18f3d0ccdc444..10fe0d6b76bf6 100644 --- a/airflow/providers/mysql/operators/mysql.py +++ b/airflow/providers/mysql/operators/mysql.py @@ -16,7 +16,7 @@ # specific language governing permissions and limitations # under the License. import ast -from typing import Dict, Iterable, Mapping, Optional, Union +from typing import Dict, Iterable, List, Mapping, Optional, Union from airflow.models import BaseOperator from airflow.providers.mysql.hooks.mysql import MySqlHook @@ -56,7 +56,7 @@ class MySqlOperator(BaseOperator): def __init__( self, *, - sql: str, + sql: Union[str, List[str]], mysql_conn_id: str = 'mysql_default', parameters: Optional[Union[Mapping, Iterable]] = None, autocommit: bool = False, diff --git a/airflow/providers/neo4j/operators/neo4j.py b/airflow/providers/neo4j/operators/neo4j.py index 5a14e46934cff..c3c3a1da9d5a8 100644 --- a/airflow/providers/neo4j/operators/neo4j.py +++ b/airflow/providers/neo4j/operators/neo4j.py @@ -30,8 +30,8 @@ class Neo4jOperator(BaseOperator): :ref:`howto/operator:Neo4jOperator` :param sql: the sql code to be executed. Can receive a str representing a - sql statement, a list of str (sql statements) - :type sql: str or list[str] + sql statement + :type sql: str :param neo4j_conn_id: Reference to :ref:`Neo4j connection id `. :type neo4j_conn_id: str """ diff --git a/airflow/providers/postgres/operators/postgres.py b/airflow/providers/postgres/operators/postgres.py index e4ef23fce23e6..51733cd03e47e 100644 --- a/airflow/providers/postgres/operators/postgres.py +++ b/airflow/providers/postgres/operators/postgres.py @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from typing import Iterable, Mapping, Optional, Union +from typing import Iterable, List, Mapping, Optional, Union from airflow.models import BaseOperator from airflow.providers.postgres.hooks.postgres import PostgresHook @@ -49,7 +49,7 @@ class PostgresOperator(BaseOperator): def __init__( self, *, - sql: str, + sql: Union[str, List[str]], postgres_conn_id: str = 'postgres_default', autocommit: bool = False, parameters: Optional[Union[Mapping, Iterable]] = None, diff --git a/airflow/providers/samba/hooks/samba.py b/airflow/providers/samba/hooks/samba.py index f8234ebe3df2b..cde873f771a5a 100644 --- a/airflow/providers/samba/hooks/samba.py +++ b/airflow/providers/samba/hooks/samba.py @@ -16,67 +16,232 @@ # specific language governing permissions and limitations # under the License. -import os +import posixpath +from functools import wraps +from shutil import copyfileobj +from typing import Optional -from smbclient import SambaClient +import smbclient from airflow.hooks.base import BaseHook class SambaHook(BaseHook): - """Allows for interaction with an samba server.""" + """Allows for interaction with a Samba server. + + The hook should be used as a context manager in order to correctly + set up a session and disconnect open connections upon exit. + + :param samba_conn_id: The connection id reference. + :type samba_conn_id: str + :param share: + An optional share name. If this is unset then the "schema" field of + the connection is used in its place. + :type share: str + """ conn_name_attr = 'samba_conn_id' default_conn_name = 'samba_default' conn_type = 'samba' hook_name = 'Samba' - def __init__(self, samba_conn_id: str = default_conn_name) -> None: + def __init__(self, samba_conn_id: str = default_conn_name, share: Optional[str] = None) -> None: super().__init__() - self.conn = self.get_connection(samba_conn_id) + conn = self.get_connection(samba_conn_id) + + if not conn.login: + self.log.info("Login not provided") + + if not conn.password: + self.log.info("Password not provided") + + self._host = conn.host + self._share = share or conn.schema + self._connection_cache = connection_cache = {} + self._conn_kwargs = { + "username": conn.login, + "password": conn.password, + "port": conn.port or 445, + "connection_cache": connection_cache, + } + + def __enter__(self): + # This immediately connects to the host (which can be + # perceived as a benefit), but also help work around an issue: + # + # https://github.com/jborean93/smbprotocol/issues/109. + smbclient.register_session(self._host, **self._conn_kwargs) + return self + + def __exit__(self, exc_type, exc_value, traceback): + for host, connection in self._connection_cache.items(): + self.log.info("Disconnecting from %s", host) + connection.disconnect() + self._connection_cache.clear() + + def _join_path(self, path): + return f"//{posixpath.join(self._host, self._share, path)}" + + @wraps(smbclient.link) + def link(self, src, dst, follow_symlinks=True): + return smbclient.link( + self._join_path(src), + self._join_path(dst), + follow_symlinks=follow_symlinks, + **self._conn_kwargs, + ) + + @wraps(smbclient.listdir) + def listdir(self, path): + return smbclient.listdir(self._join_path(path), **self._conn_kwargs) + + @wraps(smbclient.lstat) + def lstat(self, path): + return smbclient.lstat(self._join_path(path), **self._conn_kwargs) + + @wraps(smbclient.makedirs) + def makedirs(self, path, exist_ok=False): + return smbclient.makedirs(self._join_path(path), exist_ok=exist_ok, **self._conn_kwargs) + + @wraps(smbclient.mkdir) + def mkdir(self, path): + return smbclient.mkdir(self._join_path(path), **self._conn_kwargs) + + @wraps(smbclient.open_file) + def open_file( + self, + path, + mode="r", + buffering=-1, + encoding=None, + errors=None, + newline=None, + share_access=None, + desired_access=None, + file_attributes=None, + file_type="file", + ): + return smbclient.open_file( + self._join_path(path), + mode=mode, + buffering=buffering, + encoding=encoding, + errors=errors, + newline=newline, + share_access=share_access, + desired_access=desired_access, + file_attributes=file_attributes, + file_type=file_type, + **self._conn_kwargs, + ) + + @wraps(smbclient.readlink) + def readlink(self, path): + return smbclient.readlink(self._join_path(path), **self._conn_kwargs) + + @wraps(smbclient.remove) + def remove(self, path): + return smbclient.remove(self._join_path(path), **self._conn_kwargs) + + @wraps(smbclient.removedirs) + def removedirs(self, path): + return smbclient.removedirs(self._join_path(path), **self._conn_kwargs) + + @wraps(smbclient.rename) + def rename(self, src, dst): + return smbclient.rename(self._join_path(src), self._join_path(dst), **self._conn_kwargs) - def get_conn(self) -> SambaClient: - """ - Return a samba client object. + @wraps(smbclient.replace) + def replace(self, src, dst): + return smbclient.replace(self._join_path(src), self._join_path(dst), **self._conn_kwargs) - You can provide optional parameters in the extra fields of - your connection. + @wraps(smbclient.rmdir) + def rmdir(self, path): + return smbclient.rmdir(self._join_path(path), **self._conn_kwargs) - Below is an inexhaustive list of these parameters: + @wraps(smbclient.scandir) + def scandir(self, path, search_pattern="*"): + return smbclient.scandir( + self._join_path(path), + search_pattern=search_pattern, + **self._conn_kwargs, + ) - `logdir` - Base directory name for log/debug files. + @wraps(smbclient.stat) + def stat(self, path, follow_symlinks=True): + return smbclient.stat(self._join_path(path), follow_symlinks=follow_symlinks, **self._conn_kwargs) + + @wraps(smbclient.stat_volume) + def stat_volume(self, path): + return smbclient.stat_volume(self._join_path(path), **self._conn_kwargs) + + @wraps(smbclient.symlink) + def symlink(self, src, dst, target_is_directory=False): + return smbclient.symlink( + self._join_path(src), + self._join_path(dst), + target_is_directory=target_is_directory, + **self._conn_kwargs, + ) - `kerberos` - Try to authenticate with kerberos. + @wraps(smbclient.truncate) + def truncate(self, path, length): + return smbclient.truncate(self._join_path(path), length, **self._conn_kwargs) + + @wraps(smbclient.unlink) + def unlink(self, path): + return smbclient.unlink(self._join_path(path), **self._conn_kwargs) + + @wraps(smbclient.utime) + def utime(self, path, times=None, ns=None, follow_symlinks=True): + return smbclient.utime( + self._join_path(path), + times=times, + ns=ns, + follow_symlinks=follow_symlinks, + **self._conn_kwargs, + ) - `workgroup` - Set the SMB domain of the username. + @wraps(smbclient.walk) + def walk(self, path, topdown=True, onerror=None, follow_symlinks=False): + return smbclient.walk( + self._join_path(path), + topdown=topdown, + onerror=onerror, + follow_symlinks=follow_symlinks, + **self._conn_kwargs, + ) - `netbios_name` - This option allows you to override the NetBIOS name that - Samba uses for itself. + @wraps(smbclient.getxattr) + def getxattr(self, path, attribute, follow_symlinks=True): + return smbclient.getxattr( + self._join_path(path), attribute, follow_symlinks=follow_symlinks, **self._conn_kwargs + ) + + @wraps(smbclient.listxattr) + def listxattr(self, path, follow_symlinks=True): + return smbclient.listxattr( + self._join_path(path), follow_symlinks=follow_symlinks, **self._conn_kwargs + ) + + @wraps(smbclient.removexattr) + def removexattr(self, path, attribute, follow_symlinks=True): + return smbclient.removexattr( + self._join_path(path), attribute, follow_symlinks=follow_symlinks, **self._conn_kwargs + ) - For additional details, see `smbclient.SambaClient`. - """ - samba = SambaClient( - server=self.conn.host, - share=self.conn.schema, - username=self.conn.login, - ip=self.conn.host, - password=self.conn.password, - **self.conn.extra_dejson, + @wraps(smbclient.setxattr) + def setxattr(self, path, attribute, value, flags=0, follow_symlinks=True): + return smbclient.setxattr( + self._join_path(path), + attribute, + value, + flags=flags, + follow_symlinks=follow_symlinks, + **self._conn_kwargs, ) - return samba - def push_from_local(self, destination_filepath: str, local_filepath: str) -> None: + def push_from_local(self, destination_filepath: str, local_filepath: str): """Push local file to samba server""" - samba = self.get_conn() - if samba.exists(destination_filepath): - if samba.isfile(destination_filepath): - samba.remove(destination_filepath) - else: - folder = os.path.dirname(destination_filepath) - if not samba.exists(folder): - samba.mkdir(folder) - samba.upload(local_filepath, destination_filepath) + with open(local_filepath, "rb") as f, self.open_file(destination_filepath, mode="w") as g: + copyfileobj(f, g) diff --git a/airflow/providers/slack/operators/slack.py b/airflow/providers/slack/operators/slack.py index 96a8ea09582d0..053ce3b38e8b2 100644 --- a/airflow/providers/slack/operators/slack.py +++ b/airflow/providers/slack/operators/slack.py @@ -161,6 +161,7 @@ class SlackAPIFileOperator(SlackAPIOperator): .. code-block:: python + # Send file with filename and filetype slack = SlackAPIFileOperator( task_id="slack_file_upload", dag=dag, @@ -169,7 +170,16 @@ class SlackAPIFileOperator(SlackAPIOperator): initial_comment="Hello World!", filename="hello_world.csv", filetype="csv", - content="hello,world,csv,file", + ) + + # Send file content + slack = SlackAPIFileOperator( + task_id="slack_file_upload", + dag=dag, + slack_conn_id="slack", + channel="#general", + initial_comment="Hello World!", + content="file content in txt", ) :param channel: channel in which to sent file on slack name (templated) @@ -192,9 +202,9 @@ def __init__( self, channel: str = '#general', initial_comment: str = 'No message has been set!', - filename: str = 'default_name.csv', - filetype: str = 'csv', - content: str = 'default,content,csv,file', + filename: str = None, + filetype: str = None, + content: str = None, **kwargs, ) -> None: self.method = 'files.upload' @@ -203,13 +213,31 @@ def __init__( self.filename = filename self.filetype = filetype self.content = content + self.file_params = {} super().__init__(method=self.method, **kwargs) def construct_api_call_params(self) -> Any: - self.api_params = { - 'channels': self.channel, - 'content': self.content, - 'filename': self.filename, - 'filetype': self.filetype, - 'initial_comment': self.initial_comment, - } + if self.content is not None: + self.api_params = { + 'channels': self.channel, + 'content': self.content, + 'initial_comment': self.initial_comment, + } + elif self.filename is not None: + self.api_params = { + 'channels': self.channel, + 'filename': self.filename, + 'filetype': self.filetype, + 'initial_comment': self.initial_comment, + } + self.file_params = {'file': self.filename} + + def execute(self, **kwargs): + """ + The SlackAPIOperator calls will not fail even if the call is not unsuccessful. + It should not prevent a DAG from completing in success + """ + if not self.api_params: + self.construct_api_call_params() + slack = SlackHook(token=self.token, slack_conn_id=self.slack_conn_id) + slack.call(self.method, data=self.api_params, files=self.file_params) diff --git a/airflow/providers/snowflake/CHANGELOG.rst b/airflow/providers/snowflake/CHANGELOG.rst index 0f0c2fe41c0fd..fecc953ab5c71 100644 --- a/airflow/providers/snowflake/CHANGELOG.rst +++ b/airflow/providers/snowflake/CHANGELOG.rst @@ -22,19 +22,11 @@ Changelog 2.1.0 ..... -Breaking changes -~~~~~~~~~~~~~~~~ - - Features ~~~~~~~~ * ``Adding: Snowflake Role in snowflake provider hook (#16735)`` -Bug Fixes -~~~~~~~~~ - - .. Below changes are excluded from the changelog. Move them to appropriate section above if needed. Do not delete the lines(!): * ``Logging and returning info about query execution SnowflakeHook (#15736)`` diff --git a/airflow/providers/yandex/example_dags/example_yandexcloud_dataproc.py b/airflow/providers/yandex/example_dags/example_yandexcloud_dataproc.py index 42071a6e6aa58..6d9a384fb8140 100644 --- a/airflow/providers/yandex/example_dags/example_yandexcloud_dataproc.py +++ b/airflow/providers/yandex/example_dags/example_yandexcloud_dataproc.py @@ -50,6 +50,8 @@ zone=AVAILABILITY_ZONE_ID, connection_id=CONNECTION_ID, s3_bucket=S3_BUCKET_NAME_FOR_JOB_LOGS, + computenode_count=1, + computenode_max_hosts_count=5, ) create_hive_query = DataprocCreateHiveJobOperator( diff --git a/airflow/providers/yandex/hooks/yandex.py b/airflow/providers/yandex/hooks/yandex.py index bb4763e5404ef..ee1ae0dffe5ab 100644 --- a/airflow/providers/yandex/hooks/yandex.py +++ b/airflow/providers/yandex/hooks/yandex.py @@ -29,8 +29,8 @@ class YandexCloudBaseHook(BaseHook): """ A base hook for Yandex.Cloud related tasks. - :param connection_id: The connection ID to use when fetching connection info. - :type connection_id: str + :param yandex_conn_id: The connection ID to use when fetching connection info. + :type yandex_conn_id: str """ conn_name_attr = 'yandex_conn_id' diff --git a/airflow/providers/yandex/hooks/yandexcloud_dataproc.py b/airflow/providers/yandex/hooks/yandexcloud_dataproc.py index cff2a99417b1c..4cb52155f5ae6 100644 --- a/airflow/providers/yandex/hooks/yandexcloud_dataproc.py +++ b/airflow/providers/yandex/hooks/yandexcloud_dataproc.py @@ -23,8 +23,8 @@ class DataprocHook(YandexCloudBaseHook): """ A base hook for Yandex.Cloud Data Proc. - :param connection_id: The connection ID to use when fetching connection info. - :type connection_id: str + :param yandex_conn_id: The connection ID to use when fetching connection info. + :type yandex_conn_id: str """ def __init__(self, *args, **kwargs) -> None: diff --git a/airflow/providers/yandex/operators/yandexcloud_dataproc.py b/airflow/providers/yandex/operators/yandexcloud_dataproc.py index 0ab6245a7d628..7150e145ccf63 100644 --- a/airflow/providers/yandex/operators/yandexcloud_dataproc.py +++ b/airflow/providers/yandex/operators/yandexcloud_dataproc.py @@ -71,6 +71,28 @@ class DataprocCreateClusterOperator(BaseOperator): :type computenode_disk_type: str :param connection_id: ID of the Yandex.Cloud Airflow connection. :type connection_id: Optional[str] + :type computenode_max_count: int + :param computenode_max_count: Maximum number of nodes of compute autoscaling subcluster. + :param computenode_warmup_duration: The warmup time of the instance in seconds. During this time, + traffic is sent to the instance, + but instance metrics are not collected. In seconds. + :type computenode_warmup_duration: int + :param computenode_stabilization_duration: Minimum amount of time in seconds for monitoring before + Instance Groups can reduce the number of instances in the group. + During this time, the group size doesn't decrease, + even if the new metric values indicate that it should. In seconds. + :type computenode_stabilization_duration: int + :param computenode_preemptible: Preemptible instances are stopped at least once every 24 hours, + and can be stopped at any time if their resources are needed by Compute. + :type computenode_preemptible: bool + :param computenode_cpu_utilization_target: Defines an autoscaling rule + based on the average CPU utilization of the instance group. + in percents. 10-100. + By default is not set and default autoscaling strategy is used. + :type computenode_cpu_utilization_target: int + :param computenode_decommission_timeout: Timeout to gracefully decommission nodes during downscaling. + In seconds. + :type computenode_decommission_timeout: int """ def __init__( @@ -78,31 +100,38 @@ def __init__( *, folder_id: Optional[str] = None, cluster_name: Optional[str] = None, - cluster_description: str = '', - cluster_image_version: str = '1.1', + cluster_description: Optional[str] = '', + cluster_image_version: Optional[str] = None, ssh_public_keys: Optional[Union[str, Iterable[str]]] = None, subnet_id: Optional[str] = None, services: Iterable[str] = ('HDFS', 'YARN', 'MAPREDUCE', 'HIVE', 'SPARK'), s3_bucket: Optional[str] = None, zone: str = 'ru-central1-b', service_account_id: Optional[str] = None, - masternode_resource_preset: str = 's2.small', - masternode_disk_size: int = 15, - masternode_disk_type: str = 'network-ssd', - datanode_resource_preset: str = 's2.small', - datanode_disk_size: int = 15, - datanode_disk_type: str = 'network-ssd', - datanode_count: int = 2, - computenode_resource_preset: str = 's2.small', - computenode_disk_size: int = 15, - computenode_disk_type: str = 'network-ssd', + masternode_resource_preset: Optional[str] = None, + masternode_disk_size: Optional[int] = None, + masternode_disk_type: Optional[str] = None, + datanode_resource_preset: Optional[str] = None, + datanode_disk_size: Optional[int] = None, + datanode_disk_type: Optional[str] = None, + datanode_count: int = 1, + computenode_resource_preset: Optional[str] = None, + computenode_disk_size: Optional[int] = None, + computenode_disk_type: Optional[str] = None, computenode_count: int = 0, + computenode_max_hosts_count: Optional[int] = None, + computenode_measurement_duration: Optional[int] = None, + computenode_warmup_duration: Optional[int] = None, + computenode_stabilization_duration: Optional[int] = None, + computenode_preemptible: bool = False, + computenode_cpu_utilization_target: Optional[int] = None, + computenode_decommission_timeout: Optional[int] = None, connection_id: Optional[str] = None, **kwargs, ) -> None: super().__init__(**kwargs) self.folder_id = folder_id - self.connection_id = connection_id + self.yandex_conn_id = connection_id self.cluster_name = cluster_name self.cluster_description = cluster_description self.cluster_image_version = cluster_image_version @@ -123,11 +152,19 @@ def __init__( self.computenode_disk_size = computenode_disk_size self.computenode_disk_type = computenode_disk_type self.computenode_count = computenode_count + self.computenode_max_hosts_count = computenode_max_hosts_count + self.computenode_measurement_duration = computenode_measurement_duration + self.computenode_warmup_duration = computenode_warmup_duration + self.computenode_stabilization_duration = computenode_stabilization_duration + self.computenode_preemptible = computenode_preemptible + self.computenode_cpu_utilization_target = computenode_cpu_utilization_target + self.computenode_decommission_timeout = computenode_decommission_timeout + self.hook: Optional[DataprocHook] = None def execute(self, context) -> None: self.hook = DataprocHook( - connection_id=self.connection_id, + yandex_conn_id=self.yandex_conn_id, ) operation_result = self.hook.client.create_cluster( folder_id=self.folder_id, @@ -151,9 +188,16 @@ def execute(self, context) -> None: computenode_disk_size=self.computenode_disk_size, computenode_disk_type=self.computenode_disk_type, computenode_count=self.computenode_count, + computenode_max_hosts_count=self.computenode_max_hosts_count, + computenode_measurement_duration=self.computenode_measurement_duration, + computenode_warmup_duration=self.computenode_warmup_duration, + computenode_stabilization_duration=self.computenode_stabilization_duration, + computenode_preemptible=self.computenode_preemptible, + computenode_cpu_utilization_target=self.computenode_cpu_utilization_target, + computenode_decommission_timeout=self.computenode_decommission_timeout, ) context['task_instance'].xcom_push(key='cluster_id', value=operation_result.response.id) - context['task_instance'].xcom_push(key='yandexcloud_connection_id', value=self.connection_id) + context['task_instance'].xcom_push(key='yandexcloud_connection_id', value=self.yandex_conn_id) class DataprocDeleteClusterOperator(BaseOperator): @@ -171,17 +215,17 @@ def __init__( self, *, connection_id: Optional[str] = None, cluster_id: Optional[str] = None, **kwargs ) -> None: super().__init__(**kwargs) - self.connection_id = connection_id + self.yandex_conn_id = connection_id self.cluster_id = cluster_id self.hook: Optional[DataprocHook] = None def execute(self, context) -> None: cluster_id = self.cluster_id or context['task_instance'].xcom_pull(key='cluster_id') - connection_id = self.connection_id or context['task_instance'].xcom_pull( + yandex_conn_id = self.yandex_conn_id or context['task_instance'].xcom_pull( key='yandexcloud_connection_id' ) self.hook = DataprocHook( - connection_id=connection_id, + yandex_conn_id=yandex_conn_id, ) self.hook.client.delete_cluster(cluster_id) @@ -236,11 +280,11 @@ def __init__( def execute(self, context) -> None: cluster_id = self.cluster_id or context['task_instance'].xcom_pull(key='cluster_id') - connection_id = self.connection_id or context['task_instance'].xcom_pull( + yandex_conn_id = self.connection_id or context['task_instance'].xcom_pull( key='yandexcloud_connection_id' ) self.hook = DataprocHook( - connection_id=connection_id, + yandex_conn_id=yandex_conn_id, ) self.hook.client.create_hive_job( query=self.query, @@ -312,11 +356,11 @@ def __init__( def execute(self, context) -> None: cluster_id = self.cluster_id or context['task_instance'].xcom_pull(key='cluster_id') - connection_id = self.connection_id or context['task_instance'].xcom_pull( + yandex_conn_id = self.connection_id or context['task_instance'].xcom_pull( key='yandexcloud_connection_id' ) self.hook = DataprocHook( - connection_id=connection_id, + yandex_conn_id=yandex_conn_id, ) self.hook.client.create_mapreduce_job( main_class=self.main_class, @@ -389,11 +433,11 @@ def __init__( def execute(self, context) -> None: cluster_id = self.cluster_id or context['task_instance'].xcom_pull(key='cluster_id') - connection_id = self.connection_id or context['task_instance'].xcom_pull( + yandex_conn_id = self.connection_id or context['task_instance'].xcom_pull( key='yandexcloud_connection_id' ) self.hook = DataprocHook( - connection_id=connection_id, + yandex_conn_id=yandex_conn_id, ) self.hook.client.create_spark_job( main_class=self.main_class, @@ -466,11 +510,11 @@ def __init__( def execute(self, context) -> None: cluster_id = self.cluster_id or context['task_instance'].xcom_pull(key='cluster_id') - connection_id = self.connection_id or context['task_instance'].xcom_pull( + yandex_conn_id = self.connection_id or context['task_instance'].xcom_pull( key='yandexcloud_connection_id' ) self.hook = DataprocHook( - connection_id=connection_id, + yandex_conn_id=yandex_conn_id, ) self.hook.client.create_pyspark_job( main_python_file_uri=self.main_python_file_uri, diff --git a/airflow/providers_manager.py b/airflow/providers_manager.py index 5e67e06389e94..17bcf15ce583c 100644 --- a/airflow/providers_manager.py +++ b/airflow/providers_manager.py @@ -22,6 +22,7 @@ import logging import os from collections import OrderedDict +from time import perf_counter from typing import Any, Dict, NamedTuple, Set import jsonschema @@ -29,6 +30,7 @@ from airflow.utils import yaml from airflow.utils.entry_points import entry_points_with_dist +from airflow.utils.log.logging_mixin import LoggingMixin try: import importlib.resources as importlib_resources @@ -83,7 +85,7 @@ class ConnectionFormWidgetInfo(NamedTuple): field: Field -class ProvidersManager: +class ProvidersManager(LoggingMixin): """ Manages all provider packages. This is a Singleton class. The first time it is instantiated, it discovers all available providers in installed packages and @@ -99,6 +101,7 @@ def __new__(cls): return cls._instance def __init__(self): + """Initializes the manager.""" # Keeps dict of providers keyed by module name self._provider_dict: Dict[str, ProviderInfo] = {} # Keeps dict of hooks keyed by connection type @@ -112,30 +115,61 @@ def __init__(self): self._customized_form_fields_schema_validator = ( _create_customized_form_field_behaviours_schema_validator() ) - self._initialized = False + self._providers_list_initialized = False + self._providers_hooks_initialized = False + self._providers_extra_links_initialized = False - def initialize_providers_manager(self): - """Lazy initialization of provider data.""" + def initialize_providers_list(self): + """Lazy initialization of providers list.""" # We cannot use @cache here because it does not work during pytest, apparently each test # runs it it's own namespace and ProvidersManager is a different object in each namespace - # even if it is singleton but @cache on the initialize_providers_manager message still works in the + # even if it is singleton but @cache on the initialize_providers_* still works in the # way that it is called only once for one of the objects (at least this is how it looks like # from running tests) - if self._initialized: + if self._providers_list_initialized: return + start_time = perf_counter() + self.log.debug("Initializing Providers Manager list") # Local source folders are loaded first. They should take precedence over the package ones for # Development purpose. In production provider.yaml files are not present in the 'airflow" directory # So there is no risk we are going to override package provider accidentally. This can only happen # in case of local development self._discover_all_airflow_builtin_providers_from_local_sources() self._discover_all_providers_from_packages() - self._discover_hooks() self._provider_dict = OrderedDict(sorted(self._provider_dict.items())) + self.log.debug( + "Initialization of Providers Manager list took %.2f seconds", perf_counter() - start_time + ) + self._providers_list_initialized = True + + def initialize_providers_hooks(self): + """Lazy initialization of providers hooks.""" + if self._providers_hooks_initialized: + return + self.initialize_providers_list() + start_time = perf_counter() + self.log.debug("Initializing Providers Hooks") + self._discover_hooks() self._hooks_dict = OrderedDict(sorted(self._hooks_dict.items())) self._connection_form_widgets = OrderedDict(sorted(self._connection_form_widgets.items())) self._field_behaviours = OrderedDict(sorted(self._field_behaviours.items())) + self.log.debug( + "Initialization of Providers Manager hooks took %.2f seconds", perf_counter() - start_time + ) + self._providers_hooks_initialized = True + + def initialize_providers_extra_links(self): + """Lazy initialization of providers extra links.""" + if self._providers_extra_links_initialized: + return + self.initialize_providers_list() + start_time = perf_counter() + self.log.debug("Initializing Providers Extra Links") self._discover_extra_links() - self._initialized = True + self.log.debug( + "Initialization of Providers Manager extra links took %.2f seconds", perf_counter() - start_time + ) + self._providers_extra_links_initialized = True def _discover_all_providers_from_packages(self) -> None: """ @@ -397,29 +431,29 @@ def _add_extra_link(self, extra_link_class_name, provider_package) -> None: @property def providers(self) -> Dict[str, ProviderInfo]: """Returns information about available providers.""" - self.initialize_providers_manager() + self.initialize_providers_list() return self._provider_dict @property def hooks(self) -> Dict[str, HookInfo]: """Returns dictionary of connection_type-to-hook mapping""" - self.initialize_providers_manager() + self.initialize_providers_hooks() return self._hooks_dict @property - def extra_links_class_names(self): + def extra_links_class_names(self) -> Set[str]: """Returns set of extra link class names.""" - self.initialize_providers_manager() + self.initialize_providers_extra_links() return sorted(self._extra_link_class_name_set) @property def connection_form_widgets(self) -> Dict[str, ConnectionFormWidgetInfo]: """Returns widgets for connection forms.""" - self.initialize_providers_manager() + self.initialize_providers_hooks() return self._connection_form_widgets @property def field_behaviours(self) -> Dict[str, Dict]: """Returns dictionary with field behaviours for connection types.""" - self.initialize_providers_manager() + self.initialize_providers_hooks() return self._field_behaviours diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py index a62a66d930fc8..e1fd91341802d 100644 --- a/airflow/serialization/serialized_objects.py +++ b/airflow/serialization/serialized_objects.py @@ -752,7 +752,6 @@ def deserialize_dag(cls, encoded_dag: Dict[str, Any]) -> 'SerializedDAG': for k in keys_to_set_none: setattr(dag, k, None) - setattr(dag, 'full_filepath', dag.fileloc) for task in dag.task_dict.values(): task.dag = dag serializable_task: BaseOperator = task diff --git a/airflow/ti_deps/deps/trigger_rule_dep.py b/airflow/ti_deps/deps/trigger_rule_dep.py index b65a84955532d..cc1672ff31f5a 100644 --- a/airflow/ti_deps/deps/trigger_rule_dep.py +++ b/airflow/ti_deps/deps/trigger_rule_dep.py @@ -61,8 +61,8 @@ def _get_dep_statuses(self, ti, session, dep_context): yield self._passing_status(reason="The task instance did not have any upstream tasks.") return - if ti.task.trigger_rule == TR.DUMMY: - yield self._passing_status(reason="The task had a dummy trigger rule set.") + if ti.task.trigger_rule == TR.ALWAYS: + yield self._passing_status(reason="The task had a always trigger rule set.") return # see if the task name is in the task upstream for our task successes, skipped, failed, upstream_failed, done = self._get_states_count_upstream_ti( diff --git a/airflow/utils/log/file_task_handler.py b/airflow/utils/log/file_task_handler.py index 2dc9beb57b0ac..e15e375dd75c1 100644 --- a/airflow/utils/log/file_task_handler.py +++ b/airflow/utils/log/file_task_handler.py @@ -186,6 +186,16 @@ def _read(self, ti, try_number, metadata=None): ) response.encoding = "utf-8" + if response.status_code == 403: + log += ( + "*** !!!! Please make sure that all your Airflow components (e.g. " + "schedulers, webservers and workers) have" + " the same 'secret_key' configured in 'webserver' section !!!!!\n***" + ) + log += ( + "*** See more at https://airflow.apache.org/docs/apache-airflow/" + "stable/configurations-ref.html#secret-key\n***" + ) # Check if the resource was properly fetched response.raise_for_status() diff --git a/airflow/utils/state.py b/airflow/utils/state.py index e95b4095da2b2..f408c94b0d3df 100644 --- a/airflow/utils/state.py +++ b/airflow/utils/state.py @@ -39,7 +39,8 @@ class TaskInstanceState(str, Enum): QUEUED = "queued" # Executor has enqueued the task RUNNING = "running" # Task is executing SUCCESS = "success" # Task completed - SHUTDOWN = "shutdown" # External request to shut down + SHUTDOWN = "shutdown" # External request to shut down (e.g. marked failed when running) + RESTARTING = "restarting" # External request to restart (e.g. cleared when running) FAILED = "failed" # Task errored out UP_FOR_RETRY = "up_for_retry" # Task failed but has retries left UP_FOR_RESCHEDULE = "up_for_reschedule" # A waiting `reschedule` sensor @@ -84,6 +85,7 @@ class State: SCHEDULED = TaskInstanceState.SCHEDULED QUEUED = TaskInstanceState.QUEUED SHUTDOWN = TaskInstanceState.SHUTDOWN + RESTARTING = TaskInstanceState.RESTARTING UP_FOR_RETRY = TaskInstanceState.UP_FOR_RETRY UP_FOR_RESCHEDULE = TaskInstanceState.UP_FOR_RESCHEDULE UPSTREAM_FAILED = TaskInstanceState.UPSTREAM_FAILED @@ -105,6 +107,7 @@ class State: TaskInstanceState.RUNNING: 'lime', TaskInstanceState.SUCCESS: 'green', TaskInstanceState.SHUTDOWN: 'blue', + TaskInstanceState.RESTARTING: 'violet', TaskInstanceState.FAILED: 'red', TaskInstanceState.UP_FOR_RETRY: 'gold', TaskInstanceState.UP_FOR_RESCHEDULE: 'turquoise', @@ -159,6 +162,7 @@ def color_fg(cls, state): TaskInstanceState.RUNNING, TaskInstanceState.SENSING, TaskInstanceState.SHUTDOWN, + TaskInstanceState.RESTARTING, TaskInstanceState.UP_FOR_RETRY, TaskInstanceState.UP_FOR_RESCHEDULE, ] @@ -182,6 +186,11 @@ def color_fg(cls, state): A list of states indicating that a task or dag is a success state. """ + terminating_states = frozenset([TaskInstanceState.SHUTDOWN, TaskInstanceState.RESTARTING]) + """ + A list of states indicating that a task has been terminated. + """ + class PokeState: """Static class with poke states constants used in smart operator.""" diff --git a/airflow/utils/trigger_rule.py b/airflow/utils/trigger_rule.py index 0ce820bf26448..c9ac00abc195f 100644 --- a/airflow/utils/trigger_rule.py +++ b/airflow/utils/trigger_rule.py @@ -31,6 +31,7 @@ class TriggerRule: NONE_FAILED_OR_SKIPPED = 'none_failed_or_skipped' NONE_SKIPPED = 'none_skipped' DUMMY = 'dummy' + ALWAYS = 'always' _ALL_TRIGGER_RULES: Set[str] = set() diff --git a/airflow/www/forms.py b/airflow/www/forms.py index 45c9dd5f58cc2..8a28e546ceece 100644 --- a/airflow/www/forms.py +++ b/airflow/www/forms.py @@ -30,21 +30,12 @@ from flask_babel import lazy_gettext from flask_wtf import FlaskForm from wtforms import widgets -from wtforms.fields import ( - BooleanField, - Field, - IntegerField, - PasswordField, - SelectField, - StringField, - TextAreaField, -) +from wtforms.fields import Field, IntegerField, PasswordField, SelectField, StringField, TextAreaField from wtforms.validators import InputRequired, Optional from airflow.configuration import conf from airflow.utils import timezone from airflow.utils.types import DagRunType -from airflow.www.validators import ValidJson from airflow.www.widgets import ( AirflowDateTimePickerROWidget, AirflowDateTimePickerWidget, @@ -130,32 +121,22 @@ class DateTimeWithNumRunsWithDagRunsForm(DateTimeWithNumRunsForm): execution_date = SelectField("DAG run") -class DagRunForm(DynamicForm): - """Form for adding DAG Run""" +class DagRunEditForm(DynamicForm): + """Form for editing DAG Run. - dag_id = StringField(lazy_gettext('Dag Id'), validators=[InputRequired()], widget=BS3TextFieldWidget()) - start_date = DateTimeWithTimezoneField(lazy_gettext('Start Date'), widget=AirflowDateTimePickerWidget()) - end_date = DateTimeWithTimezoneField(lazy_gettext('End Date'), widget=AirflowDateTimePickerWidget()) - run_id = StringField(lazy_gettext('Run Id'), validators=[InputRequired()], widget=BS3TextFieldWidget()) - state = SelectField( - lazy_gettext('State'), - choices=( - ('success', 'success'), - ('running', 'running'), - ('failed', 'failed'), - ), - widget=Select2Widget(), - validators=[InputRequired()], - ) + We don't actually want to allow editing, so everything is read-only here. + """ + + dag_id = StringField(lazy_gettext('Dag Id'), widget=BS3TextFieldROWidget()) + start_date = DateTimeWithTimezoneField(lazy_gettext('Start Date'), widget=AirflowDateTimePickerROWidget()) + end_date = DateTimeWithTimezoneField(lazy_gettext('End Date'), widget=AirflowDateTimePickerROWidget()) + run_id = StringField(lazy_gettext('Run Id'), widget=BS3TextFieldROWidget()) + state = StringField(lazy_gettext('State'), widget=BS3TextFieldROWidget()) execution_date = DateTimeWithTimezoneField( lazy_gettext('Execution Date'), - widget=AirflowDateTimePickerWidget(), - validators=[InputRequired()], - ) - external_trigger = BooleanField(lazy_gettext('External Trigger')) - conf = TextAreaField( - lazy_gettext('Conf'), validators=[ValidJson(), Optional()], widget=BS3TextAreaFieldWidget() + widget=AirflowDateTimePickerROWidget(), ) + conf = TextAreaField(lazy_gettext('Conf'), widget=BS3TextAreaROWidget()) def populate_obj(self, item): """Populates the attributes of the passed obj with data from the form’s fields.""" @@ -165,23 +146,6 @@ def populate_obj(self, item): item.conf = json.loads(item.conf) -class DagRunEditForm(DagRunForm): - """Form for editing DAG Run""" - - dag_id = StringField(lazy_gettext('Dag Id'), validators=[InputRequired()], widget=BS3TextFieldROWidget()) - start_date = DateTimeWithTimezoneField(lazy_gettext('Start Date'), widget=AirflowDateTimePickerROWidget()) - end_date = DateTimeWithTimezoneField(lazy_gettext('End Date'), widget=AirflowDateTimePickerROWidget()) - run_id = StringField(lazy_gettext('Run Id'), validators=[InputRequired()], widget=BS3TextFieldROWidget()) - execution_date = DateTimeWithTimezoneField( - lazy_gettext('Execution Date'), - widget=AirflowDateTimePickerROWidget(), - validators=[InputRequired()], - ) - conf = TextAreaField( - lazy_gettext('Conf'), validators=[ValidJson(), Optional()], widget=BS3TextAreaROWidget() - ) - - class TaskInstanceEditForm(DynamicForm): """Form for editing TaskInstance""" diff --git a/airflow/www/static/css/graph.css b/airflow/www/static/css/graph.css index f4c7b942c5a77..ce76df7dfcc7b 100644 --- a/airflow/www/static/css/graph.css +++ b/airflow/www/static/css/graph.css @@ -148,6 +148,10 @@ g.node.shutdown rect { stroke: blue; } +g.node.restarting rect { + stroke: violet; +} + g.node.upstream_failed rect { stroke: orange; } diff --git a/airflow/www/static/css/tree.css b/airflow/www/static/css/tree.css index 05b2c81de9973..c17cf0a6eb65b 100644 --- a/airflow/www/static/css/tree.css +++ b/airflow/www/static/css/tree.css @@ -67,6 +67,10 @@ rect.shutdown { fill: blue; } +rect.restarting { + fill: violet; +} + rect.upstream_failed { fill: orange; } diff --git a/airflow/www/static/js/graph.js b/airflow/www/static/js/graph.js index d3d913b6dd27a..5b1ffee5406e3 100644 --- a/airflow/www/static/js/graph.js +++ b/airflow/www/static/js/graph.js @@ -536,7 +536,7 @@ function getNodeState(nodeId, tis) { // In this order, if any of these states appeared in childrenStates, return it as // the group state. const priority = ['failed', 'upstream_failed', 'up_for_retry', 'up_for_reschedule', - 'queued', 'scheduled', 'sensing', 'running', 'shutdown', 'removed', + 'queued', 'scheduled', 'sensing', 'running', 'shutdown', 'restarting', 'removed', 'no_status', 'success', 'skipped']; return priority.find((state) => childrenStates.has(state)) || 'no_status'; diff --git a/airflow/www/templates/airflow/dag.html b/airflow/www/templates/airflow/dag.html index 322d19face731..2c4c25cb6082e 100644 --- a/airflow/www/templates/airflow/dag.html +++ b/airflow/www/templates/airflow/dag.html @@ -144,9 +144,6 @@

  • Trigger DAG w/ config
    • - - - diff --git a/airflow/www/templates/airflow/dag_details.html b/airflow/www/templates/airflow/dag_details.html index af03e3dae8574..4d8fe0eb4b264 100644 --- a/airflow/www/templates/airflow/dag_details.html +++ b/airflow/www/templates/airflow/dag_details.html @@ -88,8 +88,8 @@

    {{ title }}

    {{ dag.task_ids }} - Filepath - {{ dag.filepath }} + Relative file location + {{ dag.relative_fileloc }} Owner diff --git a/airflow/www/templates/airflow/dags.html b/airflow/www/templates/airflow/dags.html index 290675c839e95..bdc1190c985f2 100644 --- a/airflow/www/templates/airflow/dags.html +++ b/airflow/www/templates/airflow/dags.html @@ -195,9 +195,6 @@

    {{ page_title }}

  • Trigger DAG w/ config
    • - - - {% endif %} {# Use dag_id instead of dag.dag_id, because the DAG might not exist in the webserver's DagBag #} diff --git a/airflow/www/views.py b/airflow/www/views.py index 117b173f84fc3..9d1d4ad4f1f82 100644 --- a/airflow/www/views.py +++ b/airflow/www/views.py @@ -121,7 +121,6 @@ from airflow.www.forms import ( ConnectionForm, DagRunEditForm, - DagRunForm, DateTimeForm, DateTimeWithNumRunsForm, DateTimeWithNumRunsWithDagRunsForm, @@ -279,6 +278,29 @@ def task_group_to_dict(task_group): } +def get_key_paths(input_dict): + """Return a list of dot-separated dictionary paths""" + for key, value in input_dict.items(): + if isinstance(value, dict): + for sub_key in get_key_paths(value): + yield '.'.join((key, sub_key)) + else: + yield key + + +def get_value_from_path(key_path, content): + """Return the value from a dictionary based on dot-separated path of keys""" + elem = content + for x in key_path.strip(".").split("."): + try: + x = int(x) + elem = elem[x] + except ValueError: + elem = elem.get(x) + + return elem + + def dag_edges(dag): """ Create the list of edges needed to construct the Graph view. @@ -995,11 +1017,31 @@ def rendered_templates(self): renderer = task.template_fields_renderers.get(template_field, template_field) if renderer in renderers: if isinstance(content, (dict, list)): - content = json.dumps(content, sort_keys=True, indent=4) - html_dict[template_field] = renderers[renderer](content) + json_content = json.dumps(content, sort_keys=True, indent=4) + html_dict[template_field] = renderers[renderer](json_content) + else: + html_dict[template_field] = renderers[renderer](content) else: html_dict[template_field] = Markup("
    {}
    ").format(pformat(content)) + if isinstance(content, dict): + if template_field == 'op_kwargs': + for key, value in content.items(): + renderer = task.template_fields_renderers.get(key, key) + if renderer in renderers: + html_dict['.'.join([template_field, key])] = renderers[renderer](value) + else: + html_dict['.'.join([template_field, key])] = Markup( + "
    {}
    " + ).format(pformat(value)) + else: + for dict_keys in get_key_paths(content): + template_path = '.'.join((template_field, dict_keys)) + renderer = task.template_fields_renderers.get(template_path, template_path) + if renderer in renderers: + content_value = get_value_from_path(dict_keys, content) + html_dict[template_path] = renderers[renderer](content_value) + return self.render_template( 'airflow/ti_code.html', html_dict=html_dict, @@ -1346,9 +1388,8 @@ def xcom(self, session=None): dttm = timezone.parse(execution_date) form = DateTimeForm(data={'execution_date': dttm}) root = request.args.get('root', '') - dm_db = models.DagModel ti_db = models.TaskInstance - dag = session.query(dm_db).filter(dm_db.dag_id == dag_id).first() + dag = DagModel.get_dagmodel(dag_id) ti = session.query(ti_db).filter(and_(ti_db.dag_id == dag_id, ti_db.task_id == task_id)).first() if not ti: @@ -2617,48 +2658,6 @@ def paused(self): models.DagModel.get_dagmodel(dag_id).set_is_paused(is_paused=is_paused) return "OK" - @expose('/refresh', methods=['POST']) - @auth.has_access( - [ - (permissions.ACTION_CAN_EDIT, permissions.RESOURCE_DAG), - ] - ) - @action_logging - @provide_session - def refresh(self, session=None): - """Refresh DAG.""" - dag_id = request.values.get('dag_id') - orm_dag = session.query(DagModel).filter(DagModel.dag_id == dag_id).first() - - if orm_dag: - orm_dag.last_expired = timezone.utcnow() - session.merge(orm_dag) - session.commit() - - dag = current_app.dag_bag.get_dag(dag_id) - # sync dag permission - current_app.appbuilder.sm.sync_perm_for_dag(dag_id, dag.access_control) - - flash(f"DAG [{dag_id}] is now fresh as a daisy") - return redirect(request.referrer) - - @expose('/refresh_all', methods=['POST']) - @auth.has_access( - [ - (permissions.ACTION_CAN_EDIT, permissions.RESOURCE_DAG), - ] - ) - @action_logging - def refresh_all(self): - """Refresh everything""" - current_app.dag_bag.collect_dags_from_db() - - # sync permissions for all dags - for dag_id, dag in current_app.dag_bag.dags.items(): - current_app.appbuilder.sm.sync_perm_for_dag(dag_id, dag.access_control) - flash("All DAGs are now up to date") - return redirect(url_for('Airflow.index')) - @expose('/gantt') @auth.has_access( [ @@ -3221,11 +3220,33 @@ def action_mulduplicate(self, connections, session=None): def process_form(self, form, is_created): """Process form data.""" conn_type = form.data['conn_type'] + conn_id = form.data["conn_id"] extra = { key: form.data[key] for key in self.extra_fields if key in form.data and key.startswith(f"extra__{conn_type}__") } + + # If parameters are added to the classic `Extra` field, include these values along with + # custom-field extras. + extra_conn_params = form.data.get("extra") + + if extra_conn_params: + try: + extra.update(json.loads(extra_conn_params)) + except (JSONDecodeError, TypeError): + flash( + Markup( + "

    The Extra connection field contained an invalid value for Conn ID: " + f"{conn_id}.

    " + "

    If connection parameters need to be added to Extra, " + "please make sure they are in the form of a single, valid JSON object.


    " + "The following Extra parameters were not added to the connection:
    " + f"{extra_conn_params}", + ), + category="error", + ) + if extra.keys(): form.extra.data = json.dumps(extra) @@ -3562,7 +3583,6 @@ class DagRunModelView(AirflowModelView): class_permission_name = permissions.RESOURCE_DAG_RUN method_permission_name = { - 'add': 'create', 'list': 'read', 'action_clear': 'delete', 'action_muldelete': 'delete', @@ -3571,14 +3591,12 @@ class DagRunModelView(AirflowModelView): 'action_set_success': 'edit', } base_permissions = [ - permissions.ACTION_CAN_CREATE, permissions.ACTION_CAN_READ, permissions.ACTION_CAN_EDIT, permissions.ACTION_CAN_DELETE, permissions.ACTION_CAN_ACCESS_MENU, ] - add_columns = ['state', 'dag_id', 'execution_date', 'run_id', 'external_trigger', 'conf'] list_columns = [ 'state', 'dag_id', @@ -3607,7 +3625,6 @@ class DagRunModelView(AirflowModelView): base_filters = [['dag_id', DagFilter, lambda: []]] - add_form = DagRunForm edit_form = DagRunEditForm formatters_columns = { diff --git a/breeze b/breeze index 7aa42952fb7cb..420861eefb53d 100755 --- a/breeze +++ b/breeze @@ -1199,12 +1199,6 @@ function breeze::parse_arguments() { echo shift ;; - --continue-on-pip-check-failure) - export CONTINUE_ON_PIP_CHECK_FAILURE="true" - echo "Skip PIP check failure." - echo - shift - ;; --package-format) export PACKAGE_FORMAT="${2}" echo "Selected package type: ${PACKAGE_FORMAT}" @@ -2439,9 +2433,6 @@ ${FORMATTED_INSTALLATION_METHOD} --upgrade-to-newer-dependencies Upgrades PIP packages to latest versions available without looking at the constraints. ---continue-on-pip-check-failure - Continue even if 'pip check' fails. - " } diff --git a/breeze-complete b/breeze-complete index e5b5d3f254d7f..01efc7c79cc6c 100644 --- a/breeze-complete +++ b/breeze-complete @@ -29,7 +29,7 @@ _breeze_allowed_integrations="cassandra kerberos mongo openldap pinot rabbitmq r _breeze_allowed_generate_constraints_modes="source-providers pypi-providers no-providers" _breeze_allowed_kubernetes_modes="image" _breeze_allowed_kubernetes_versions="v1.20.2 v1.19.7 v1.18.15" -_breeze_allowed_helm_versions="v3.2.4" +_breeze_allowed_helm_versions="v3.6.3" _breeze_allowed_kind_versions="v0.11.1" _breeze_allowed_mysql_versions="5.7 8" _breeze_allowed_mssql_versions="2017-latest 2019-latest" diff --git a/chart/UPDATING.rst b/chart/UPDATING.rst index f9e2ac29f3106..9413403ea209b 100644 --- a/chart/UPDATING.rst +++ b/chart/UPDATING.rst @@ -36,6 +36,8 @@ assists users migrating to a new version. Airflow Helm Chart 1.1.0 ------------------------ +Run ``helm repo update`` before upgrading the chart to the latest version. + Default Airflow version is updated to ``2.1.2`` """"""""""""""""""""""""""""""""""""""""""""""" diff --git a/chart/templates/NOTES.txt b/chart/templates/NOTES.txt index 454088f43e95e..9bc5783406471 100644 --- a/chart/templates/NOTES.txt +++ b/chart/templates/NOTES.txt @@ -87,7 +87,7 @@ You are using ssh authentication for your gitsync repo, however you currently ha making you susceptible to man-in-the-middle attacks! Information on how to set knownHosts can be found here: -https://airflow.apache.org/docs/helm-chart/latest/production-guide.html#knownhosts +https://airflow.apache.org/docs/helm-chart/stable/production-guide.html#knownhosts {{- end }} diff --git a/confirm b/confirm index e796737cbf3b0..42316da01a2fd 100755 --- a/confirm +++ b/confirm @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -set -euo pipefail +set -uo pipefail if [[ -n "${FORCE_ANSWER_TO_QUESTIONS=}" ]]; then RESPONSE=${FORCE_ANSWER_TO_QUESTIONS} @@ -31,8 +31,8 @@ if [[ -n "${FORCE_ANSWER_TO_QUESTIONS=}" ]]; then esac else echo - echo "Please confirm ${1}. Are you sure? [y/N/q]" - read -r RESPONSE + echo "Please confirm ${1} (or wait 4 seconds to skip it). Are you sure? [y/N/q]" + read -t 4 -r RESPONSE fi case "${RESPONSE}" in diff --git a/dev/README_RELEASE_HELM_CHART.md b/dev/README_RELEASE_HELM_CHART.md index 2955e8908ab95..e5aea88e14dd9 100644 --- a/dev/README_RELEASE_HELM_CHART.md +++ b/dev/README_RELEASE_HELM_CHART.md @@ -37,6 +37,7 @@ - [Update `index.yaml` to Airflow Website](#update-indexyaml-to-airflow-website) - [Notify developers of release](#notify-developers-of-release) - [Update Announcements page](#update-announcements-page) + - [Remove old releases](#remove-old-releases) @@ -490,12 +491,6 @@ for f in ../../../airflow-dev/helm-chart/$RC/*; do svn cp $f ${$(basename $f)/}; svn rm index.yaml svn commit -m "Release Airflow Helm Chart Check ${VERSION} from ${RC}" -# Remove old release -# http://www.apache.org/legal/release-policy.html#when-to-archive -cd .. -export PREVIOUS_VERSION=1.0.0 -svn rm ${PREVIOUS_VERSION} -svn commit -m "Remove old Helm Chart release: ${PREVIOUS_VERSION}" ``` Verify that the packages appear in [Airflow Helm Chart](https://dist.apache.org/repos/dist/release/airflow/helm-chart/). @@ -590,3 +585,18 @@ EOF ## Update Announcements page Update "Announcements" page at the [Official Airflow website](https://airflow.apache.org/announcements/) + +## Remove old releases + +We should keep the old version a little longer than a day or at least until the updated +``index.yaml`` is published. This is to avoid errors for users who haven't run ``helm repo update``. + +It is probably ok if we leave last 2 versions on release svn repo too. + +```shell +# http://www.apache.org/legal/release-policy.html#when-to-archive +cd airflow-release/helm-chart +export PREVIOUS_VERSION=1.0.0 +svn rm ${PREVIOUS_VERSION} +svn commit -m "Remove old Helm Chart release: ${PREVIOUS_VERSION}" +``` diff --git a/dev/provider_packages/prepare_provider_packages.py b/dev/provider_packages/prepare_provider_packages.py index e36fcf87e1af7..23c51c74b2f01 100755 --- a/dev/provider_packages/prepare_provider_packages.py +++ b/dev/provider_packages/prepare_provider_packages.py @@ -1694,8 +1694,10 @@ def replace_content(file_path, old_text, new_text, provider_package_id): os.remove(temp_file_path) +AUTOMATICALLY_GENERATED_MARKER = "AUTOMATICALLY GENERATED" AUTOMATICALLY_GENERATED_CONTENT = ( - ".. THE REMINDER OF THE FILE IS AUTOMATICALLY GENERATED. IT WILL BE OVERWRITTEN AT RELEASE TIME!" + f".. THE REMAINDER OF THE FILE IS {AUTOMATICALLY_GENERATED_MARKER}. " + f"IT WILL BE OVERWRITTEN AT RELEASE TIME!" ) @@ -1715,7 +1717,7 @@ def update_index_rst( new_text = deepcopy(old_text) lines = old_text.splitlines(keepends=False) for index, line in enumerate(lines): - if line == AUTOMATICALLY_GENERATED_CONTENT: + if AUTOMATICALLY_GENERATED_MARKER in line: new_text = "\n".join(lines[:index]) new_text += "\n" + AUTOMATICALLY_GENERATED_CONTENT + "\n" new_text += index_update diff --git a/docs/apache-airflow-providers-amazon/connections/aws.rst b/docs/apache-airflow-providers-amazon/connections/aws.rst index 0d7b04cdbc17e..c73b12039fb74 100644 --- a/docs/apache-airflow-providers-amazon/connections/aws.rst +++ b/docs/apache-airflow-providers-amazon/connections/aws.rst @@ -415,3 +415,19 @@ You can configure connection, also using environmental variable :envvar:`AIRFLOW assume_role_method=assume_role_with_web_identity&\ assume_role_with_web_identity_federation=google&\ assume_role_with_web_identity_federation_audience=aaa.polidea.com" + +Using IAM Roles for Service Accounts (IRSA) on EKS +---------------------------------------------------------------- + +If you are running Airflow on Amazon EKS, you can grant AWS related permission (such as S3 Read/Write for remote logging) to the Airflow service by granting the IAM role to it's service account. To activate this, the following steps must be followed: + +1. Create an IAM OIDC Provider on EKS cluster. +2. Create an IAM Role and Policy to attach to the Airflow service account with web identity provider created at 1. +3. Add the corresponding IAM Role to the Airflow service account as an annotation. + +.. seealso:: + https://docs.aws.amazon.com/eks/latest/userguide/iam-roles-for-service-accounts.html + +Then you can find ``AWS_ROLE_ARN`` and ``AWS_WEB_IDENTITY_TOKEN_FILE`` in environment variables of appropriate pods that `Amazon EKS Pod Identity Web Hook `__ added. Then `boto3 `__ will configure credentials using those variables. + +In order to use IRSA in Airflow, you have to create an aws connection with all fields empty. If a field such as ``role-arn`` is set, Airflow does not follow the boto3 default flow because it manually create a session using connection fields. If you did not change the default connection ID, an empty AWS connection named ``aws_default`` would be enough. diff --git a/docs/apache-airflow-providers-amazon/operators/emr.rst b/docs/apache-airflow-providers-amazon/operators/emr.rst index ca3b91bcd7a65..84fda3514dfd0 100644 --- a/docs/apache-airflow-providers-amazon/operators/emr.rst +++ b/docs/apache-airflow-providers-amazon/operators/emr.rst @@ -58,7 +58,7 @@ Create EMR Job Flow with automatic steps Purpose """"""" -This example dag ``example_emr_job_flow_automatic_steps.py`` use ``EmrCreateJobFlowOperator`` to create a new EMR job flow calculating the mathematical constant ``Pi``, and monitor the progress +This example dag ``example_emr_job_flow_automatic_steps.py`` uses ``EmrCreateJobFlowOperator`` to create a new EMR job flow calculating the mathematical constant ``Pi``, and monitors the progress with ``EmrJobFlowSensor``. The cluster will be terminated automatically after finishing the steps. JobFlow configuration diff --git a/docs/apache-airflow-providers-hashicorp/secrets-backends/hashicorp-vault.rst b/docs/apache-airflow-providers-hashicorp/secrets-backends/hashicorp-vault.rst index 09d3576ffe076..b71cd0d41e352 100644 --- a/docs/apache-airflow-providers-hashicorp/secrets-backends/hashicorp-vault.rst +++ b/docs/apache-airflow-providers-hashicorp/secrets-backends/hashicorp-vault.rst @@ -130,3 +130,41 @@ Verify that you can get the secret from ``vault``: Note that the secret ``Key`` is ``value``, and secret ``Value`` is ``world`` and ``mount_point`` is ``airflow``. + +Storing and Retrieving Config +"""""""""""""""""""""""""""""""" + +If you have set ``config_path`` as ``config`` and ``mount_point`` as ``airflow``, then for config ``sql_alchemy_conn_secret`` with +``sql_alchemy_conn_value`` as value, you would want to store your secret as: + +.. code-block:: bash + + vault kv put airflow/config/sql_alchemy_conn_value value=postgres://user:pass@host:5432/db?ssl_mode=disable + +Verify that you can get the secret from ``vault``: + +.. code-block:: console + + ❯ vault kv get airflow/config/sql_alchemy_conn_value + ====== Metadata ====== + Key Value + --- ----- + created_time 2020-03-28T02:10:54.301784Z + deletion_time n/a + destroyed false + version 1 + + ==== Data ==== + Key Value + --- ----- + value postgres://user:pass@host:5432/db?ssl_mode=disable + +Then you can use above secret for ``sql_alchemy_conn_secret`` in your configuration file. + +.. code-block:: ini + + [core] + sql_alchemy_conn_secret: "sql_alchemy_conn_value" + +Note that the secret ``Key`` is ``value``, and secret ``Value`` is ``postgres://user:pass@host:5432/db?ssl_mode=disable`` and +``mount_point`` is ``airflow``. diff --git a/docs/apache-airflow-providers-samba/index.rst b/docs/apache-airflow-providers-samba/index.rst index bb67f4440efcb..1ca77f7379223 100644 --- a/docs/apache-airflow-providers-samba/index.rst +++ b/docs/apache-airflow-providers-samba/index.rst @@ -71,7 +71,7 @@ PIP requirements PIP package Version required ================== ================== ``apache-airflow`` ``>=2.1.0`` -``pysmbclient`` ``>=0.1.3`` +``smbprotocol`` ``>=1.5.0`` ================== ================== .. include:: ../../airflow/providers/samba/CHANGELOG.rst diff --git a/docs/apache-airflow-providers/index.rst b/docs/apache-airflow-providers/index.rst index 7329b7fd780a4..71c5132acdec8 100644 --- a/docs/apache-airflow-providers/index.rst +++ b/docs/apache-airflow-providers/index.rst @@ -21,6 +21,8 @@ Provider packages .. contents:: :local: +.. _providers:community-maintained-providers: + Community maintained providers '''''''''''''''''''''''''''''' @@ -31,6 +33,9 @@ Those provider packages are separated per-provider (for example ``amazon``, ``go etc.). Those packages are available as ``apache-airflow-providers`` packages - separately per each provider (for example there is an ``apache-airflow-providers-amazon`` or ``apache-airflow-providers-google`` package). +The full list of community managed providers is available at +`Providers Index `_. + You can install those provider packages separately in order to interface with a given service. For those providers that have corresponding extras, the provider packages (latest version from PyPI) are installed automatically when Airflow is installed with the extra. diff --git a/docs/apache-airflow/best-practices.rst b/docs/apache-airflow/best-practices.rst index cb40f4c35a1a2..2857ec401e4cd 100644 --- a/docs/apache-airflow/best-practices.rst +++ b/docs/apache-airflow/best-practices.rst @@ -62,7 +62,9 @@ Some of the ways you can avoid producing a different result - .. tip:: You should define repetitive parameters such as ``connection_id`` or S3 paths in ``default_args`` rather than declaring them for each task. - The ``default_args`` help to avoid mistakes such as typographical errors. + The ``default_args`` help to avoid mistakes such as typographical errors. Also, most connection types have unique parameter names in + tasks, so you can declare a connection only once in ``default_args`` (for example ``gcp_conn_id``) and it is automatically + used by all operators that use this connection type. Deleting a task ---------------- @@ -91,13 +93,16 @@ Variables --------- You should avoid usage of Variables outside an operator's ``execute()`` method or Jinja templates if possible, -as Variables create a connection to metadata DB of Airflow to fetch the value, which can slow down parsing and place extra load on the DB. +as Variables create a connection to metadata DB of Airflow to fetch the value, which can slow down parsing and +place extra load on the DB. Airflow parses all the DAGs in the background at a specific period. -The default period is set using ``processor_poll_interval`` config, which is by default 1 second. During parsing, Airflow creates a new connection to the metadata DB for each DAG. +The default period is set using the ``processor_poll_interval`` config, which is 1 second by default. +During parsing, Airflow creates a new connection to the metadata DB for each DAG. This can result in a lot of open connections. -The best way of using variables is via a Jinja template, which will delay reading the value until the task execution. The template syntax to do this is: +The best way of using variables is via a Jinja template, which will delay reading the value until the task execution. +The template syntax to do this is: .. code-block:: @@ -109,8 +114,12 @@ or if you need to deserialize a json object from the variable : {{ var.json. }} -An alternative option is to use environment variables in the top-level python code or use Environment Variables to create and manage Airflow variables. to manage Airflow Variables. This will avoid new connections to Airflow metadata DB every time Airflow parses the python file. For more information, see: :ref:`managing_variables`. +For security purpose, you're recommended to use the :ref:`Secrets Backend` +for any variable that contains sensitive data. +An alternative option is to use environment variables in the top-level Python code or use environment variables to +create and manage Airflow variables. This will avoid new connections to Airflow metadata DB every time +Airflow parses the Python file. For more information, see: :ref:`managing_variables`. Top level Python Code --------------------- diff --git a/docs/apache-airflow/concepts/dags.rst b/docs/apache-airflow/concepts/dags.rst index 89bde8bfcf589..c564ef85b7efc 100644 --- a/docs/apache-airflow/concepts/dags.rst +++ b/docs/apache-airflow/concepts/dags.rst @@ -337,7 +337,7 @@ However, this is just the default behaviour, and you can control it using the `` * ``none_failed``: All upstream tasks have not ``failed`` or ``upstream_failed`` - that is, all upstream tasks have succeeded or been skipped * ``none_failed_or_skipped``: All upstream tasks have not ``failed`` or ``upstream_failed``, and at least one upstream task has succeeded. * ``none_skipped``: No upstream task is in a ``skipped`` state - that is, all upstream tasks are in a ``success``, ``failed``, or ``upstream_failed`` state -* ``dummy``: No dependencies at all, run this task at any time +* ``always``: No dependencies at all, run this task at any time You can also combine this with the :ref:`concepts:depends-on-past` functionality if you wish. diff --git a/docs/apache-airflow/concepts/scheduler.rst b/docs/apache-airflow/concepts/scheduler.rst index 6ea5ff236f75c..0a1079e4a290c 100644 --- a/docs/apache-airflow/concepts/scheduler.rst +++ b/docs/apache-airflow/concepts/scheduler.rst @@ -123,13 +123,14 @@ The following databases are fully supported and provide an "optimal" experience: .. warning:: - MariaDB does not implement the ``SKIP LOCKED`` or ``NOWAIT`` SQL clauses (see `MDEV-13115 - `_). Without these features running multiple schedulers is not - supported and deadlock errors have been reported. + MariaDB did not implement the ``SKIP LOCKED`` or ``NOWAIT`` SQL clauses until version + `10.6.0 `_. + Without these features, running multiple schedulers is not supported and deadlock errors have been reported. MariaDB + 10.6.0 and following may work appropriately with multiple schedulers, but this has not been tested. .. warning:: - MySQL 5.x also does not support ``SKIP LOCKED`` or ``NOWAIT``, and additionally is more prone to deciding + MySQL 5.x does not support ``SKIP LOCKED`` or ``NOWAIT``, and additionally is more prone to deciding queries are deadlocked, so running with more than a single scheduler on MySQL 5.x is not supported or recommended. diff --git a/docs/apache-airflow/concepts/tasks.rst b/docs/apache-airflow/concepts/tasks.rst index d4a6608c7c539..f481baba60485 100644 --- a/docs/apache-airflow/concepts/tasks.rst +++ b/docs/apache-airflow/concepts/tasks.rst @@ -70,6 +70,8 @@ The possible states for a Task Instance are: * ``queued``: The task has been assigned to an Executor and is awaiting a worker * ``running``: The task is running on a worker (or on a local/synchronous executor) * ``success``: The task finished running without errors +* ``shutdown``: The task was externally requested to shut down when it was running +* ``restarting``: The task was externally requested to restart when it was running * ``failed``: The task had an error during execution and failed to run * ``skipped``: The task was skipped due to branching, LatestOnly, or similar. * ``upstream_failed``: An upstream task failed and the :ref:`Trigger Rule ` says we needed it diff --git a/docs/apache-airflow/howto/custom-operator.rst b/docs/apache-airflow/howto/custom-operator.rst index 8074bf534eb0a..25558f76458f5 100644 --- a/docs/apache-airflow/howto/custom-operator.rst +++ b/docs/apache-airflow/howto/custom-operator.rst @@ -195,7 +195,7 @@ with actual value. Note that Jinja substitutes the operator attributes and not t In the example, the ``template_fields`` should be ``['guest_name']`` and not ``['name']`` -Additionally you may provide ``template_fields_renderers`` dictionary which defines in what style the value +Additionally you may provide ``template_fields_renderers`` a dictionary which defines in what style the value from template field renders in Web UI. For example: .. code-block:: python @@ -208,12 +208,48 @@ from template field renders in Web UI. For example: super().__init__(**kwargs) self.request_body = request_body +In the situation where ``template_field`` is itself a dictionary, it is also possible to specify a +dot-separated key path to extract and render individual elements appropriately. For example: + +.. code-block:: python + + class MyConfigOperator(BaseOperator): + template_fields = ["configuration"] + template_fields_renderers = { + "configuration": "json", + "configuration.query.sql": "sql", + } + + def __init__(self, configuration: dict, **kwargs) -> None: + super().__init__(**kwargs) + self.configuration = configuration + +Then using this template as follows: + +.. code-block:: python + + with dag: + config_task = MyConfigOperator( + task_id="task_id_1", + configuration={"query": {"job_id": "123", "sql": "select * from my_table"}}, + dag=dag, + ) + +This will result in the UI rendering ``configuration`` as json in addition to the value contained in the +configuration at ``query.sql`` to be rendered with the SQL lexer. + +.. image:: ../img/template_field_renderer_path.png + Currently available lexers: - bash - doc + - hql + - html + - jinja - json - md + - powershell - py - rst - sql diff --git a/docs/apache-airflow/img/task_lifecycle_diagram.png b/docs/apache-airflow/img/task_lifecycle_diagram.png index ad0bd9ecf49ec..810942fc74001 100644 Binary files a/docs/apache-airflow/img/task_lifecycle_diagram.png and b/docs/apache-airflow/img/task_lifecycle_diagram.png differ diff --git a/docs/apache-airflow/img/template_field_renderer_path.png b/docs/apache-airflow/img/template_field_renderer_path.png new file mode 100644 index 0000000000000..bbecf61ea5da0 Binary files /dev/null and b/docs/apache-airflow/img/template_field_renderer_path.png differ diff --git a/docs/apache-airflow/index.rst b/docs/apache-airflow/index.rst index ba9db5d5c317b..18b2f7a973666 100644 --- a/docs/apache-airflow/index.rst +++ b/docs/apache-airflow/index.rst @@ -97,7 +97,7 @@ unit of work and continuity. lineage dag-serialization modules_management - Release policies + Release Policies changelog best-practices production-deployment diff --git a/docs/apache-airflow/logging-monitoring/check-health.rst b/docs/apache-airflow/logging-monitoring/check-health.rst index a5f86644baf4b..4468deb1d3684 100644 --- a/docs/apache-airflow/logging-monitoring/check-health.rst +++ b/docs/apache-airflow/logging-monitoring/check-health.rst @@ -20,9 +20,13 @@ Checking Airflow Health Status ============================== -Airflow has two methods to check the health of components - HTTP checks and CLI checks. Their choice depends on the role of the component as well as what tools it uses to monitor the deployment. +Airflow has two methods to check the health of components - HTTP checks and CLI checks. All available checks are +accessible through the CLI, but only some are accessible through HTTP due to the role of the component being checked +and the tools being used to monitor the deployment. -For example, when running on Kubernetes, use `a Liveness probes `__ (``livenessProbe`` property) with :ref:`CLI checks ` on the scheduler deployment to restart it when it fail. For the webserver, you can configure the readiness probe (``readinessProbe`` property) using :ref:`check-health/http-endpoint`. +For example, when running on Kubernetes, use `a Liveness probes `__ (``livenessProbe`` property) +with :ref:`CLI checks ` on the scheduler deployment to restart it when it fails. +For the webserver, you can configure the readiness probe (``readinessProbe`` property) using :ref:`check-health/http-endpoint`. For an example for a Docker Compose environment, see the ``docker-compose.yaml`` file available in the :doc:`/start/docker`. diff --git a/docs/apache-airflow/logging-monitoring/errors.rst b/docs/apache-airflow/logging-monitoring/errors.rst index 37ed307ba49ab..578666b078d77 100644 --- a/docs/apache-airflow/logging-monitoring/errors.rst +++ b/docs/apache-airflow/logging-monitoring/errors.rst @@ -41,7 +41,7 @@ Add your ``SENTRY_DSN`` to your configuration file e.g. ``airflow.cfg`` in ``[se .. note:: If this value is not provided, the SDK will try to read it from the ``SENTRY_DSN`` environment variable. -You can supply `additional configuration options `__ based on the Python platform via ``[sentry]`` section. +You can supply `additional configuration options `__ based on the Python platform via ``[sentry]`` section. Unsupported options: ``integrations``, ``in_app_include``, ``in_app_exclude``, ``ignore_errors``, ``before_breadcrumb``, ``before_send``, ``transport``. Tags @@ -60,7 +60,7 @@ Breadcrumbs ------------ -When a task fails with an error `breadcrumbs `__ will be added for the other tasks in the current dag run. +When a task fails with an error `breadcrumbs `__ will be added for the other tasks in the current dag run. ======================================= ============================================================== Name Description diff --git a/docs/apache-airflow/logging-monitoring/metrics.rst b/docs/apache-airflow/logging-monitoring/metrics.rst index d410261725ffd..b55ad929155f2 100644 --- a/docs/apache-airflow/logging-monitoring/metrics.rst +++ b/docs/apache-airflow/logging-monitoring/metrics.rst @@ -50,15 +50,15 @@ the metrics that start with the elements of the list: statsd_allow_list = scheduler,executor,dagrun If you want to redirect metrics to different name, you can configure ``stat_name_handler`` option -in ``[scheduler]`` section. It should point to a function that validate the statsd stat name, apply changes -to the stat name if necessary and return the transformed stat name. The function may looks as follow: +in ``[scheduler]`` section. It should point to a function that validates the statsd stat name, applies changes +to the stat name if necessary, and returns the transformed stat name. The function may looks as follow: .. code-block:: python def my_custom_stat_name_handler(stat_name: str) -> str: return stat_name.lower()[:32] -If you want to use a custom Statsd client outwith the default one provided by Airflow the following key must be added +If you want to use a custom Statsd client instead of the default one provided by Airflow, the following key must be added to the configuration file alongside the module path of your custom Statsd client. This module must be available on your :envvar:`PYTHONPATH`. diff --git a/docs/apache-airflow/security/access-control.rst b/docs/apache-airflow/security/access-control.rst index 26157256abd5b..084043594dfa7 100644 --- a/docs/apache-airflow/security/access-control.rst +++ b/docs/apache-airflow/security/access-control.rst @@ -223,8 +223,6 @@ Get DAG as duration graph DAGs.can_read, Task Instances.can_read Show all tries DAGs.can_read, Task Instances.can_read Viewer Show landing times DAGs.can_read, Task Instances.can_read Viewer Toggle DAG paused status DAGs.can_edit User -Refresh DAG DAGs.can_edit User -Refresh all DAGs DAGs.can_edit User Show Gantt Chart DAGs.can_read, Task Instances.can_read Viewer Get external links DAGs.can_read, Task Instances.can_read Viewer Show Task Instances DAGs.can_read, Task Instances.can_read Viewer diff --git a/docs/apache-airflow/security/index.rst b/docs/apache-airflow/security/index.rst index 2ebc01b71a118..65d02f7f536bf 100644 --- a/docs/apache-airflow/security/index.rst +++ b/docs/apache-airflow/security/index.rst @@ -15,8 +15,6 @@ specific language governing permissions and limitations under the License. - - Security ======== diff --git a/docs/apache-airflow/security/secrets/secrets-backend/index.rst b/docs/apache-airflow/security/secrets/secrets-backend/index.rst index 7afc0fea5043e..272a2a5834d4d 100644 --- a/docs/apache-airflow/security/secrets/secrets-backend/index.rst +++ b/docs/apache-airflow/security/secrets/secrets-backend/index.rst @@ -15,8 +15,7 @@ specific language governing permissions and limitations under the License. - -Secrets backend +Secrets Backend --------------- .. versionadded:: 1.10.10 diff --git a/docs/apache-airflow/security/webserver.rst b/docs/apache-airflow/security/webserver.rst index 5fb03c5230167..3906abf62c18e 100644 --- a/docs/apache-airflow/security/webserver.rst +++ b/docs/apache-airflow/security/webserver.rst @@ -90,7 +90,7 @@ Other Methods ''''''''''''' Since the Airflow 2.0, the default UI is the Flask App Builder RBAC. A ``webserver_config.py`` configuration file -it's automatically generated and can be used to configure the Airflow to support authentication +is automatically generated and can be used to configure the Airflow to support authentication methods like OAuth, OpenID, LDAP, REMOTE_USER. For previous versions from Airflow, the ``$AIRFLOW_HOME/airflow.cfg`` following entry needs to be set to enable @@ -107,7 +107,7 @@ with the following entry in the ``$AIRFLOW_HOME/webserver_config.py``. AUTH_TYPE = AUTH_DB -Another way to create users it's in the UI login page, allowing user self registration through a "Register" button. +Another way to create users is in the UI login page, allowing user self registration through a "Register" button. The following entries in the ``$AIRFLOW_HOME/webserver_config.py`` can be edited to make it possible: .. code-block:: ini diff --git a/docs/apache-airflow/start/docker-compose.yaml b/docs/apache-airflow/start/docker-compose.yaml index 95796715fcc85..832092eec9e95 100644 --- a/docs/apache-airflow/start/docker-compose.yaml +++ b/docs/apache-airflow/start/docker-compose.yaml @@ -44,7 +44,11 @@ version: '3' x-airflow-common: &airflow-common + # In order to add custom dependencies or upgrade provider packages you can use your extended image. + # Comment the image line, place your Dockerfile in the directory where you placed the docker-compose.yaml + # and uncomment the "build" line below, Then run `docker-compose build` to build the images. image: ${AIRFLOW_IMAGE_NAME:-apache/airflow:|version|} + # build: . environment: &airflow-common-env AIRFLOW__CORE__EXECUTOR: CeleryExecutor @@ -60,7 +64,7 @@ x-airflow-common: - ./dags:/opt/airflow/dags - ./logs:/opt/airflow/logs - ./plugins:/opt/airflow/plugins - user: "${AIRFLOW_UID:-50000}:${AIRFLOW_GID:-50000}" + user: "${AIRFLOW_UID:-50000}:${AIRFLOW_GID:-0}" depends_on: &airflow-common-depends-on redis: diff --git a/docs/apache-airflow/start/docker.rst b/docs/apache-airflow/start/docker.rst index e05ec799209c2..bc3470d285983 100644 --- a/docs/apache-airflow/start/docker.rst +++ b/docs/apache-airflow/start/docker.rst @@ -81,6 +81,18 @@ If you need install a new Python library or system library, you can :doc:`build .. _initializing_docker_compose_environment: +Using custom images +=================== + +When you want to run Airflow locally, you might want to use an extended image, containing some additional dependencies - for +example you might add new python packages, or upgrade airflow providers to a later version. This can be done very easily +by placing a custom Dockerfile alongside your ``docker-compose.yaml``. Then you can use ``docker-compose build`` command +to build your image (you need to do it only once). You can also add the ``--build`` flag to your ``docker-compose`` commands +to rebuild the images on-the-fly when you run other ``docker-compose`` commands. + +Examples of how you can extend the image with custom providers, python packages, +apt packages and more can be found in :doc:`Building the image `. + Initializing Environment ======================== @@ -93,7 +105,7 @@ On **Linux**, the mounted volumes in container use the native Linux filesystem u mkdir ./dags ./logs ./plugins echo -e "AIRFLOW_UID=$(id -u)\nAIRFLOW_GID=0" > .env -See:ref:`Docker Compose environment variables ` +See :ref:`Docker Compose environment variables ` On **all operating systems**, you need to run database migrations and create the first user account. To do it, run. diff --git a/docs/apache-airflow/start/index.rst b/docs/apache-airflow/start/index.rst index c86ef8380e437..b8f0c0bcf016c 100644 --- a/docs/apache-airflow/start/index.rst +++ b/docs/apache-airflow/start/index.rst @@ -15,7 +15,7 @@ specific language governing permissions and limitations under the License. -Quick start +Quick Start =========== This section contains quick start guides to help you get up and running with Apache Airflow. diff --git a/docs/apache-airflow/timezone.rst b/docs/apache-airflow/timezone.rst index 449393cec6506..2a191f2dc70c2 100644 --- a/docs/apache-airflow/timezone.rst +++ b/docs/apache-airflow/timezone.rst @@ -17,22 +17,22 @@ -Time zones +Time Zones ========== Support for time zones is enabled by default. Airflow stores datetime information in UTC internally and in the database. -It allows you to run your DAGs with time zone dependent schedules. At the moment Airflow does not convert them to the -end user’s time zone in the user interface. There it will always be displayed in UTC. Also templates used in Operators -are not converted. Time zone information is exposed and it is up to the writer of DAG what do with it. +It allows you to run your DAGs with time zone dependent schedules. At the moment, Airflow does not convert them to the +end user’s time zone in the user interface. It will always be displayed in UTC there. Also, templates used in Operators +are not converted. Time zone information is exposed and it is up to the writer of DAG to decide what do with it. This is handy if your users live in more than one time zone and you want to display datetime information according to each user’s wall clock. -Even if you are running Airflow in only one time zone it is still good practice to store data in UTC in your database -(also before Airflow became time zone aware this was also to recommended or even required setup). The main reason is -Daylight Saving Time (DST). Many countries have a system of DST, where clocks are moved forward in spring and backward +Even if you are running Airflow in only one time zone, it is still good practice to store data in UTC in your database +(also before Airflow became time zone aware this was also the recommended or even required setup). The main reason is +that many countries use Daylight Saving Time (DST), where clocks are moved forward in spring and backward in autumn. If you’re working in local time, you’re likely to encounter errors twice a year, when the transitions -happen. (The pendulum and pytz documentation discusses these issues in greater detail.) This probably doesn’t matter +happen. (The pendulum and pytz documentation discuss these issues in greater detail.) This probably doesn’t matter for a simple DAG, but it’s a problem if you are in, for example, financial services where you have end of day deadlines to meet. @@ -68,7 +68,7 @@ a datetime object is aware. Otherwise, it’s naive. You can use ``timezone.is_localized()`` and ``timezone.is_naive()`` to determine whether datetimes are aware or naive. -Because Airflow uses time-zone-aware datetime objects. If your code creates datetime objects they need to be aware too. +Because Airflow uses time zone aware datetime objects. If your code creates datetime objects they need to be aware too. .. code-block:: python @@ -100,7 +100,7 @@ Unfortunately, during DST transitions, some datetimes don’t exist or are ambig In such situations, pendulum raises an exception. That’s why you should always create aware datetime objects when time zone support is enabled. -In practice, this is rarely an issue. Airflow gives you aware datetime objects in the models and DAGs, and most often, +In practice, this is rarely an issue. Airflow gives you time zone aware datetime objects in the models and DAGs, and most often, new datetime objects are created from existing ones through timedelta arithmetic. The only datetime that’s often created in application code is the current time, and ``timezone.utcnow()`` automatically does the right thing. diff --git a/docs/docker-stack/build-arg-ref.rst b/docs/docker-stack/build-arg-ref.rst index 8780970f613c5..f2507e05db423 100644 --- a/docs/docker-stack/build-arg-ref.rst +++ b/docs/docker-stack/build-arg-ref.rst @@ -79,12 +79,6 @@ for examples of using those arguments. +------------------------------------------+------------------------------------------+------------------------------------------+ | Build argument | Default value | Description | +==========================================+==========================================+==========================================+ -| ``CONTINUE_ON_PIP_CHECK_FAILURE`` | ``false`` | By default the image build fails if pip | -| | | check fails for it. This is good for | -| | | interactive building but on CI the | -| | | image should be built regardless - we | -| | | have a separate step to verify image. | -+------------------------------------------+------------------------------------------+------------------------------------------+ | ``UPGRADE_TO_NEWER_DEPENDENCIES`` | ``false`` | If set to true, the dependencies are | | | | upgraded to newer versions matching | | | | setup.py before installation. | diff --git a/docs/docker-stack/build.rst b/docs/docker-stack/build.rst index 7d89f2fa23adf..c469a35dfc653 100644 --- a/docs/docker-stack/build.rst +++ b/docs/docker-stack/build.rst @@ -250,6 +250,19 @@ You should be aware, about a few things: Examples of image extending --------------------------- +Example of upgrading Airflow Provider packages +.............................................. + +The :ref:`Airflow Providers ` are released independently of core +Airflow and sometimes you might want to upgrade specific providers only to fix some problems or +use features available in that provider version. Here is an example of how you can do it + +.. exampleinclude:: docker-examples/extending/add-providers/Dockerfile + :language: Dockerfile + :start-after: [START Dockerfile] + :end-before: [END Dockerfile] + + Example of adding ``apt`` package ................................. diff --git a/docs/docker-stack/docker-examples/extending/add-apt-packages/Dockerfile b/docs/docker-stack/docker-examples/extending/add-apt-packages/Dockerfile index 62de197973833..f11e87adc666d 100644 --- a/docs/docker-stack/docker-examples/extending/add-apt-packages/Dockerfile +++ b/docs/docker-stack/docker-examples/extending/add-apt-packages/Dockerfile @@ -15,7 +15,7 @@ # This is an example Dockerfile. It is not intended for PRODUCTION use # [START Dockerfile] -FROM apache/airflow +FROM apache/airflow:2.1.2 USER root RUN apt-get update \ && apt-get install -y --no-install-recommends \ diff --git a/docs/docker-stack/docker-examples/extending/add-build-essential-extend/Dockerfile b/docs/docker-stack/docker-examples/extending/add-build-essential-extend/Dockerfile index b34fdc9ab3cf8..47ac51ffbe07c 100644 --- a/docs/docker-stack/docker-examples/extending/add-build-essential-extend/Dockerfile +++ b/docs/docker-stack/docker-examples/extending/add-build-essential-extend/Dockerfile @@ -15,7 +15,7 @@ # This is an example Dockerfile. It is not intended for PRODUCTION use # [START Dockerfile] -FROM apache/airflow +FROM apache/airflow:2.1.2 USER root RUN apt-get update \ && apt-get install -y --no-install-recommends \ diff --git a/docs/docker-stack/docker-examples/extending/add-providers/Dockerfile b/docs/docker-stack/docker-examples/extending/add-providers/Dockerfile new file mode 100644 index 0000000000000..b65262fe41d9d --- /dev/null +++ b/docs/docker-stack/docker-examples/extending/add-providers/Dockerfile @@ -0,0 +1,20 @@ +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# This is an example Dockerfile. It is not intended for PRODUCTION use +# [START Dockerfile] +FROM apache/airflow:2.1.2 +RUN pip install --no-cache-dir apache-airflow-providers-docker==2.1.0 +# [END Dockerfile] diff --git a/docs/docker-stack/docker-examples/extending/add-pypi-packages/Dockerfile b/docs/docker-stack/docker-examples/extending/add-pypi-packages/Dockerfile index cc2559f79d062..310f84cf68dc8 100644 --- a/docs/docker-stack/docker-examples/extending/add-pypi-packages/Dockerfile +++ b/docs/docker-stack/docker-examples/extending/add-pypi-packages/Dockerfile @@ -15,6 +15,6 @@ # This is an example Dockerfile. It is not intended for PRODUCTION use # [START Dockerfile] -FROM apache/airflow +FROM apache/airflow:2.1.2 RUN pip install --no-cache-dir lxml # [END Dockerfile] diff --git a/docs/docker-stack/docker-examples/extending/embedding-dags/Dockerfile b/docs/docker-stack/docker-examples/extending/embedding-dags/Dockerfile index c849697859fb2..48701aae23d87 100644 --- a/docs/docker-stack/docker-examples/extending/embedding-dags/Dockerfile +++ b/docs/docker-stack/docker-examples/extending/embedding-dags/Dockerfile @@ -15,7 +15,7 @@ # This is an example Dockerfile. It is not intended for PRODUCTION use # [START Dockerfile] -FROM apache/airflow +FROM apache/airflow:2.1.2 COPY --chown=airflow:root test_dag.py /opt/airflow/dags diff --git a/docs/docker-stack/docker-examples/extending/writable-directory/Dockerfile b/docs/docker-stack/docker-examples/extending/writable-directory/Dockerfile index ba07f6816888f..8fbb98dfef00c 100644 --- a/docs/docker-stack/docker-examples/extending/writable-directory/Dockerfile +++ b/docs/docker-stack/docker-examples/extending/writable-directory/Dockerfile @@ -15,7 +15,7 @@ # This is an example Dockerfile. It is not intended for PRODUCTION use # [START Dockerfile] -FROM apache/airflow +FROM apache/airflow:2.1.2 RUN umask 0002; \ mkdir -p ~/writeable-directory # [END Dockerfile] diff --git a/docs/exts/airflow_intersphinx.py b/docs/exts/airflow_intersphinx.py index 1ae91bb9f35f7..2c0e0eb13b549 100644 --- a/docs/exts/airflow_intersphinx.py +++ b/docs/exts/airflow_intersphinx.py @@ -70,7 +70,7 @@ def _generate_provider_intersphinx_mapping(): airflow_mapping[pkg_name] = ( # base URI - f'/docs/{pkg_name}/latest/', + f'/docs/{pkg_name}/{"stable" if for_production else "latest"}/', (doc_inventory if os.path.exists(doc_inventory) else cache_inventory,), ) for pkg_name in ['apache-airflow-providers', 'docker-stack']: diff --git a/docs/helm-chart/index.rst b/docs/helm-chart/index.rst index aa2e2050b675f..25318cc9bdb9d 100644 --- a/docs/helm-chart/index.rst +++ b/docs/helm-chart/index.rst @@ -103,6 +103,9 @@ To upgrade the chart with the release name ``airflow``: helm upgrade airflow apache-airflow/airflow --namespace airflow +.. note:: + To upgrade to a new version of the chart, run ``helm repo update`` first. + Uninstalling the Chart ---------------------- diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 403750f7a1906..1acce0b1996be 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -1227,6 +1227,7 @@ stringified subchart subclasses subclassing +subcluster subcommand subcommands subdag @@ -1363,6 +1364,7 @@ videointelligence virtualenv vm volumeMounts +warmup wasb webProperty webhdfs diff --git a/scripts/ci/kubernetes/ci_run_kubernetes_tests.sh b/scripts/ci/kubernetes/ci_run_kubernetes_tests.sh index 6cab719cdffc1..a97f6929e1716 100755 --- a/scripts/ci/kubernetes/ci_run_kubernetes_tests.sh +++ b/scripts/ci/kubernetes/ci_run_kubernetes_tests.sh @@ -89,11 +89,19 @@ function create_virtualenv() { pip install --upgrade "pip==${AIRFLOW_PIP_VERSION}" "wheel==${WHEEL_VERSION}" - pip install pytest freezegun \ - --constraint "https://raw.githubusercontent.com/${CONSTRAINTS_GITHUB_REPOSITORY}/${DEFAULT_CONSTRAINTS_BRANCH}/constraints-${HOST_PYTHON_VERSION}.txt" + local constraints=( + --constraint + "https://raw.githubusercontent.com/${CONSTRAINTS_GITHUB_REPOSITORY}/${DEFAULT_CONSTRAINTS_BRANCH}/constraints-${HOST_PYTHON_VERSION}.txt" + ) + if [[ -n ${GITHUB_REGISTRY_PULL_IMAGE_TAG=} ]]; then + # Disable constraints when building in CI with specific version of sources + # In case there will be conflicting constraints + constraints=() + fi + + pip install pytest freezegun "${constraints[@]}" - pip install -e ".[cncf.kubernetes,postgres]" \ - --constraint "https://raw.githubusercontent.com/${CONSTRAINTS_GITHUB_REPOSITORY}/${DEFAULT_CONSTRAINTS_BRANCH}/constraints-${HOST_PYTHON_VERSION}.txt" + pip install -e ".[cncf.kubernetes,postgres]" "${constraints[@]}" } function run_tests() { diff --git a/scripts/ci/libraries/_build_images.sh b/scripts/ci/libraries/_build_images.sh index 94f5c8eaf8125..4ead442d9d9de 100644 --- a/scripts/ci/libraries/_build_images.sh +++ b/scripts/ci/libraries/_build_images.sh @@ -193,7 +193,7 @@ function build_images::confirm_image_rebuild() { echo "${COLOR_RED}ERROR: The ${THE_IMAGE_TYPE} needs to be rebuilt - it is outdated. ${COLOR_RESET}" echo """ - Make sure you build the images bu running + Make sure you build the images by running: ./breeze --python ${PYTHON_MAJOR_MINOR_VERSION} build-image @@ -669,7 +669,6 @@ Docker building ${AIRFLOW_CI_IMAGE}. --build-arg ADDITIONAL_RUNTIME_APT_DEPS="${ADDITIONAL_RUNTIME_APT_DEPS}" \ --build-arg ADDITIONAL_RUNTIME_APT_ENV="${ADDITIONAL_RUNTIME_APT_ENV}" \ --build-arg UPGRADE_TO_NEWER_DEPENDENCIES="${UPGRADE_TO_NEWER_DEPENDENCIES}" \ - --build-arg CONTINUE_ON_PIP_CHECK_FAILURE="${CONTINUE_ON_PIP_CHECK_FAILURE}" \ --build-arg CONSTRAINTS_GITHUB_REPOSITORY="${CONSTRAINTS_GITHUB_REPOSITORY}" \ --build-arg AIRFLOW_CONSTRAINTS_REFERENCE="${DEFAULT_CONSTRAINTS_BRANCH}" \ --build-arg AIRFLOW_CONSTRAINTS="${AIRFLOW_CONSTRAINTS}" \ @@ -810,7 +809,6 @@ function build_images::build_prod_images() { --build-arg INSTALL_FROM_PYPI="${INSTALL_FROM_PYPI}" \ --build-arg INSTALL_FROM_DOCKER_CONTEXT_FILES="${INSTALL_FROM_DOCKER_CONTEXT_FILES}" \ --build-arg UPGRADE_TO_NEWER_DEPENDENCIES="${UPGRADE_TO_NEWER_DEPENDENCIES}" \ - --build-arg CONTINUE_ON_PIP_CHECK_FAILURE="${CONTINUE_ON_PIP_CHECK_FAILURE}" \ --build-arg BUILD_ID="${CI_BUILD_ID}" \ --build-arg COMMIT_SHA="${COMMIT_SHA}" \ --build-arg CONSTRAINTS_GITHUB_REPOSITORY="${CONSTRAINTS_GITHUB_REPOSITORY}" \ @@ -845,7 +843,6 @@ function build_images::build_prod_images() { --build-arg INSTALL_FROM_PYPI="${INSTALL_FROM_PYPI}" \ --build-arg INSTALL_FROM_DOCKER_CONTEXT_FILES="${INSTALL_FROM_DOCKER_CONTEXT_FILES}" \ --build-arg UPGRADE_TO_NEWER_DEPENDENCIES="${UPGRADE_TO_NEWER_DEPENDENCIES}" \ - --build-arg CONTINUE_ON_PIP_CHECK_FAILURE="${CONTINUE_ON_PIP_CHECK_FAILURE}" \ --build-arg AIRFLOW_VERSION="${AIRFLOW_VERSION}" \ --build-arg AIRFLOW_BRANCH="${AIRFLOW_BRANCH_FOR_PYPI_PRELOADING}" \ --build-arg AIRFLOW_EXTRAS="${AIRFLOW_EXTRAS}" \ @@ -1014,23 +1011,6 @@ ${COLOR_BLUE} ./breeze build-image --production-image --upgrade-to-newer-dependencies --python 3.6 ${COLOR_RESET} -* If you want to build the image regardless if 'pip check' fails for it, you can add - --continue-on-pip-check-failure flag and enter the image and inspect dependencies. - -CI image: - -${COLOR_BLUE} - ./breeze build-image --upgrade-to-newer-dependencies --python 3.6 --continue-on-pip-check-failure - docker run -it apache/airflow:main-3.6-ci bash -${COLOR_RESET} - -Production image: - -${COLOR_BLUE} - ./breeze build-image --production-image --upgrade-to-newer-dependencies --python 3.6 --continue-on-pip-check-failure - docker run -it apache/airflow:main-3.6 bash -${COLOR_RESET} - * You will see error messages there telling which requirements are conflicting and which packages caused the conflict. Add the limitation that caused the conflict to EAGER_UPGRADE_ADDITIONAL_REQUIREMENTS variable in Dockerfile.ci. Note that the limitations might be different for Dockerfile.ci and Dockerfile diff --git a/scripts/ci/libraries/_initialization.sh b/scripts/ci/libraries/_initialization.sh index fefcb1a5b4040..223d9ea999343 100644 --- a/scripts/ci/libraries/_initialization.sh +++ b/scripts/ci/libraries/_initialization.sh @@ -221,7 +221,7 @@ function initialization::initialize_files_for_rebuild_check() { "scripts/docker/common.sh" "scripts/docker/install_additional_dependencies.sh" "scripts/docker/install_airflow.sh" - "scripts/docker/install_airflow_from_branch_tip.sh" + "scripts/docker/install_airflow_dependencies_from_branch_tip.sh" "scripts/docker/install_from_docker_context_files.sh" "scripts/docker/install_mysql.sh" "airflow/www/package.json" @@ -422,7 +422,7 @@ function initialization::initialize_image_build_variables() { export INSTALLED_PROVIDERS export INSTALLED_EXTRAS="async,amazon,celery,cncf.kubernetes,docker,dask,elasticsearch,ftp,grpc,hashicorp,http,imap,ldap,google,microsoft.azure,mysql,postgres,redis,sendgrid,sftp,slack,ssh,statsd,virtualenv" - AIRFLOW_PIP_VERSION=${AIRFLOW_PIP_VERSION:="21.1"} + AIRFLOW_PIP_VERSION=${AIRFLOW_PIP_VERSION:="21.2.2"} export AIRFLOW_PIP_VERSION # We also pin version of wheel used to get consistent builds @@ -446,9 +446,6 @@ function initialization::initialize_image_build_variables() { # Installs different airflow version than current from the sources export INSTALL_AIRFLOW_VERSION=${INSTALL_AIRFLOW_VERSION:=""} - # Continue on PIP CHECK failure - export CONTINUE_ON_PIP_CHECK_FAILURE=${CONTINUE_ON_PIP_CHECK_FAILURE:="false"} - # Determines if airflow should be installed from a specified reference in GitHub export INSTALL_AIRFLOW_REFERENCE=${INSTALL_AIRFLOW_REFERENCE:=""} @@ -503,7 +500,7 @@ function initialization::initialize_kubernetes_variables() { CURRENT_KIND_VERSIONS+=("v0.11.1") export CURRENT_KIND_VERSIONS # Currently supported versions of Helm - CURRENT_HELM_VERSIONS+=("v3.2.4") + CURRENT_HELM_VERSIONS+=("v3.6.3") export CURRENT_HELM_VERSIONS # Current executor in chart CURRENT_EXECUTOR+=("KubernetesExecutor") @@ -684,7 +681,6 @@ Common image build variables: INSTALL_FROM_PYPI: '${INSTALL_FROM_PYPI}' AIRFLOW_PRE_CACHED_PIP_PACKAGES: '${AIRFLOW_PRE_CACHED_PIP_PACKAGES}' UPGRADE_TO_NEWER_DEPENDENCIES: '${UPGRADE_TO_NEWER_DEPENDENCIES}' - CONTINUE_ON_PIP_CHECK_FAILURE: '${CONTINUE_ON_PIP_CHECK_FAILURE}' CHECK_IMAGE_FOR_REBUILD: '${CHECK_IMAGE_FOR_REBUILD}' AIRFLOW_CONSTRAINTS_LOCATION: '${AIRFLOW_CONSTRAINTS_LOCATION}' AIRFLOW_CONSTRAINTS_REFERENCE: '${AIRFLOW_CONSTRAINTS_REFERENCE}' diff --git a/scripts/ci/libraries/_kind.sh b/scripts/ci/libraries/_kind.sh index 970a6d36c23cd..d4910d92c8e07 100644 --- a/scripts/ci/libraries/_kind.sh +++ b/scripts/ci/libraries/_kind.sh @@ -258,15 +258,20 @@ function kind::check_cluster_ready_for_airflow() { function kind::build_image_for_kubernetes_tests() { cd "${AIRFLOW_SOURCES}" || exit 1 + local image_tag="latest" + if [[ -n ${GITHUB_REGISTRY_PULL_IMAGE_TAG=} ]]; then + image_tag="${GITHUB_REGISTRY_PULL_IMAGE_TAG}" + fi + echo "Building ${AIRFLOW_PROD_IMAGE_KUBERNETES}:latest from ${AIRFLOW_PROD_IMAGE}:${image_tag}" docker_v build --tag "${AIRFLOW_PROD_IMAGE_KUBERNETES}:latest" . -f - < /dev/null || true) if [[ -z "${python_image_hash=}" || "${FORCE_PULL_IMAGES}" == "true" || \ ${CHECK_IF_BASE_PYTHON_IMAGE_UPDATED} == "true" ]]; then - push_pull_remove_images::pull_base_python_image + if [[ ${GITHUB_REGISTRY_PULL_IMAGE_TAG} == "latest" ]]; then + # Pull base python image when building latest image + push_pull_remove_images::pull_base_python_image + fi fi if [[ "${DOCKER_CACHE}" == "pulled" ]]; then push_pull_remove_images::pull_image_if_not_present_or_forced \ @@ -160,7 +163,10 @@ function push_pull_remove_images::pull_prod_images_if_needed() { python_image_hash=$(docker images -q "${AIRFLOW_PYTHON_BASE_IMAGE}" 2> /dev/null || true) if [[ -z "${python_image_hash=}" || "${FORCE_PULL_IMAGES}" == "true" || \ ${CHECK_IF_BASE_PYTHON_IMAGE_UPDATED} == "true" ]]; then - push_pull_remove_images::pull_base_python_image + if [[ ${GITHUB_REGISTRY_PULL_IMAGE_TAG} == "latest" ]]; then + # Pull base python image when building latest image + push_pull_remove_images::pull_base_python_image + fi fi if [[ "${DOCKER_CACHE}" == "pulled" ]]; then # "Build" segment of production image diff --git a/scripts/ci/pre_commit/pre_commit_check_provider_yaml_files.py b/scripts/ci/pre_commit/pre_commit_check_provider_yaml_files.py index 24d963bedef2b..c6c0584021c35 100755 --- a/scripts/ci/pre_commit/pre_commit_check_provider_yaml_files.py +++ b/scripts/ci/pre_commit/pre_commit_check_provider_yaml_files.py @@ -119,13 +119,13 @@ def assert_sets_equal(set1, set2): lines = [] if difference1: - lines.append('Items in the first set but not the second:') + lines.append(' -- Items in the left set but not the right:') for item in sorted(difference1): - lines.append(repr(item)) + lines.append(f' {item!r}') if difference2: - lines.append('Items in the second set but not the first:') + lines.append(' -- Items in the right set but not the left:') for item in sorted(difference2): - lines.append(repr(item)) + lines.append(f' {item!r}') standard_msg = '\n'.join(lines) raise AssertionError(standard_msg) @@ -155,6 +155,7 @@ def parse_module_data(provider_data, resource_type, yaml_file_path): def check_completeness_of_list_of_hooks_sensors_hooks(yaml_files: Dict[str, Dict]): print("Checking completeness of list of {sensors, hooks, operators}") + print(" -- {sensors, hooks, operators} - Expected modules(Left): Current Modules(Right)") for (yaml_file_path, provider_data), resource_type in product( yaml_files.items(), ["sensors", "operators", "hooks"] ): @@ -193,6 +194,8 @@ def check_duplicates_in_integrations_names_of_hooks_sensors_operators(yaml_files def check_completeness_of_list_of_transfers(yaml_files: Dict[str, Dict]): print("Checking completeness of list of transfers") resource_type = 'transfers' + + print(" -- Expected transfers modules(Left): Current transfers Modules(Right)") for yaml_file_path, provider_data in yaml_files.items(): expected_modules, provider_package, resource_data = parse_module_data( provider_data, resource_type, yaml_file_path @@ -309,7 +312,10 @@ def check_doc_files(yaml_files: Dict[str, Dict]): } try: + print(" -- Checking document urls: expected(left), current(right)") assert_sets_equal(set(expected_doc_urls), set(current_doc_urls)) + + print(" -- Checking logo urls: expected(left), current(right)") assert_sets_equal(set(expected_logo_urls), set(current_logo_urls)) except AssertionError as ex: print(ex) diff --git a/scripts/ci/tools/fix_ownership.sh b/scripts/ci/tools/fix_ownership.sh index 6ed1161be747f..de1562122a779 100755 --- a/scripts/ci/tools/fix_ownership.sh +++ b/scripts/ci/tools/fix_ownership.sh @@ -33,8 +33,12 @@ sanity_checks::sanitize_mounted_files read -r -a EXTRA_DOCKER_FLAGS <<<"$(local_mounts::convert_local_mounts_to_docker_params)" -docker_v run --entrypoint /bin/bash "${EXTRA_DOCKER_FLAGS[@]}" \ - --rm \ - --env-file "${AIRFLOW_SOURCES}/scripts/ci/docker-compose/_docker.env" \ - "${AIRFLOW_CI_IMAGE}" \ - -c /opt/airflow/scripts/in_container/run_fix_ownership.sh || true +if docker image inspect "${AIRFLOW_CI_IMAGE}" >/dev/null 2>&1; then + docker_v run --entrypoint /bin/bash "${EXTRA_DOCKER_FLAGS[@]}" \ + --rm \ + --env-file "${AIRFLOW_SOURCES}/scripts/ci/docker-compose/_docker.env" \ + "${AIRFLOW_CI_IMAGE}" \ + -c /opt/airflow/scripts/in_container/run_fix_ownership.sh || true +else + echo "Skip fixing ownership as seems that you do not have the ${AIRFLOW_CI_IMAGE} image yet" +fi diff --git a/scripts/docker/compile_www_assets.sh b/scripts/docker/compile_www_assets.sh index 59a7017fd157f..50e1318c548a4 100755 --- a/scripts/docker/compile_www_assets.sh +++ b/scripts/docker/compile_www_assets.sh @@ -35,7 +35,7 @@ function compile_www_assets() { www_dir="$(python -m site --user-site)/airflow/www" fi pushd ${www_dir} || exit 1 - yarn install --frozen-lockfile --no-cache + yarn install --frozen-lockfile --no-cache --network-concurrency=1 yarn run prod find package.json yarn.lock static/css static/js -type f | sort | xargs md5sum > "${md5sum_file}" rm -rf "${www_dir}/node_modules" diff --git a/scripts/docker/install_additional_dependencies.sh b/scripts/docker/install_additional_dependencies.sh index 6c035ae5def09..4f9c05f6b7680 100755 --- a/scripts/docker/install_additional_dependencies.sh +++ b/scripts/docker/install_additional_dependencies.sh @@ -23,7 +23,6 @@ test -v ADDITIONAL_PYTHON_DEPS test -v EAGER_UPGRADE_ADDITIONAL_REQUIREMENTS test -v AIRFLOW_INSTALL_USER_FLAG test -v AIRFLOW_PIP_VERSION -test -v CONTINUE_ON_PIP_CHECK_FAILURE # shellcheck source=scripts/docker/common.sh . "$( dirname "${BASH_SOURCE[0]}" )/common.sh" @@ -41,7 +40,7 @@ function install_additional_dependencies() { ${ADDITIONAL_PYTHON_DEPS} ${EAGER_UPGRADE_ADDITIONAL_REQUIREMENTS} # make sure correct PIP version is used pip install ${AIRFLOW_INSTALL_USER_FLAG} --upgrade "pip==${AIRFLOW_PIP_VERSION}" - pip check || ${CONTINUE_ON_PIP_CHECK_FAILURE} + pip check else echo echo Installing additional dependencies upgrading only if needed @@ -51,7 +50,7 @@ function install_additional_dependencies() { ${ADDITIONAL_PYTHON_DEPS} # make sure correct PIP version is used pip install ${AIRFLOW_INSTALL_USER_FLAG} --upgrade "pip==${AIRFLOW_PIP_VERSION}" - pip check || ${CONTINUE_ON_PIP_CHECK_FAILURE} + pip check fi } diff --git a/scripts/docker/install_airflow.sh b/scripts/docker/install_airflow.sh index 49040275a04ac..e2bca4fc839a0 100755 --- a/scripts/docker/install_airflow.sh +++ b/scripts/docker/install_airflow.sh @@ -60,7 +60,7 @@ function install_airflow() { # make sure correct PIP version is used pip install ${AIRFLOW_INSTALL_USER_FLAG} --upgrade "pip==${AIRFLOW_PIP_VERSION}" - pip check || ${CONTINUE_ON_PIP_CHECK_FAILURE} + pip check else \ echo echo Installing all packages with constraints and upgrade if needed @@ -76,7 +76,7 @@ function install_airflow() { "${AIRFLOW_INSTALLATION_METHOD}[${AIRFLOW_EXTRAS}]${AIRFLOW_VERSION_SPECIFICATION}" \ # make sure correct PIP version is used pip install ${AIRFLOW_INSTALL_USER_FLAG} --upgrade "pip==${AIRFLOW_PIP_VERSION}" - pip check || ${CONTINUE_ON_PIP_CHECK_FAILURE} + pip check fi } diff --git a/scripts/docker/install_airflow_from_branch_tip.sh b/scripts/docker/install_airflow_dependencies_from_branch_tip.sh similarity index 85% rename from scripts/docker/install_airflow_from_branch_tip.sh rename to scripts/docker/install_airflow_dependencies_from_branch_tip.sh index 925a872fa50ea..61aaa13ef467f 100755 --- a/scripts/docker/install_airflow_from_branch_tip.sh +++ b/scripts/docker/install_airflow_dependencies_from_branch_tip.sh @@ -30,28 +30,29 @@ . "$( dirname "${BASH_SOURCE[0]}" )/common.sh" -function install_airflow_from_branch_tip() { +function install_airflow_dependencies_from_branch_tip() { echo echo "Installing airflow from ${AIRFLOW_BRANCH}. It is used to cache dependencies" echo if [[ ${INSTALL_MYSQL_CLIENT} != "true" ]]; then AIRFLOW_EXTRAS=${AIRFLOW_EXTRAS/mysql,} fi - # Install latest set of dependencies using constraints + # Install latest set of dependencies using constraints. In case constraints were upgraded and there + # are conflicts, this might fail, but it should be fixed in the following installation steps pip install ${AIRFLOW_INSTALL_USER_FLAG} \ "https://github.com/${AIRFLOW_REPO}/archive/${AIRFLOW_BRANCH}.tar.gz#egg=apache-airflow[${AIRFLOW_EXTRAS}]" \ - --constraint "${AIRFLOW_CONSTRAINTS_LOCATION}" + --constraint "${AIRFLOW_CONSTRAINTS_LOCATION}" || true # make sure correct PIP version is used pip install ${AIRFLOW_INSTALL_USER_FLAG} --upgrade "pip==${AIRFLOW_PIP_VERSION}" pip freeze | grep apache-airflow-providers | xargs pip uninstall --yes || true echo echo Uninstalling just airflow. Dependencies remain. echo - pip uninstall --yes apache-airflow + pip uninstall --yes apache-airflow || true } common::get_airflow_version_specification common::override_pip_version_if_needed common::get_constraints_location -install_airflow_from_branch_tip +install_airflow_dependencies_from_branch_tip diff --git a/scripts/docker/install_from_docker_context_files.sh b/scripts/docker/install_from_docker_context_files.sh index 813d1b02df7e0..d8ed6bc72bd9a 100755 --- a/scripts/docker/install_from_docker_context_files.sh +++ b/scripts/docker/install_from_docker_context_files.sh @@ -96,7 +96,7 @@ function install_airflow_and_providers_from_docker_context_files(){ # make sure correct PIP version is left installed pip install ${AIRFLOW_INSTALL_USER_FLAG} --upgrade "pip==${AIRFLOW_PIP_VERSION}" - pip check || ${CONTINUE_ON_PIP_CHECK_FAILURE} + pip check } diff --git a/scripts/in_container/prod/airflow_scheduler_autorestart.sh b/scripts/in_container/prod/airflow_scheduler_autorestart.sh index 09e03447fef8a..0fdc6445d0681 100755 --- a/scripts/in_container/prod/airflow_scheduler_autorestart.sh +++ b/scripts/in_container/prod/airflow_scheduler_autorestart.sh @@ -18,7 +18,11 @@ while echo "Running"; do airflow scheduler -n 5 - echo "Scheduler crashed with exit code $?. Respawning.." >&2 - date >> /tmp/airflow_scheduler_errors.txt + return_code=$? + if (( return_code != 0 )); then + echo "Scheduler crashed with exit code $return_code. Respawning.." >&2 + date >> /tmp/airflow_scheduler_errors.txt + fi + sleep 1 done diff --git a/setup.cfg b/setup.cfg index fbe58cb24f29f..d3c5f574c0b7d 100644 --- a/setup.cfg +++ b/setup.cfg @@ -100,7 +100,7 @@ install_requires = # https://github.com/readthedocs/sphinx_rtd_theme/issues/1115 docutils<0.17 flask>=1.1.0, <2.0 - flask-appbuilder~=3.3 + flask-appbuilder>=3.3.2, <4.0.0 flask-caching>=1.5.0, <2.0.0 flask-login>=0.3, <0.5 flask-wtf>=0.14.3, <0.15 diff --git a/setup.py b/setup.py index 83ed5e7a49422..3b2389c2c147c 100644 --- a/setup.py +++ b/setup.py @@ -182,6 +182,7 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version amazon = [ 'boto3>=1.15.0,<1.18.0', 'watchtower~=1.0.6', + 'jsonpath_ng>=1.5.3', ] apache_beam = [ 'apache-beam>=2.20.0', @@ -270,10 +271,8 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version facebook = [ 'facebook-business>=6.0.2', ] -flask_oauth = [ - 'Flask-OAuthlib>=0.9.1,<0.9.6', # Flask OAuthLib 0.9.6 requires Flask-Login 0.5.0 - breaks FAB - 'oauthlib!=2.0.3,!=2.0.4,!=2.0.5,<3.0.0,>=1.1.2', - 'requests-oauthlib<1.2.0', +flask_appbuilder_authlib = [ + 'authlib', ] google = [ 'PyOpenSSL', @@ -427,7 +426,7 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version 'tableauserverclient', ] samba = [ - 'pysmbclient>=0.1.3', + 'smbprotocol>=1.5.0', ] segment = [ 'analytics-python>=1.2.9', @@ -478,7 +477,7 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version 'pywinrm~=0.4', ] yandex = [ - 'yandexcloud>=0.22.0', + 'yandexcloud>=0.97.0', ] zendesk = [ 'zdesk', @@ -622,8 +621,8 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version 'cncf.kubernetes': kubernetes, # also has provider, but it extends the core with the KubernetesExecutor 'dask': dask, 'deprecated_api': deprecated_api, - 'github_enterprise': flask_oauth, - 'google_auth': flask_oauth, + 'github_enterprise': flask_appbuilder_authlib, + 'google_auth': flask_appbuilder_authlib, 'kerberos': kerberos, 'ldap': ldap, 'leveldb': leveldb, diff --git a/tests/cli/commands/test_celery_command.py b/tests/cli/commands/test_celery_command.py index e2c16685979c1..b3d0365c087f7 100644 --- a/tests/cli/commands/test_celery_command.py +++ b/tests/cli/commands/test_celery_command.py @@ -142,6 +142,44 @@ def test_same_pid_file_is_used_in_start_and_stop( celery_command.stop_worker(stop_args) mock_read_pid_from_pidfile.assert_called_once_with(pid_file) + @mock.patch("airflow.cli.commands.celery_command.remove_existing_pidfile") + @mock.patch("airflow.cli.commands.celery_command.read_pid_from_pidfile") + @mock.patch("airflow.cli.commands.celery_command.worker_bin.worker") + @mock.patch("airflow.cli.commands.celery_command.psutil.Process") + @mock.patch("airflow.cli.commands.celery_command.setup_locations") + @conf_vars({("core", "executor"): "CeleryExecutor"}) + def test_custom_pid_file_is_used_in_start_and_stop( + self, + mock_setup_locations, + mock_process, + mock_celery_worker, + mock_read_pid_from_pidfile, + mock_remove_existing_pidfile, + ): + pid_file = "custom_test_pid_file" + mock_setup_locations.return_value = (pid_file, None, None, None) + # Call worker + worker_args = self.parser.parse_args(['celery', 'worker', '--skip-serve-logs', '--pid', pid_file]) + celery_command.worker(worker_args) + run_mock = mock_celery_worker.return_value.run + assert run_mock.call_args + args, kwargs = run_mock.call_args + assert 'pidfile' in kwargs + assert kwargs['pidfile'] == pid_file + assert not args + stop_args = self.parser.parse_args(['celery', 'stop', '--pid', pid_file]) + celery_command.stop_worker(stop_args) + run_mock = mock_celery_worker.return_value.run + assert run_mock.call_args + args, kwargs = run_mock.call_args + assert 'pidfile' in kwargs + assert kwargs['pidfile'] == pid_file + assert not args + + mock_read_pid_from_pidfile.assert_called_once_with(pid_file) + mock_process.return_value.terminate.assert_called() + mock_remove_existing_pidfile.assert_called_once_with(pid_file) + @pytest.mark.backend("mysql", "postgres") class TestWorkerStart(unittest.TestCase): diff --git a/tests/cli/commands/test_kubernetes_command.py b/tests/cli/commands/test_kubernetes_command.py index f2a8605a6ce91..490c7fafc5b66 100644 --- a/tests/cli/commands/test_kubernetes_command.py +++ b/tests/cli/commands/test_kubernetes_command.py @@ -55,12 +55,8 @@ def test_generate_dag_yaml(self): class TestCleanUpPodsCommand(unittest.TestCase): - label_selector = kubernetes.client.V1LabelSelector( - match_expressions=[ - kubernetes.client.V1LabelSelectorRequirement(key=label, operator="Exists") - for label in ['dag_id', 'task_id', 'execution_date', 'try_number', 'airflow_version'] - ] - ) + + label_selector = ','.join(['dag_id', 'task_id', 'execution_date', 'try_number', 'airflow_version']) @classmethod def setUpClass(cls): diff --git a/tests/cluster_policies/__init__.py b/tests/cluster_policies/__init__.py index d395ec0982c04..521a8522908b8 100644 --- a/tests/cluster_policies/__init__.py +++ b/tests/cluster_policies/__init__.py @@ -52,7 +52,7 @@ def _check_task_rules(current_task: BaseOperator): if notices: notices_list = " * " + "\n * ".join(notices) raise AirflowClusterPolicyViolation( - f"DAG policy violation (DAG ID: {current_task.dag_id}, Path: {current_task.dag.filepath}):\n" + f"DAG policy violation (DAG ID: {current_task.dag_id}, Path: {current_task.dag.fileloc}):\n" f"Notices:\n" f"{notices_list}" ) @@ -70,7 +70,7 @@ def dag_policy(dag: DAG): """Ensure that DAG has at least one tag""" if not dag.tags: raise AirflowClusterPolicyViolation( - f"DAG {dag.dag_id} has no tags. At least one tag required. File path: {dag.filepath}" + f"DAG {dag.dag_id} has no tags. At least one tag required. File path: {dag.fileloc}" ) diff --git a/tests/conftest.py b/tests/conftest.py index 896e32ac1ab06..4bcab64ba695e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -428,8 +428,9 @@ def app(): @pytest.fixture def dag_maker(request): - from airflow.models import DAG + from airflow.models import DAG, DagModel from airflow.utils import timezone + from airflow.utils.session import provide_session from airflow.utils.state import State DEFAULT_DATE = timezone.datetime(2016, 1, 1) @@ -444,33 +445,41 @@ def __exit__(self, type, value, traceback): dag.__exit__(type, value, traceback) if type is None: dag.clear() - self.dag_run = dag.create_dagrun( - run_id=self.kwargs.get("run_id", "test"), - state=self.kwargs.get('state', State.RUNNING), - execution_date=self.kwargs.get('execution_date', self.kwargs['start_date']), - start_date=self.kwargs['start_date'], + self.dag_model = DagModel( + dag_id=dag.dag_id, + next_dagrun=dag.start_date, + is_active=True, + is_paused=False, + max_active_tasks=dag.max_active_tasks, + has_task_concurrency_limits=False, ) + self.session.add(self.dag_model) + self.session.flush() - def __call__(self, dag_id='test_dag', **kwargs): + def create_dagrun(self, **kwargs): + dag = self.dag + defaults = dict( + run_id='test', + state=State.RUNNING, + execution_date=self.start_date, + start_date=self.start_date, + ) + kwargs = {**defaults, **kwargs} + self.dag_run = dag.create_dagrun(**kwargs) + return self.dag_run + + @provide_session + def __call__(self, dag_id='test_dag', session=None, **kwargs): self.kwargs = kwargs - if "start_date" not in kwargs: + self.session = session + self.start_date = self.kwargs.get('start_date', None) + if not self.start_date: if hasattr(request.module, 'DEFAULT_DATE'): - kwargs['start_date'] = getattr(request.module, 'DEFAULT_DATE') + self.start_date = getattr(request.module, 'DEFAULT_DATE') else: - kwargs['start_date'] = DEFAULT_DATE - dagrun_fields_not_in_dag = [ - 'state', - 'execution_date', - 'run_type', - 'queued_at', - "run_id", - "creating_job_id", - "external_trigger", - "last_scheduling_decision", - "dag_hash", - ] - kwargs = {k: v for k, v in kwargs.items() if k not in dagrun_fields_not_in_dag} - self.dag = DAG(dag_id, **kwargs) + self.start_date = DEFAULT_DATE + self.kwargs['start_date'] = self.start_date + self.dag = DAG(dag_id, **self.kwargs) return self return DagFactory() diff --git a/tests/dag_processing/test_manager.py b/tests/dag_processing/test_manager.py index d34eca23585b3..52d7487cd9f57 100644 --- a/tests/dag_processing/test_manager.py +++ b/tests/dag_processing/test_manager.py @@ -407,9 +407,9 @@ def test_find_zombies(self): seconds=manager._zombie_threshold_secs + 1 ) manager._find_zombies() - requests = manager._callback_to_execute[dag.full_filepath] + requests = manager._callback_to_execute[dag.fileloc] assert 1 == len(requests) - assert requests[0].full_filepath == dag.full_filepath + assert requests[0].full_filepath == dag.fileloc assert requests[0].msg == "Detected as zombie" assert requests[0].is_failure_callback is True assert isinstance(requests[0].simple_task_instance, SimpleTaskInstance) @@ -451,7 +451,7 @@ def test_handle_failure_callback_with_zombies_are_correctly_passed_to_dag_file_p expected_failure_callback_requests = [ TaskCallbackRequest( - full_filepath=dag.full_filepath, + full_filepath=dag.fileloc, simple_task_instance=SimpleTaskInstance(ti), msg="Message", ) diff --git a/tests/dag_processing/test_processor.py b/tests/dag_processing/test_processor.py index 90c2baea85327..3de9952d155bf 100644 --- a/tests/dag_processing/test_processor.py +++ b/tests/dag_processing/test_processor.py @@ -115,6 +115,10 @@ def setUpClass(cls): non_serialized_dagbag.sync_to_db() cls.dagbag = DagBag(read_dags_from_db=True) + @staticmethod + def assert_scheduled_ti_count(session, count): + assert count == session.query(TaskInstance).filter_by(state=State.SCHEDULED).count() + def test_dag_file_processor_sla_miss_callback(self): """ Test that the dag file processor calls the sla miss callback @@ -387,8 +391,8 @@ def test_dag_file_processor_process_task_instances(self, state, start_date, end_ ti.start_date = start_date ti.end_date = end_date - count = self.scheduler_job._schedule_dag_run(dr, session) - assert count == 1 + self.scheduler_job._schedule_dag_run(dr, session) + self.assert_scheduled_ti_count(session, 1) session.refresh(ti) assert ti.state == State.SCHEDULED @@ -444,8 +448,8 @@ def test_dag_file_processor_process_task_instances_with_task_concurrency( ti.start_date = start_date ti.end_date = end_date - count = self.scheduler_job._schedule_dag_run(dr, session) - assert count == 1 + self.scheduler_job._schedule_dag_run(dr, session) + self.assert_scheduled_ti_count(session, 1) session.refresh(ti) assert ti.state == State.SCHEDULED @@ -504,8 +508,8 @@ def test_dag_file_processor_process_task_instances_depends_on_past(self, state, ti.start_date = start_date ti.end_date = end_date - count = self.scheduler_job._schedule_dag_run(dr, session) - assert count == 2 + self.scheduler_job._schedule_dag_run(dr, session) + self.assert_scheduled_ti_count(session, 2) session.refresh(tis[0]) session.refresh(tis[1]) @@ -547,9 +551,9 @@ def test_scheduler_job_add_new_task(self): BashOperator(task_id='dummy2', dag=dag, owner='airflow', bash_command='echo test') SerializedDagModel.write_dag(dag=dag) - scheduled_tis = self.scheduler_job._schedule_dag_run(dr, session) + self.scheduler_job._schedule_dag_run(dr, session) + self.assert_scheduled_ti_count(session, 2) session.flush() - assert scheduled_tis == 2 drs = DagRun.find(dag_id=dag.dag_id, session=session) assert len(drs) == 1 @@ -681,7 +685,7 @@ def test_process_file_should_failure_callback(self): requests = [ TaskCallbackRequest( - full_filepath=dag.full_filepath, + full_filepath=dag.fileloc, simple_task_instance=SimpleTaskInstance(ti), msg="Message", ) diff --git a/tests/jobs/test_backfill_job.py b/tests/jobs/test_backfill_job.py index c110e632a5388..d70606ac1928c 100644 --- a/tests/jobs/test_backfill_job.py +++ b/tests/jobs/test_backfill_job.py @@ -46,7 +46,7 @@ from airflow.utils.state import State from airflow.utils.timeout import timeout from airflow.utils.types import DagRunType -from tests.test_utils.db import clear_db_pools, clear_db_runs, set_default_pool_slots +from tests.test_utils.db import clear_db_dags, clear_db_pools, clear_db_runs, set_default_pool_slots from tests.test_utils.mock_executor import MockExecutor logger = logging.getLogger(__name__) @@ -59,44 +59,10 @@ def dag_bag(): return DagBag(include_examples=True) -@pytest.fixture -def get_dummy_dag_and_run(dag_maker): - def _get_dummy_dag_and_run( - dag_id='test_dag', pool=Pool.DEFAULT_POOL_NAME, task_concurrency=None, task_id='op', **kwargs - ): - with dag_maker(dag_id=dag_id, schedule_interval='@daily', **kwargs) as dag: - DummyOperator(task_id=task_id, pool=pool, task_concurrency=task_concurrency) - - return dag, dag_maker.dag_run - - return _get_dummy_dag_and_run - - -@pytest.fixture -def get_dag_test_max_active_limits(dag_maker): - def _get_dag_test_max_active_limits(dag_id='test_dag', max_active_runs=1, **kwargs): - with dag_maker( - dag_id=dag_id, - start_date=DEFAULT_DATE, - schedule_interval="@hourly", - max_active_runs=max_active_runs, - **kwargs, - ) as dag: - op1 = DummyOperator(task_id='leave1') - op2 = DummyOperator(task_id='leave2') - op3 = DummyOperator(task_id='upstream_level_1') - op4 = DummyOperator(task_id='upstream_level_2') - - op1 >> op2 >> op3 - op4 >> op3 - return dag, dag_maker.dag_run - - return _get_dag_test_max_active_limits - - class TestBackfillJob: @staticmethod def clean_db(): + clear_db_dags() clear_db_runs() clear_db_pools() @@ -106,6 +72,20 @@ def set_instance_attrs(self, dag_bag): self.parser = cli_parser.get_parser() self.dagbag = dag_bag + def _get_dummy_dag( + self, + dag_maker_fixture, + dag_id='test_dag', + pool=Pool.DEFAULT_POOL_NAME, + task_concurrency=None, + task_id='op', + **kwargs, + ): + with dag_maker_fixture(dag_id=dag_id, schedule_interval='@daily', **kwargs) as dag: + DummyOperator(task_id=task_id, pool=pool, task_concurrency=task_concurrency) + + return dag + def _times_called_with(self, method, class_): count = 0 for args in method.call_args_list: @@ -113,8 +93,9 @@ def _times_called_with(self, method, class_): count += 1 return count - def test_unfinished_dag_runs_set_to_failed(self, get_dummy_dag_and_run): - dag, dag_run = get_dummy_dag_and_run(dag_id='dummy_dag') + def test_unfinished_dag_runs_set_to_failed(self, dag_maker): + dag = self._get_dummy_dag(dag_maker) + dag_run = dag_maker.create_dagrun() job = BackfillJob( dag=dag, @@ -129,8 +110,9 @@ def test_unfinished_dag_runs_set_to_failed(self, get_dummy_dag_and_run): assert State.FAILED == dag_run.state - def test_dag_run_with_finished_tasks_set_to_success(self, get_dummy_dag_and_run): - dag, dag_run = get_dummy_dag_and_run(dag_id='dummy_dag') + def test_dag_run_with_finished_tasks_set_to_success(self, dag_maker): + dag = self._get_dummy_dag(dag_maker) + dag_run = dag_maker.create_dagrun() for ti in dag_run.get_task_instances(): ti.set_state(State.SUCCESS) @@ -289,8 +271,9 @@ def test_backfill_examples(self, dag_id, expected_execution_order): for task_id in expected_execution_order ] == executor.sorted_tasks - def test_backfill_conf(self, get_dummy_dag_and_run): - dag, _ = get_dummy_dag_and_run(dag_id='test_backfill_conf') + def test_backfill_conf(self, dag_maker): + dag = self._get_dummy_dag(dag_maker, dag_id='test_backfill_conf') + dag_maker.create_dagrun() executor = MockExecutor() @@ -312,12 +295,14 @@ def test_backfill_conf(self, get_dummy_dag_and_run): assert conf_ == dr[0].conf @patch('airflow.jobs.backfill_job.BackfillJob.log') - def test_backfill_respect_task_concurrency_limit(self, mock_log, get_dummy_dag_and_run): + def test_backfill_respect_task_concurrency_limit(self, mock_log, dag_maker): task_concurrency = 2 - dag, _ = get_dummy_dag_and_run( + dag = self._get_dummy_dag( + dag_maker, dag_id='test_backfill_respect_task_concurrency_limit', task_concurrency=task_concurrency, ) + dag_maker.create_dagrun() executor = MockExecutor() @@ -364,9 +349,9 @@ def test_backfill_respect_task_concurrency_limit(self, mock_log, get_dummy_dag_a assert times_task_concurrency_limit_reached_in_debug > 0 @patch('airflow.jobs.backfill_job.BackfillJob.log') - def test_backfill_respect_dag_concurrency_limit(self, mock_log, get_dummy_dag_and_run): - - dag, _ = get_dummy_dag_and_run(dag_id='test_backfill_respect_concurrency_limit') + def test_backfill_respect_dag_concurrency_limit(self, mock_log, dag_maker): + dag = self._get_dummy_dag(dag_maker, dag_id='test_backfill_respect_concurrency_limit') + dag_maker.create_dagrun() dag.max_active_tasks = 2 executor = MockExecutor() @@ -415,11 +400,12 @@ def test_backfill_respect_dag_concurrency_limit(self, mock_log, get_dummy_dag_an assert times_dag_concurrency_limit_reached_in_debug > 0 @patch('airflow.jobs.backfill_job.BackfillJob.log') - def test_backfill_respect_default_pool_limit(self, mock_log, get_dummy_dag_and_run): + def test_backfill_respect_default_pool_limit(self, mock_log, dag_maker): default_pool_slots = 2 set_default_pool_slots(default_pool_slots) - dag, _ = get_dummy_dag_and_run(dag_id='test_backfill_with_no_pool_limit') + dag = self._get_dummy_dag(dag_maker, dag_id='test_backfill_with_no_pool_limit') + dag_maker.create_dagrun() executor = MockExecutor() @@ -469,11 +455,13 @@ def test_backfill_respect_default_pool_limit(self, mock_log, get_dummy_dag_and_r assert 0 == times_task_concurrency_limit_reached_in_debug assert times_pool_limit_reached_in_debug > 0 - def test_backfill_pool_not_found(self, get_dummy_dag_and_run): - dag, _ = get_dummy_dag_and_run( + def test_backfill_pool_not_found(self, dag_maker): + dag = self._get_dummy_dag( + dag_maker, dag_id='test_backfill_pool_not_found', pool='king_pool', ) + dag_maker.create_dagrun() executor = MockExecutor() @@ -490,7 +478,7 @@ def test_backfill_pool_not_found(self, get_dummy_dag_and_run): return @patch('airflow.jobs.backfill_job.BackfillJob.log') - def test_backfill_respect_pool_limit(self, mock_log, get_dummy_dag_and_run): + def test_backfill_respect_pool_limit(self, mock_log, dag_maker): session = settings.Session() slots = 2 @@ -501,10 +489,12 @@ def test_backfill_respect_pool_limit(self, mock_log, get_dummy_dag_and_run): session.add(pool) session.commit() - dag, _ = get_dummy_dag_and_run( + dag = self._get_dummy_dag( + dag_maker, dag_id='test_backfill_respect_pool_limit', pool=pool.pool, ) + dag_maker.create_dagrun() executor = MockExecutor() @@ -550,10 +540,11 @@ def test_backfill_respect_pool_limit(self, mock_log, get_dummy_dag_and_run): assert 0 == times_dag_concurrency_limit_reached_in_debug assert times_pool_limit_reached_in_debug > 0 - def test_backfill_run_rescheduled(self, get_dummy_dag_and_run): - dag, _ = get_dummy_dag_and_run( - dag_id="test_backfill_run_rescheduled", task_id="test_backfill_run_rescheduled_task-1" + def test_backfill_run_rescheduled(self, dag_maker): + dag = self._get_dummy_dag( + dag_maker, dag_id="test_backfill_run_rescheduled", task_id="test_backfill_run_rescheduled_task-1" ) + dag_maker.create_dagrun() executor = MockExecutor() @@ -581,10 +572,11 @@ def test_backfill_run_rescheduled(self, get_dummy_dag_and_run): ti.refresh_from_db() assert ti.state == State.SUCCESS - def test_backfill_rerun_failed_tasks(self, get_dummy_dag_and_run): - dag, _ = get_dummy_dag_and_run( - dag_id="test_backfill_rerun_failed", task_id="test_backfill_rerun_failed_task-1" + def test_backfill_rerun_failed_tasks(self, dag_maker): + dag = self._get_dummy_dag( + dag_maker, dag_id="test_backfill_rerun_failed", task_id="test_backfill_rerun_failed_task-1" ) + dag_maker.create_dagrun() executor = MockExecutor() @@ -614,12 +606,11 @@ def test_backfill_rerun_failed_tasks(self, get_dummy_dag_and_run): def test_backfill_rerun_upstream_failed_tasks(self, dag_maker): - with dag_maker( - dag_id='test_backfill_rerun_upstream_failed', start_date=DEFAULT_DATE, schedule_interval='@daily' - ) as dag: + with dag_maker(dag_id='test_backfill_rerun_upstream_failed', schedule_interval='@daily') as dag: op1 = DummyOperator(task_id='test_backfill_rerun_upstream_failed_task-1') op2 = DummyOperator(task_id='test_backfill_rerun_upstream_failed_task-2') op1.set_upstream(op2) + dag_maker.create_dagrun() executor = MockExecutor() @@ -647,10 +638,11 @@ def test_backfill_rerun_upstream_failed_tasks(self, dag_maker): ti.refresh_from_db() assert ti.state == State.SUCCESS - def test_backfill_rerun_failed_tasks_without_flag(self, get_dummy_dag_and_run): - dag, _ = get_dummy_dag_and_run( - dag_id='test_backfill_rerun_failed', task_id='test_backfill_rerun_failed_task-1' + def test_backfill_rerun_failed_tasks_without_flag(self, dag_maker): + dag = self._get_dummy_dag( + dag_maker, dag_id='test_backfill_rerun_failed', task_id='test_backfill_rerun_failed_task-1' ) + dag_maker.create_dagrun() executor = MockExecutor() @@ -680,7 +672,6 @@ def test_backfill_rerun_failed_tasks_without_flag(self, get_dummy_dag_and_run): def test_backfill_retry_intermittent_failed_task(self, dag_maker): with dag_maker( dag_id='test_intermittent_failure_job', - start_date=DEFAULT_DATE, schedule_interval="@daily", default_args={ 'retries': 2, @@ -688,6 +679,7 @@ def test_backfill_retry_intermittent_failed_task(self, dag_maker): }, ) as dag: task1 = DummyOperator(task_id="task1") + dag_maker.create_dagrun() executor = MockExecutor(parallelism=16) executor.mock_task_results[ @@ -707,7 +699,6 @@ def test_backfill_retry_intermittent_failed_task(self, dag_maker): def test_backfill_retry_always_failed_task(self, dag_maker): with dag_maker( dag_id='test_always_failure_job', - start_date=DEFAULT_DATE, schedule_interval="@daily", default_args={ 'retries': 1, @@ -715,6 +706,7 @@ def test_backfill_retry_always_failed_task(self, dag_maker): }, ) as dag: task1 = DummyOperator(task_id="task1") + dag_maker.create_dagrun() executor = MockExecutor(parallelism=16) executor.mock_task_results[ @@ -734,7 +726,6 @@ def test_backfill_ordered_concurrent_execute(self, dag_maker): with dag_maker( dag_id='test_backfill_ordered_concurrent_execute', - start_date=DEFAULT_DATE, schedule_interval="@daily", ) as dag: op1 = DummyOperator(task_id='leave1') @@ -747,6 +738,7 @@ def test_backfill_ordered_concurrent_execute(self, dag_maker): op1.set_downstream(op3) op4.set_downstream(op5) op3.set_downstream(op4) + dag_maker.create_dagrun() executor = MockExecutor(parallelism=16) job = BackfillJob( @@ -881,10 +873,29 @@ def test_cli_receives_delay_arg(self): parsed_args = self.parser.parse_args(args) assert 0.5 == parsed_args.delay_on_limit - def test_backfill_max_limit_check_within_limit(self, get_dag_test_max_active_limits): - dag, _ = get_dag_test_max_active_limits( - dag_id='test_backfill_max_limit_check_within_limit', max_active_runs=16 + def _get_dag_test_max_active_limits( + self, dag_maker_fixture, dag_id='test_dag', max_active_runs=1, **kwargs + ): + with dag_maker_fixture( + dag_id=dag_id, + schedule_interval="@hourly", + max_active_runs=max_active_runs, + **kwargs, + ) as dag: + op1 = DummyOperator(task_id='leave1') + op2 = DummyOperator(task_id='leave2') + op3 = DummyOperator(task_id='upstream_level_1') + op4 = DummyOperator(task_id='upstream_level_2') + + op1 >> op2 >> op3 + op4 >> op3 + return dag + + def test_backfill_max_limit_check_within_limit(self, dag_maker): + dag = self._get_dag_test_max_active_limits( + dag_maker, dag_id='test_backfill_max_limit_check_within_limit', max_active_runs=16 ) + dag_maker.create_dagrun() start_date = DEFAULT_DATE - datetime.timedelta(hours=1) end_date = DEFAULT_DATE @@ -898,7 +909,7 @@ def test_backfill_max_limit_check_within_limit(self, get_dag_test_max_active_lim assert 2 == len(dagruns) assert all(run.state == State.SUCCESS for run in dagruns) - def test_backfill_max_limit_check(self, get_dag_test_max_active_limits): + def test_backfill_max_limit_check(self, dag_maker): dag_id = 'test_backfill_max_limit_check' run_id = 'test_dag_run' start_date = DEFAULT_DATE - datetime.timedelta(hours=1) @@ -911,9 +922,12 @@ def run_backfill(cond): # this session object is different than the one in the main thread with create_session() as thread_session: try: - dag, _ = get_dag_test_max_active_limits( - # Existing dagrun that is not within the backfill range + dag = self._get_dag_test_max_active_limits( + dag_maker, dag_id=dag_id, + ) + dag_maker.create_dagrun( + # Existing dagrun that is not within the backfill range run_id=run_id, execution_date=DEFAULT_DATE + datetime.timedelta(hours=1), ) @@ -960,11 +974,14 @@ def run_backfill(cond): finally: dag_run_created_cond.release() - def test_backfill_max_limit_check_no_count_existing(self, get_dag_test_max_active_limits): + def test_backfill_max_limit_check_no_count_existing(self, dag_maker): start_date = DEFAULT_DATE end_date = DEFAULT_DATE # Existing dagrun that is within the backfill range - dag, _ = get_dag_test_max_active_limits(dag_id='test_backfill_max_limit_check_no_count_existing') + dag = self._get_dag_test_max_active_limits( + dag_maker, dag_id='test_backfill_max_limit_check_no_count_existing' + ) + dag_maker.create_dagrun() executor = MockExecutor() job = BackfillJob( @@ -980,8 +997,11 @@ def test_backfill_max_limit_check_no_count_existing(self, get_dag_test_max_activ assert 1 == len(dagruns) assert State.SUCCESS == dagruns[0].state - def test_backfill_max_limit_check_complete_loop(self, get_dag_test_max_active_limits): - dag, _ = get_dag_test_max_active_limits(dag_id='test_backfill_max_limit_check_complete_loop') + def test_backfill_max_limit_check_complete_loop(self, dag_maker): + dag = self._get_dag_test_max_active_limits( + dag_maker, dag_id='test_backfill_max_limit_check_complete_loop' + ) + dag_maker.create_dagrun() start_date = DEFAULT_DATE - datetime.timedelta(hours=1) end_date = DEFAULT_DATE @@ -1003,9 +1023,6 @@ def test_sub_set_subdag(self, dag_maker): with dag_maker( 'test_sub_set_subdag', - start_date=DEFAULT_DATE, - default_args={'owner': 'owner1'}, - execution_date=DEFAULT_DATE, ) as dag: op1 = DummyOperator(task_id='leave1') op2 = DummyOperator(task_id='leave2') @@ -1018,7 +1035,7 @@ def test_sub_set_subdag(self, dag_maker): op4.set_downstream(op5) op3.set_downstream(op4) - dr = dag_maker.dag_run + dr = dag_maker.create_dagrun() executor = MockExecutor() sub_dag = dag.partial_subset( @@ -1043,9 +1060,6 @@ def test_sub_set_subdag(self, dag_maker): def test_backfill_fill_blanks(self, dag_maker): with dag_maker( 'test_backfill_fill_blanks', - start_date=DEFAULT_DATE, - default_args={'owner': 'owner1'}, - execution_date=DEFAULT_DATE, ) as dag: op1 = DummyOperator(task_id='op1') op2 = DummyOperator(task_id='op2') @@ -1054,7 +1068,7 @@ def test_backfill_fill_blanks(self, dag_maker): op5 = DummyOperator(task_id='op5') op6 = DummyOperator(task_id='op6') - dr = dag_maker.dag_run + dr = dag_maker.create_dagrun() executor = MockExecutor() @@ -1231,11 +1245,9 @@ def test_backfill_execute_subdag_with_removed_task(self): dag.clear() def test_update_counters(self, dag_maker): - with dag_maker( - dag_id='test_manage_executor_state', start_date=DEFAULT_DATE, execution_date=DEFAULT_DATE - ) as dag: - task1 = DummyOperator(task_id='dummy', dag=dag, owner='airflow') - dr = dag_maker.dag_run + with dag_maker(dag_id='test_manage_executor_state', start_date=DEFAULT_DATE) as dag: + task1 = DummyOperator(task_id='dummy', owner='airflow') + dr = dag_maker.create_dagrun() job = BackfillJob(dag=dag) session = settings.Session() @@ -1380,9 +1392,7 @@ def test_reset_orphaned_tasks_with_orphans(self, dag_maker): states_to_reset = [State.QUEUED, State.SCHEDULED, State.NONE] tasks = [] - with dag_maker( - dag_id=prefix, start_date=DEFAULT_DATE, schedule_interval="@daily", run_id='test1' - ) as dag: + with dag_maker(dag_id=prefix, start_date=DEFAULT_DATE, schedule_interval="@daily") as dag: for i in range(len(states)): task_id = f"{prefix}_task_{i}" task = DummyOperator(task_id=task_id) @@ -1392,7 +1402,7 @@ def test_reset_orphaned_tasks_with_orphans(self, dag_maker): job = BackfillJob(dag=dag) # create dagruns - dr1 = dag_maker.dag_run + dr1 = dag_maker.create_dagrun() dr2 = dag.create_dagrun(run_id='test2', state=State.SUCCESS) # create taskinstances and set states @@ -1445,15 +1455,13 @@ def test_reset_orphaned_tasks_specified_dagrun(self, dag_maker): dag_id=dag_id, start_date=DEFAULT_DATE, schedule_interval='@daily', - state=State.SUCCESS, - run_id='test1', ) as dag: DummyOperator(task_id=task_id, dag=dag) job = BackfillJob(dag=dag) session = settings.Session() # make two dagruns, only reset for one - dr1 = dag_maker.dag_run # Already created in dag_maker with state=SUCCESS + dr1 = dag_maker.create_dagrun(state=State.SUCCESS) dr2 = dag.create_dagrun(run_id='test2', state=State.RUNNING) ti1 = dr1.get_task_instances(session=session)[0] ti2 = dr2.get_task_instances(session=session)[0] diff --git a/tests/jobs/test_local_task_job.py b/tests/jobs/test_local_task_job.py index b7910cb5749b5..4cf991ae9c07e 100644 --- a/tests/jobs/test_local_task_job.py +++ b/tests/jobs/test_local_task_job.py @@ -21,19 +21,18 @@ import signal import time import uuid +from datetime import timedelta from multiprocessing import Lock, Value from unittest import mock from unittest.mock import patch import pytest -from parameterized import parameterized from airflow import settings from airflow.exceptions import AirflowException, AirflowFailException from airflow.executors.sequential_executor import SequentialExecutor from airflow.jobs.local_task_job import LocalTaskJob from airflow.jobs.scheduler_job import SchedulerJob -from airflow.models.dag import DAG, DagModel from airflow.models.dagbag import DagBag from airflow.models.taskinstance import TaskInstance from airflow.operators.dummy import DummyOperator @@ -72,10 +71,19 @@ def clear_db_class(): db.clear_db_task_fail() +@pytest.fixture(scope='module') +def dagbag(): + return DagBag( + dag_folder=TEST_DAG_FOLDER, + include_examples=False, + ) + + @pytest.mark.usefixtures('clear_db_class', 'clear_db') class TestLocalTaskJob: @pytest.fixture(autouse=True) - def set_instance_attrs(self): + def set_instance_attrs(self, dagbag): + self.dagbag = dagbag with patch('airflow.jobs.base_job.sleep') as self.mock_base_job_sleep: yield @@ -91,12 +99,10 @@ def test_localtaskjob_essential_attr(self, dag_maker): of LocalTaskJob can be assigned with proper values without intervention """ - with dag_maker( - 'test_localtaskjob_essential_attr', start_date=DEFAULT_DATE, default_args={'owner': 'owner1'} - ): + with dag_maker('test_localtaskjob_essential_attr'): op1 = DummyOperator(task_id='op1') - dr = dag_maker.dag_run + dr = dag_maker.create_dagrun() ti = dr.get_task_instance(task_id=op1.task_id) @@ -115,7 +121,7 @@ def test_localtaskjob_heartbeat(self, dag_maker): with dag_maker('test_localtaskjob_heartbeat'): op1 = DummyOperator(task_id='op1') - dr = dag_maker.dag_run + dr = dag_maker.create_dagrun() ti = dr.get_task_instance(task_id=op1.task_id, session=session) ti.state = State.RUNNING ti.hostname = "blablabla" @@ -147,7 +153,7 @@ def test_localtaskjob_heartbeat_with_run_as_user(self, psutil_mock, dag_maker): session = settings.Session() with dag_maker('test_localtaskjob_heartbeat'): op1 = DummyOperator(task_id='op1', run_as_user='myuser') - dr = dag_maker.dag_run + dr = dag_maker.create_dagrun() ti = dr.get_task_instance(task_id=op1.task_id, session=session) ti.state = State.RUNNING ti.pid = 2 @@ -183,6 +189,48 @@ def test_localtaskjob_heartbeat_with_run_as_user(self, psutil_mock, dag_maker): with pytest.raises(AirflowException, match='PID of job runner does not match'): job1.heartbeat_callback() + @conf_vars({('core', 'default_impersonation'): 'testuser'}) + @mock.patch('airflow.jobs.local_task_job.psutil') + def test_localtaskjob_heartbeat_with_default_impersonation(self, psutil_mock, dag_maker): + session = settings.Session() + with dag_maker('test_localtaskjob_heartbeat'): + op1 = DummyOperator(task_id='op1') + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance(task_id=op1.task_id, session=session) + ti.state = State.RUNNING + ti.pid = 2 + ti.hostname = get_hostname() + session.commit() + + job1 = LocalTaskJob(task_instance=ti, ignore_ti_state=True, executor=SequentialExecutor()) + ti.task = op1 + ti.refresh_from_task(op1) + job1.task_runner = StandardTaskRunner(job1) + job1.task_runner.process = mock.Mock() + job1.task_runner.process.pid = 2 + # Here, ti.pid is 2, the parent process of ti.pid is a mock(different). + # And task_runner process is 2. Should fail + with pytest.raises(AirflowException, match='PID of job runner does not match'): + job1.heartbeat_callback() + + job1.task_runner.process.pid = 1 + # We make the parent process of ti.pid to equal the task_runner process id + psutil_mock.Process.return_value.ppid.return_value = 1 + ti.state = State.RUNNING + ti.pid = 2 + # The task_runner process id is 1, same as the parent process of ti.pid + # as seen above + assert job1.task_runner.run_as_user == 'testuser' + session.merge(ti) + session.commit() + job1.heartbeat_callback(session=None) + + # Here the task_runner process id is changed to 2 + # while parent process of ti.pid is kept at 1, which is different + job1.task_runner.process.pid = 2 + with pytest.raises(AirflowException, match='PID of job runner does not match'): + job1.heartbeat_callback() + def test_heartbeat_failed_fast(self): """ Test that task heartbeat will sleep when it fails fast @@ -191,13 +239,10 @@ def test_heartbeat_failed_fast(self): dag_id = 'test_heartbeat_failed_fast' task_id = 'test_heartbeat_failed_fast_op' with create_session() as session: - dagbag = DagBag( - dag_folder=TEST_DAG_FOLDER, - include_examples=False, - ) + dag_id = 'test_heartbeat_failed_fast' task_id = 'test_heartbeat_failed_fast_op' - dag = dagbag.get_dag(dag_id) + dag = self.dagbag.get_dag(dag_id) task = dag.get_task(task_id) dag.create_dagrun( @@ -228,17 +273,12 @@ def test_heartbeat_failed_fast(self): delta = (time2 - time1).total_seconds() assert abs(delta - job.heartrate) < 0.5 - @pytest.mark.quarantined def test_mark_success_no_kill(self): """ Test that ensures that mark_success in the UI doesn't cause the task to fail, and that the task exits """ - dagbag = DagBag( - dag_folder=TEST_DAG_FOLDER, - include_examples=False, - ) - dag = dagbag.dags.get('test_mark_success') + dag = self.dagbag.dags.get('test_mark_success') task = dag.get_task('task1') session = settings.Session() @@ -254,9 +294,9 @@ def test_mark_success_no_kill(self): ti = TaskInstance(task=task, execution_date=DEFAULT_DATE) ti.refresh_from_db() job1 = LocalTaskJob(task_instance=ti, ignore_ti_state=True) + settings.engine.dispose() process = multiprocessing.Process(target=job1.run) process.start() - ti.refresh_from_db() for _ in range(0, 50): if ti.state == State.RUNNING: break @@ -266,19 +306,13 @@ def test_mark_success_no_kill(self): ti.state = State.SUCCESS session.merge(ti) session.commit() - process.join(timeout=10) - assert not process.is_alive() ti.refresh_from_db() assert State.SUCCESS == ti.state def test_localtaskjob_double_trigger(self): - dagbag = DagBag( - dag_folder=TEST_DAG_FOLDER, - include_examples=False, - ) - dag = dagbag.dags.get('test_localtaskjob_double_trigger') + dag = self.dagbag.dags.get('test_localtaskjob_double_trigger') task = dag.get_task('test_localtaskjob_double_trigger_task') session = settings.Session() @@ -314,11 +348,8 @@ def test_localtaskjob_double_trigger(self): @pytest.mark.quarantined def test_localtaskjob_maintain_heart_rate(self): - dagbag = DagBag( - dag_folder=TEST_DAG_FOLDER, - include_examples=False, - ) - dag = dagbag.dags.get('test_localtaskjob_double_trigger') + + dag = self.dagbag.dags.get('test_localtaskjob_double_trigger') task = dag.get_task('test_localtaskjob_double_trigger_task') session = settings.Session() @@ -397,6 +428,7 @@ def task_function(ti): python_callable=task_function, on_failure_callback=check_failure, ) + dag_maker.create_dagrun() ti = TaskInstance(task=task, execution_date=DEFAULT_DATE) ti.refresh_from_db() @@ -438,6 +470,7 @@ def task_function(ti): python_callable=task_function, on_failure_callback=failure_callback, ) + dag_maker.create_dagrun() ti = TaskInstance(task=task, execution_date=DEFAULT_DATE) ti.refresh_from_db() @@ -466,7 +499,6 @@ def dummy_return_code(*args, **kwargs): assert ti.state == State.FAILED # task exits with failure state assert failure_callback_called.value == 1 - @pytest.mark.quarantined def test_mark_success_on_success_callback(self, dag_maker): """ Test that ensures that where a task is marked success in the UI @@ -481,11 +513,12 @@ def test_mark_success_on_success_callback(self, dag_maker): def success_callback(context): with shared_mem_lock: success_callback_called.value += 1 + assert context['dag_run'].dag_id == 'test_mark_success' def task_function(ti): - time.sleep(60) + # This should not happen -- the state change should be noticed and the task should get killed with shared_mem_lock: task_terminated_externally.value = 0 @@ -517,22 +550,14 @@ def task_function(ti): ti.state = State.SUCCESS session.merge(ti) session.commit() - + ti.refresh_from_db() process.join(timeout=10) assert success_callback_called.value == 1 assert task_terminated_externally.value == 1 - assert not process.is_alive() - @parameterized.expand( - [ - (signal.SIGTERM,), - (signal.SIGKILL,), - ] - ) - @pytest.mark.quarantined - def test_process_kill_calls_on_failure_callback(self, signal_type, dag_maker): + def test_task_sigkill_calls_on_failure_callback(self, dag_maker): """ - Test that ensures that when a task is killed with sigterm or sigkill + Test that ensures that when a task is killed with sigkill on_failure_callback gets executed """ # use shared memory value so we can properly track value change even if @@ -544,10 +569,49 @@ def test_process_kill_calls_on_failure_callback(self, signal_type, dag_maker): def failure_callback(context): with shared_mem_lock: failure_callback_called.value += 1 - assert context['dag_run'].dag_id == 'test_mark_failure' + assert context['dag_run'].dag_id == 'test_send_sigkill' def task_function(ti): + os.kill(os.getpid(), signal.SIGKILL) + # This should not happen -- the state change should be noticed and the task should get killed + with shared_mem_lock: + task_terminated_externally.value = 0 + + with dag_maker(dag_id='test_send_sigkill'): + task = PythonOperator( + task_id='test_on_failure', + python_callable=task_function, + on_failure_callback=failure_callback, + ) + ti = TaskInstance(task=task, execution_date=DEFAULT_DATE) + ti.refresh_from_db() + job1 = LocalTaskJob(task_instance=ti, ignore_ti_state=True, executor=SequentialExecutor()) + settings.engine.dispose() + process = multiprocessing.Process(target=job1.run) + process.start() + time.sleep(0.3) + process.join(timeout=10) + assert failure_callback_called.value == 1 + assert task_terminated_externally.value == 1 + + def test_process_sigterm_calls_on_failure_callback(self, dag_maker): + """ + Test that ensures that when a task runner is killed with sigterm + on_failure_callback gets executed + """ + # use shared memory value so we can properly track value change even if + # it's been updated across processes. + failure_callback_called = Value('i', 0) + task_terminated_externally = Value('i', 1) + shared_mem_lock = Lock() + + def failure_callback(context): + with shared_mem_lock: + failure_callback_called.value += 1 + assert context['dag_run'].dag_id == 'test_mark_failure' + + def task_function(ti): time.sleep(60) # This should not happen -- the state change should be noticed and the task should get killed with shared_mem_lock: @@ -562,26 +626,22 @@ def task_function(ti): ti = TaskInstance(task=task, execution_date=DEFAULT_DATE) ti.refresh_from_db() job1 = LocalTaskJob(task_instance=ti, ignore_ti_state=True, executor=SequentialExecutor()) - job1.task_runner = StandardTaskRunner(job1) - settings.engine.dispose() process = multiprocessing.Process(target=job1.run) process.start() - - for _ in range(0, 20): + for _ in range(0, 25): ti.refresh_from_db() - if ti.state == State.RUNNING and ti.pid is not None: + if ti.state == State.RUNNING: break time.sleep(0.2) - assert ti.pid is not None - assert ti.state == State.RUNNING - os.kill(ti.pid, signal_type) + os.kill(process.pid, signal.SIGTERM) + ti.refresh_from_db() process.join(timeout=10) assert failure_callback_called.value == 1 assert task_terminated_externally.value == 1 - assert not process.is_alive() - @parameterized.expand( + @pytest.mark.parametrize( + "conf, dependencies, init_state, first_run_state, second_run_state, error_message", [ ( {('scheduler', 'schedule_after_task_execution'): 'True'}, @@ -615,27 +675,17 @@ def task_function(ti): None, "A -> C & B -> C, when A is QUEUED but B has FAILED, C is marked UPSTREAM_FAILED.", ), - ] + ], ) def test_fast_follow( - self, conf, dependencies, init_state, first_run_state, second_run_state, error_message + self, conf, dependencies, init_state, first_run_state, second_run_state, error_message, dag_maker ): with conf_vars(conf): session = settings.Session() - dag = DAG('test_dagrun_fast_follow', start_date=DEFAULT_DATE) - - dag_model = DagModel( - dag_id=dag.dag_id, - next_dagrun=dag.start_date, - is_active=True, - ) - session.add(dag_model) - session.flush() - python_callable = lambda: True - with dag: + with dag_maker('test_dagrun_fast_follow') as dag: task_a = PythonOperator(task_id='A', python_callable=python_callable) task_b = PythonOperator(task_id='B', python_callable=python_callable) task_c = PythonOperator(task_id='C', python_callable=python_callable) @@ -683,34 +733,119 @@ def test_fast_follow( if scheduler_job.processor_agent: scheduler_job.processor_agent.end() + def test_task_sigkill_works_with_retries(self, dag_maker): + """ + Test that ensures that tasks are retried when they receive sigkill + """ + # use shared memory value so we can properly track value change even if + # it's been updated across processes. + retry_callback_called = Value('i', 0) + task_terminated_externally = Value('i', 1) + shared_mem_lock = Lock() + + def retry_callback(context): + with shared_mem_lock: + retry_callback_called.value += 1 + assert context['dag_run'].dag_id == 'test_mark_failure_2' + + def task_function(ti): + os.kill(os.getpid(), signal.SIGKILL) + # This should not happen -- the state change should be noticed and the task should get killed + with shared_mem_lock: + task_terminated_externally.value = 0 + + with dag_maker( + dag_id='test_mark_failure_2', start_date=DEFAULT_DATE, default_args={'owner': 'owner1'} + ): + task = PythonOperator( + task_id='test_on_failure', + python_callable=task_function, + retries=1, + retry_delay=timedelta(seconds=2), + on_retry_callback=retry_callback, + ) + ti = TaskInstance(task=task, execution_date=DEFAULT_DATE) + ti.refresh_from_db() + job1 = LocalTaskJob(task_instance=ti, ignore_ti_state=True, executor=SequentialExecutor()) + job1.task_runner = StandardTaskRunner(job1) + job1.task_runner.start() + settings.engine.dispose() + process = multiprocessing.Process(target=job1.run) + process.start() + time.sleep(0.4) + process.join(timeout=10) + ti.refresh_from_db() + assert ti.state == State.UP_FOR_RETRY + assert retry_callback_called.value == 1 + assert task_terminated_externally.value == 1 + + def test_process_sigterm_works_with_retries(self, dag_maker): + """ + Test that ensures that task runner sets tasks to retry when they(task runner) + receive sigterm + """ + # use shared memory value so we can properly track value change even if + # it's been updated across processes. + retry_callback_called = Value('i', 0) + task_terminated_externally = Value('i', 1) + shared_mem_lock = Lock() + + def retry_callback(context): + with shared_mem_lock: + retry_callback_called.value += 1 + assert context['dag_run'].dag_id == 'test_mark_failure_2' + + def task_function(ti): + time.sleep(60) + # This should not happen -- the state change should be noticed and the task should get killed + with shared_mem_lock: + task_terminated_externally.value = 0 + + with dag_maker(dag_id='test_mark_failure_2'): + task = PythonOperator( + task_id='test_on_failure', + python_callable=task_function, + retries=1, + retry_delay=timedelta(seconds=2), + on_retry_callback=retry_callback, + ) + ti = TaskInstance(task=task, execution_date=DEFAULT_DATE) + ti.refresh_from_db() + job1 = LocalTaskJob(task_instance=ti, ignore_ti_state=True, executor=SequentialExecutor()) + job1.task_runner = StandardTaskRunner(job1) + job1.task_runner.start() + settings.engine.dispose() + process = multiprocessing.Process(target=job1.run) + process.start() + for _ in range(0, 25): + ti.refresh_from_db() + if ti.state == State.RUNNING and ti.pid is not None: + break + time.sleep(0.2) + os.kill(process.pid, signal.SIGTERM) + process.join(timeout=10) + ti.refresh_from_db() + assert ti.state == State.UP_FOR_RETRY + assert retry_callback_called.value == 1 + assert task_terminated_externally.value == 1 + def test_task_exit_should_update_state_of_finished_dagruns_with_dag_paused(self, dag_maker): """Test that with DAG paused, DagRun state will update when the tasks finishes the run""" - dag = DAG(dag_id='test_dags', start_date=DEFAULT_DATE) - op1 = PythonOperator(task_id='dummy', dag=dag, owner='airflow', python_callable=lambda: True) + with dag_maker(dag_id='test_dags') as dag: + op1 = PythonOperator(task_id='dummy', python_callable=lambda: True) session = settings.Session() - orm_dag = DagModel( - dag_id=dag.dag_id, - has_task_concurrency_limits=False, - next_dagrun=dag.start_date, - next_dagrun_create_after=dag.following_schedule(DEFAULT_DATE), - is_active=True, - is_paused=True, - ) - session.add(orm_dag) + dagmodel = dag_maker.dag_model + dagmodel.next_dagrun_create_after = dag.following_schedule(DEFAULT_DATE) + dagmodel.is_paused = True + session.merge(dagmodel) session.flush() # Write Dag to DB dagbag = DagBag(dag_folder="/dev/null", include_examples=False, read_dags_from_db=False) dagbag.bag_dag(dag, root_dag=dag) dagbag.sync_to_db() - dr = dag.create_dagrun( - run_type=DagRunType.SCHEDULED, - state=State.RUNNING, - execution_date=DEFAULT_DATE, - start_date=DEFAULT_DATE, - session=session, - ) + dr = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED) assert dr.state == State.RUNNING ti = TaskInstance(op1, dr.execution_date) @@ -733,18 +868,17 @@ def clean_db_helper(): class TestLocalTaskJobPerformance: @pytest.mark.parametrize("return_codes", [[0], 9 * [None] + [0]]) # type: ignore @mock.patch("airflow.jobs.local_task_job.get_task_runner") - def test_number_of_queries_single_loop(self, mock_get_task_runner, return_codes): + def test_number_of_queries_single_loop(self, mock_get_task_runner, return_codes, dag_maker): unique_prefix = str(uuid.uuid4()) - dag = DAG(dag_id=f'{unique_prefix}_test_number_of_queries', start_date=DEFAULT_DATE) - task = DummyOperator(task_id='test_state_succeeded1', dag=dag) + with dag_maker(dag_id=f'{unique_prefix}_test_number_of_queries'): + task = DummyOperator(task_id='test_state_succeeded1') - dag.clear() - dag.create_dagrun(run_id=unique_prefix, execution_date=DEFAULT_DATE, state=State.NONE) + dag_maker.create_dagrun(run_id=unique_prefix, state=State.NONE) ti = TaskInstance(task=task, execution_date=DEFAULT_DATE) mock_get_task_runner.return_value.return_code.side_effects = return_codes job = LocalTaskJob(task_instance=ti, executor=MockExecutor()) - with assert_queries_count(16): + with assert_queries_count(18): job.run() diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py index 9bd2b9b8e77b6..7b1f1fbbcb028 100644 --- a/tests/jobs/test_scheduler_job.py +++ b/tests/jobs/test_scheduler_job.py @@ -200,8 +200,8 @@ def test_process_executor_events(self, mock_stats_incr, mock_task_callback): dag_id2 = "test_process_executor_events_2" task_id_1 = 'dummy_task' - dag = DAG(dag_id=dag_id, start_date=DEFAULT_DATE, full_filepath="/test_path1/") - dag2 = DAG(dag_id=dag_id2, start_date=DEFAULT_DATE, full_filepath="/test_path1/") + dag = DAG(dag_id=dag_id, start_date=DEFAULT_DATE) + dag2 = DAG(dag_id=dag_id2, start_date=DEFAULT_DATE) task1 = DummyOperator(dag=dag, task_id=task_id_1) DummyOperator(dag=dag2, task_id=task_id_1) dag.fileloc = "/test_path1/" @@ -1670,10 +1670,11 @@ def test_dagrun_callbacks_are_called(self, state, expected_callback_msg): ti = dr.get_task_instance('dummy') ti.set_state(state, session) - self.scheduler_job._schedule_dag_run(dr, session) + with mock.patch.object(settings, "USE_JOB_SCHEDULE", False): + self.scheduler_job._do_scheduling(session) expected_callback = DagCallbackRequest( - full_filepath=dr.dag.fileloc, + full_filepath=dag.fileloc, dag_id=dr.dag_id, is_failure_callback=bool(state == State.FAILED), execution_date=dr.execution_date, @@ -1689,6 +1690,64 @@ def test_dagrun_callbacks_are_called(self, state, expected_callback_msg): session.rollback() session.close() + def test_dagrun_callbacks_commited_before_sent(self): + """ + Tests that before any callbacks are sent to the processor, the session is committed. This ensures + that the dagrun details are up to date when the callbacks are run. + """ + dag = DAG(dag_id='test_dagrun_callbacks_commited_before_sent', start_date=DEFAULT_DATE) + DummyOperator(task_id='dummy', dag=dag, owner='airflow') + + self.scheduler_job = SchedulerJob(subdir=os.devnull) + self.scheduler_job.processor_agent = mock.Mock() + self.scheduler_job._send_dag_callbacks_to_processor = mock.Mock() + self.scheduler_job._schedule_dag_run = mock.Mock() + + # Sync DAG into DB + with mock.patch.object(settings, "STORE_DAG_CODE", False): + self.scheduler_job.dagbag.bag_dag(dag, root_dag=dag) + self.scheduler_job.dagbag.sync_to_db() + + session = settings.Session() + orm_dag = session.query(DagModel).get(dag.dag_id) + assert orm_dag is not None + + # Create DagRun + self.scheduler_job._create_dag_runs([orm_dag], session) + + drs = DagRun.find(dag_id=dag.dag_id, session=session) + assert len(drs) == 1 + dr = drs[0] + + ti = dr.get_task_instance('dummy') + ti.set_state(State.SUCCESS, session) + + with mock.patch.object(settings, "USE_JOB_SCHEDULE", False), mock.patch( + "airflow.jobs.scheduler_job.prohibit_commit" + ) as mock_gaurd: + mock_gaurd.return_value.__enter__.return_value.commit.side_effect = session.commit + + def mock_schedule_dag_run(*args, **kwargs): + mock_gaurd.reset_mock() + return None + + def mock_send_dag_callbacks_to_processor(*args, **kwargs): + mock_gaurd.return_value.__enter__.return_value.commit.assert_called_once() + + self.scheduler_job._send_dag_callbacks_to_processor.side_effect = ( + mock_send_dag_callbacks_to_processor + ) + self.scheduler_job._schedule_dag_run.side_effect = mock_schedule_dag_run + + self.scheduler_job._do_scheduling(session) + + # Verify dag failure callback request is sent to file processor + self.scheduler_job._send_dag_callbacks_to_processor.assert_called_once() + # and mock_send_dag_callbacks_to_processor has asserted the callback was sent after a commit + + session.rollback() + session.close() + @parameterized.expand([(State.SUCCESS,), (State.FAILED,)]) def test_dagrun_callbacks_are_not_added_when_callbacks_are_not_defined(self, state): """ @@ -1725,10 +1784,15 @@ def test_dagrun_callbacks_are_not_added_when_callbacks_are_not_defined(self, sta ti = dr.get_task_instance('test_task') ti.set_state(state, session) - self.scheduler_job._schedule_dag_run(dr, session) + 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 - self.scheduler_job._send_dag_callbacks_to_processor.assert_called_once_with(dr, None) + 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() session.close() @@ -2372,12 +2436,10 @@ def test_verify_integrity_if_dag_not_changed(self): # Verify that DagRun.verify_integrity is not called with mock.patch('airflow.jobs.scheduler_job.DagRun.verify_integrity') as mock_verify_integrity: - scheduled_tis = self.scheduler_job._schedule_dag_run(dr, session) + self.scheduler_job._schedule_dag_run(dr, session) mock_verify_integrity.assert_not_called() session.flush() - assert scheduled_tis == 1 - tis_count = ( session.query(func.count(TaskInstance.task_id)) .filter( @@ -2436,11 +2498,9 @@ def test_verify_integrity_if_dag_changed(self): dag_version_2 = SerializedDagModel.get_latest_version_hash(dr.dag_id, session=session) assert dag_version_2 != dag_version_1 - scheduled_tis = self.scheduler_job._schedule_dag_run(dr, session) + self.scheduler_job._schedule_dag_run(dr, session) session.flush() - assert scheduled_tis == 2 - drs = DagRun.find(dag_id=dag.dag_id, session=session) assert len(drs) == 1 dr = drs[0] @@ -2465,6 +2525,7 @@ def test_verify_integrity_if_dag_changed(self): session.rollback() session.close() + @pytest.mark.quarantined def test_retry_still_in_executor(self): """ Checks if the scheduler does not put a task in limbo, when a task is retried diff --git a/tests/models/__init__.py b/tests/models/__init__.py index 7fba62e42e300..2d4a0d9a430de 100644 --- a/tests/models/__init__.py +++ b/tests/models/__init__.py @@ -21,4 +21,4 @@ from airflow.utils import timezone DEFAULT_DATE = timezone.datetime(2016, 1, 1) -TEST_DAGS_FOLDER = os.path.join(os.path.dirname(os.path.realpath(__file__)), '../dags') +TEST_DAGS_FOLDER = os.path.normpath(os.path.join(os.path.dirname(os.path.realpath(__file__)), '..', 'dags')) diff --git a/tests/models/test_baseoperator.py b/tests/models/test_baseoperator.py index becb970b1f140..ce848db0ed4c7 100644 --- a/tests/models/test_baseoperator.py +++ b/tests/models/test_baseoperator.py @@ -31,6 +31,7 @@ from airflow.models.baseoperator import BaseOperatorMeta, chain, cross_downstream from airflow.operators.dummy import DummyOperator from airflow.utils.edgemodifier import Label +from airflow.utils.trigger_rule import TriggerRule from tests.models import DEFAULT_DATE from tests.test_utils.mock_operators import DeprecatedOperator, MockNamedTuple, MockOperator @@ -595,6 +596,25 @@ def test_set_xcomargs_dependencies_error_when_outside_dag(self): op1 = DummyOperator(task_id="op1") CustomOp(task_id="op2", field=op1.output) + def test_invalid_trigger_rule(self): + with pytest.raises( + AirflowException, + match=( + f"The trigger_rule must be one of {TriggerRule.all_triggers()}," + "'.op1'; received 'some_rule'." + ), + ): + DummyOperator(task_id="op1", trigger_rule="some_rule") + + @parameterized.expand((("string", "dummy"), ("enum", TriggerRule.DUMMY))) + def test_replace_dummy_trigger_rule(self, name, rule): + with pytest.warns( + DeprecationWarning, match="dummy Trigger Rule is deprecated. Please use `TriggerRule.ALWAYS`." + ): + op1 = DummyOperator(task_id="op1", trigger_rule=rule) + + assert op1.trigger_rule == TriggerRule.ALWAYS + class InitSubclassOp(DummyOperator): def __init_subclass__(cls, class_arg=None, **kwargs) -> None: diff --git a/tests/models/test_dag.py b/tests/models/test_dag.py index 7cbda6690c8b6..0b673b9d95c6c 100644 --- a/tests/models/test_dag.py +++ b/tests/models/test_dag.py @@ -25,6 +25,7 @@ import unittest from contextlib import redirect_stdout from datetime import timedelta +from pathlib import Path from tempfile import NamedTemporaryFile from typing import Optional from unittest import mock @@ -1432,7 +1433,7 @@ def test_clear_set_dagrun_state_for_parent_dag(self, dag_run_state): @parameterized.expand( [(state, State.NONE) for state in State.task_states if state != State.RUNNING] - + [(State.RUNNING, State.SHUTDOWN)] + + [(State.RUNNING, State.RESTARTING)] ) # type: ignore def test_clear_dag(self, ti_state_begin, ti_state_end: Optional[str]): dag_id = 'test_clear_dag' @@ -1770,6 +1771,19 @@ def test_dags_needing_dagruns_only_unpaused(self): session.rollback() session.close() + @pytest.mark.parametrize( + ('fileloc', 'expected_relative'), + [ + (os.path.join(settings.DAGS_FOLDER, 'a.py'), Path('a.py')), + ('/tmp/foo.py', Path('/tmp/foo.py')), + ], + ) + def test_relative_fileloc(self, fileloc, expected_relative): + dag = DAG(dag_id='test') + dag.fileloc = fileloc + + assert dag.relative_fileloc == expected_relative + class TestQueries(unittest.TestCase): def setUp(self) -> None: diff --git a/tests/models/test_dagbag.py b/tests/models/test_dagbag.py index b4edc0c7a4f0e..37bb0106da5a3 100644 --- a/tests/models/test_dagbag.py +++ b/tests/models/test_dagbag.py @@ -246,7 +246,7 @@ def test_get_dag_fileloc(self): expected = { 'example_bash_operator': 'airflow/example_dags/example_bash_operator.py', 'example_subdag_operator': 'airflow/example_dags/example_subdag_operator.py', - 'example_subdag_operator.section-1': 'airflow/example_dags/subdags/subdag.py', + 'example_subdag_operator.section-1': 'airflow/example_dags/example_subdag_operator.py', 'test_zip_dag': 'dags/test_zip.zip/test_zip.py', } @@ -507,12 +507,15 @@ def subdag_1(): assert len(test_dag.subdags) == 6 # Perform processing dag - dagbag, found_dags, _ = self.process_dag(nested_subdags) + dagbag, found_dags, filename = self.process_dag(nested_subdags) # Validate correctness # all dags from test_dag should be listed self.validate_dags(test_dag, found_dags, dagbag) + for dag in dagbag.dags.values(): + assert dag.fileloc == filename + def test_skip_cycle_dags(self): """ Don't crash when loading an invalid (contains a cycle) DAG file. @@ -703,7 +706,9 @@ def test_serialized_dag_errors_are_import_errors(self, mock_serialize): ) assert dagbag.import_errors == {} - dagbag.sync_to_db(session=session) + with self.assertLogs(level="ERROR") as cm: + dagbag.sync_to_db(session=session) + self.assertIn("SerializationError", "\n".join(cm.output)) assert path in dagbag.import_errors err = dagbag.import_errors[path] diff --git a/tests/models/test_renderedtifields.py b/tests/models/test_renderedtifields.py index f76078c008cb5..02c72be994cb0 100644 --- a/tests/models/test_renderedtifields.py +++ b/tests/models/test_renderedtifields.py @@ -26,6 +26,7 @@ from parameterized import parameterized from airflow import settings +from airflow.configuration import TEST_DAGS_FOLDER from airflow.models import Variable from airflow.models.dag import DAG from airflow.models.renderedtifields import RenderedTaskInstanceFields as RTIF @@ -244,6 +245,7 @@ def test_get_k8s_pod_yaml(self, redact): dag = DAG("test_get_k8s_pod_yaml", start_date=START_DATE) with dag: task = BashOperator(task_id="test", bash_command="echo hi") + dag.fileloc = TEST_DAGS_FOLDER + '/test_get_k8s_pod_yaml.py' ti = TI(task=task, execution_date=EXECUTION_DATE) diff --git a/tests/models/test_serialized_dag.py b/tests/models/test_serialized_dag.py index 3e68ddc2cc413..aea143fc6a24a 100644 --- a/tests/models/test_serialized_dag.py +++ b/tests/models/test_serialized_dag.py @@ -71,7 +71,7 @@ def test_write_dag(self): assert SDM.has_dag(dag.dag_id) result = session.query(SDM.fileloc, SDM.data).filter(SDM.dag_id == dag.dag_id).one() - assert result.fileloc == dag.full_filepath + assert result.fileloc == dag.fileloc # Verifies JSON schema. SerializedDAG.validate_schema(result.data) @@ -138,8 +138,8 @@ def test_remove_dags_by_filepath(self): # Tests removing by file path. dag_removed_by_file = filtered_example_dags_list[0] # remove repeated files for those DAGs that define multiple dags in the same file (set comprehension) - example_dag_files = list({dag.full_filepath for dag in filtered_example_dags_list}) - example_dag_files.remove(dag_removed_by_file.full_filepath) + example_dag_files = list({dag.fileloc for dag in filtered_example_dags_list}) + example_dag_files.remove(dag_removed_by_file.fileloc) SDM.remove_deleted_dags(example_dag_files) assert not SDM.has_dag(dag_removed_by_file.dag_id) diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index 39c99dacd82b4..7d64c76c5a92f 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -18,6 +18,7 @@ import datetime import os +import signal import time import unittest import urllib @@ -67,7 +68,7 @@ from airflow.utils.state import State from airflow.utils.types import DagRunType from airflow.version import version -from tests.models import DEFAULT_DATE +from tests.models import DEFAULT_DATE, TEST_DAGS_FOLDER from tests.test_utils import db from tests.test_utils.asserts import assert_queries_count from tests.test_utils.config import conf_vars @@ -530,6 +531,37 @@ def raise_skip_exception(): ti.run() assert State.SKIPPED == ti.state + def test_task_sigterm_works_with_retries(self): + """ + Test that ensures that tasks are retried when they receive sigterm + """ + dag = DAG(dag_id='test_mark_failure_2', start_date=DEFAULT_DATE, default_args={'owner': 'owner1'}) + + def task_function(ti): + # pylint: disable=unused-argument + os.kill(ti.pid, signal.SIGTERM) + + task = PythonOperator( + task_id='test_on_failure', + python_callable=task_function, + retries=1, + retry_delay=datetime.timedelta(seconds=2), + dag=dag, + ) + + dag.create_dagrun( + run_id="test", + state=State.RUNNING, + execution_date=DEFAULT_DATE, + start_date=DEFAULT_DATE, + ) + ti = TI(task=task, execution_date=DEFAULT_DATE) + ti.refresh_from_db() + with self.assertRaises(AirflowException): + ti.run() + ti.refresh_from_db() + assert ti.state == State.UP_FOR_RETRY + def test_retry_delay(self): """ Test that retry delays are respected @@ -1867,6 +1899,26 @@ def test_task_stats(self, stats_mock): assert call(f'ti.start.{dag.dag_id}.{op.task_id}') in stats_mock.mock_calls assert stats_mock.call_count == 5 + def test_command_as_list(self): + dag = DAG( + 'test_dag', + start_date=DEFAULT_DATE, + end_date=DEFAULT_DATE + datetime.timedelta(days=10), + ) + dag.fileloc = os.path.join(TEST_DAGS_FOLDER, 'x.py') + op = DummyOperator(task_id='dummy_op', dag=dag) + ti = TI(task=op, execution_date=DEFAULT_DATE) + assert ti.command_as_list() == [ + 'airflow', + 'tasks', + 'run', + dag.dag_id, + op.task_id, + DEFAULT_DATE.isoformat(), + '--subdir', + 'DAGS_FOLDER/x.py', + ] + def test_generate_command_default_param(self): dag_id = 'test_generate_command_default_param' task_id = 'task' @@ -1893,8 +1945,9 @@ def test_generate_command_specific_param(self): def test_get_rendered_template_fields(self): - with DAG('test-dag', start_date=DEFAULT_DATE): + with DAG('test-dag', start_date=DEFAULT_DATE) as dag: task = BashOperator(task_id='op1', bash_command="{{ task.task_id }}") + dag.fileloc = TEST_DAGS_FOLDER + '/test_get_k8s_pod_yaml.py' ti = TI(task=task, execution_date=DEFAULT_DATE) @@ -1952,6 +2005,8 @@ def test_render_k8s_pod_yaml(self, pod_mutation_hook): 'test_get_rendered_k8s_spec', 'op1', '2016-01-01T00:00:00+00:00', + '--subdir', + __file__, ], 'image': ':', 'name': 'base', diff --git a/tests/models/test_timestamp.py b/tests/models/test_timestamp.py index 692e1a77919ad..cb50cb1377b6a 100644 --- a/tests/models/test_timestamp.py +++ b/tests/models/test_timestamp.py @@ -24,17 +24,14 @@ from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import State +from tests.test_utils.db import clear_db_logs, clear_db_runs -@pytest.yield_fixture(name="clear_db_fixture") -@provide_session -def clear_db(session=None): - session.query(Log).delete() - session.query(TaskInstance).delete() - yield session - session.query(Log).delete() - session.query(TaskInstance).delete() - session.commit() +@pytest.fixture(autouse=True) +def clear_db(): + clear_db_logs() + clear_db_runs() + yield def add_log(execdate, session, timezone_override=None): @@ -50,28 +47,28 @@ def add_log(execdate, session, timezone_override=None): return log -def test_timestamp_behaviour(clear_db_fixture): - for session in clear_db_fixture: - execdate = timezone.utcnow() - with freeze_time(execdate): - current_time = timezone.utcnow() - old_log = add_log(execdate, session) - session.expunge(old_log) - log_time = session.query(Log).one().dttm - assert log_time == current_time - assert log_time.tzinfo.name == 'UTC' +@provide_session +def test_timestamp_behaviour(session=None): + execdate = timezone.utcnow() + with freeze_time(execdate): + current_time = timezone.utcnow() + old_log = add_log(execdate, session) + session.expunge(old_log) + log_time = session.query(Log).one().dttm + assert log_time == current_time + assert log_time.tzinfo.name == 'UTC' -def test_timestamp_behaviour_with_timezone(clear_db_fixture): - for session in clear_db_fixture: - execdate = timezone.utcnow() - with freeze_time(execdate): - current_time = timezone.utcnow() - old_log = add_log(execdate, session, timezone_override=pendulum.timezone('Europe/Warsaw')) - session.expunge(old_log) - # No matter what timezone we set - we should always get back UTC - log_time = session.query(Log).one().dttm - assert log_time == current_time - assert old_log.dttm.tzinfo.name != 'UTC' - assert log_time.tzinfo.name == 'UTC' - assert old_log.dttm.astimezone(pendulum.timezone('UTC')) == log_time +@provide_session +def test_timestamp_behaviour_with_timezone(session=None): + execdate = timezone.utcnow() + with freeze_time(execdate): + current_time = timezone.utcnow() + old_log = add_log(execdate, session, timezone_override=pendulum.timezone('Europe/Warsaw')) + session.expunge(old_log) + # No matter what timezone we set - we should always get back UTC + log_time = session.query(Log).one().dttm + assert log_time == current_time + assert old_log.dttm.tzinfo.name != 'UTC' + assert log_time.tzinfo.name == 'UTC' + assert old_log.dttm.astimezone(pendulum.timezone('UTC')) == log_time diff --git a/tests/providers/amazon/aws/hooks/test_base_aws.py b/tests/providers/amazon/aws/hooks/test_base_aws.py index abb6c587b1462..c9343878c3f74 100644 --- a/tests/providers/amazon/aws/hooks/test_base_aws.py +++ b/tests/providers/amazon/aws/hooks/test_base_aws.py @@ -230,6 +230,62 @@ def test_get_credentials_from_extra_with_s3_config_and_profile( hook._get_credentials(region_name=None) mock_parse_s3_config.assert_called_once_with('aws-credentials', 'aws', 'test') + @unittest.skipIf(mock_sts is None, 'mock_sts package not present') + @mock.patch.object(AwsBaseHook, 'get_connection') + @mock_sts + def test_assume_role(self, mock_get_connection): + aws_conn_id = 'aws/test' + role_arn = 'arn:aws:iam::123456:role/role_arn' + slugified_role_session_name = 'airflow_aws-test' + + mock_connection = Connection( + conn_id=aws_conn_id, + extra=json.dumps( + { + "role_arn": role_arn, + } + ), + ) + mock_get_connection.return_value = mock_connection + + def mock_assume_role(**kwargs): + assert kwargs['RoleArn'] == role_arn + # The role session name gets invalid characters removed/replaced with hyphens + # (e.g. / is replaced with -) + assert kwargs['RoleSessionName'] == slugified_role_session_name + sts_response = { + 'ResponseMetadata': {'HTTPStatusCode': 200}, + 'Credentials': { + 'Expiration': datetime.now(), + 'AccessKeyId': 1, + 'SecretAccessKey': 1, + 'SessionToken': 1, + }, + } + return sts_response + + with mock.patch( + 'airflow.providers.amazon.aws.hooks.base_aws.requests.Session.get' + ) as mock_get, mock.patch('airflow.providers.amazon.aws.hooks.base_aws.boto3') as mock_boto3: + mock_get.return_value.ok = True + + mock_client = mock_boto3.session.Session.return_value.client + mock_client.return_value.assume_role.side_effect = mock_assume_role + + hook = AwsBaseHook(aws_conn_id=aws_conn_id, client_type='s3') + hook.get_client_type('s3') + + calls_assume_role = [ + mock.call.session.Session().client('sts', config=None), + mock.call.session.Session() + .client() + .assume_role( + RoleArn=role_arn, + RoleSessionName=slugified_role_session_name, + ), + ] + mock_boto3.assert_has_calls(calls_assume_role) + @unittest.skipIf(mock_sts is None, 'mock_sts package not present') @mock.patch.object(AwsBaseHook, 'get_connection') @mock_sts diff --git a/tests/providers/amazon/aws/operators/test_ecs.py b/tests/providers/amazon/aws/operators/test_ecs.py index 4013450ec9c55..8cf38c5859b00 100644 --- a/tests/providers/amazon/aws/operators/test_ecs.py +++ b/tests/providers/amazon/aws/operators/test_ecs.py @@ -97,23 +97,38 @@ def test_template_fields_overrides(self): @parameterized.expand( [ - ['EC2', None, None, {'launchType': 'EC2'}], - ['FARGATE', None, None, {'launchType': 'FARGATE', 'platformVersion': 'LATEST'}], [ 'EC2', None, + None, + None, + {'launchType': 'EC2'}, + ], + [ + 'FARGATE', + None, + 'LATEST', + None, + {'launchType': 'FARGATE', 'platformVersion': 'LATEST'}, + ], + [ + 'EC2', + None, + None, {'testTagKey': 'testTagValue'}, {'launchType': 'EC2', 'tags': [{'key': 'testTagKey', 'value': 'testTagValue'}]}, ], [ '', None, + None, {'testTagKey': 'testTagValue'}, {'tags': [{'key': 'testTagKey', 'value': 'testTagValue'}]}, ], [ None, {'capacityProvider': 'FARGATE_SPOT'}, + 'LATEST', None, { 'capacityProviderStrategy': {'capacityProvider': 'FARGATE_SPOT'}, @@ -123,6 +138,7 @@ def test_template_fields_overrides(self): [ 'FARGATE', {'capacityProvider': 'FARGATE_SPOT', 'weight': 123, 'base': 123}, + 'LATEST', None, { 'capacityProviderStrategy': { @@ -136,6 +152,7 @@ def test_template_fields_overrides(self): [ 'EC2', {'capacityProvider': 'FARGATE_SPOT'}, + 'LATEST', None, { 'capacityProviderStrategy': {'capacityProvider': 'FARGATE_SPOT'}, @@ -147,11 +164,21 @@ def test_template_fields_overrides(self): @mock.patch.object(ECSOperator, '_wait_for_task_ended') @mock.patch.object(ECSOperator, '_check_success_task') def test_execute_without_failures( - self, launch_type, capacity_provider_strategy, tags, expected_args, check_mock, wait_mock + self, + launch_type, + capacity_provider_strategy, + platform_version, + tags, + expected_args, + check_mock, + wait_mock, ): self.set_up_operator( - launch_type=launch_type, capacity_provider_strategy=capacity_provider_strategy, tags=tags + launch_type=launch_type, + capacity_provider_strategy=capacity_provider_strategy, + platform_version=platform_version, + tags=tags, ) client_mock = self.aws_hook_mock.return_value.get_conn.return_value client_mock.run_task.return_value = RESPONSE_WITHOUT_FAILURES diff --git a/tests/providers/amazon/aws/sensors/test_sqs.py b/tests/providers/amazon/aws/sensors/test_sqs.py index 90349a321c1e9..82a1aacb2a4a2 100644 --- a/tests/providers/amazon/aws/sensors/test_sqs.py +++ b/tests/providers/amazon/aws/sensors/test_sqs.py @@ -17,6 +17,7 @@ # under the License. +import json import unittest from unittest import mock @@ -107,3 +108,180 @@ def test_poke_receive_raise_exception(self, mock_conn): self.sensor.poke(self.mock_context) assert 'test exception' in ctx.value.args[0] + + @mock.patch.object(SQSHook, 'get_conn') + def test_poke_visibility_timeout(self, mock_conn): + # Check without visibility_timeout parameter + self.sqs_hook.create_queue('test') + self.sqs_hook.send_message(queue_url='test', message_body='hello') + + self.sensor.poke(self.mock_context) + + calls_receive_message = [ + mock.call().receive_message(QueueUrl='test', MaxNumberOfMessages=5, WaitTimeSeconds=1) + ] + mock_conn.assert_has_calls(calls_receive_message) + # Check with visibility_timeout parameter + self.sensor = SQSSensor( + task_id='test_task2', + dag=self.dag, + sqs_queue='test', + aws_conn_id='aws_default', + visibility_timeout=42, + ) + self.sensor.poke(self.mock_context) + + calls_receive_message = [ + mock.call().receive_message( + QueueUrl='test', MaxNumberOfMessages=5, WaitTimeSeconds=1, VisibilityTimeout=42 + ) + ] + mock_conn.assert_has_calls(calls_receive_message) + + @mock_sqs + def test_poke_message_invalid_filtering(self): + self.sqs_hook.create_queue('test') + self.sqs_hook.send_message(queue_url='test', message_body='hello') + sensor = SQSSensor( + task_id='test_task2', + dag=self.dag, + sqs_queue='test', + aws_conn_id='aws_default', + message_filtering='invalid_option', + ) + with pytest.raises(NotImplementedError) as ctx: + sensor.poke(self.mock_context) + assert 'Override this method to define custom filters' in ctx.value.args[0] + + @mock.patch.object(SQSHook, "get_conn") + def test_poke_message_filtering_literal_values(self, mock_conn): + self.sqs_hook.create_queue('test') + matching = [{"id": 11, "body": "a matching message"}] + non_matching = [{"id": 12, "body": "a non-matching message"}] + all = matching + non_matching + + def mock_receive_message(**kwargs): + messages = [] + for message in all: + messages.append( + { + 'MessageId': message['id'], + 'ReceiptHandle': 100 + message['id'], + 'Body': message['body'], + } + ) + return {'Messages': messages} + + mock_conn.return_value.receive_message.side_effect = mock_receive_message + + def mock_delete_message_batch(**kwargs): + return {'Successful'} + + mock_conn.return_value.delete_message_batch.side_effect = mock_delete_message_batch + + # Test that messages are filtered + self.sensor.message_filtering = 'literal' + self.sensor.message_filtering_match_values = ["a matching message"] + result = self.sensor.poke(self.mock_context) + assert result + + # Test that only filtered messages are deleted + delete_entries = [{'Id': x['id'], 'ReceiptHandle': 100 + x['id']} for x in matching] + calls_delete_message_batch = [ + mock.call().delete_message_batch(QueueUrl='test', Entries=delete_entries) + ] + mock_conn.assert_has_calls(calls_delete_message_batch) + + @mock.patch.object(SQSHook, "get_conn") + def test_poke_message_filtering_jsonpath(self, mock_conn): + self.sqs_hook.create_queue('test') + matching = [ + {"id": 11, "key": {"matches": [1, 2]}}, + {"id": 12, "key": {"matches": [3, 4, 5]}}, + {"id": 13, "key": {"matches": [10]}}, + ] + non_matching = [ + {"id": 14, "key": {"nope": [5, 6]}}, + {"id": 15, "key": {"nope": [7, 8]}}, + ] + all = matching + non_matching + + def mock_receive_message(**kwargs): + messages = [] + for message in all: + messages.append( + { + 'MessageId': message['id'], + 'ReceiptHandle': 100 + message['id'], + 'Body': json.dumps(message), + } + ) + return {'Messages': messages} + + mock_conn.return_value.receive_message.side_effect = mock_receive_message + + def mock_delete_message_batch(**kwargs): + return {'Successful'} + + mock_conn.return_value.delete_message_batch.side_effect = mock_delete_message_batch + + # Test that messages are filtered + self.sensor.message_filtering = 'jsonpath' + self.sensor.message_filtering_config = 'key.matches[*]' + result = self.sensor.poke(self.mock_context) + assert result + + # Test that only filtered messages are deleted + delete_entries = [{'Id': x['id'], 'ReceiptHandle': 100 + x['id']} for x in matching] + calls_delete_message_batch = [ + mock.call().delete_message_batch(QueueUrl='test', Entries=delete_entries) + ] + mock_conn.assert_has_calls(calls_delete_message_batch) + + @mock.patch.object(SQSHook, "get_conn") + def test_poke_message_filtering_jsonpath_values(self, mock_conn): + self.sqs_hook.create_queue('test') + matching = [ + {"id": 11, "key": {"matches": [1, 2]}}, + {"id": 12, "key": {"matches": [1, 4, 5]}}, + {"id": 13, "key": {"matches": [4, 5]}}, + ] + non_matching = [ + {"id": 21, "key": {"matches": [10]}}, + {"id": 22, "key": {"nope": [5, 6]}}, + {"id": 23, "key": {"nope": [7, 8]}}, + ] + all = matching + non_matching + + def mock_receive_message(**kwargs): + messages = [] + for message in all: + messages.append( + { + 'MessageId': message['id'], + 'ReceiptHandle': 100 + message['id'], + 'Body': json.dumps(message), + } + ) + return {'Messages': messages} + + mock_conn.return_value.receive_message.side_effect = mock_receive_message + + def mock_delete_message_batch(**kwargs): + return {'Successful'} + + mock_conn.return_value.delete_message_batch.side_effect = mock_delete_message_batch + + # Test that messages are filtered + self.sensor.message_filtering = 'jsonpath' + self.sensor.message_filtering_config = 'key.matches[*]' + self.sensor.message_filtering_match_values = [1, 4] + result = self.sensor.poke(self.mock_context) + assert result + + # Test that only filtered messages are deleted + delete_entries = [{'Id': x['id'], 'ReceiptHandle': 100 + x['id']} for x in matching] + calls_delete_message_batch = [ + mock.call().delete_message_batch(QueueUrl='test', Entries=delete_entries) + ] + mock_conn.assert_has_calls(calls_delete_message_batch) diff --git a/tests/providers/google/cloud/operators/test_cloud_storage_transfer_service.py b/tests/providers/google/cloud/operators/test_cloud_storage_transfer_service.py index 3ea08bb6df384..27e71c7151ecf 100644 --- a/tests/providers/google/cloud/operators/test_cloud_storage_transfer_service.py +++ b/tests/providers/google/cloud/operators/test_cloud_storage_transfer_service.py @@ -40,6 +40,7 @@ HTTP_DATA_SOURCE, LIST_URL, NAME, + PATH, SCHEDULE, SCHEDULE_END_DATE, SCHEDULE_START_DATE, @@ -77,6 +78,7 @@ OPERATION_NAME = "operation-name" AWS_BUCKET_NAME = "aws-bucket-name" GCS_BUCKET_NAME = "gcp-bucket-name" +GCS_PATH = None DESCRIPTION = "description" DEFAULT_DATE = timezone.datetime(2017, 1, 1) @@ -112,7 +114,7 @@ DESCRIPTION: DESCRIPTION, STATUS: 'ENABLED', SCHEDULE: SCHEDULE_DICT, - TRANSFER_SPEC: {GCS_DATA_SINK: {BUCKET_NAME: GCS_BUCKET_NAME}}, + TRANSFER_SPEC: {GCS_DATA_SINK: {BUCKET_NAME: GCS_BUCKET_NAME, PATH: GCS_PATH}}, } # type: Dict VALID_TRANSFER_JOB_GCS = deepcopy(VALID_TRANSFER_JOB_BASE) VALID_TRANSFER_JOB_GCS[TRANSFER_SPEC].update(deepcopy(SOURCE_GCS)) @@ -126,7 +128,7 @@ SCHEDULE: SCHEDULE_NATIVE, TRANSFER_SPEC: { GCS_DATA_SOURCE: {BUCKET_NAME: GCS_BUCKET_NAME}, - GCS_DATA_SINK: {BUCKET_NAME: GCS_BUCKET_NAME}, + GCS_DATA_SINK: {BUCKET_NAME: GCS_BUCKET_NAME, PATH: GCS_PATH}, }, } @@ -134,7 +136,7 @@ DESCRIPTION: DESCRIPTION, STATUS: 'ENABLED', SCHEDULE: SCHEDULE_DICT, - TRANSFER_SPEC: {GCS_DATA_SINK: {BUCKET_NAME: GCS_BUCKET_NAME}}, + TRANSFER_SPEC: {GCS_DATA_SINK: {BUCKET_NAME: GCS_BUCKET_NAME, PATH: GCS_PATH}}, } # type: Dict VALID_TRANSFER_JOB_GCS_RAW = deepcopy(VALID_TRANSFER_JOB_RAW) diff --git a/tests/providers/samba/hooks/test_samba.py b/tests/providers/samba/hooks/test_samba.py index f384e9e77cb70..4fb36d8ca9f94 100644 --- a/tests/providers/samba/hooks/test_samba.py +++ b/tests/providers/samba/hooks/test_samba.py @@ -16,23 +16,24 @@ # specific language governing permissions and limitations # under the License. -import json import unittest +from inspect import getfullargspec from unittest import mock -from unittest.mock import call import pytest +from parameterized import parameterized from airflow.exceptions import AirflowException from airflow.models import Connection from airflow.providers.samba.hooks.samba import SambaHook -connection = Connection( +PATH_PARAMETER_NAMES = {"path", "src", "dst"} + +CONNECTION = Connection( host='ip', schema='share', login='username', password='password', - extra=json.dumps({'workgroup': 'workgroup'}), ) @@ -41,90 +42,91 @@ def test_get_conn_should_fail_if_conn_id_does_not_exist(self): with pytest.raises(AirflowException): SambaHook('conn') - @mock.patch('airflow.providers.samba.hooks.samba.SambaClient') - @mock.patch('airflow.hooks.base.BaseHook.get_connection') - def test_get_conn(self, get_conn_mock, get_client_mock): - get_conn_mock.return_value = connection - hook = SambaHook('samba_default') - conn = hook.get_conn() - assert str(get_client_mock.mock_calls[0].workgroup) == "workgroup" - assert conn is get_client_mock() - get_conn_mock.assert_called_once_with('samba_default') - - @mock.patch('airflow.providers.samba.hooks.samba.SambaHook.get_conn') - @mock.patch('airflow.hooks.base.BaseHook.get_connection') - def test_push_from_local_should_succeed_if_destination_has_same_name_but_not_a_file( - self, base_conn_mock, samba_hook_mock - ): - base_conn_mock.return_value = connection - samba_hook_mock.get_conn.return_value = mock.Mock() - - samba_hook_mock.return_value.exists.return_value = True - samba_hook_mock.return_value.isfile.return_value = False - samba_hook_mock.return_value.exists.return_value = True - - hook = SambaHook('samba_default') - destination_filepath = "/path/to/dest/file" - local_filepath = "/path/to/local/file" - hook.push_from_local(destination_filepath=destination_filepath, local_filepath=local_filepath) - - base_conn_mock.assert_called_once_with('samba_default') - samba_hook_mock.assert_called_once() - samba_hook_mock.return_value.exists.assert_called_once_with(destination_filepath) - samba_hook_mock.return_value.isfile.assert_called_once_with(destination_filepath) - samba_hook_mock.return_value.remove.assert_not_called() - samba_hook_mock.return_value.upload.assert_called_once_with(local_filepath, destination_filepath) - - @mock.patch('airflow.providers.samba.hooks.samba.SambaHook.get_conn') + @mock.patch('smbclient.register_session') @mock.patch('airflow.hooks.base.BaseHook.get_connection') - def test_push_from_local_should_delete_file_if_exists_and_save_file( - self, base_conn_mock, samba_hook_mock - ): - base_conn_mock.return_value = connection - samba_hook_mock.get_conn.return_value = mock.Mock() - - samba_hook_mock.return_value.exists.return_value = False - samba_hook_mock.return_value.exists.return_value = False - - hook = SambaHook('samba_default') - destination_folder = "/path/to/dest" - destination_filepath = destination_folder + "/file" - local_filepath = "/path/to/local/file" - hook.push_from_local(destination_filepath=destination_filepath, local_filepath=local_filepath) - - base_conn_mock.assert_called_once_with('samba_default') - samba_hook_mock.assert_called_once() - samba_hook_mock.return_value.exists.assert_has_calls( - [call(destination_filepath), call(destination_folder)] - ) - samba_hook_mock.return_value.isfile.assert_not_called() - samba_hook_mock.return_value.remove.assert_not_called() - samba_hook_mock.return_value.mkdir.assert_called_once_with(destination_folder) - samba_hook_mock.return_value.upload.assert_called_once_with(local_filepath, destination_filepath) - - @mock.patch('airflow.providers.samba.hooks.samba.SambaHook.get_conn') + def test_context_manager(self, get_conn_mock, register_session): + get_conn_mock.return_value = CONNECTION + register_session.return_value = None + with SambaHook('samba_default'): + args, kwargs = tuple(register_session.call_args_list[0]) + assert args == (CONNECTION.host,) + assert kwargs == { + "username": CONNECTION.login, + "password": CONNECTION.password, + "port": 445, + "connection_cache": {}, + } + cache = kwargs.get("connection_cache") + mock_connection = mock.Mock() + mock_connection.disconnect.return_value = None + cache["foo"] = mock_connection + + # Test that the connection was disconnected upon exit. + assert len(mock_connection.disconnect.mock_calls) == 1 + + @parameterized.expand( + [ + "getxattr", + "link", + "listdir", + "listxattr", + "lstat", + "makedirs", + "mkdir", + "open_file", + "readlink", + "remove", + "removedirs", + "removexattr", + "rename", + "replace", + "rmdir", + "scandir", + "setxattr", + "stat", + "stat_volume", + "symlink", + "truncate", + "unlink", + "utime", + "walk", + ], + ) @mock.patch('airflow.hooks.base.BaseHook.get_connection') - def test_push_from_local_should_create_directory_if_not_exist_and_save_file( - self, base_conn_mock, samba_hook_mock - ): - base_conn_mock.return_value = connection - samba_hook_mock.get_conn.return_value = mock.Mock() - - samba_hook_mock.return_value.exists.return_value = False - samba_hook_mock.return_value.exists.return_value = False - + def test_method(self, name, get_conn_mock): + get_conn_mock.return_value = CONNECTION hook = SambaHook('samba_default') - destination_folder = "/path/to/dest" - destination_filepath = destination_folder + "/file" - local_filepath = "/path/to/local/file" - hook.push_from_local(destination_filepath=destination_filepath, local_filepath=local_filepath) - - base_conn_mock.assert_called_once_with('samba_default') - samba_hook_mock.assert_called_once() - samba_hook_mock.return_value.exists.assert_has_calls( - [call(destination_filepath), call(destination_folder)] - ) - samba_hook_mock.return_value.isfile.assert_not_called() - samba_hook_mock.return_value.remove.assert_not_called() - samba_hook_mock.return_value.mkdir.assert_called_once_with(destination_folder) - samba_hook_mock.return_value.upload.assert_called_once_with(local_filepath, destination_filepath) + connection_settings = { + 'connection_cache': {}, + 'username': CONNECTION.login, + 'password': CONNECTION.password, + 'port': 445, + } + with mock.patch('smbclient.' + name) as p: + kwargs = {} + method = getattr(hook, name) + spec = getfullargspec(method) + + if spec.defaults: + for default in reversed(spec.defaults): + arg = spec.args.pop() + kwargs[arg] = default + + # Ignore "self" argument. + args = spec.args[1:] + + method(*args, **kwargs) + assert len(p.mock_calls) == 1 + + # Verify positional arguments. If the argument is a path parameter, then we expect + # the hook implementation to fully qualify the path. + p_args, p_kwargs = tuple(p.call_args_list[0]) + for arg, provided in zip(args, p_args): + if arg in PATH_PARAMETER_NAMES: + expected = "//" + CONNECTION.host + "/" + CONNECTION.schema + "/" + arg + else: + expected = arg + assert expected == provided + + # We expect keyword arguments to include the connection settings. + assert dict(kwargs, **connection_settings) == p_kwargs diff --git a/tests/providers/slack/operators/test_slack.py b/tests/providers/slack/operators/test_slack.py index a60c193341081..964bb9cbef7d5 100644 --- a/tests/providers/slack/operators/test_slack.py +++ b/tests/providers/slack/operators/test_slack.py @@ -188,12 +188,28 @@ def test_init_with_valid_params(self): assert slack_api_post_operator.slack_conn_id == test_slack_conn_id @mock.patch('airflow.providers.slack.operators.slack.SlackHook') - def test_api_call_params_with_default_args(self, mock_hook): + def test_api_call_params_with_content_args(self, mock_hook): test_slack_conn_id = 'test_slack_conn_id' slack_api_post_operator = SlackAPIFileOperator( - task_id='slack', - slack_conn_id=test_slack_conn_id, + task_id='slack', slack_conn_id=test_slack_conn_id, content='test-content' + ) + + slack_api_post_operator.execute() + + expected_api_params = { + 'channels': '#general', + 'initial_comment': 'No message has been set!', + 'content': 'test-content', + } + assert expected_api_params == slack_api_post_operator.api_params + + @mock.patch('airflow.providers.slack.operators.slack.SlackHook') + def test_api_call_params_with_file_args(self, mock_hook): + test_slack_conn_id = 'test_slack_conn_id' + + slack_api_post_operator = SlackAPIFileOperator( + task_id='slack', slack_conn_id=test_slack_conn_id, filename='test.csv', filetype='csv' ) slack_api_post_operator.execute() @@ -201,8 +217,10 @@ def test_api_call_params_with_default_args(self, mock_hook): expected_api_params = { 'channels': '#general', 'initial_comment': 'No message has been set!', - 'filename': 'default_name.csv', + 'filename': 'test.csv', 'filetype': 'csv', - 'content': 'default,content,csv,file', } + + expected_file_params = {'file': 'test.csv'} assert expected_api_params == slack_api_post_operator.api_params + assert expected_file_params == slack_api_post_operator.file_params diff --git a/tests/providers/yandex/hooks/test_yandexcloud_dataproc.py b/tests/providers/yandex/hooks/test_yandexcloud_dataproc.py index fd6defddae95a..0d93a4a56367c 100644 --- a/tests/providers/yandex/hooks/test_yandexcloud_dataproc.py +++ b/tests/providers/yandex/hooks/test_yandexcloud_dataproc.py @@ -35,7 +35,7 @@ AVAILABILITY_ZONE_ID = 'ru-central1-c' CLUSTER_NAME = 'dataproc_cluster' -CLUSTER_IMAGE_VERSION = '1.1' +CLUSTER_IMAGE_VERSION = '1.4' # https://cloud.yandex.com/docs/resource-manager/operations/folder/get-id FOLDER_ID = 'my_folder_id' diff --git a/tests/providers/yandex/operators/test_yandexcloud_dataproc.py b/tests/providers/yandex/operators/test_yandexcloud_dataproc.py index 66c531f4ef34b..d52607ceb8739 100644 --- a/tests/providers/yandex/operators/test_yandexcloud_dataproc.py +++ b/tests/providers/yandex/operators/test_yandexcloud_dataproc.py @@ -37,7 +37,7 @@ AVAILABILITY_ZONE_ID = 'ru-central1-c' CLUSTER_NAME = 'dataproc_cluster' -CLUSTER_IMAGE_VERSION = '1.1' +CLUSTER_IMAGE_VERSION = '1.4' # https://cloud.yandex.com/docs/resource-manager/operations/folder/get-id FOLDER_ID = 'my_folder_id' @@ -92,20 +92,27 @@ def test_create_cluster(self, create_cluster_mock, *_): operator.execute(context) create_cluster_mock.assert_called_once_with( cluster_description='', - cluster_image_version='1.1', + cluster_image_version='1.4', cluster_name=None, computenode_count=0, - computenode_disk_size=15, - computenode_disk_type='network-ssd', - computenode_resource_preset='s2.small', - datanode_count=2, - datanode_disk_size=15, - datanode_disk_type='network-ssd', - datanode_resource_preset='s2.small', + computenode_disk_size=None, + computenode_disk_type=None, + computenode_resource_preset=None, + computenode_max_hosts_count=None, + computenode_measurement_duration=None, + computenode_warmup_duration=None, + computenode_stabilization_duration=None, + computenode_preemptible=False, + computenode_cpu_utilization_target=None, + computenode_decommission_timeout=None, + datanode_count=1, + datanode_disk_size=None, + datanode_disk_type=None, + datanode_resource_preset=None, folder_id='my_folder_id', - masternode_disk_size=15, - masternode_disk_type='network-ssd', - masternode_resource_preset='s2.small', + masternode_disk_size=None, + masternode_disk_type=None, + masternode_resource_preset=None, s3_bucket='my_bucket_name', service_account_id=None, services=('HDFS', 'YARN', 'MAPREDUCE', 'HIVE', 'SPARK'), diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py index 40d9bc077c441..7e2798c4cbda4 100644 --- a/tests/serialization/test_dag_serialization.py +++ b/tests/serialization/test_dag_serialization.py @@ -380,10 +380,6 @@ def validate_deserialized_dag(self, serialized_dag, dag): for task_id in dag.task_ids: self.validate_deserialized_task(serialized_dag.get_task(task_id), dag.get_task(task_id)) - # Verify that the DAG object has 'full_filepath' attribute - # and is equal to fileloc - assert serialized_dag.full_filepath == dag.fileloc - def validate_deserialized_task( self, serialized_task, diff --git a/tests/ti_deps/deps/test_trigger_rule_dep.py b/tests/ti_deps/deps/test_trigger_rule_dep.py index 1443e92ca9794..8a1247b92f81e 100644 --- a/tests/ti_deps/deps/test_trigger_rule_dep.py +++ b/tests/ti_deps/deps/test_trigger_rule_dep.py @@ -48,11 +48,11 @@ def test_no_upstream_tasks(self): ti = self._get_task_instance(TriggerRule.ALL_DONE, State.UP_FOR_RETRY) assert TriggerRuleDep().is_met(ti=ti) - def test_dummy_tr(self): + def test_always_tr(self): """ - The dummy trigger rule should always pass this dep + The always trigger rule should always pass this dep """ - ti = self._get_task_instance(TriggerRule.DUMMY, State.UP_FOR_RETRY) + ti = self._get_task_instance(TriggerRule.ALWAYS, State.UP_FOR_RETRY) assert TriggerRuleDep().is_met(ti=ti) def test_one_success_tr_success(self): diff --git a/tests/utils/test_trigger_rule.py b/tests/utils/test_trigger_rule.py index be7a903afb32b..9a03808ac56fb 100644 --- a/tests/utils/test_trigger_rule.py +++ b/tests/utils/test_trigger_rule.py @@ -32,4 +32,5 @@ def test_valid_trigger_rules(self): assert TriggerRule.is_valid(TriggerRule.NONE_FAILED_OR_SKIPPED) assert TriggerRule.is_valid(TriggerRule.NONE_SKIPPED) assert TriggerRule.is_valid(TriggerRule.DUMMY) - assert len(TriggerRule.all_triggers()) == 9 + assert TriggerRule.is_valid(TriggerRule.ALWAYS) + assert len(TriggerRule.all_triggers()) == 10 diff --git a/tests/www/views/test_views.py b/tests/www/views/test_views.py index ebb4f462a2985..a15f708e5f354 100644 --- a/tests/www/views/test_views.py +++ b/tests/www/views/test_views.py @@ -22,7 +22,7 @@ from airflow.configuration import initialize_config from airflow.plugins_manager import AirflowPlugin, EntryPointSource -from airflow.www.views import get_safe_url, truncate_task_duration +from airflow.www.views import get_key_paths, get_safe_url, get_value_from_path, truncate_task_duration from tests.test_utils.config import conf_vars from tests.test_utils.mock_plugins import mock_plugin_manager from tests.test_utils.www import check_content_in_response, check_content_not_in_response @@ -243,3 +243,26 @@ def get_task_instance(session, task): dagrun.refresh_from_db(session=session) # dagrun should be set to QUEUED assert dagrun.get_state() == State.QUEUED + + +TEST_CONTENT_DICT = {"key1": {"key2": "val2", "key3": "val3", "key4": {"key5": "val5"}}} + + +@pytest.mark.parametrize( + "test_content_dict, expected_paths", [(TEST_CONTENT_DICT, ("key1.key2", "key1.key3", "key1.key4.key5"))] +) +def test_generate_key_paths(test_content_dict, expected_paths): + for key_path in get_key_paths(test_content_dict): + assert key_path in expected_paths + + +@pytest.mark.parametrize( + "test_content_dict, test_key_path, expected_value", + [ + (TEST_CONTENT_DICT, "key1.key2", "val2"), + (TEST_CONTENT_DICT, "key1.key3", "val3"), + (TEST_CONTENT_DICT, "key1.key4.key5", "val5"), + ], +) +def test_get_value_from_path(test_content_dict, test_key_path, expected_value): + assert expected_value == get_value_from_path(test_key_path, test_content_dict) diff --git a/tests/www/views/test_views_acl.py b/tests/www/views/test_views_acl.py index 8fe011b22e23c..4dab3860e2c91 100644 --- a/tests/www/views/test_views_acl.py +++ b/tests/www/views/test_views_acl.py @@ -694,21 +694,9 @@ def test_failed_success(client_all_dags_edit_tis): check_content_in_response('Marked failed on 1 task instances', resp) -@pytest.mark.parametrize( - "url, expected_content", - [ - ("paused?dag_id=example_bash_operator&is_paused=false", "OK"), - ("refresh?dag_id=example_bash_operator", ""), - ], - ids=[ - "paused", - "refresh", - ], -) -def test_post_success(dag_test_client, url, expected_content): - # post request failure won't test - resp = dag_test_client.post(url, follow_redirects=True) - check_content_in_response(expected_content, resp) +def test_paused_post_success(dag_test_client): + resp = dag_test_client.post("paused?dag_id=example_bash_operator&is_paused=false", follow_redirects=True) + check_content_in_response("OK", resp) @pytest.fixture(scope="module") @@ -771,9 +759,3 @@ def test_get_logs_with_metadata_failure(dag_faker_client): ) check_content_not_in_response('"message":', resp) check_content_not_in_response('"metadata":', resp) - - -def test_refresh_failure_for_viewer(viewer_client): - # viewer role can't refresh - resp = viewer_client.post('refresh?dag_id=example_bash_operator') - check_content_in_response('Redirecting', resp, resp_code=302) diff --git a/tests/www/views/test_views_base.py b/tests/www/views/test_views_base.py index 84a1ea6b13231..537f661e588c5 100644 --- a/tests/www/views/test_views_base.py +++ b/tests/www/views/test_views_base.py @@ -58,7 +58,7 @@ def test_home(capture_templates, admin_client): val_state_color_mapping = ( 'const STATE_COLOR = {"failed": "red", ' '"null": "lightblue", "queued": "gray", ' - '"removed": "lightgrey", "running": "lime", ' + '"removed": "lightgrey", "restarting": "violet", "running": "lime", ' '"scheduled": "tan", "sensing": "lightseagreen", ' '"shutdown": "blue", "skipped": "pink", ' '"success": "green", "up_for_reschedule": "turquoise", ' diff --git a/tests/www/views/test_views_connection.py b/tests/www/views/test_views_connection.py index 557697729a1be..249bf2a28e8f8 100644 --- a/tests/www/views/test_views_connection.py +++ b/tests/www/views/test_views_connection.py @@ -15,6 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import json from unittest import mock import pytest @@ -56,6 +57,57 @@ def test_prefill_form_null_extra(): cmv.prefill_form(form=mock_form, pk=1) +def test_process_form_extras(): + """ + Test the handling of connection parameters set with the classic `Extra` field as well as custom fields. + """ + + # Testing parameters set in both `Extra` and custom fields. + mock_form = mock.Mock() + mock_form.data = { + "conn_type": "test", + "conn_id": "extras_test", + "extra": '{"param1": "param1_val"}', + "extra__test__custom_field": "custom_field_val", + } + + cmv = ConnectionModelView() + cmv.extra_fields = ["extra__test__custom_field"] # Custom field + cmv.process_form(form=mock_form, is_created=True) + + assert json.loads(mock_form.extra.data) == { + "extra__test__custom_field": "custom_field_val", + "param1": "param1_val", + } + + # Testing parameters set in `Extra` field only. + mock_form = mock.Mock() + mock_form.data = { + "conn_type": "test2", + "conn_id": "extras_test2", + "extra": '{"param2": "param2_val"}', + } + + cmv = ConnectionModelView() + cmv.process_form(form=mock_form, is_created=True) + + assert json.loads(mock_form.extra.data) == {"param2": "param2_val"} + + # Testing parameters set in custom fields only. + mock_form = mock.Mock() + mock_form.data = { + "conn_type": "test3", + "conn_id": "extras_test3", + "extra__test3__custom_field": "custom_field_val3", + } + + cmv = ConnectionModelView() + cmv.extra_fields = ["extra__test3__custom_field"] # Custom field + cmv.process_form(form=mock_form, is_created=True) + + assert json.loads(mock_form.extra.data) == {"extra__test3__custom_field": "custom_field_val3"} + + def test_duplicate_connection(admin_client): """Test Duplicate multiple connection with suffix""" conn1 = Connection( diff --git a/tests/www/views/test_views_dagrun.py b/tests/www/views/test_views_dagrun.py index 94e44395a0cf2..a57a59528feb4 100644 --- a/tests/www/views/test_views_dagrun.py +++ b/tests/www/views/test_views_dagrun.py @@ -15,14 +15,11 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -import json - import pytest from airflow.models import DagBag, DagRun, TaskInstance from airflow.utils import timezone from airflow.utils.session import create_session -from tests.test_utils.config import conf_vars from tests.test_utils.www import check_content_in_response @@ -45,97 +42,6 @@ def reset_dagrun(): session.query(TaskInstance).delete() -@pytest.mark.parametrize( - "post_date, expected", - [ - ("2018-07-06 05:04:03Z", timezone.datetime(2018, 7, 6, 5, 4, 3)), - ("2018-07-06 05:04:03-04:00", timezone.datetime(2018, 7, 6, 9, 4, 3)), - ("2018-07-06 05:04:03-08:00", timezone.datetime(2018, 7, 6, 13, 4, 3)), - ("2018-07-06 05:04:03", timezone.datetime(2018, 7, 6, 5, 4, 3)), - ], - ids=["UTC", "EDT", "PST", "naive"], -) -def test_create_dagrun(session, admin_client, post_date, expected): - data = { - "state": "running", - "dag_id": "example_bash_operator", - "execution_date": post_date, - "run_id": "test_create_dagrun", - } - resp = admin_client.post('/dagrun/add', data=data, follow_redirects=True) - check_content_in_response('Added Row', resp) - - dr = session.query(DagRun).one() - - assert dr.execution_date == expected - - -@conf_vars({("core", "default_timezone"): "America/Toronto"}) -def test_create_dagrun_without_timezone_default(session, admin_client): - data = { - "state": "running", - "dag_id": "example_bash_operator", - "execution_date": "2018-07-06 05:04:03", - "run_id": "test_create_dagrun", - } - resp = admin_client.post('/dagrun/add', data=data, follow_redirects=True) - check_content_in_response('Added Row', resp) - - dr = session.query(DagRun).one() - - assert dr.execution_date == timezone.datetime(2018, 7, 6, 9, 4, 3) - - -def test_create_dagrun_valid_conf(session, admin_client): - conf_value = dict(Valid=True) - data = { - "state": "running", - "dag_id": "example_bash_operator", - "execution_date": "2018-07-06 05:05:03-02:00", - "run_id": "test_create_dagrun_valid_conf", - "conf": json.dumps(conf_value), - } - - resp = admin_client.post('/dagrun/add', data=data, follow_redirects=True) - check_content_in_response('Added Row', resp) - dr = session.query(DagRun).one() - assert dr.conf == conf_value - - -def test_create_dagrun_invalid_conf(session, admin_client): - data = { - "state": "running", - "dag_id": "example_bash_operator", - "execution_date": "2018-07-06 05:06:03", - "run_id": "test_create_dagrun_invalid_conf", - "conf": "INVALID: [JSON", - } - - resp = admin_client.post('/dagrun/add', data=data, follow_redirects=True) - check_content_in_response('JSON Validation Error:', resp) - dr = session.query(DagRun).all() - assert not dr - - -def test_list_dagrun_includes_conf(session, admin_client): - data = { - "state": "running", - "dag_id": "example_bash_operator", - "execution_date": "2018-07-06 05:06:03", - "run_id": "test_list_dagrun_includes_conf", - "conf": '{"include": "me"}', - } - admin_client.post('/dagrun/add', data=data, follow_redirects=True) - dr = session.query(DagRun).one() - - expect_date = timezone.convert_to_utc(timezone.datetime(2018, 7, 6, 5, 6, 3)) - assert dr.execution_date == expect_date - assert dr.conf == {"include": "me"} - - resp = admin_client.get('/dagrun/list', follow_redirects=True) - check_content_in_response("{"include": "me"}", resp) - - @pytest.fixture() def running_dag_run(session): dag = DagBag().get_dag("example_bash_operator") diff --git a/tests/www/views/test_views_tasks.py b/tests/www/views/test_views_tasks.py index cf9b95bc3e53d..04db6fcdf71e3 100644 --- a/tests/www/views/test_views_tasks.py +++ b/tests/www/views/test_views_tasks.py @@ -502,18 +502,6 @@ def test_run_with_not_runnable_states(_, admin_client, session, state): assert re.search(msg, resp.get_data(as_text=True)) -def test_refresh(admin_client): - resp = admin_client.post('refresh?dag_id=example_bash_operator') - check_content_in_response('', resp, resp_code=302) - - -def test_refresh_all(app, admin_client): - with unittest.mock.patch.object(app.dag_bag, 'collect_dags_from_db') as collect_dags_from_db: - resp = admin_client.post("/refresh_all", follow_redirects=True) - check_content_in_response('', resp) - collect_dags_from_db.assert_called_once_with() - - @pytest.fixture() def new_id_example_bash_operator(): dag_id = 'example_bash_operator'