From 3464be7f700911e7897d036ed9f0127d8a672bae Mon Sep 17 00:00:00 2001
From: Apoorv Mehrotra <apoorvmehrotra3@gmail.com>
Date: Wed, 27 Nov 2024 00:51:46 +0530
Subject: [PATCH] Fixes dbt retry does not respect --threads (#10591)

---
 .../unreleased/Fixes-20240822-122132.yaml     |  6 ++
 core/dbt/task/retry.py                        |  2 +-
 tests/functional/retry/fixtures.py            | 11 ++++
 tests/functional/retry/test_retry_threads.py  | 56 +++++++++++++++++++
 4 files changed, 74 insertions(+), 1 deletion(-)
 create mode 100644 .changes/unreleased/Fixes-20240822-122132.yaml
 create mode 100644 tests/functional/retry/test_retry_threads.py

diff --git a/.changes/unreleased/Fixes-20240822-122132.yaml b/.changes/unreleased/Fixes-20240822-122132.yaml
new file mode 100644
index 00000000000..d169520ea99
--- /dev/null
+++ b/.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"
diff --git a/core/dbt/task/retry.py b/core/dbt/task/retry.py
index 9b3c3874718..fd2b4362c94 100644
--- a/core/dbt/task/retry.py
+++ b/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,
diff --git a/tests/functional/retry/fixtures.py b/tests/functional/retry/fixtures.py
index 64adf9c684f..52c9d8afdfc 100644
--- a/tests/functional/retry/fixtures.py
+++ b/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
+
+"""
diff --git a/tests/functional/retry/test_retry_threads.py b/tests/functional/retry/test_retry_threads.py
new file mode 100644
index 00000000000..8e9fae8135f
--- /dev/null
+++ b/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