nucliadb-utils 6.9.1.post5229__py3-none-any.whl → 6.10.0.post5732__py3-none-any.whl

nucliadb_utils/fastapi/versioning.py

@@ -20,9 +20,11 @@
 # This code is inspired by fastapi_versioning 1/3/2022 with MIT licence
 
 from collections import defaultdict
-from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, cast
+from collections.abc import Callable
+from typing import Any, Sequence, TypeVar, cast
 
 from fastapi import FastAPI
+from fastapi.middleware import Middleware
 from fastapi.routing import APIRoute
 from starlette.routing import BaseRoute
 
@@ -39,8 +41,8 @@ def version(major: int, minor: int = 0) -> Callable[[CallableT], CallableT]:  #
 
 def version_to_route(
     route: BaseRoute,
-    default_version: Tuple[int, int],
-) -> Tuple[Tuple[int, int], APIRoute]:  # pragma: no cover
+    default_version: tuple[int, int],
+) -> tuple[tuple[int, int], APIRoute]:  # pragma: no cover
     api_route = cast(APIRoute, route)
     version = getattr(api_route.endpoint, "_api_version", default_version)
     return version, api_route
@@ -50,17 +52,19 @@ def VersionedFastAPI(
     app: FastAPI,
     version_format: str = "{major}.{minor}",
     prefix_format: str = "/v{major}_{minor}",
-    default_version: Tuple[int, int] = (1, 0),
+    default_version: tuple[int, int] = (1, 0),
     enable_latest: bool = False,
-    kwargs: Optional[Dict[str, object]] = None,
+    middleware: Sequence[Middleware] | None = None,
+    kwargs: dict[str, object] | None = None,
 ) -> FastAPI:  # pragma: no cover
     kwargs = kwargs or {}
 
     parent_app = FastAPI(
         title=app.title,
+        middleware=middleware,
         **kwargs,  # type: ignore
     )
-    version_route_mapping: Dict[Tuple[int, int], List[APIRoute]] = defaultdict(list)
+    version_route_mapping: dict[tuple[int, int], list[APIRoute]] = defaultdict(list)
     version_routes = [version_to_route(route, default_version) for route in app.routes]
 
     for version, route in version_routes:

nucliadb_utils/featureflagging.py

@@ -18,7 +18,7 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 import json
-from typing import Any, Optional
+from typing import Any
 
 import mrflagly
 import pydantic_settings
@@ -28,7 +28,10 @@ from nucliadb_utils.settings import nuclia_settings, running_settings
 
 
 class Settings(pydantic_settings.BaseSettings):
-    flag_settings_url: Optional[str] = None
+    flag_settings_url: str | None = None
+
+    # temporary flag to test this FF enabled/disabled easily
+    disable_ask_decoupled_ff: bool = False
 
 
 DEFAULT_FLAG_DATA: dict[str, Any] = {
@@ -45,7 +48,10 @@ DEFAULT_FLAG_DATA: dict[str, Any] = {
         "rollout": 0,
         "variants": {"environment": ["local"]},
     },
-    const.Features.REBALANCE_ENABLED: {"rollout": 0, "variants": {"environment": ["local"]}},
+    const.Features.ASK_DECOUPLED: {
+        "rollout": 0,
+        "variants": {"environment": [] if Settings().disable_ask_decoupled_ff else ["local"]},
+    },
 }
 
 
@@ -57,7 +63,7 @@ class FlagService:
         else:
             self.flag_service = mrflagly.FlagService(url=settings.flag_settings_url)
 
-    def enabled(self, flag_key: str, default: bool = False, context: Optional[dict] = None) -> bool:
+    def enabled(self, flag_key: str, default: bool = False, context: dict | None = None) -> bool:
         if context is None:
             context = {}
         context["environment"] = running_settings.running_environment

nucliadb_utils/grpc.py

@@ -19,7 +19,6 @@
 
 import json
 import logging
-from typing import Optional
 
 from grpc import ChannelCredentials, aio
 
