Fix job attachment 401s for submitted jobs

azure-quantum/azure/quantum/job/base_job.py

@@ -344,12 +344,8 @@ def upload_attachment(
         :rtype: str
         """
 
-        # Use Job's default container if not specified
         if container_uri is None:
-            if self._details.container_uri is None:
-                container_uri = self.workspace.get_container_uri(job_id=self.id)
-            else:
-                container_uri = self._details.container_uri
+            container_uri = self._get_attachment_container_uri(required_permission="w")
 
         uploaded_blob_uri = self.upload_input_data(
             container_uri = container_uri,
@@ -377,13 +373,9 @@ def download_attachment(
         :rtype: bytes
         """
 
-        # Use Job's default container if not specified
         if container_uri is None:
-            if self._details.container_uri is None:
-                container_uri = self.workspace.get_container_uri(job_id=self.id)
-            else:
-                container_uri = self._details.container_uri
-
+            container_uri = self._get_attachment_container_uri(required_permission="r")
+
         container_client = ContainerClient.from_container_url(container_uri)
         blob_client = container_client.get_blob_client(name)
         response = blob_client.download_blob().readall()
@@ -399,16 +391,51 @@ def list_attachments(self) -> list[BlobProperties]:
         :rtype: list[~azure.storage.blob.BlobProperties]
         """
 
-        # Use the job's linked storage container.
-        if self._details.container_uri is None:
-            container_uri = self.workspace.get_container_uri(job_id=self.id)
-        else:
-            container_uri = self._details.container_uri
+        container_uri = self._get_attachment_container_uri(required_permission="l")
 
         container_client = ContainerClient.from_container_url(container_uri)
         return list(container_client.list_blobs())
 
 
+    def _get_attachment_container_uri(self, required_permission: str) -> str:
+        container_uri = self._details.container_uri
+        if container_uri is None:
+            return self.workspace.get_container_uri(job_id=self.id)
+
+        query_params = parse_qs(urlparse(container_uri).query)
+        token_expire_query_param = query_params.get("se")
+        token_permissions = query_params.get("sp", [""])[0]
+        if (
+            query_params.get("sig")
+            and token_expire_query_param
+            and required_permission in token_permissions
+        ):
+            try:
+                token_expire_time = datetime.fromisoformat(
+                    token_expire_query_param[0].replace("Z", "+00:00")
+                )
+                if token_expire_time.tzinfo is None:
+                    token_expire_time = token_expire_time.replace(tzinfo=timezone.utc)
+                if datetime.now(tz=timezone.utc) < token_expire_time - timedelta(minutes=5):
+                    return container_uri
+            except ValueError:
+                pass
+
+        refreshed_container_uri = self.workspace.get_container_uri(
+            job_id=self.id,
+            container_name=self.container_name,
+        )
+        stored_hostname = urlparse(container_uri).hostname
+        refreshed_hostname = urlparse(refreshed_container_uri).hostname
+        if stored_hostname != refreshed_hostname:
+            raise ValueError(
+                "Refreshed attachment container hostname "
+                f"'{refreshed_hostname}' does not match job container hostname "
+                f"'{stored_hostname}'."
+            )
+        return refreshed_container_uri
+
+
     def _get_blob_uri_with_sas_token(self, blob_uri: str) -> str:
         """Get Blob URI with SAS-token if one was not specified in blob_uri parameter
         :param blob_uri: Blob URI

azure-quantum/azure/quantum/storage.py

@@ -10,6 +10,7 @@
     ContainerClient,
     BlobClient,
     BlobSasPermissions,
+    ContainerSasPermissions,
     ContentSettings,
     generate_blob_sas,
     generate_container_sas,
@@ -68,8 +69,8 @@ def get_container_uri(connection_string: str, container_name: str) -> str:
         container.account_name,
         container.container_name,
         account_key=container.credential.account_key,
-        permission=BlobSasPermissions(
-            read=True, add=True, write=True, create=True
+        permission=ContainerSasPermissions(
+            read=True, add=True, write=True, create=True, list=True
         ),
         expiry=datetime.utcnow() + timedelta(days=14),
     )

azure-quantum/tests/test_job_attachments.py

@@ -3,16 +3,30 @@
 # Licensed under the MIT License.
 ##
 
-from unittest.mock import Mock, patch
+from unittest.mock import Mock, call, patch
+
+import pytest
+
 from azure.quantum import Job, JobDetails
 
 
-CONTAINER_URI = "https://acct.blob.core.windows.net/job-id?sas"
+JOB_ID = "job-id"
+DEFAULT_CONTAINER_NAME = f"job-{JOB_ID}"
+UNSIGNED_CONTAINER_URI = f"https://acct.blob.core.windows.net/{DEFAULT_CONTAINER_NAME}"
+SIGNED_CONTAINER_URI = (
+    f"{UNSIGNED_CONTAINER_URI}?sp=racwdl&se=2099-01-01T00%3A00%3A00Z&sig=signature"
+)
+READ_ONLY_CONTAINER_URI = (
+    f"{UNSIGNED_CONTAINER_URI}?sp=rl&se=2099-01-01T00%3A00%3A00Z&sig=signature"
+)
+EXPIRED_CONTAINER_URI = (
+    f"{UNSIGNED_CONTAINER_URI}?sp=racwdl&se=2000-01-01T00%3A00%3A00Z&sig=signature"
+)
 
 
-def _job_with_container(container_uri=CONTAINER_URI, workspace=None) -> Job:
+def _job_with_container(container_uri=UNSIGNED_CONTAINER_URI, workspace=None) -> Job:
     job_details = JobDetails(
-        id="job-id",
+        id=JOB_ID,
         name="",
         provider_id="",
         target="",
@@ -25,7 +39,9 @@ def _job_with_container(container_uri=CONTAINER_URI, workspace=None) -> Job:
 
 @patch("azure.quantum.job.base_job.ContainerClient")
 def test_list_attachments_returns_container_blobs(mock_container_client):
-    job = _job_with_container()
+    workspace = Mock()
+    workspace.get_container_uri.return_value = SIGNED_CONTAINER_URI
+    job = _job_with_container(workspace=workspace)
 
     blob_a = Mock()
     blob_b = Mock()
@@ -34,21 +50,218 @@ def test_list_attachments_returns_container_blobs(mock_container_client):
 
     result = job.list_attachments()
 
-    mock_container_client.from_container_url.assert_called_once_with(CONTAINER_URI)
+    workspace.get_container_uri.assert_called_once_with(
+        job_id=JOB_ID,
+        container_name=DEFAULT_CONTAINER_NAME,
+    )
+    mock_container_client.from_container_url.assert_called_once_with(SIGNED_CONTAINER_URI)
     assert result == [blob_a, blob_b]
 
 
 @patch("azure.quantum.job.base_job.ContainerClient")
 def test_list_attachments_uses_workspace_container_when_unset(mock_container_client):
     workspace = Mock()
-    workspace.get_container_uri.return_value = CONTAINER_URI
+    workspace.get_container_uri.return_value = SIGNED_CONTAINER_URI
     job = _job_with_container(container_uri=None, workspace=workspace)
 
     container = mock_container_client.from_container_url.return_value
     container.list_blobs.return_value = []
 
     result = job.list_attachments()
 
-    workspace.get_container_uri.assert_called_once_with(job_id="job-id")
-    mock_container_client.from_container_url.assert_called_once_with(CONTAINER_URI)
+    workspace.get_container_uri.assert_called_once_with(job_id=JOB_ID)
+    mock_container_client.from_container_url.assert_called_once_with(SIGNED_CONTAINER_URI)
     assert result == []
+
+
+def test_upload_attachment_uses_fresh_workspace_container_uri():
+    workspace = Mock()
+    workspace.get_container_uri.return_value = SIGNED_CONTAINER_URI
+    job = _job_with_container(workspace=workspace)
+    job.upload_input_data = Mock(return_value="uploaded-uri")
+
+    result = job.upload_attachment("attachment", b"data")
+
+    workspace.get_container_uri.assert_called_once_with(
+        job_id=JOB_ID,
+        container_name=DEFAULT_CONTAINER_NAME,
+    )
+    job.upload_input_data.assert_called_once_with(
+        container_uri=SIGNED_CONTAINER_URI,
+        blob_name="attachment",
+        input_data=b"data",
+    )
+    assert result == "uploaded-uri"
+
+
+@patch("azure.quantum.job.base_job.ContainerClient")
+def test_download_attachment_uses_fresh_workspace_container_uri(mock_container_client):
+    workspace = Mock()
+    workspace.get_container_uri.return_value = SIGNED_CONTAINER_URI
+    job = _job_with_container(workspace=workspace)
+    blob_client = mock_container_client.from_container_url.return_value.get_blob_client.return_value
+    blob_client.download_blob.return_value.readall.return_value = b"data"
+
+    result = job.download_attachment("attachment")
+
+    workspace.get_container_uri.assert_called_once_with(
+        job_id=JOB_ID,
+        container_name=DEFAULT_CONTAINER_NAME,
+    )
+    mock_container_client.from_container_url.assert_called_once_with(SIGNED_CONTAINER_URI)
+    assert result == b"data"
+
+
+@patch("azure.quantum.job.base_job.ContainerClient")
+def test_attachment_methods_honor_explicit_container_uri(mock_container_client):
+    workspace = Mock()
+    job = _job_with_container(workspace=workspace)
+    job.upload_input_data = Mock(return_value="uploaded-uri")
+    blob_client = mock_container_client.from_container_url.return_value.get_blob_client.return_value
+    blob_client.download_blob.return_value.readall.return_value = b"data"
+    explicit_uri = "https://custom.blob.core.windows.net/container?sas"
+
+    job.upload_attachment("upload", b"data", container_uri=explicit_uri)
+    job.download_attachment("download", container_uri=explicit_uri)
+
+    workspace.get_container_uri.assert_not_called()
+    job.upload_input_data.assert_called_once_with(
+        container_uri=explicit_uri,
+        blob_name="upload",
+        input_data=b"data",
+    )
+    mock_container_client.from_container_url.assert_called_once_with(explicit_uri)
+
+
+@patch("azure.quantum.job.base_job.ContainerClient")
+def test_attachment_methods_reuse_valid_signed_job_uri(mock_container_client):
+    workspace = Mock()
+    job = _job_with_container(container_uri=SIGNED_CONTAINER_URI, workspace=workspace)
+    job.upload_input_data = Mock(return_value="uploaded-uri")
+    container_client = mock_container_client.from_container_url.return_value
+    container_client.list_blobs.return_value = []
+    container_client.get_blob_client.return_value.download_blob.return_value.readall.return_value = b"data"
+
+    job.upload_attachment("upload", b"data")
+    attachments = job.list_attachments()
+    downloaded = job.download_attachment("download")
+
+    workspace.get_container_uri.assert_not_called()
+    job.upload_input_data.assert_called_once_with(
+        container_uri=SIGNED_CONTAINER_URI,
+        blob_name="upload",
+        input_data=b"data",
+    )
+    assert mock_container_client.from_container_url.call_args_list == [
+        call(SIGNED_CONTAINER_URI),
+        call(SIGNED_CONTAINER_URI),
+    ]
+    assert attachments == []
+    assert downloaded == b"data"
+
+
+def test_upload_attachment_refreshes_expired_job_uri():
+    workspace = Mock()
+    workspace.get_container_uri.return_value = SIGNED_CONTAINER_URI
+    job = _job_with_container(container_uri=EXPIRED_CONTAINER_URI, workspace=workspace)
+    job.upload_input_data = Mock(return_value="uploaded-uri")
+
+    job.upload_attachment("attachment", b"data")
+
+    workspace.get_container_uri.assert_called_once_with(
+        job_id=JOB_ID,
+        container_name=DEFAULT_CONTAINER_NAME,
+    )
+    job.upload_input_data.assert_called_once_with(
+        container_uri=SIGNED_CONTAINER_URI,
+        blob_name="attachment",
+        input_data=b"data",
+    )
+
+
+def test_upload_attachment_refreshes_job_uri_without_write_permission():
+    workspace = Mock()
+    workspace.get_container_uri.return_value = SIGNED_CONTAINER_URI
+    job = _job_with_container(container_uri=READ_ONLY_CONTAINER_URI, workspace=workspace)
+    job.upload_input_data = Mock(return_value="uploaded-uri")
+
+    job.upload_attachment("attachment", b"data")
+
+    workspace.get_container_uri.assert_called_once_with(
+        job_id=JOB_ID,
+        container_name=DEFAULT_CONTAINER_NAME,
+    )
+    job.upload_input_data.assert_called_once_with(
+        container_uri=SIGNED_CONTAINER_URI,
+        blob_name="attachment",
+        input_data=b"data",
+    )
+
+
+@patch("azure.quantum.job.base_job.ContainerClient")
+def test_read_only_job_uri_is_reused_for_list_and_download(mock_container_client):
+    workspace = Mock()
+    job = _job_with_container(container_uri=READ_ONLY_CONTAINER_URI, workspace=workspace)
+    container_client = mock_container_client.from_container_url.return_value
+    container_client.list_blobs.return_value = []
+    container_client.get_blob_client.return_value.download_blob.return_value.readall.return_value = b"data"
+
+    attachments = job.list_attachments()
+    downloaded = job.download_attachment("download")
+
+    workspace.get_container_uri.assert_not_called()
+    assert mock_container_client.from_container_url.call_args_list == [
+        call(READ_ONLY_CONTAINER_URI),
+        call(READ_ONLY_CONTAINER_URI),
+    ]
+    assert attachments == []
+    assert downloaded == b"data"
+
+
+def test_upload_attachment_rejects_refreshed_storage_hostname_mismatch():
+    workspace = Mock()
+    workspace.get_container_uri.return_value = (
+        f"https://other-acct.blob.core.windows.net/{DEFAULT_CONTAINER_NAME}?sas"
+    )
+    job = _job_with_container(workspace=workspace)
+
+    with pytest.raises(ValueError, match="does not match job container hostname"):
+        job.upload_attachment("attachment", b"data")
+
+
+@patch("azure.quantum.job.base_job.ContainerClient")
+def test_attachment_methods_preserve_custom_container_name(mock_container_client):
+    custom_container_name = "custom-container"
+    custom_unsigned_uri = f"https://acct.blob.core.windows.net/{custom_container_name}"
+    custom_signed_uri = f"{custom_unsigned_uri}?sas"
+    workspace = Mock()
+    workspace.get_container_uri.return_value = custom_signed_uri
+    job = _job_with_container(container_uri=custom_unsigned_uri, workspace=workspace)
+    job.upload_input_data = Mock(return_value="uploaded-uri")
+    container_client = mock_container_client.from_container_url.return_value
+    container_client.list_blobs.return_value = []
+    container_client.get_blob_client.return_value.download_blob.return_value.readall.return_value = b"data"
+
+    job.upload_attachment("upload", b"data")
+    attachments = job.list_attachments()
+    downloaded = job.download_attachment("download")
+
+    workspace.get_container_uri.assert_has_calls(
+        [
+            call(job_id=JOB_ID, container_name=custom_container_name),
+            call(job_id=JOB_ID, container_name=custom_container_name),
+            call(job_id=JOB_ID, container_name=custom_container_name),
+        ]
+    )
+    assert workspace.get_container_uri.call_count == 3
+    job.upload_input_data.assert_called_once_with(
+        container_uri=custom_signed_uri,
+        blob_name="upload",
+        input_data=b"data",
+    )
+    assert mock_container_client.from_container_url.call_args_list == [
+        call(custom_signed_uri),
+        call(custom_signed_uri),
+    ]
+    assert attachments == []
+    assert downloaded == b"data"