Fixes dbt retry does not respect --threads (#10591) (#11447)

Co-authored-by: Apoorv Mehrotra <[email protected]>

Author: dbeatty10

.changes/unreleased/Fixes-20240822-122132.yaml

@@ -0,0 +1,6 @@
+kind: Fixes
+body: dbt retry does not respect --threads
+time: 2024-08-22T12:21:32.358066+05:30
+custom:
+    Author: donjin-master
+    Issue: "10584"

core/dbt/task/retry.py

@@ -42,7 +42,7 @@
     "warn_error",
 }
 
-ALLOW_CLI_OVERRIDE_FLAGS = {"vars"}
+ALLOW_CLI_OVERRIDE_FLAGS = {"vars", "threads"}
 
 TASK_DICT = {
     "build": BuildTask,

tests/functional/retry/fixtures.py

@@ -1,5 +1,6 @@
 models__sample_model = """select 1 as id, baz as foo"""
 models__second_model = """select 1 as id, 2 as bar"""
+models__thread_model = """select idx as id"""
 
 models__union_model = """
 select foo + bar as sum3 from {{ ref('sample_model') }}
@@ -58,3 +59,13 @@
         data_tests:
           - not_null
 """
+
+schema_test_thread_yml = """
+models:
+  - name: thread_model
+    columns:
+      - name: id
+        data_tests:
+          - not_null
+
+"""

tests/functional/retry/test_retry_threads.py

@@ -0,0 +1,56 @@
+import pytest
+
+from dbt.contracts.results import RunStatus, TestStatus
+from dbt.tests.util import run_dbt, write_file
+from tests.functional.retry.fixtures import models__thread_model, schema_test_thread_yml
+
+
+class TestCustomThreadRetry:
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "thread_model.sql": models__thread_model,
+            "schema.yml": schema_test_thread_yml,
+        }
+
+    def test_thread_target(self, project):
+        # Passing Threads to check
+        results = run_dbt(
+            ["build", "--select", "thread_model", "--threads", "3"], expect_pass=False
+        )
+        expected_statuses = {
+            "thread_model": RunStatus.Error,
+            "not_null_thread_model_id": TestStatus.Skipped,
+        }
+        assert {n.node.name: n.status for n in results.results} == expected_statuses
+
+        # Retry Running the Dbt with simple Retry
+        results = run_dbt(["retry", "--threads", "2"], expect_pass=False)
+        expected_statuses = {
+            "thread_model": RunStatus.Error,
+            "not_null_thread_model_id": TestStatus.Skipped,
+        }
+        assert {n.node.name: n.status for n in results.results} == expected_statuses
+        assert results.args["threads"] == 2
+
+        # running with retry withour threads
+        results = run_dbt(["retry"], expect_pass=False)
+        expected_statuses = {
+            "thread_model": RunStatus.Error,
+            "not_null_thread_model_id": TestStatus.Skipped,
+        }
+        assert {n.node.name: n.status for n in results.results} == expected_statuses
+        assert results.args["threads"] == 2
+
+        # Retry with fixing the model and running with --threads 1
+        fixed_sql = "select 1 as id"
+        write_file(fixed_sql, "models", "thread_model.sql")
+
+        results = run_dbt(["retry", "--threads", "1"])
+        expected_statuses = {
+            "thread_model": RunStatus.Success,
+            "not_null_thread_model_id": TestStatus.Pass,
+        }
+
+        assert {n.node.name: n.status for n in results.results} == expected_statuses
+        assert results.args["threads"] == 1