@@ -58,7 +57,7 @@ RETRY_OPTIONS = [
 def get_traced_grpc_channel(
     address: str,
     service_name: str,
-    credentials: Optional[ChannelCredentials] = None,
+    credentials: ChannelCredentials | None = None,
     variant: str = "",
     max_send_message: int = 100,
 ) -> aio.Channel:
@@ -75,7 +74,8 @@ def get_traced_grpc_channel(
         options = [
             ("grpc.max_receive_message_length", max_send_message * 1024 * 1024),
             ("grpc.max_send_message_length", max_send_message * 1024 * 1024),
-        ] + RETRY_OPTIONS
+            *RETRY_OPTIONS,
+        ]
         channel = aio.insecure_channel(address, options=options)
     return channel
 

nucliadb_utils/helpers.py

@@ -19,7 +19,7 @@
 #
 import asyncio
 import logging
-from typing import AsyncGenerator, Awaitable, Callable
+from collections.abc import AsyncGenerator, Awaitable, Callable
 
 from nucliadb_telemetry.errors import capture_exception
 

nucliadb_utils/nats.py

@@ -21,8 +21,9 @@ import asyncio
 import logging
 import sys
 import time
+from collections.abc import Awaitable, Callable
 from functools import cached_property, partial
-from typing import Any, Awaitable, Callable, Optional, Union
+from typing import Any
 
 import nats
 import nats.errors
@@ -67,7 +68,7 @@ class NatsMessageProgressUpdater(MessageProgressUpdater):
 
 
 class NatsConnectionManager:
-    _nc: Union[NATSClient, NatsClientTelemetry]
+    _nc: NATSClient | NatsClientTelemetry
     _subscriptions: list[tuple[Subscription, Callable[[], Awaitable[None]]]]
     _pull_subscriptions: list[
         tuple[
@@ -81,8 +82,8 @@ class NatsConnectionManager:
         *,
         service_name: str,
         nats_servers: list[str],
-        nats_creds: Optional[str] = None,
-        pull_utilization_metrics: Optional[Counter] = None,
+        nats_creds: str | None = None,
+        pull_utilization_metrics: Counter | None = None,
     ):
         self._service_name = service_name
         self._nats_servers = nats_servers
@@ -91,9 +92,9 @@ class NatsConnectionManager:
         self._pull_subscriptions = []
         self._lock = asyncio.Lock()
         self._healthy = True
-        self._last_unhealthy: Optional[float] = None
+        self._last_unhealthy: float | None = None
         self._needs_reconnection = False
-        self._reconnect_task: Optional[asyncio.Task] = None
+        self._reconnect_task: asyncio.Task | None = None
         self._expected_subscriptions: set[str] = set()
         self._initialized = False
         self.pull_utilization_metrics = pull_utilization_metrics
@@ -274,11 +275,11 @@ class NatsConnectionManager:
         logger.info("Connection is closed on NATS")
 
     @property
-    def nc(self) -> Union[NATSClient, NatsClientTelemetry]:
+    def nc(self) -> NATSClient | NatsClientTelemetry:
         return self._nc
 
     @cached_property
-    def js(self) -> Union[JetStreamContext, JetStreamContextTelemetry]:
+    def js(self) -> JetStreamContext | JetStreamContextTelemetry:
         return get_traced_jetstream(self._nc, self._service_name)
 
     async def subscribe(
@@ -291,7 +292,7 @@ class NatsConnectionManager:
         subscription_lost_cb: Callable[[], Awaitable[None]],
         flow_control: bool = False,
         manual_ack: bool = True,
-        config: Optional[nats.js.api.ConsumerConfig] = None,
+        config: nats.js.api.ConsumerConfig | None = None,
     ) -> Subscription:
         sub = await self.js.subscribe(
             subject=subject,
@@ -314,8 +315,8 @@ class NatsConnectionManager:
         stream: str,
         cb: Callable[[Msg], Awaitable[None]],
         subscription_lost_cb: Callable[[], Awaitable[None]],
-        durable: Optional[str] = None,
-        config: Optional[nats.js.api.ConsumerConfig] = None,
+        durable: str | None = None,
+        config: nats.js.api.ConsumerConfig | None = None,
     ) -> JetStreamContext.PullSubscription:
         wrapped_cb: Callable[[Msg], Awaitable[None]]
         if isinstance(self.js, JetStreamContextTelemetry):
@@ -370,9 +371,7 @@ class NatsConnectionManager:
 
         return psub
 
-    async def _remove_subscription(
-        self, subscription: Union[Subscription, JetStreamContext.PullSubscription]
-    ):
+    async def _remove_subscription(self, subscription: Subscription | JetStreamContext.PullSubscription):
         async with self._lock:
             for index, (sub, _) in enumerate(self._subscriptions):
                 if sub is not subscription:
@@ -391,7 +390,7 @@ class NatsConnectionManager:
                     pass
                 return
 
-    async def unsubscribe(self, subscription: Union[Subscription, JetStreamContext.PullSubscription]):
+    async def unsubscribe(self, subscription: Subscription | JetStreamContext.PullSubscription):
         await subscription.unsubscribe()
         await self._remove_subscription(subscription)
 
@@ -403,7 +402,7 @@ class NatsConnectionManager:
         while True:
             await asyncio.sleep(30)
 
-            existing_subs = set(sub._consumer for sub, _, _, _ in self._pull_subscriptions)
+            existing_subs = {sub._consumer for sub, _, _, _ in self._pull_subscriptions}
             missing_subs = self._expected_subscriptions - existing_subs
             if missing_subs:
                 logger.warning(f"Some NATS subscriptions are missing {missing_subs}")

nucliadb_utils/nuclia_usage/utils/kb_usage_report.py

@@ -22,7 +22,6 @@ import logging
 from collections.abc import Iterable
 from contextlib import suppress
 from datetime import datetime, timezone
-from typing import Optional
 
 from nats.js.client import JetStreamContext
 
@@ -90,14 +89,14 @@ class KbUsageReportUtility:
     def send_kb_usage(
         self,
         service: Service,
-        account_id: Optional[str],
-        kb_id: Optional[str],
+        account_id: str | None,
+        kb_id: str | None,
         kb_source: KBSource,
         processes: Iterable[Process] = (),
         predicts: Iterable[Predict] = (),
         searches: Iterable[Search] = (),
-        storage: Optional[Storage] = None,
-        activity_log_match: Optional[ActivityLogMatch] = None,
+        storage: Storage | None = None,
+        activity_log_match: ActivityLogMatch | None = None,
     ):
         usage = KbUsage()
         usage.service = service  # type: ignore

nucliadb_utils/run.py

@@ -21,7 +21,7 @@ import asyncio
 import inspect
 import logging
 import signal
-from typing import Awaitable, Callable
+from collections.abc import Awaitable, Callable
 
 logger = logging.getLogger(__name__)
 

nucliadb_utils/settings.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 from enum import Enum
-from typing import Dict, List, Optional
 
 from pydantic import AliasChoices, Field, model_validator
 from pydantic_settings import BaseSettings
@@ -26,7 +25,7 @@ from pydantic_settings import BaseSettings
 
 class RunningSettings(BaseSettings):
     debug: bool = False
-    sentry_url: Optional[str] = None
+    sentry_url: str | None = None
     running_environment: str = Field(
         default="local",
         validation_alias=AliasChoices("environment", "running_environment"),
@@ -42,7 +41,7 @@ running_settings = RunningSettings()
 
 
 class HTTPSettings(BaseSettings):
-    cors_origins: List[str] = ["*"]
+    cors_origins: list[str] = ["*"]
 
 
 http_settings = HTTPSettings()
@@ -70,47 +69,47 @@ class StorageSettings(BaseSettings):
         default=FileBackendConfig.NOT_SET, description="File backend storage type"
     )
 
-    gcs_base64_creds: Optional[str] = Field(
+    gcs_base64_creds: str | None = Field(
         default=None,
-        description="GCS JSON credentials of a service account encoded in Base64: https://cloud.google.com/iam/docs/service-account-overview",  # noqa
+        description="GCS JSON credentials of a service account encoded in Base64: https://cloud.google.com/iam/docs/service-account-overview",
     )
-    gcs_bucket: Optional[str] = Field(
+    gcs_bucket: str | None = Field(
         default=None,
         description="GCS Bucket name where files are stored: https://cloud.google.com/storage/docs/buckets",
     )
-    gcs_location: Optional[str] = Field(
+    gcs_location: str | None = Field(
         default=None,
         description="GCS Bucket location: https://cloud.google.com/storage/docs/locations",
     )
-    gcs_project: Optional[str] = Field(
+    gcs_project: str | None = Field(
         default=None,
-        description="Google Cloud Project ID: https://cloud.google.com/resource-manager/docs/creating-managing-projects",  # noqa
+        description="Google Cloud Project ID: https://cloud.google.com/resource-manager/docs/creating-managing-projects",
     )
-    gcs_bucket_labels: Dict[str, str] = Field(
+    gcs_bucket_labels: dict[str, str] = Field(
         default={},
-        description="Map of labels with which GCS buckets will be labeled with: https://cloud.google.com/storage/docs/tags-and-labels",  # noqa
+        description="Map of labels with which GCS buckets will be labeled with: https://cloud.google.com/storage/docs/tags-and-labels",
     )
     gcs_endpoint_url: str = "https://www.googleapis.com"
 
-    s3_client_id: Optional[str] = None
-    s3_client_secret: Optional[str] = None
+    s3_client_id: str | None = None
+    s3_client_secret: str | None = None
     s3_ssl: bool = True
     s3_verify_ssl: bool = True
     s3_max_pool_connections: int = 30
-    s3_endpoint: Optional[str] = None
-    s3_region_name: Optional[str] = None
-    s3_kms_key_id: Optional[str] = None
-    s3_bucket: Optional[str] = Field(default=None, description="KnowledgeBox S3 bucket name template")
-    s3_bucket_tags: Dict[str, str] = Field(
+    s3_endpoint: str | None = None
+    s3_region_name: str | None = None
+    s3_kms_key_id: str | None = None
+    s3_bucket: str | None = Field(default=None, description="KnowledgeBox S3 bucket name template")
+    s3_bucket_tags: dict[str, str] = Field(
         default={},
-        description="Map of tags with which S3 buckets will be tagged with: https://docs.aws.amazon.com/AmazonS3/latest/API/API_PutBucketTagging.html",  # noqa
+        description="Map of tags with which S3 buckets will be tagged with: https://docs.aws.amazon.com/AmazonS3/latest/API/API_PutBucketTagging.html",
     )
 
-    local_files: Optional[str] = Field(
+    local_files: str | None = Field(
         default=None,
         description="If using LOCAL `file_backend` storage, directory where files should be stored",
     )
-    local_indexing_bucket: Optional[str] = Field(
+    local_indexing_bucket: str | None = Field(
         default="indexer",
         description="If using LOCAL `file_backend` storage, subdirectory where indexing data is stored",
     )
@@ -119,28 +118,28 @@ class StorageSettings(BaseSettings):
         description="Number of days that uploaded files are kept in Nulia's processing engine",
     )
 
-    azure_account_url: Optional[str] = Field(
+    azure_account_url: str | None = Field(
         default=None,
-        description="Azure Account URL. The driver implementation uses Azure's default credential authentication method: https://learn.microsoft.com/en-us/python/api/azure-identity/azure.identity.defaultazurecredential?view=azure-python",  # noqa
+        description="Azure Account URL. The driver implementation uses Azure's default credential authentication method: https://learn.microsoft.com/en-us/python/api/azure-identity/azure.identity.defaultazurecredential?view=azure-python",
         examples=["https://<storageaccountname>.blob.core.windows.net"],
     )
 
-    azure_kb_account_url: Optional[str] = Field(
+    azure_kb_account_url: str | None = Field(
         default=None,
-        description="Azure Account URL for KB containers. If unspecified, uses `azure_account_url`",  # noqa
+        description="Azure Account URL for KB containers. If unspecified, uses `azure_account_url`",
         examples=["https://<storageaccountname>.blob.core.windows.net"],
     )
 
     # For testing purposes: Azurite docker image requires a connection string as it
     # doesn't support Azure's default credential authentication method
-    azure_connection_string: Optional[str] = None
+    azure_connection_string: str | None = None
 
 
 storage_settings = StorageSettings()
 
 
 class NucliaSettings(BaseSettings):
-    nuclia_service_account: Optional[str] = None
+    nuclia_service_account: str | None = None
     nuclia_public_url: str = "https://{zone}.nuclia.cloud"
     nuclia_processing_cluster_url: str = "http://processing-api.processing.svc.cluster.local:8080"
     nuclia_inner_predict_url: str = "http://predict.learning.svc.cluster.local:8080"
@@ -149,7 +148,7 @@ class NucliaSettings(BaseSettings):
     nuclia_zone: str = "europe-1"
     onprem: bool = True
 
-    nuclia_jwt_key: Optional[str] = None
+    nuclia_jwt_key: str | None = None
     nuclia_hash_seed: int = 42
     nuclia_partitions: int = 1
 
@@ -157,7 +156,7 @@ class NucliaSettings(BaseSettings):
     dummy_predict: bool = False
     dummy_learning_services: bool = False
     local_predict: bool = False
-    local_predict_headers: Dict[str, str] = {}
+    local_predict_headers: dict[str, str] = {}
 
     @model_validator(mode="before")
     @classmethod
@@ -171,15 +170,15 @@ nuclia_settings = NucliaSettings()
 
 
 class NucliaDBSettings(BaseSettings):
-    nucliadb_ingest: Optional[str] = "ingest-orm-grpc.nucliadb.svc.cluster.local:8030"
+    nucliadb_ingest: str | None = "ingest-orm-grpc.nucliadb.svc.cluster.local:8030"
 
 
 nucliadb_settings = NucliaDBSettings()
 
 
 class TransactionSettings(BaseSettings):
-    transaction_jetstream_auth: Optional[str] = None
-    transaction_jetstream_servers: List[str] = ["nats://localhost:4222"]
+    transaction_jetstream_auth: str | None = None
+    transaction_jetstream_servers: list[str] = ["nats://localhost:4222"]
     transaction_local: bool = False
     transaction_commit_timeout: int = Field(
         default=60, description="Transaction commit timeout in seconds"
@@ -190,10 +189,10 @@ transaction_settings = TransactionSettings()
 
 
 class IndexingSettings(BaseSettings):
-    index_jetstream_servers: List[str] = []
-    index_jetstream_auth: Optional[str] = None
+    index_jetstream_servers: list[str] = []
+    index_jetstream_auth: str | None = None
     index_local: bool = False
-    index_nidx_subject: Optional[str] = None
+    index_nidx_subject: str | None = None
     index_searcher_refresh_interval: float = 1.0
 
 
@@ -202,9 +201,9 @@ indexing_settings = IndexingSettings()
 
 class AuditSettings(BaseSettings):
     audit_driver: str = "basic"
-    audit_jetstream_target: Optional[str] = "audit.{partition}.{type}"
-    audit_jetstream_servers: List[str] = []
-    audit_jetstream_auth: Optional[str] = None
+    audit_jetstream_target: str | None = "audit.{partition}.{type}"
+    audit_jetstream_servers: list[str] = []
+    audit_jetstream_auth: str | None = None
     audit_partitions: int = 3
     audit_stream: str = "audit"
     audit_hash_seed: int = 1234
@@ -214,9 +213,9 @@ audit_settings = AuditSettings()
 
 
 class UsageSettings(BaseSettings):
-    usage_jetstream_subject: Optional[str] = "kb-usage.nuclia_db"
-    usage_jetstream_servers: List[str] = []
-    usage_jetstream_auth: Optional[str] = None
+    usage_jetstream_subject: str | None = "kb-usage.nuclia_db"
+    usage_jetstream_servers: list[str] = []
+    usage_jetstream_auth: str | None = None
     usage_stream: str = "kb-usage"
 
 

nucliadb_utils/signals.py

@@ -18,10 +18,10 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 import asyncio
-from collections.abc import Awaitable
+from collections.abc import Awaitable, Callable
 from enum import Enum
 from inspect import iscoroutinefunction
-from typing import Any, Callable, Type
+from typing import Any
 
 from nucliadb_telemetry.errors import capture_exception
 from nucliadb_utils import logger
@@ -33,7 +33,7 @@ class ListenerPriority(Enum):
 
 
 class Signal:
-    def __init__(self, payload_model: Type):
+    def __init__(self, payload_model: type):
         self.payload_model_type = payload_model
         self.callbacks: dict[str, tuple[Callable[..., Awaitable], int]] = {}
 

nucliadb_utils/storages/azure.py

@@ -20,9 +20,10 @@
 
 from __future__ import annotations
 
+import base64
 import logging
+from collections.abc import AsyncGenerator, AsyncIterator
 from datetime import datetime
-from typing import AsyncGenerator, AsyncIterator, Optional, Union
 
 from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError
 from azure.identity import DefaultAzureCredential
@@ -65,7 +66,7 @@ class AzureStorageField(StorageField):
             origin_bucket_name, origin_uri, destination_bucket_name, destination_uri
         )
 
-    async def iter_data(self, range: Optional[Range] = None) -> AsyncGenerator[bytes, None]:
+    async def iter_data(self, range: Range | None = None) -> AsyncGenerator[bytes]:
         if self.field is not None:
             bucket = self.field.bucket_name
             key = self.field.uri
@@ -132,7 +133,7 @@ class AzureStorageField(StorageField):
         self.field.ClearField("upload_uri")
         self.field.ClearField("parts")
 
-    async def exists(self) -> Optional[ObjectMetadata]:
+    async def exists(self) -> ObjectMetadata | None:
         key = None
         bucket = None
         if self.field is not None and self.field.uri != "":
@@ -170,9 +171,9 @@ class AzureStorage(Storage):
         self,
         account_url: str,
         kb_account_url: str,
-        deadletter_bucket: Optional[str] = "deadletter",
-        indexing_bucket: Optional[str] = "indexing",
-        connection_string: Optional[str] = None,
+        deadletter_bucket: str | None = "deadletter",
+        indexing_bucket: str | None = "indexing",
+        connection_string: str | None = None,
     ):
         self.object_store = AzureObjectStore(account_url, connection_string=connection_string)
         self.kb_object_store = AzureObjectStore(kb_account_url, connection_string=connection_string)
@@ -185,7 +186,7 @@ class AzureStorage(Storage):
         else:
             return self.kb_object_store
 
-    async def initialize(self, service_name: Optional[str] = None):
+    async def initialize(self, service_name: str | None = None):
         await self.object_store.initialize()
         await self.kb_object_store.initialize()
         for bucket in [
@@ -208,7 +209,7 @@ class AzureStorage(Storage):
         except KeyError:
             pass
 
-    async def create_bucket(self, bucket_name: str, kbid: Optional[str] = None):
+    async def create_bucket(self, bucket_name: str, kbid: str | None = None):
         if await self.object_store_for_bucket(bucket_name).bucket_exists(bucket_name):
             return
         await self.object_store_for_bucket(bucket_name).bucket_create(bucket_name)
@@ -230,8 +231,8 @@ class AzureStorage(Storage):
         return await self.kb_object_store.bucket_delete(bucket_name)
 
     async def iterate_objects(
-        self, bucket: str, prefix: str, start: Optional[str] = None
-    ) -> AsyncGenerator[ObjectInfo, None]:
+        self, bucket: str, prefix: str, start: str | None = None
+    ) -> AsyncGenerator[ObjectInfo]:
         async for obj in self.object_store_for_bucket(bucket).iterate(bucket, prefix, start):
             yield obj
 
@@ -240,10 +241,10 @@ class AzureStorage(Storage):
 
 
 class AzureObjectStore(ObjectStore):
-    def __init__(self, account_url: str, connection_string: Optional[str] = None):
+    def __init__(self, account_url: str, connection_string: str | None = None):
         self.account_url = account_url
         self.connection_string = connection_string
-        self._service_client: Optional[BlobServiceClient] = None
+        self._service_client: BlobServiceClient | None = None
 
     @property
     def service_client(self) -> BlobServiceClient:
@@ -344,11 +345,11 @@ class AzureObjectStore(ObjectStore):
         self,
         bucket: str,
         key: str,
-        data: Union[bytes, AsyncGenerator[bytes, None]],
+        data: bytes | AsyncGenerator[bytes],
         metadata: ObjectMetadata,
     ) -> None:
         container_client = self.service_client.get_container_client(bucket)
-        length: Optional[int] = None
+        length: int | None = None
         if isinstance(data, bytes):
             length = len(data)
             metadata.size = length
@@ -384,8 +385,8 @@ class AzureObjectStore(ObjectStore):
         return await downloader.readall()
 
     async def download_stream(
-        self, bucket: str, key: str, range: Optional[Range] = None
-    ) -> AsyncGenerator[bytes, None]:
+        self, bucket: str, key: str, range: Range | None = None
+    ) -> AsyncGenerator[bytes]:
         range = range or Range()
         container_client = self.service_client.get_container_client(bucket)
         blob_client = container_client.get_blob_client(key)
@@ -405,8 +406,8 @@ class AzureObjectStore(ObjectStore):
             yield chunk
 
     async def iterate(
-        self, bucket: str, prefix: str, start: Optional[str] = None
-    ) -> AsyncGenerator[ObjectInfo, None]:
+        self, bucket: str, prefix: str, start: str | None = None
+    ) -> AsyncGenerator[ObjectInfo]:
         container_client = self.service_client.get_container_client(bucket)
         async for blob in container_client.list_blobs(name_starts_with=prefix):
             if start and blob.name <= start:
@@ -426,13 +427,20 @@ class AzureObjectStore(ObjectStore):
     @ops_observer.wrap({"type": "multipart_start"})
     async def upload_multipart_start(self, bucket: str, key: str, metadata: ObjectMetadata) -> None:
         container_client = self.service_client.get_container_client(bucket)
-        custom_metadata = {key: str(value) for key, value in metadata.model_dump().items()}
+        custom_metadata = {
+            "base64_filename": base64.b64encode(metadata.filename.encode()).decode(),
+            "content_type": metadata.content_type,
+            "size": str(metadata.size),
+        }
         blob_client = container_client.get_blob_client(key)
+        safe_filename = (
+            metadata.filename.encode("ascii", "replace").decode().replace('"', "").replace("\n", "")
+        )
         await blob_client.create_append_blob(
             metadata=custom_metadata,
             content_settings=ContentSettings(
                 content_type=metadata.content_type,
-                content_disposition=f"attachment; filename={metadata.filename}",
+                content_disposition=f'attachment; filename="{safe_filename}"',
             ),
         )
 
@@ -460,7 +468,12 @@ def parse_object_metadata(properties: BlobProperties, key: str) -> ObjectMetadat
         size = int(custom_metadata_size)
     else:
         size = properties.size
-    filename = custom_metadata.get("filename") or key.split("/")[-1]
+
+    b64_filename = custom_metadata.get("base64_filename")
+    if b64_filename:
+        filename = base64.b64decode(b64_filename.encode()).decode()
+    else:
+        filename = key.split("/")[-1]
     content_type = custom_metadata.get("content_type") or properties.content_settings.content_type or ""
     return ObjectMetadata(
         filename=filename,