Skip to content

Commit

Permalink
add unpause option
Browse files Browse the repository at this point in the history
  • Loading branch information
astro-anand committed Apr 25, 2024
1 parent c28cec3 commit 7d93914
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions astronomer_starship/compat/starship_hook.py
Original file line number Diff line number Diff line change
Expand Up @@ -165,6 +165,7 @@ def __init__(
target_auth: Union[tuple, list] = None,
source_headers: dict = None,
target_headers: dict = None,
unpause_dags_in_target=False,
logger_name: str | None = None,
):
super().__init__(logger_name)
Expand All @@ -175,6 +176,7 @@ def __init__(
self.target_api_hook = StarshipAPIHook(
webserver_url=target_webserver_url, auth=target_auth, headers=target_headers
)
self.unpause_dags_in_target = unpause_dags_in_target

def load_dagruns_to_target(
self,
Expand All @@ -200,6 +202,9 @@ def load_dagruns_to_target(
self.get_and_set_dagruns(dag_id)
self.get_and_set_task_instances(dag_id)

if self.unpause_dags_in_target:
self.target_api_hook.set_dag_state(dag_id=dag_id, action="unpause")

def get_and_set_dagruns(self, dag_id: str) -> None:
dag_runs = self.source_api_hook.get_dag_runs(
dag_id=dag_id,
Expand Down

0 comments on commit 7d93914

Please sign in to comment.