Repository: PrefectHQ/prefect

Issue:
Azure storage fails with NoneType object has no attribute rstrip
## Description

I am running a basic flow on a KubernetesAgent on Azure and having issues with Prefect’s Azure storage.

```python
storage = Azure("prefect", connection_string=connection_string)
```

Where `”prefect”` is a container in a storage account described by my `connection_string`. Unfortunately, during the flow’s execution with this storage type yields a `AttributeError("'NoneType' object has no attribute 'rstrip'"). 

Specifying `blob_name`, or `stored_as_script` parameters had no effect. The connection string provided was in the format: `BlobEndpoint=https://<>.blob.core.windows.net/……`


## Expected Behavior

I expected the Azure blob storage class to work without causing the flow run to fail. 


## Reproduction

```python
from prefect import task, Flow
from prefect.storage import Azure
from prefect.run_configs import KubernetesRun

connection_string = 'connection_string'

@task
def say_hello():
   print("Hello, world!")

# Prefect cloud config settings -----------------------------------------------------------

run_config = KubernetesRun(
   cpu_request=2,
   memory_request="2Gi",
   image="gcr.io/carbonplan/hub-notebook:b2419ff",
)

storage = Azure("prefect", connection_string=connection_string)

with Flow("debug_flow", storage=storage, run_config=run_config) as flow:
   h = say_hello()
```



Result:
```
└── 19:56:01 | INFO    | Entered state <Failed>: Failed to load and execute Flow's environment: AttributeError("'NoneType' object has no attribute 'rstrip'")
Flow run failed!
```

Full traceback:
<details>




```python-traceback
'NoneType' object has no attribute 'rstrip'
Traceback (most recent call last):
  File "/srv/conda/envs/notebook/bin/prefect", line 10, in <module>
    sys.exit(cli())
  File "/srv/conda/envs/notebook/lib/python3.8/site-packages/click/core.py", line 829, in __call__
    return self.main(*args, **kwargs)
  File "/srv/conda/envs/notebook/lib/python3.8/site-packages/click/core.py", line 782, in main
    rv = self.invoke(ctx)
  File "/srv/conda/envs/notebook/lib/python3.8/site-packages/click/core.py", line 1259, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/srv/conda/envs/notebook/lib/python3.8/site-packages/click/core.py", line 1259, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/srv/conda/envs/notebook/lib/python3.8/site-packages/click/core.py", line 1066, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/srv/conda/envs/notebook/lib/python3.8/site-packages/click/core.py", line 610, in invoke
    return callback(*args, **kwargs)
  File "/srv/conda/envs/notebook/lib/python3.8/site-packages/prefect/cli/execute.py", line 96, in flow_run
    raise exc
  File "/srv/conda/envs/notebook/lib/python3.8/site-packages/prefect/cli/execute.py", line 73, in flow_run
    flow = storage.get_flow(flow_data.name)
  File "/srv/conda/envs/notebook/lib/python3.8/site-packages/prefect/storage/azure.py", line 74, in get_flow
    client = self._azure_block_blob_service.get_blob_client(
  File "/srv/conda/envs/notebook/lib/python3.8/site-packages/prefect/storage/azure.py", line 154, in _azure_block_blob_service
    return azure.storage.blob.BlobServiceClient.from_connection_string(
  File "/srv/conda/envs/notebook/lib/python3.8/site-packages/azure/storage/blob/_blob_service_client.py", line 174, in from_connection_string
    account_url, secondary, credential = parse_connection_str(conn_str, credential, 'blob')
  File "/srv/conda/envs/notebook/lib/python3.8/site-packages/azure/storage/blob/_shared/base_client.py", line 363, in parse_connection_str
    conn_str = conn_str.rstrip(";")
AttributeError: 'NoneType' object has no attribute 'rstrip'

```

</details>


## Environment

Prefect diagnostics (local machine)
<details>

```json

{
  "config_overrides": {},
  "env_vars": [
    "PREFECT__CONTEXT__SECRETS__cmip6downscaling"
  ],
  "system_information": {
    "platform": "Linux-5.4.109+-x86_64-with-glibc2.10",
    "prefect_backend": "cloud",
    "prefect_version": "0.15.3",
    "python_version": "3.8.10"
  }
}
```
</details>

