-
Notifications
You must be signed in to change notification settings - Fork 1.6k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
9 changed files
with
411 additions
and
4 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,6 @@ | ||
kind: Features | ||
body: dbt retry | ||
time: 2023-06-02T08:33:02.410456-07:00 | ||
custom: | ||
Author: stu-k aranke | ||
Issue: "7299" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,113 @@ | ||
from pathlib import Path | ||
|
||
from dbt.cli.flags import Flags | ||
from dbt.cli.types import Command as CliCommand | ||
from dbt.config import RuntimeConfig | ||
from dbt.contracts.results import NodeStatus | ||
from dbt.contracts.state import PreviousState | ||
from dbt.exceptions import DbtRuntimeError | ||
from dbt.graph import GraphQueue | ||
from dbt.task.base import ConfiguredTask | ||
from dbt.task.build import BuildTask | ||
from dbt.task.compile import CompileTask | ||
from dbt.task.generate import GenerateTask | ||
from dbt.task.run import RunTask | ||
from dbt.task.run_operation import RunOperationTask | ||
from dbt.task.seed import SeedTask | ||
from dbt.task.snapshot import SnapshotTask | ||
from dbt.task.test import TestTask | ||
|
||
RETRYABLE_STATUSES = {NodeStatus.Error, NodeStatus.Fail, NodeStatus.Skipped, NodeStatus.RuntimeErr} | ||
|
||
TASK_DICT = { | ||
"build": BuildTask, | ||
"compile": CompileTask, | ||
"generate": GenerateTask, | ||
"seed": SeedTask, | ||
"snapshot": SnapshotTask, | ||
"test": TestTask, | ||
"run": RunTask, | ||
"run-operation": RunOperationTask, | ||
} | ||
|
||
CMD_DICT = { | ||
"build": CliCommand.BUILD, | ||
"compile": CliCommand.COMPILE, | ||
"generate": CliCommand.DOCS_GENERATE, | ||
"seed": CliCommand.SEED, | ||
"snapshot": CliCommand.SNAPSHOT, | ||
"test": CliCommand.TEST, | ||
"run": CliCommand.RUN, | ||
"run-operation": CliCommand.RUN_OPERATION, | ||
} | ||
|
||
|
||
class RetryTask(ConfiguredTask): | ||
def __init__(self, args, config, manifest): | ||
super().__init__(args, config, manifest) | ||
|
||
state_path = self.args.state or self.config.target_path | ||
|
||
if self.args.warn_error: | ||
RETRYABLE_STATUSES.add(NodeStatus.Warn) | ||
|
||
self.previous_state = PreviousState( | ||
state_path=Path(state_path), | ||
target_path=Path(self.config.target_path), | ||
project_root=Path(self.config.project_root), | ||
) | ||
|
||
if not self.previous_state.results: | ||
raise DbtRuntimeError( | ||
f"Could not find previous run in '{state_path}' target directory" | ||
) | ||
|
||
self.previous_args = self.previous_state.results.args | ||
self.previous_command_name = self.previous_args.get("which") | ||
self.task_class = TASK_DICT.get(self.previous_command_name) | ||
|
||
def run(self): | ||
unique_ids = set( | ||
[ | ||
result.unique_id | ||
for result in self.previous_state.results.results | ||
if result.status in RETRYABLE_STATUSES | ||
] | ||
) | ||
|
||
cli_command = CMD_DICT.get(self.previous_command_name) | ||
|
||
# Remove these args when their default values are present, otherwise they'll raise an exception | ||
args_to_remove = { | ||
"show": lambda x: True, | ||
"resource_types": lambda x: x == [], | ||
"warn_error_options": lambda x: x == {"exclude": [], "include": []}, | ||
} | ||
|
||
for k, v in args_to_remove.items(): | ||
if k in self.previous_args and v(self.previous_args[k]): | ||
del self.previous_args[k] | ||
|
||
retry_flags = Flags.from_dict(cli_command, self.previous_args) | ||
retry_config = RuntimeConfig.from_args(args=retry_flags) | ||
|
||
class TaskWrapper(self.task_class): | ||
def get_graph_queue(self): | ||
new_graph = self.graph.get_subset_graph(unique_ids) | ||
return GraphQueue( | ||
new_graph.graph, | ||
self.manifest, | ||
unique_ids, | ||
) | ||
|
||
task = TaskWrapper( | ||
retry_flags, | ||
retry_config, | ||
self.manifest, | ||
) | ||
|
||
return_value = task.run() | ||
return return_value | ||
|
||
def interpret_results(self, *args, **kwargs): | ||
return self.task_class.interpret_results(*args, **kwargs) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,47 @@ | ||
models__sample_model = """select 1 as id, baz as foo""" | ||
models__second_model = """select 1 as id, 2 as bar""" | ||
|
||
models__union_model = """ | ||
select foo + bar as sum3 from {{ ref('sample_model') }} | ||
left join {{ ref('second_model') }} on sample_model.id = second_model.id | ||
""" | ||
|
||
schema_yml = """ | ||
models: | ||
- name: sample_model | ||
columns: | ||
- name: foo | ||
tests: | ||
- accepted_values: | ||
values: [3] | ||
quote: false | ||
config: | ||
severity: warn | ||
- name: second_model | ||
columns: | ||
- name: bar | ||
tests: | ||
- accepted_values: | ||
values: [3] | ||
quote: false | ||
config: | ||
severity: warn | ||
- name: union_model | ||
columns: | ||
- name: sum3 | ||
tests: | ||
- accepted_values: | ||
values: [3] | ||
quote: false | ||
""" | ||
|
||
macros__alter_timezone_sql = """ | ||
{% macro alter_timezone(timezone='America/Los_Angeles') %} | ||
{% set sql %} | ||
SET TimeZone='{{ timezone }}'; | ||
{% endset %} | ||
{% do run_query(sql) %} | ||
{% do log("Timezone set to: " + timezone, info=True) %} | ||
{% endmacro %} | ||
""" |
Oops, something went wrong.