def test_local_secrets_auto_load_json_strings():
secret = Secret(name="test")
with set_temporary_config({"cloud.use_local_secrets": True}):
with prefect.context(secrets=dict(test='{"x": 42}')):
assert secret.get() == {"x": 42}
with pytest.raises(ValueError):
secret.get()
def test_secrets_raise_if_in_flow_context():
secret = Secret(name="test")
with set_temporary_config({"cloud.use_local_secrets": True}):
with prefect.context(secrets=dict(test=42)):
with prefect.Flow("test"):
with pytest.raises(ValueError):
secret.get()
def test_secret_value_depends_on_use_local_secrets():
secret = Secret(name="test")
with set_temporary_config({"cloud.use_local_secrets": False}):
with prefect.context(secrets=dict(test=42)):
with pytest.raises(AuthorizationError) as exc:
secret.get()
assert "Client.login" in str(exc.value)
def test_secret_value_pulled_from_context():
secret = Secret(name="test")
with set_temporary_config({"cloud.use_local_secrets": True}):
with prefect.context(secrets=dict(test=42)):
assert secret.get() == 42
with pytest.raises(ValueError):
secret.get()
def run(
self,
container_name: str = None,
datastore_name: str = None,
create_container_if_not_exists: bool = False,
overwrite_existing_datastore: bool = False,
azure_credentials_secret: str = "AZ_CREDENTIALS",
set_as_default: bool = False,
) -> AzureBlobDatastore:
"""
Task run method.
Args:
- container_name (str, optional): The name of the container.
- datastore_name (str, optional): The name of the datastore. If not defined, the
container name will be used.
- create_container_if_not_exists (bool, optional): Create a container, if one does
not exist with the given name.
- overwrite_existing_datastore (bool, optional): Overwrite an existing datastore.
If the datastore does not exist, it will be created.
- azure_credentials_secret (str, optinonal): The name of the Prefect Secret that
stores your Azure credentials; this Secret must be a JSON string with two keys:
`ACCOUNT_NAME` and either `ACCOUNT_KEY` or `SAS_TOKEN` (if both are defined
then`ACCOUNT_KEY` is used)
- set_as_default (bool optional): Set the created Datastore as the default
datastore for the Workspace.
Return:
- (azureml.data.azure_storage_datastore.AzureBlobDatastore): The registered Datastore.
"""
if container_name is None:
raise ValueError("A container name must be provided.")
if datastore_name is None:
datastore_name = container_name
# get Azure credentials
azure_credentials = Secret(azure_credentials_secret).get()
az_account_name = azure_credentials["ACCOUNT_NAME"]
az_account_key = azure_credentials.get("ACCOUNT_KEY")
az_sas_token = azure_credentials.get("SAS_TOKEN")
datastore = azureml.core.datastore.Datastore.register_azure_blob_container(
workspace=self.workspace,
datastore_name=datastore_name,
container_name=container_name,
account_name=az_account_name,
account_key=az_account_key,
sas_token=az_sas_token,
overwrite=overwrite_existing_datastore,
create_if_not_exists=create_container_if_not_exists,
)
if set_as_default:
datastore.set_as_default()
return datastore
def test_secret_value_depends_on_use_local_secrets(monkeypatch):
secret = Secret(name="test")
with set_temporary_config({
"cloud.use_local_secrets": False,
"cloud.auth_token": None
}):
with prefect.context(secrets=dict(test=42)):
with pytest.raises(ClientError):
secret.get()
def test_secret_raises_informative_error_for_server():
secret = Secret(name="test")
with set_temporary_config({
"cloud.use_local_secrets": False,
"backend": "server"
}):
with pytest.raises(ValueError) as exc:
secret.get()
assert str(exc.value) == 'Local Secret "test" was not found.'
def test_secret_value_depends_on_use_local_secrets():
secret = Secret(name="test")
with set_temporary_config({
"cloud.use_local_secrets": False,
"cloud.auth_token": None
}):
with prefect.context(secrets=dict(test=42)):
with pytest.raises(AuthorizationError, match="Client.login"):
secret.get()
def test_secret_value_depends_on_use_local_secrets(monkeypatch):
response = {"errors": "Malformed Authorization header"}
post = MagicMock(return_value=MagicMock(json=MagicMock(return_value=response)))
session = MagicMock()
session.return_value.post = post
monkeypatch.setattr("requests.Session", session)
secret = Secret(name="test")
with set_temporary_config(
{"cloud.use_local_secrets": False, "cloud.auth_token": None}
):
with prefect.context(secrets=dict()):
with pytest.raises(ClientError):
secret.get()
def run(
self,
data: str,
blob_name: str = None,
azure_credentials_secret: str = "AZ_CREDENTIALS",
container: str = None,
) -> str:
"""
Task run method.
Args:
- data (str): the data payload to upload
- blob_name (str, optional): the name to upload the data under; if not
provided, a random `uuid` will be created
- azure_credentials_secret (str, optional): the name of the Prefect Secret
that stores your Azure credentials; this Secret must be a JSON string
with two keys: `ACCOUNT_NAME` and either `ACCOUNT_KEY` or `SAS_TOKEN`
- container (str, optional): the name of the Blob Storage container to upload to
Returns:
- str: the name of the blob the data payload was uploaded to
"""
if container is None:
raise ValueError("A container name must be provided.")
## get Azure credentials
azure_credentials = Secret(azure_credentials_secret).get()
az_account_name = azure_credentials["ACCOUNT_NAME"]
az_account_key = azure_credentials.get("ACCOUNT_KEY")
az_sas_token = azure_credentials.get("SAS_TOKEN")
blob_service = azure.storage.blob.BlockBlobService(
account_name=az_account_name,
account_key=az_account_key,
sas_token=az_sas_token,
)
## create key if not provided
if blob_name is None:
blob_name = str(uuid.uuid4())
blob_service.create_blob_from_text(container_name=container,
blob_name=blob_name,
text=data)
return blob_name
def test_local_secrets_remain_plain_dictionaries():
secret = Secret(name="test")
with set_temporary_config({"cloud.use_local_secrets": True}):
with prefect.context(secrets=dict(test={"x": 42})):
assert isinstance(prefect.context.secrets["test"], dict)
val = secret.get()
assert val == {"x": 42}
assert isinstance(val, dict) and not isinstance(val, box.Box)
def test_secrets_use_client(monkeypatch):
response = {"data": {"secretValue": "1234"}}
post = MagicMock(return_value=MagicMock(json=MagicMock(return_value=response)))
monkeypatch.setattr("requests.post", post)
with set_temporary_config(
{"cloud.auth_token": "secret_token", "cloud.use_local_secrets": False}
):
my_secret = Secret(name="the-key")
val = my_secret.get()
assert val == "1234"
def initialize_service(self) -> None:
"""
Initialize a Blob service.
"""
import azure.storage.blob
kwargs = dict()
if self.connection_string:
kwargs["connection_string"] = self.connection_string
else:
azure_credentials = Secret(self.azure_credentials_secret).get()
if isinstance(azure_credentials, str):
azure_credentials = json.loads(azure_credentials)
kwargs["account_name"] = azure_credentials["ACCOUNT_NAME"]
kwargs["account_key"] = azure_credentials.get("ACCOUNT_KEY")
kwargs["sas_token"] = azure_credentials.get("SAS_TOKEN")
blob_service = azure.storage.blob.BlockBlobService(**kwargs)
self.service = blob_service
def initialize_service(self) -> None:
"""
Initialize a Blob service.
"""
import azure.storage.blob
azure_credentials = Secret(self.azure_credentials_secret).get()
if isinstance(azure_credentials, str):
azure_credentials = json.loads(azure_credentials)
az_account_name = azure_credentials["ACCOUNT_NAME"]
az_account_key = azure_credentials.get("ACCOUNT_KEY")
az_sas_token = azure_credentials.get("SAS_TOKEN")
blob_service = azure.storage.blob.BlockBlobService(
account_name=az_account_name,
account_key=az_account_key,
sas_token=az_sas_token,
)
self.service = blob_service
def test_secrets_use_client(monkeypatch, cloud_api):
response = {"data": {"secret_value": '"1234"'}}
post = MagicMock(return_value=MagicMock(json=MagicMock(return_value=response)))
session = MagicMock()
session.return_value.post = post
monkeypatch.setattr("requests.Session", session)
with set_temporary_config(
{"cloud.auth_token": "secret_token", "cloud.use_local_secrets": False}
):
my_secret = Secret(name="the-key")
val = my_secret.get()
assert val == "1234"
def run(
self,
blob_name: str,
azure_credentials_secret: str = "AZ_CREDENTIALS",
container: str = None,
) -> str:
"""
Task run method.
Args:
- blob_name (str): the name of the blob within this container to retrieve
- azure_credentials_secret (str, optional): the name of the Prefect Secret
that stores your Azure credentials; this Secret must be a JSON string
with two keys: `ACCOUNT_NAME` and either `ACCOUNT_KEY` or `SAS_TOKEN`
- container (str, optional): the name of the Blob Storage container to download from
Returns:
- str: the contents of this blob_name / container, as a string
"""
if container is None:
raise ValueError("A container name must be provided.")
# get Azure credentials
azure_credentials = Secret(azure_credentials_secret).get()
az_account_name = azure_credentials["ACCOUNT_NAME"]
az_account_key = azure_credentials.get("ACCOUNT_KEY")
az_sas_token = azure_credentials.get("SAS_TOKEN")
blob_service = azure.storage.blob.BlockBlobService(
account_name=az_account_name,
account_key=az_account_key,
sas_token=az_sas_token,
)
blob_result = blob_service.get_blob_to_text(container_name=container,
blob_name=blob_name)
content_string = blob_result.content
return content_string
def test_cloud_secrets_auto_load_json_strings(monkeypatch, cloud_api):
response = {"data": {"secret_value": '{"x": 42}'}}
post = MagicMock(return_value=MagicMock(json=MagicMock(return_value=response)))
session = MagicMock()
session.return_value.post = post
monkeypatch.setattr("requests.Session", session)
with set_temporary_config(
{"cloud.auth_token": "secret_token", "cloud.use_local_secrets": False}
):
my_secret = Secret(name="the-key")
val = my_secret.get()
assert isinstance(val, dict)
def test_cloud_secrets_use_context_first(monkeypatch):
response = {"data": {"secret_value": '"1234"'}}
post = MagicMock(return_value=MagicMock(json=MagicMock(
return_value=response)))
session = MagicMock()
session.return_value.post = post
monkeypatch.setattr("requests.Session", session)
with set_temporary_config({
"cloud.api_key": "api-key",
"cloud.use_local_secrets": False
}):
with prefect.context(secrets={"the-key": "foo"}):
my_secret = Secret(name="the-key")
val = my_secret.get()
assert val == "foo"
def test_cloud_secrets_remain_plain_dictionaries(monkeypatch, cloud_api):
response = {"data": {"secret_value": {"a": "1234", "b": [1, 2, {"c": 3}]}}}
post = MagicMock(return_value=MagicMock(json=MagicMock(return_value=response)))
session = MagicMock()
session.return_value.post = post
monkeypatch.setattr("requests.Session", session)
with set_temporary_config(
{"cloud.auth_token": "secret_token", "cloud.use_local_secrets": False}
):
my_secret = Secret(name="the-key")
val = my_secret.get()
assert val == {"a": "1234", "b": [1, 2, {"c": 3}]}
assert isinstance(val, dict) and not isinstance(val, box.Box)
val2 = val["b"]
assert isinstance(val2, list) and not isinstance(val2, box.BoxList)
val3 = val["b"][2]
assert isinstance(val3, dict) and not isinstance(val3, box.Box)
def test_secret_raises_if_doesnt_exist():
secret = Secret(name="test")
with set_temporary_config({"cloud.use_local_secrets": True}):
with pytest.raises(ValueError, match="not found"):
secret.get()
def test_secrets_dont_raise_just_because_flow_key_is_populated():
secret = Secret(name="test")
with set_temporary_config({"cloud.use_local_secrets": True}):
with prefect.context(secrets=dict(test=42), flow="not None"):
assert secret.get() == 42
def test_secret_get_none():
secret = Secret(name="test")
with set_temporary_config({"cloud.use_local_secrets": True}):
assert secret.get() is None
def test_secret_raises_if_doesnt_exist():
secret = Secret(name="test")
with set_temporary_config({"cloud.use_local_secrets": True}):
with pytest.raises(ValueError) as exc:
secret.get()
assert "not found" in str(exc.value)
