88 lines
3.6 KiB
Python
88 lines
3.6 KiB
Python
"""Materialize opaque Agent runtime file handles from the OSS exchange prefix."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from arr_processing.errors import ProcessingTransportError
|
|
from arr_storage.exchange import OutputExchangeConfig
|
|
from arr_storage.remote import CloudClientError, CloudClientPort
|
|
|
|
|
|
class PrefixRemoteFilePort:
|
|
"""Resolve an opaque handle only within one controlled exchange prefix."""
|
|
|
|
def __init__(self, client: CloudClientPort, config: OutputExchangeConfig) -> None:
|
|
self._client = client
|
|
self._config = config
|
|
|
|
def materialize(self, file_handle: str, destination: Path, max_bytes: int) -> None:
|
|
try:
|
|
object_key = self._config.object_key(file_handle)
|
|
except ValueError:
|
|
raise ProcessingTransportError(
|
|
"PROCESSING_ARTIFACT_HANDLE_INVALID", retryable=False
|
|
) from None
|
|
if not isinstance(max_bytes, int) or isinstance(max_bytes, bool) or max_bytes < 0:
|
|
raise ProcessingTransportError(
|
|
"PROCESSING_ARTIFACT_REQUEST_INVALID", retryable=False
|
|
)
|
|
try:
|
|
descriptor = self._client.stat_object(object_key)
|
|
if descriptor.byte_size > max_bytes:
|
|
raise ProcessingTransportError(
|
|
"PROCESSING_ARTIFACT_TOO_LARGE", retryable=False
|
|
)
|
|
destination.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
|
|
output = os.open(destination, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
|
|
try:
|
|
copied = 0
|
|
with self._client.stream_object(object_key) as source, os.fdopen(
|
|
output, "wb"
|
|
) as target:
|
|
output = -1
|
|
while True:
|
|
chunk = source.read(1024 * 1024)
|
|
if not chunk:
|
|
break
|
|
if not isinstance(chunk, bytes):
|
|
raise OSError("remote output stream is invalid")
|
|
copied += len(chunk)
|
|
if copied > max_bytes:
|
|
raise ProcessingTransportError(
|
|
"PROCESSING_ARTIFACT_TOO_LARGE", retryable=False
|
|
)
|
|
target.write(chunk)
|
|
target.flush()
|
|
os.fsync(target.fileno())
|
|
if copied != descriptor.byte_size:
|
|
raise OSError("remote output size changed")
|
|
finally:
|
|
if output >= 0:
|
|
os.close(output)
|
|
except ProcessingTransportError:
|
|
destination.unlink(missing_ok=True)
|
|
raise
|
|
except FileExistsError:
|
|
raise ProcessingTransportError(
|
|
"PROCESSING_ARTIFACT_DESTINATION_CONFLICT", retryable=False
|
|
) from None
|
|
except CloudClientError as error:
|
|
destination.unlink(missing_ok=True)
|
|
code = {
|
|
"not_found": "PROCESSING_ARTIFACT_NOT_FOUND",
|
|
"forbidden": "PROCESSING_ARTIFACT_FORBIDDEN",
|
|
"invalid_request": "PROCESSING_ARTIFACT_REQUEST_INVALID",
|
|
}.get(error.kind, "PROCESSING_ARTIFACT_UNAVAILABLE")
|
|
raise ProcessingTransportError(
|
|
code,
|
|
retryable=error.kind in {"not_found", "unavailable"},
|
|
) from None
|
|
except OSError:
|
|
destination.unlink(missing_ok=True)
|
|
raise ProcessingTransportError(
|
|
"PROCESSING_ARTIFACT_UNAVAILABLE", retryable=True
|
|
) from None
|
|
|