Fix RedshiftDataOperator not running in deferred mode when it should

[borismo]

Problems

• When a RedshiftDataOperator task is configured with deferrable=True and wait_for_completion=True (the default), it doesn't go in deferred state. Instead it stays in running state until the statement completes.
• Also, if wait_for_completion=False and deferrable=True, after the statement is submitted, the task will still go into deferred mode and wait for the statement to complete.

Reasons

• Currently, if deferrable=True, self.wait_for_completion is set to False in execute(), but never used after.
• execute does not check whether the task should wait for completion, only if it should be deferred.

Solution

• Overwrite wait_for_completion instead of self.wait_for_completion when deferrable. Also, remove redundant condition on self.wait_for_completion
• Before going into deferrable mode, also check that the task is supposed to wait for completion.

How I tested

Checked with this simple DAG that the operator now behaves as expected, for all 4 combinations:

from airflow.decorators import dag
from airflow.providers.amazon.aws.operators.redshift_data import RedshiftDataOperator

@dag(
    "Foo",
)
def _():
    for task_id, config in {
        "wait_defer": {"deferrable": True, "wait_for_completion": True},
        "wait_no_defer": {"deferrable": False, "wait_for_completion": True},
        "no_wait_defer": {"deferrable": True, "wait_for_completion": False},
        "no_wait_no_defer": {"deferrable": False, "wait_for_completion": False},
    }.items():
        RedshiftDataOperator(
            task_id=task_id,
            aws_conn_id="redshift_data",
            cluster_identifier="data-warehouse",
            db_user="airflow",
            database="bar",
            sql="""CREATE TEMPORARY TABLE tmp_foo AS
            SELECT *
            FROM some.big_table
            LIMIT 10;""",
            deferrable=config["deferrable"],
            wait_for_completion=config["wait_for_completion"],
        )

_()

---

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

[boring-cyborg]

Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide (https://github.com/apache/airflow/blob/main/contributing-docs/README.rst)
Here are some useful points:

• Pay attention to the quality of your code (ruff, mypy and type annotations). Our pre-commits will help you with that.
• In case of a new feature add useful documentation (in docstrings or in docs/ directory). Adding a new operator? Check this short guide Consider adding an example DAG that shows how users should use it.
• Consider using Breeze environment for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
• Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
• Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
• Be sure to read the Airflow Coding style.
• Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
  Apache Airflow is a community-driven project and together we are making it better 🚀.
  In case of doubts contact the developers at:
  Mailing List: dev@airflow.apache.org
  Slack: https://s.apache.org/airflow-slack

[eladkal]

Wasn't this already fixed in #41191 ?

[borismo]

At first I thought too, but it's a different operator. RedshiftDataOperator vs. RedshiftCreateClusterOperator. module: airflow/providers/amazon/aws/operators/redshift_cluster.py vs. airflow/providers/amazon/aws/operators/redshift_data.py.

[eladkal]

Ah OK.
Can you please add unit test to avoid regression?

[vincbeck]

Agree with Elad, a unit test would be great, but otherwise LGTM

[borismo]

Here's what I did:

• tweaked the existing deferrable_operator fixture so it waits for completion. The test_execute_defer test was failing with my fix. I think this test should use an operator configured to wait.
• added a test which verifies that if we don't want to wait for completion, the operator should not check whether the statement has completed and also shouldn't enter deferred mode.

Not sure if it's a good idea to iterate over deferrable. I did so to keep it DRY, since if separate, the 2 two tests would be almost identical. And I intend the test to really emphasize on the fact that the deferrable value doesn't matter.

Warning

I realize now that this fix will change the operator's behavior: for folks who have been using it with deferrable = True and wait_for_completion = False to make it go into deferrable mode (a workaround IMO), the task instance will now immediately be marked as success without waiting. Could break downstream tasks that expect the statement to be completed.

[vincbeck]

I realize now that this fix will change the operator's behavior: for folks who have been using it with deferrable = True and wait_for_completion = False to make it go into deferrable mode (a workaround IMO), the task instance will now immediately be marked as success without waiting. Could break downstream tasks that expect the statement to be completed.

You are correct but this is a bug fix. The previous behavior was wrong so to me we should go ahead with that change

[eladkal]

Can you add to the provider changelog a block of

Main
......

.. warning::

to the changelog where you explain how to mitigate the change. It needs to have 2-4 sentenses that just users would understand what was change and how to mitigate this.

Example:

https://github.com/apache/airflow/blob/main/airflow/providers/amazon/CHANGELOG.rst#870

[boring-cyborg]

Awesome work, congrats on your first merged pull request! You are invited to check our Issue Tracker for additional contributions.