You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "subbota19 (via GitHub)" <gi...@apache.org> on 2025/02/12 11:29:04 UTC

[PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

subbota19 opened a new pull request, #46672:
URL: https://github.com/apache/airflow/pull/46672

   <!--
    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.
    -->
   
   <!--
   Thank you for contributing! Please make sure that your code changes
   are covered with tests. And in case of new features or big changes
   remember to adjust the documentation.
   
   Feel free to ping committers for the review!
   
   In case of an existing issue, reference it using one of the following:
   
   closes: #ISSUE
   related: #ISSUE
   
   How to write a good git commit message:
   http://chris.beams.io/posts/git-commit/
   -->
   
   <!-- Please keep an empty line above the dashes. -->
   ---
   
   Fix https://github.com/apache/airflow/issues/46648.
   
   **Problem Description**
   In deferable mode, Airflow uses triggers to check the job status, so every time the deferable interval is reached, it checks the status and yields an event on [this line](https://github.com/apache/airflow/blob/main/providers/snowflake/src/airflow/providers/snowflake/triggers/snowflake_trigger.py#L89).
   
   However, when not using deferable logic, the execution follows the [else statement](https://github.com/apache/airflow/blob/main/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py#L456), which does not include a while loop to continuously monitor execution. Instead, on [this line](https://github.com/apache/airflow/blob/main/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py#L457), the task iterates through active Snowflake jobs, checks their status and then sleeps for _poll_interval_.
   
   This results in the following behavior:
   
   ✅ If a job finishes in less than 5 seconds ( _time.sleep_ using the default _poll_interval_), the logic works correctly.
   ❌ If a job takes longer than 5 seconds, there is no mechanism to keep monitoring it, and the task incorrectly finishes after the loop ends - even if the Snowflake job is still running.
   
   **Suggested Fix**
   I propose adding a running status check inside **poll_on_queries** and modifying the else statement with a loop to ensure Airflow waits until all queries finish, the logic is quick similar to what we use in [deferable trigger while loop](https://github.com/subbota19/airflow/blob/main/providers/snowflake/src/airflow/providers/snowflake/triggers/snowflake_trigger.py#L79) 
   
   With this change, the Airflow task would periodically call **poll_on_queries** and continue waiting until all jobs are complete.
   
   **Tests**
   Simulates a scenario where the Snowflake job is running for two checks (10 seconds apart, assuming default _poll_interval_): ffter the second check, the job is completed successfully and task is marked as successful.
   
   Simulates a scenario where the Snowflake job failed and ensures AirflowException is raised.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

Posted by "subbota19 (via GitHub)" <gi...@apache.org>.
subbota19 closed pull request #46672: Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active
URL: https://github.com/apache/airflow/pull/46672


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

Posted by "subbota19 (via GitHub)" <gi...@apache.org>.
subbota19 commented on code in PR #46672:
URL: https://github.com/apache/airflow/pull/46672#discussion_r1976704960


##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py:
##########
@@ -455,16 +455,21 @@ def execute(self, context: Context) -> None:
                 method_name="execute_complete",
             )
         else:
-            statement_status = self.poll_on_queries()
-            if statement_status["error"]:
-                raise AirflowException(statement_status["error"])
+            while True:
+                statement_status = self.poll_on_queries()

Review Comment:
   Hi @eladkal !
   Thanks for the suggestion. Let me add one point - inside the [poll_on_queries](https://github.com/apache/airflow/blob/main/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py#L463) method we actually do [sleep](https://github.com/apache/airflow/blob/main/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py#L482) for _poll_interval_(by default it's 5 second), so I believe the while loop already has a mechanism to handle throttling issues.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

Posted by "boring-cyborg[bot] (via GitHub)" <gi...@apache.org>.
boring-cyborg[bot] commented on PR #46672:
URL: https://github.com/apache/airflow/pull/46672#issuecomment-2704665430

   Awesome work, congrats on your first merged pull request! You are invited to check our [Issue Tracker](https://github.com/apache/airflow/issues) for additional contributions.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #46672:
URL: https://github.com/apache/airflow/pull/46672#discussion_r1976620434


##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py:
##########
@@ -455,16 +455,21 @@ def execute(self, context: Context) -> None:
                 method_name="execute_complete",
             )
         else:
-            statement_status = self.poll_on_queries()
-            if statement_status["error"]:
-                raise AirflowException(statement_status["error"])
+            while True:
+                statement_status = self.poll_on_queries()

Review Comment:
   This is endless call to the Snowflake API without any waits between call.
   This can result in throttling issues isn't it?
   
   similar but not exactly the same issue with Glue was handled with sleep https://github.com/apache/airflow/pull/43688



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

Posted by "subbota19 (via GitHub)" <gi...@apache.org>.
subbota19 commented on PR #46672:
URL: https://github.com/apache/airflow/pull/46672#issuecomment-2692052148

   > > Hi @potiuk! I've resolved the conflict on my end and added suggestions, could you please double-check it? Thanks!
   > 
   > @subbota19 , take changes from main. test failures are already fixed.
   
   Hi @rawwar, thx, I've rebased to the main branch


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

Posted by "boring-cyborg[bot] (via GitHub)" <gi...@apache.org>.
boring-cyborg[bot] commented on PR #46672:
URL: https://github.com/apache/airflow/pull/46672#issuecomment-2653439169

   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]( https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst#prerequisites-for-pre-commit-hooks) 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](https://github.com/apache/airflow/blob/main/docs/apache-airflow/howto/custom-operator.rst) Consider adding an example DAG that shows how users should use it.
   - Consider using [Breeze environment](https://github.com/apache/airflow/blob/main/dev/breeze/doc/README.rst) 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](https://www.apache.org/foundation/policies/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]( https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#coding-style-and-best-practices).
   - 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
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk merged PR #46672:
URL: https://github.com/apache/airflow/pull/46672


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

Posted by "subbota19 (via GitHub)" <gi...@apache.org>.
subbota19 commented on PR #46672:
URL: https://github.com/apache/airflow/pull/46672#issuecomment-2690221043

   Hi @potiuk!
   I've resolved the conflict on my end and added suggestions, could you please double-check it? Thanks!


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

Posted by "rawwar (via GitHub)" <gi...@apache.org>.
rawwar commented on code in PR #46672:
URL: https://github.com/apache/airflow/pull/46672#discussion_r1955604980


##########
providers/snowflake/tests/provider_tests/snowflake/operators/test_snowflake.py:
##########
@@ -391,3 +391,66 @@ def test_snowflake_sql_api_execute_operator_running_before_defer(
         operator.execute(create_context(operator))
 
         assert mock_defer.called
+
+    def test_snowflake_sql_api_execute_operator_polling_running(
+        self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output
+    ):
+        """
+        Tests that the execute method correctly loops and waits until all queries complete
+        when `deferrable=False`

Review Comment:
   ```suggestion
           when ``deferrable=False``
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

Posted by "subbota19 (via GitHub)" <gi...@apache.org>.
subbota19 commented on code in PR #46672:
URL: https://github.com/apache/airflow/pull/46672#discussion_r1955909908


##########
providers/snowflake/tests/provider_tests/snowflake/operators/test_snowflake.py:
##########
@@ -391,3 +391,66 @@ def test_snowflake_sql_api_execute_operator_running_before_defer(
         operator.execute(create_context(operator))
 
         assert mock_defer.called
+
+    def test_snowflake_sql_api_execute_operator_polling_running(
+        self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output
+    ):
+        """
+        Tests that the execute method correctly loops and waits until all queries complete
+        when `deferrable=False`
+        """
+        operator = SnowflakeSqlApiOperator(
+            task_id=TASK_ID,
+            snowflake_conn_id=CONN_ID,
+            sql=SQL_MULTIPLE_STMTS,
+            statement_count=4,
+            do_xcom_push=False,
+            deferrable=False,
+        )
+
+        mock_execute_query.return_value = ["uuid1"]
+
+        mock_get_sql_api_query_status.side_effect = [
+            # Initial get_sql_api_query_status check
+            {"status": "running"},
+            # 1st poll_on_queries check (poll_interval: 5s)
+            {"status": "running"},
+            # 2nd poll_on_queries check (poll_interval: 5s)
+            {"status": "running"},
+            # 3rd poll_on_queries check (poll_interval: 5s)
+            {"status": "success"},
+        ]
+
+        with mock.patch("time.sleep") as mock_sleep:
+            operator.execute(context=None)
+            mock_check_query_output.assert_called_once_with(["uuid1"])
+            assert mock_sleep.call_count == 3
+
+    def test_snowflake_sql_api_execute_operator_polling_failed(
+        self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output
+    ):
+        """
+        Tests that the execute method raises AirflowException if any query fails during polling
+        when `deferrable=False`

Review Comment:
   Hi @rawwar!
   Updated. I was wondering if it's possible to use a pre-commit hook to handle such docstring requirements. It could be beneficial to add something like this.



##########
providers/snowflake/tests/provider_tests/snowflake/operators/test_snowflake.py:
##########
@@ -391,3 +391,66 @@ def test_snowflake_sql_api_execute_operator_running_before_defer(
         operator.execute(create_context(operator))
 
         assert mock_defer.called
+
+    def test_snowflake_sql_api_execute_operator_polling_running(
+        self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output
+    ):
+        """
+        Tests that the execute method correctly loops and waits until all queries complete
+        when `deferrable=False`

Review Comment:
   Updated.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #46672:
URL: https://github.com/apache/airflow/pull/46672#issuecomment-2664174646

   You need to resolve conflicts now


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

Posted by "rawwar (via GitHub)" <gi...@apache.org>.
rawwar commented on PR #46672:
URL: https://github.com/apache/airflow/pull/46672#issuecomment-2691881986

   > Hi @potiuk! I've resolved the conflict on my end and added suggestions, could you please double-check it? Thanks!
   
   @subbota19 , take changes from main. CI failures are already fixed.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

Posted by "subbota19 (via GitHub)" <gi...@apache.org>.
subbota19 commented on PR #46672:
URL: https://github.com/apache/airflow/pull/46672#issuecomment-2703201538

   Hi @eladkal @potiuk !
   I believe we can go ahead and merge this PR


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

Posted by "subbota19 (via GitHub)" <gi...@apache.org>.
subbota19 commented on PR #46672:
URL: https://github.com/apache/airflow/pull/46672#issuecomment-2666562973

   Hi @potiuk! 
   I've resolved the conflict on my end, could you please double-check it? Thanks!


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

Posted by "rawwar (via GitHub)" <gi...@apache.org>.
rawwar commented on PR #46672:
URL: https://github.com/apache/airflow/pull/46672#issuecomment-2678713479

   > Hi @rawwar ! Could you please review this PR again and resolve the remaining conversations? Thanks!
   
   I think you should be able to mark the conversations as resolved. If not, a committer will be able to help. I don't have permission to mark the conversation as resolved.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

Posted by "subbota19 (via GitHub)" <gi...@apache.org>.
subbota19 commented on PR #46672:
URL: https://github.com/apache/airflow/pull/46672#issuecomment-2678049935

   Hi @rawwar !
   Could you please review this PR again and resolve the remaining conversations? Thanks!


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix: SnowflakeSqlApiOperator(without deferable flag) incorrectly marks the task as successful even when jobs are still active [airflow]

Posted by "rawwar (via GitHub)" <gi...@apache.org>.
rawwar commented on code in PR #46672:
URL: https://github.com/apache/airflow/pull/46672#discussion_r1955604794


##########
providers/snowflake/tests/provider_tests/snowflake/operators/test_snowflake.py:
##########
@@ -391,3 +391,66 @@ def test_snowflake_sql_api_execute_operator_running_before_defer(
         operator.execute(create_context(operator))
 
         assert mock_defer.called
+
+    def test_snowflake_sql_api_execute_operator_polling_running(
+        self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output
+    ):
+        """
+        Tests that the execute method correctly loops and waits until all queries complete
+        when `deferrable=False`
+        """
+        operator = SnowflakeSqlApiOperator(
+            task_id=TASK_ID,
+            snowflake_conn_id=CONN_ID,
+            sql=SQL_MULTIPLE_STMTS,
+            statement_count=4,
+            do_xcom_push=False,
+            deferrable=False,
+        )
+
+        mock_execute_query.return_value = ["uuid1"]
+
+        mock_get_sql_api_query_status.side_effect = [
+            # Initial get_sql_api_query_status check
+            {"status": "running"},
+            # 1st poll_on_queries check (poll_interval: 5s)
+            {"status": "running"},
+            # 2nd poll_on_queries check (poll_interval: 5s)
+            {"status": "running"},
+            # 3rd poll_on_queries check (poll_interval: 5s)
+            {"status": "success"},
+        ]
+
+        with mock.patch("time.sleep") as mock_sleep:
+            operator.execute(context=None)
+            mock_check_query_output.assert_called_once_with(["uuid1"])
+            assert mock_sleep.call_count == 3
+
+    def test_snowflake_sql_api_execute_operator_polling_failed(
+        self, mock_execute_query, mock_get_sql_api_query_status, mock_check_query_output
+    ):
+        """
+        Tests that the execute method raises AirflowException if any query fails during polling
+        when `deferrable=False`

Review Comment:
   ```suggestion
           when ``deferrable=False``
   ```
   we need to use 2 backticks on either sides for inline literals. refer to any other method's docstring



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org