Prefect diagnostics (docker image)
<details>

```json

{
  "config_overrides": {},
  "env_vars": [],
  "system_information": {
    "platform": "Linux-5.4.109+-x86_64-with-glibc2.10",
    "prefect_backend": "cloud",
    "prefect_version": "0.15.3",
    "python_version": "3.8.10"
  }
}
```
</details>

Prefect Agent Pod Describe
<details>

```

Name:         prefect-agent-6b69496ddc-j2p64
Namespace:    prefect
Priority:     0
Node:         aks-user-41025623-vmss000000/10.240.0.5
Start Time:   Fri, 03 Sep 2021 11:46:04 -0700
Labels:       app=prefect-agent
              pod-template-hash=6b69496ddc
Annotations:  <none>
Status:       Running
IP:           10.244.5.6
IPs:
  IP:           10.244.5.6
Controlled By:  ReplicaSet/prefect-agent-6b69496ddc
Containers:
  agent:
    Container ID:  docker://1accd37abfb466696cebc03da46bba99486d5311f6740064a815c90ec8d44ad7
    Image:         prefecthq/prefect:0.15.3-python3.6
    Image ID:      docker-pullable://prefecthq/prefect@sha256:efa99417e77e83fc17f6d69005ea26363507b11878103912f5543174af07402b
    Port:          <none>
    Host Port:     <none>
    Command:
      /bin/bash
      -c
    Args:
      prefect agent kubernetes start
    State:          Running
      Started:      Fri, 03 Sep 2021 11:46:07 -0700
    Ready:          True
    Restart Count:  0
    Liveness:       http-get http://:8080/api/health delay=40s timeout=1s period=40s #success=1 #failure=2
    Environment:
      PREFECT__CLOUD__AGENT__AUTH_TOKEN:
      PREFECT__CLOUD__API:                   https://api.prefect.io
      NAMESPACE:                             prefect
      IMAGE_PULL_SECRETS:
      PREFECT__CLOUD__AGENT__LABELS:         ['az-eu-west']
      JOB_MEM_REQUEST:
      JOB_MEM_LIMIT:
      JOB_CPU_REQUEST:
      JOB_CPU_LIMIT:
      IMAGE_PULL_POLICY:
      SERVICE_ACCOUNT_NAME:
      PREFECT__BACKEND:                      cloud
      PREFECT__CLOUD__AGENT__AGENT_ADDRESS:  http://:8080
      PREFECT__CLOUD__API_KEY:               SECRET
      PREFECT__CLOUD__TENANT_ID:
      PREFECT__CLOUD__AGENT__ENV_VARS:       {"PREFECT__CONTEXT__SECRETS__MYSECRET": "SECRET"}
    Mounts:
      /var/run/secrets/kubernetes.io/serviceaccount from default-token-z5pck (ro)
Conditions:
  Type              Status
  Initialized       True
  Ready             True
  ContainersReady   True
  PodScheduled      True
Volumes:
  default-token-z5pck:
    Type:        Secret (a volume populated by a Secret)
    SecretName:  default-token-z5pck
    Optional:    false
QoS Class:       BestEffort
Node-Selectors:  <none>
Tolerations:     node.kubernetes.io/not-ready:NoExecute op=Exists for 300s
                 node.kubernetes.io/unreachable:NoExecute op=Exists for 300s
Events:          <none>

```
</details>



Discussion:
Hi! Thanks for this really well written issue. I can't see a way that the `connection_string` would be `None` here from looking at the `Azure` storage implementation. If you print the `storage.connection_string` does it have a value? What does `storage._azure_block_blob_service` return?
Thanks!
Printing the following:
`connection_string` : <Storage: Azure>
`storage.connection_string` : 'https://<>.blob.core.windows.net/;QueueEndpoint=……'
`storage._azure_block_blob_service` : 

