Skip to content

Commit 0acac85

Browse files
bradheCopilot
andauthored
Add some new utility and info functions to Tower SDK (#168)
* Add some new utility and info functions to Tower SDK * chore: Add run_number() helper function * Shoudl be team, not team, in the info module * Fix linting errors * Update src/tower/info/__init__.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
1 parent 8617d78 commit 0acac85

3 files changed

Lines changed: 244 additions & 0 deletions

File tree

src/tower/__init__.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,25 @@
1919
wait_for_runs,
2020
)
2121

22+
from ._utils import (
23+
param,
24+
parameter,
25+
secret,
26+
)
27+
28+
2229
from ._features import override_get_attr, get_available_features, is_feature_enabled
2330

2431
if TYPE_CHECKING:
2532
from ._tables import tables as tables
2633
from ._llms import llms as llms
2734
from ._dbt import dbt as dbt
2835

36+
#
37+
# Sub-packages to expose
38+
#
39+
from . import info
40+
2941

3042
def __getattr__(name: str):
3143
"""

src/tower/_utils.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
def secret(name: str, default: str = ""):
2+
"""
3+
Retrieve a secret value from environment variables.
4+
5+
Args:
6+
name (str): The name of the secret (environment variable).
7+
default (str, optional): The default value to return if the secret is not found. Defaults to "".
8+
9+
Returns:
10+
str: The value of the secret or the default value.
11+
"""
12+
import os
13+
14+
return os.getenv(name, default)
15+
16+
17+
def param(name: str, default: str = ""):
18+
"""
19+
Retrieve a parameter value from this invocation. In this implementation,
20+
it fetches the value from environment variables.
21+
22+
Args:
23+
name (str): The name of the parameter (environment variable).
24+
default (str, optional): The default value to return if the parameter is not found. Defaults to "".
25+
26+
Returns:
27+
str: The value of the parameter or the default value.
28+
"""
29+
import os
30+
31+
return os.getenv(name, default)
32+
33+
34+
def parameter(name: str, default: str = ""):
35+
"""
36+
Alias for param function to retrieve a parameter value from environment variables.
37+
38+
Args:
39+
name (str): The name of the parameter (environment variable).
40+
default (str, optional): The default value to return if the parameter is not found. Defaults to "".
41+
42+
Returns:
43+
str: The value of the parameter or the default value.
44+
"""
45+
return param(name, default)

src/tower/info/__init__.py

Lines changed: 187 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,187 @@
1+
from typing import Optional
2+
3+
4+
def schedule_name() -> Optional[str]:
5+
"""
6+
Retrieve the name of the schedule that invoked this run from the runtime
7+
environment.
8+
9+
Returns:
10+
Optional[str]: The name of the schedule if set, otherwise None.
11+
"""
12+
return _get_runtime_env_variable("SCHEDULE_NAME", None)
13+
14+
15+
def is_scheduled_run() -> bool:
16+
"""
17+
Check if the current run is a scheduled run based on environment variables.
18+
Returns:
19+
bool: True if it is a scheduled run, otherwise False.
20+
"""
21+
return schedule_name() is not None
22+
23+
24+
def schedule_id() -> Optional[str]:
25+
"""
26+
Retrieve the ID of the schedule that invoked this run from the runtime
27+
environment.
28+
29+
Returns:
30+
Optional[str]: The ID of the schedule if set, otherwise None.
31+
"""
32+
return _get_runtime_env_variable("SCHEDULE_ID", None)
33+
34+
35+
def run_id() -> Optional[str]:
36+
"""
37+
Retrieve the ID of the current run from the runtime environment.
38+
39+
Returns:
40+
Optional[str]: The ID of the run if set, otherwise None.
41+
"""
42+
return _get_runtime_env_variable("RUN_ID", None)
43+
44+
45+
def run_number() -> Optional[int]:
46+
"""
47+
Retrieve the number of the current run from the runtime environment.
48+
49+
Returns:
50+
Optional[int]: The run number if set, otherwise None.
51+
"""
52+
run_number_str = _get_runtime_env_variable("RUN_NUMBER", None)
53+
return int(run_number_str) if run_number_str is not None else None
54+
55+
56+
def hostname() -> Optional[str]:
57+
"""
58+
Retrieve the hostname of the current run assigned by the runtime
59+
environment.
60+
61+
Returns:
62+
Optional[str]: The hostname if set, otherwise None.
63+
"""
64+
return _get_env_variable("TOWER__HOST", None)
65+
66+
67+
def port() -> Optional[int]:
68+
"""
69+
Retrieve the port number assigned to this run by the current runtime
70+
environment.
71+
72+
Returns:
73+
Optional[int]: The port number if set, otherwise None.
74+
"""
75+
port_str = _get_env_variable("TOWER__PORT", None)
76+
return int(port_str) if port_str is not None else None
77+
78+
79+
def is_cloud_run() -> bool:
80+
"""
81+
Check if the current run is executing in the Tower cloud environment.
82+
83+
Returns:
84+
bool: True if running in Tower cloud, otherwise False.
85+
"""
86+
val = _get_runtime_env_variable("IS_TOWER_MANAGED", "")
87+
return _strtobool(val)
88+
89+
90+
def runner_name() -> str:
91+
"""Retrieve the name of the runner executing this run from the runtime. If the
92+
name is unknown, an empty string is returned.
93+
94+
Returns:
95+
str: The name of the runner or an empty string if unknown.
96+
"""
97+
return _get_runtime_env_variable("RUNNER_NAME", "")
98+
99+
100+
def runner_id() -> str:
101+
"""
102+
Retrieve the ID of the runner executing this run from the runtime. If the
103+
ID is unknown, an empty string is returned.
104+
105+
Returns:
106+
str: The ID of the runner or an empty string if unknown.
107+
"""
108+
return _get_runtime_env_variable("RUNNER_ID", "")
109+
110+
111+
def app_name() -> str:
112+
"""
113+
Retrieve the name of the app being executed in this run from the runtime.
114+
115+
Returns:
116+
str: The name of the app or an empty string if unknown.
117+
"""
118+
return _get_runtime_env_variable("APP_NAME")
119+
120+
121+
def team_name() -> str:
122+
"""
123+
Retrieve the name of the team associated with this run from the runtime.
124+
125+
Returns:
126+
str: The name of the team or an empty string if unknown.
127+
"""
128+
return _get_runtime_env_variable("TEAM_NAME", "")
129+
130+
131+
def environment() -> str:
132+
"""
133+
Retrieve the name of the environment this run is running in.
134+
135+
Returns:
136+
str: The name of the environment or an empty string if unknown.
137+
"""
138+
return _get_runtime_env_variable("ENVIRONMENT_NAME", "")
139+
140+
141+
def is_manual_run() -> bool:
142+
"""
143+
Check if the current run was manually triggered.
144+
145+
Returns:
146+
bool: True if the run was manually triggered, otherwise False.
147+
"""
148+
val = _get_runtime_env_variable("IS_MANUAL_RUN", "false")
149+
return _strtobool(val)
150+
151+
152+
def _get_runtime_env_variable(name: str, default: Optional[str] = "") -> Optional[str]:
153+
"""
154+
Helper function to retrieve a runtime environment variable.
155+
156+
Args:
157+
name (str): The name of the runtime environment variable.
158+
159+
Returns:
160+
Optional[str]: The value of the runtime environment variable if set, otherwise None.
161+
"""
162+
return _get_env_variable(f"TOWER__RUNTIME__{name}", default)
163+
164+
165+
def _get_env_variable(var_name: str, default: Optional[str] = "") -> Optional[str]:
166+
"""
167+
Helper function to retrieve an environment variable.
168+
169+
Args:
170+
var_name (str): The name of the environment variable.
171+
172+
Returns:
173+
Optional[str]: The value of the environment variable if set, otherwise None.
174+
"""
175+
import os
176+
177+
return os.getenv(var_name, default)
178+
179+
180+
def _strtobool(val: str) -> bool:
181+
val = val.lower()
182+
if val in ("y", "yes", "t", "true", "on", "1"):
183+
return True
184+
elif val in ("n", "no", "f", "false", "off", "0"):
185+
return False
186+
else:
187+
raise ValueError(f"invalid truth value {val!r}")

0 commit comments

Comments
 (0)