/
base.py
66 lines (52 loc) · 2.45 KB
/
base.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
from prefect.client.secrets import Secret as _Secret
from prefect.core.task import Task
from prefect.engine.results import SecretResult
class SecretBase(Task):
"""
Base Secrets Task. This task does not perform any action but rather serves as the base
task class which should be inherited from when writing new Secret Tasks.
Users should subclass this Task and override its `run` method for plugging into other
Secret stores, as it is handled differently during execution to ensure the underlying
secret value is not accidentally persisted in a non-safe location.
Args:
- **kwargs (Any, optional): additional keyword arguments to pass to the Task constructor
Raises:
- ValueError: if a `result` keyword is passed
"""
def __init__(self, **kwargs):
if kwargs.get("result"):
raise ValueError("Result types for Secrets are not configurable.")
kwargs["checkpoint"] = False
super().__init__(**kwargs)
self.result = SecretResult(secret_task=self)
class PrefectSecret(SecretBase):
"""
Prefect Secrets Task. This task retrieves the underlying secret through
the Prefect Secrets API (which has the ability to toggle between local vs. Cloud secrets).
Args:
- name (str, optional): The name of the underlying secret
- **kwargs (Any, optional): additional keyword arguments to pass to the Task constructor
Raises:
- ValueError: if a `result` keyword is passed
"""
def __init__(self, name=None, **kwargs):
self.secret_name = name
super().__init__(name=name, **kwargs)
def run(self, name: str = None):
"""
The run method for Secret Tasks. This method actually retrieves and returns the
underlying secret value using the `Secret.get()` method. Note that this method first
checks context for the secret value, and if not found either raises an error or queries
Prefect Cloud, depending on whether `config.cloud.use_local_secrets` is `True` or
`False`.
Args:
- name (str, optional): the name of the underlying Secret to retrieve. Defaults
to the name provided at initialization.
Returns:
- Any: the underlying value of the Prefect Secret
"""
if name is None:
name = self.secret_name
if name is None:
raise ValueError("A secret name must be provided.")
return _Secret(name).get()