```python
ValueError                                Traceback (most recent call last)
<ipython-input-6-bd406a774b26> in <module>
----> 1 storage._azure_block_blob_service

/srv/conda/envs/notebook/lib/python3.8/site-packages/prefect/storage/azure.py in _azure_block_blob_service(self)
    152         import azure.storage.blob
    153 
--> 154         return azure.storage.blob.BlobServiceClient.from_connection_string(
    155             conn_str=self.connection_string
    156         )

/srv/conda/envs/notebook/lib/python3.8/site-packages/azure/storage/blob/_blob_service_client.py in from_connection_string(cls, conn_str, credential, **kwargs)
    172                 :caption: Creating the BlobServiceClient from a connection string.
    173         """
--> 174         account_url, secondary, credential = parse_connection_str(conn_str, credential, 'blob')
    175         if 'secondary_hostname' not in kwargs:
    176             kwargs['secondary_hostname'] = secondary

/srv/conda/envs/notebook/lib/python3.8/site-packages/azure/storage/blob/_shared/base_client.py in parse_connection_str(conn_str, credential, service)
    364     conn_settings = [s.split("=", 1) for s in conn_str.split(";")]
    365     if any(len(tup) != 2 for tup in conn_settings):
--> 366         raise ValueError("Connection string is either blank or malformed.")
    367     conn_settings = dict((key.upper(), val) for key, val in conn_settings)
    368     endpoints = _SERVICE_PARAMS[service]

ValueError: Connection string is either blank or malformed.
```

It looks like it's upset by the format of your connection string. From a quick look, it seems that it must be formatted as `KEY=VALUE;OTHER_KEY=OTHER_VALUE` and that you are missing the `BlobEndpoint=` form the start?

See https://docs.microsoft.com/en-us/azure/storage/common/storage-configure-connection-string

I could be wrong though, let me know what you discover!
Thanks @madkinsz, I must have missed adding 'BlobEndpoint=' in my example. Adding it gives:

`storage._azure_block_blob_service`: <azure.storage.blob._blob_service_client.BlobServiceClient at 0x7fd83c0edd00>

Unfortunately, when registering and running this flow, I still get the same error.


```
└── 21:13:33 | INFO    | Entered state <Failed>: Failed to load and execute Flow's environment: AttributeError("'NoneType' object has no attribute 'rstrip'")
Flow run failed!
```

I tried switching the connection string from a `BlobEndpoint` to a `DefaultEndpointsProtocol `. This had no change in the error message.

```
DefaultEndpointsProtocol=https;AccountName=<account_name>;AccountKey=<account_key>;EndpointSuffix=core.windows.net"
```
One more thing to add here, just to rule out the theory of a malformed connection string, 

```python
bc = BlobClient.from_connection_string(conn_str, 'prefect', 'path/to/blob')
```

works fine. So The `BlobClient` is able to parse the connection string when passed directly. Just not when passed via the storage class.
I was able to get the example above to work by specifying the `AZURE_STORAGE_CONNECTION_STRING` environment variable on both the client and agent, forgoing the explicit `connection_string` argument to the `Azure` storage class.

I have a theory of what is going on here but I could use some help filling in some missing understanding of how things are intended to work. The basic theory is that the agent is not receiving the `connection_string` when instantiating the flow on the agent. My basic question is about how the Storage class is intended to communicate credentials (i.e. `Azure.connection_string`) from the client to the agent. My working model was that the storage class was pickled and sent to the agent but I'm no longer sure this is what happens. 

A few snippits to help illustrate the code path here. 

Inside of `Azure.get_flow`, a blob client is created:

https://github.com/PrefectHQ/prefect/blob/5221e6c4ef68eb8c7659c212ea937d1856cee8a4/src/prefect/storage/azure.py#L74-L76

this accesses the `connection_string`:

https://github.com/PrefectHQ/prefect/blob/5221e6c4ef68eb8c7659c212ea937d1856cee8a4/src/prefect/storage/azure.py#L154-L156

which on the agent is `None` (unless the above environment variable is set on the agent):
 
https://github.com/PrefectHQ/prefect/blob/5221e6c4ef68eb8c7659c212ea937d1856cee8a4/src/prefect/storage/azure.py#L48-L50

Resolving patch:
diff --git a/src/prefect/engine/results/azure_result.py b/src/prefect/engine/results/azure_result.py
--- a/src/prefect/engine/results/azure_result.py
+++ b/src/prefect/engine/results/azure_result.py
@@ -24,7 +24,7 @@ class AzureResult(Result):
             Blob storage. If not provided the value set in the environment as
             `AZURE_STORAGE_CONNECTION_STRING` will be used
         - connection_string_secret (str, optional): the name of a Prefect Secret
-            which stores your Azure connection tring
+            which stores your Azure connection string
         - **kwargs (Any, optional): any additional `Result` initialization options
     """
 
@@ -156,17 +156,7 @@ def exists(self, location: str, **kwargs: Any) -> bool:
         Returns:
             - bool: whether or not the target result exists.
         """
-        from azure.core.exceptions import ResourceNotFoundError
 
-        # initialize client and download
-        client = self.service.get_blob_client(
+        return self.service.get_blob_client(
             container=self.container, blob=location.format(**kwargs)
-        )
-
-        # Catch exception because Azure python bindings do not yet have an exists method
-        # https://github.com/Azure/azure-sdk-for-python/issues/9507
-        try:
-            client.get_blob_properties()
-            return True
-        except ResourceNotFoundError:
-            return False
+        ).exists()
diff --git a/src/prefect/serialization/storage.py b/src/prefect/serialization/storage.py
--- a/src/prefect/serialization/storage.py
+++ b/src/prefect/serialization/storage.py
@@ -40,6 +40,7 @@ class Meta:
         object_class = Azure
 
     container = fields.String(allow_none=False)
+    connection_string_secret = fields.String(allow_none=True)
     blob_name = fields.String(allow_none=True)
     stored_as_script = fields.Bool(allow_none=True)
 
diff --git a/src/prefect/storage/azure.py b/src/prefect/storage/azure.py
--- a/src/prefect/storage/azure.py
+++ b/src/prefect/storage/azure.py
@@ -4,6 +4,8 @@
 import pendulum
 from slugify import slugify
 
+import prefect
+from prefect.client import Secret
 from prefect.engine.results import AzureResult
 from prefect.storage import Storage
 from prefect.utilities.storage import (
@@ -27,7 +29,8 @@ class Azure(Storage):
 
     Args:
         - container (str): the name of the Azure Blob Container to store the Flow
-        - connection_string (str, optional): an Azure connection string for communicating with
+        - connection_string_secret (str, optional): the name of a Prefect secret
+            that contains an Azure connection string for communicating with
             Blob storage. If not provided the value set in the environment as
             `AZURE_STORAGE_CONNECTION_STRING` will be used
         - blob_name (str, optional): a unique key to use for uploading this Flow to Azure. This
@@ -40,20 +43,18 @@ class Azure(Storage):
     def __init__(
         self,
         container: str,
-        connection_string: str = None,
+        connection_string_secret: str = None,
         blob_name: str = None,
         stored_as_script: bool = False,
         **kwargs: Any
     ) -> None:
-        self.connection_string = connection_string or os.getenv(
-            "AZURE_STORAGE_CONNECTION_STRING"
-        )
-
         self.container = container
+        self.connection_string_secret = connection_string_secret
         self.blob_name = blob_name
 
         result = AzureResult(
-            connection_string=self.connection_string, container=container
+            connection_string_secret=self.connection_string_secret,
+            container=container,
         )
         super().__init__(result=result, stored_as_script=stored_as_script, **kwargs)
 
@@ -70,15 +71,19 @@ def get_flow(self, flow_name: str) -> "Flow":
         if flow_name not in self.flows:
             raise ValueError("Flow is not contained in this Storage")
         flow_location = self.flows[flow_name]
+        try:
+            client = self._azure_block_blob_service.get_blob_client(
+                container=self.container, blob=flow_location
+            )
 
-        client = self._azure_block_blob_service.get_blob_client(
-            container=self.container, blob=flow_location
-        )
-
-        self.logger.info("Downloading {} from {}".format(flow_location, self.container))
-
-        content = client.download_blob().content_as_bytes()
+            self.logger.info(
+                "Downloading {} from {}".format(flow_location, self.container)
+            )
 
+            content = client.download_blob().content_as_bytes()
+        except Exception as err:
+            self.logger.error("Error downloading Flow from Azure: {}".format(err))
+            raise
         if self.stored_as_script:
             return extract_flow_from_file(file_contents=content, flow_name=flow_name)  # type: ignore
 
@@ -147,6 +152,20 @@ def build(self) -> "Storage":
 
         return self
 
+    @property
+    def connection_string(self):  # type: ignore
+        if self.connection_string_secret is not None:
+            return Secret(self.connection_string_secret).get()
+        conn_string = prefect.context.get("secrets", {}).get(
+            "AZURE_STORAGE_CONNECTION_STRING"
+        ) or os.getenv("AZURE_STORAGE_CONNECTION_STRING")
+        if conn_string is None:
+            raise Exception(
+                "Azure connection string not provided. Set `AZURE_STORAGE_CONNECTION_STRING` environment"
+                " variable or save connection string as Prefect secret."
+            )
+        return conn_string
+
     @property
     def _azure_block_blob_service(self):  # type: ignore
         import azure.storage.blob


Test patch:
diff --git a/tests/engine/results/test_azure_result.py b/tests/engine/results/test_azure_result.py
--- a/tests/engine/results/test_azure_result.py
+++ b/tests/engine/results/test_azure_result.py
@@ -91,26 +91,13 @@ def test_azure_result_is_pickleable(self, azure_client):
         res = cloudpickle.loads(cloudpickle.dumps(result))
         assert isinstance(res, AzureResult)
 
-    def test_azure_exists(self, monkeypatch):
-        client = MagicMock(get_blob_properties=MagicMock())
+    @pytest.mark.parametrize("exists", [True, False])
+    def test_azure_exists(self, monkeypatch, exists):
+        client = MagicMock(exists=MagicMock(return_value=exists))
         service = MagicMock(get_blob_client=MagicMock(return_value=client))
         monkeypatch.setattr(
             "prefect.engine.results.azure_result.AzureResult.service", service
         )
 
         result = AzureResult(container="foo", location="{thing}/here.txt")
-        assert result.exists("44.txt") is True
-
-    def test_azure_does_not_exists(self, monkeypatch):
-        from azure.core.exceptions import ResourceNotFoundError
-
-        client = MagicMock(
-            get_blob_properties=MagicMock(side_effect=ResourceNotFoundError)
-        )
-        service = MagicMock(get_blob_client=MagicMock(return_value=client))
-        monkeypatch.setattr(
-            "prefect.engine.results.azure_result.AzureResult.service", service
-        )
-
-        result = AzureResult(container="foo", location="{thing}/here.txt")
-        assert result.exists("44.txt") is False
+        assert result.exists("44.txt") is exists
diff --git a/tests/serialization/test_storage.py b/tests/serialization/test_storage.py
--- a/tests/serialization/test_storage.py
+++ b/tests/serialization/test_storage.py
@@ -168,7 +168,7 @@ def test_azure_empty_serialize():
 def test_azure_full_serialize():
     azure = storage.Azure(
         container="container",
-        connection_string="conn",
+        connection_string_secret="conn",
         blob_name="name",
         secrets=["foo"],
         labels=["bar", "baz"],
@@ -179,13 +179,14 @@ def test_azure_full_serialize():
     assert serialized
     assert serialized["__version__"] == prefect.__version__
     assert serialized["container"] == "container"
+    assert serialized["connection_string_secret"] == "conn"
     assert serialized["blob_name"] == "name"
     assert serialized["secrets"] == ["foo"]
 
 
 def test_azure_creds_not_serialized():
     azure = storage.Azure(
-        container="container", connection_string="conn", blob_name="name"
+        container="container", connection_string_secret="conn", blob_name="name"
     )
     serialized = AzureSchema().dump(azure)
 
@@ -193,13 +194,14 @@ def test_azure_creds_not_serialized():
     assert serialized["__version__"] == prefect.__version__
     assert serialized["container"] == "container"
     assert serialized["blob_name"] == "name"
+    assert serialized["connection_string_secret"] == "conn"
     assert serialized.get("connection_string") is None
 
 
 def test_azure_serialize_with_flows():
     azure = storage.Azure(
         container="container",
-        connection_string="conn",
+        connection_string_secret="conn",
         blob_name="name",
         secrets=["foo"],
     )
@@ -211,7 +213,7 @@ def test_azure_serialize_with_flows():
     assert serialized["__version__"] == prefect.__version__
     assert serialized["container"] == "container"
     assert serialized["blob_name"] == "name"
-    assert serialized.get("connection_string") is None
+    assert serialized["connection_string_secret"] == "conn"
     assert serialized["flows"] == {"test": "key"}
 
     deserialized = AzureSchema().load(serialized)
diff --git a/tests/storage/test_azure_storage.py b/tests/storage/test_azure_storage.py
--- a/tests/storage/test_azure_storage.py
+++ b/tests/storage/test_azure_storage.py
@@ -2,7 +2,7 @@
 
 import pytest
 
-from prefect import Flow
+from prefect import Flow, context
 from prefect.storage import Azure
 
 pytest.importorskip("azure.storage.blob")
@@ -17,14 +17,14 @@ def test_create_azure_storage():
 def test_create_azure_storage_init_args():
     storage = Azure(
         container="test",
-        connection_string="conn",
+        connection_string_secret="conn",
         blob_name="name",
         secrets=["foo"],
     )
     assert storage
     assert storage.flows == dict()
     assert storage.container == "test"
-    assert storage.connection_string == "conn"
+    assert storage.connection_string_secret == "conn"
     assert storage.blob_name == "name"
     assert storage.secrets == ["foo"]
 
@@ -36,16 +36,39 @@ def test_serialize_azure_storage():
     assert serialized_storage["type"] == "Azure"
 
 
-def test_blob_service_client_property(monkeypatch):
+@pytest.mark.parametrize(
+    "secret_name,secret_arg",
+    [
+        ("SECRET_NAME", "conn_string_value_one"),
+        ("AZURE_STORAGE_CONNECTION_STRING", "conn_string_value_two"),
+    ],
+)
+def test_blob_service_client_property(monkeypatch, secret_name, secret_arg):
     connection = MagicMock()
     azure = MagicMock(from_connection_string=connection)
     monkeypatch.setattr("azure.storage.blob.BlobServiceClient", azure)
 
-    storage = Azure(container="test", connection_string="conn")
+    with context(secrets={secret_name: secret_arg}):
+        storage = Azure(container="test", connection_string_secret=secret_name)
+        client = storage._azure_block_blob_service()
+        azure_client = storage._azure_block_blob_service
+        assert storage.connection_string == secret_arg
 
-    azure_client = storage._azure_block_blob_service
     assert azure_client
-    connection.assert_called_with(conn_str="conn")
+    connection.assert_called_with(conn_str=secret_arg)
+
+
+@pytest.mark.parametrize(
+    "secret_name,secret_arg",
+    [
+        ("SECRET_NAME", "conn_string_value_one"),
+        ("AZURE_STORAGE_CONNECTION_STRING", "conn_string_value_two"),
+    ],
+)
+def test_connection_string_property(secret_name, secret_arg):
+    with context(secrets={secret_name: secret_arg}):
+        storage = Azure(container="test", connection_string_secret=secret_name)
+        assert storage.connection_string == secret_arg
 
 
 def test_add_flow_to_azure():
diff --git a/tests/tasks/jupyter/sample_notebook.ipynb b/tests/tasks/jupyter/sample_notebook.ipynb
--- a/tests/tasks/jupyter/sample_notebook.ipynb
+++ b/tests/tasks/jupyter/sample_notebook.ipynb
@@ -110,7 +110,8 @@
    "outputs": [],
    "source": [
     "import time\n",
-    "for i in range(1,3):\n",
+    "\n",
+    "for i in range(1, 3):\n",
     "    print(i)\n",
     "    time.sleep(1)"
    ]