# Copyright (C) 2016-2026 The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information
import asyncio
import contextlib
from datetime import datetime, timedelta, timezone
from itertools import product
import logging
import string
from typing import Dict, Iterable, Iterator, Mapping, Optional, Self, Tuple, Union
from urllib.parse import parse_qs, urlparse
import warnings
import aiohttp
from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError
from azure.core.pipeline.transport import AioHttpTransport
from azure.storage.blob import (
BlobSasPermissions,
ContainerClient,
ContainerSasPermissions,
generate_blob_sas,
generate_container_sas,
)
from azure.storage.blob.aio import BlobClient as AsyncBlobClient
from azure.storage.blob.aio import ContainerClient as AsyncContainerClient
from swh.objstorage.constants import LiteralPrimaryHash
from swh.objstorage.exc import ObjNotFoundError
from swh.objstorage.interface import HashDict
from swh.objstorage.objstorage import CompressionFormat, ObjStorage, timed
logger = logging.getLogger(__name__)
[docs]
def get_container_url(
account_name: str,
account_key: str,
container_name: str,
access_policy: str = "read_only",
expiry: timedelta = timedelta(days=365),
container_url_template: str = (
"https://{account_name}.blob.core.windows.net/{container_name}?{signature}"
),
**kwargs,
) -> str:
"""Get the full url, for the given container on the given account, with a
Shared Access Signature granting the specified access policy.
Args:
account_name: name of the storage account for which to generate the URL
account_key: shared account key of the storage account used to generate the SAS
container_name: name of the container for which to grant access in the storage
account
access_policy: one of ``read_only``, ``append_only``, ``full``
expiry: the interval in the future with which the signature will expire
Returns:
the full URL of the container, with the shared access signature.
"""
access_policies = {
"read_only": ContainerSasPermissions(
read=True, list=True, delete=False, write=False
),
"append_only": ContainerSasPermissions(
read=True, list=True, delete=False, write=True
),
"full": ContainerSasPermissions(read=True, list=True, delete=True, write=True),
}
current_time = datetime.utcnow()
signature = generate_container_sas(
account_name,
container_name,
account_key=account_key,
permission=access_policies[access_policy],
start=current_time + timedelta(minutes=-1),
expiry=current_time + expiry,
)
return container_url_template.format(
account_name=account_name,
container_name=container_name,
signature=signature,
)
class _BaseAzureCloudObjStorage(ObjStorage):
primary_hash: LiteralPrimaryHash = "sha1"
def __init__(
self,
*,
compression: CompressionFormat | None = None,
use_secondary_endpoint_for_downloads=False,
**kwargs,
):
super().__init__(**kwargs)
if compression is None:
logger.warning(
"Deprecated: compression is undefined. "
"Defaulting to gzip, but please set it explicitly."
)
compression = "gzip"
self.compression = compression
self.use_secondary = use_secondary_endpoint_for_downloads
self._entered = False
self._exit_stack = contextlib.ExitStack()
self._async_loop = asyncio.new_event_loop()
self._async_exit_stack = contextlib.AsyncExitStack()
def __enter__(self) -> Self:
if self._entered:
raise RuntimeError(
f"{self.__class__.__name__} is not a re-entrant context manager"
)
self._entered = True
return self
def __exit__(self, *exc_details):
self.close()
def __del__(self):
self.close()
def close(self):
try:
async_loop = self._async_loop
except AttributeError:
# We can't clean up resources. The interpreter is probably shutting down anyway,
# so it's not a big deal.
pass
else:
if not async_loop.is_closed():
async_loop.run_until_complete(self._async_exit_stack.aclose())
async_loop.run_until_complete(self._async_loop.shutdown_asyncgens())
async_loop.close()
try:
exit_stack = self._exit_stack
except AttributeError:
# ditto
pass
else:
exit_stack.close()
def get_async_container_clients(self) -> Iterable[AsyncContainerClient]:
"""Returns a collection of container clients, to be passed to
``get_async_blob_client``.
Each container may not be used in more than one asyncio loop."""
raise NotImplementedError(
f"{self.__class__.__name__}.get_async_container_client"
)
def get_async_blob_client(self, hex_obj_id) -> AsyncBlobClient:
"""Get the azure blob client for the given hex obj id and a collection
yielded by ``get_async_container_clients``."""
raise NotImplementedError(f"{self.__class__.__name__}.get_async_blob_client")
def _internal_id(self, obj_id: HashDict) -> str:
"""Internal id is the hex version in objstorage."""
primary_hash = obj_id[self.primary_hash]
return primary_hash.hex()
def check_config(self, *, check_write):
"""Check the configuration for this object storage"""
return self._async_loop.run_until_complete(
self._check_config_async(check_write=check_write)
)
async def _check_config_async(self, *, check_write):
"""Check the configuration for this object storage"""
now = datetime.now(tz=timezone.utc).isoformat()
for container_client in self.get_async_container_clients():
parsed = urlparse(container_client.url)
# The query string for the container_client url contains a bunch of
# fields that are associated with the permissions given by the
# shared access signature. We treat all fields as optional, this is
# not a 100% bullet-proof set of checks, for instance if a
# connection_string is used
qs = parse_qs(parsed.query)
# st is the "signed start" (the signature is only valid after that
# date) as an ISO-8601 encoded datetime
if "st" in qs and qs["st"][0] > now:
raise ValueError(
"Shared access signature is not valid yet: %s" % qs["st"][0]
)
# se is the "signed expiry" (the signature is invalid after that
# date) as an ISO-8601 encoded datetime
if "se" in qs and qs["se"][0] < now:
raise ValueError(
"Shared access signature has expired: %s" % qs["se"][0]
)
# sr is the "signed resource". "c" means container.
if "sr" in qs and "c" not in qs["sr"][0]:
raise ValueError(
"Shared access signature is not for a container service"
)
# sp is the "signed permissions"
if (
"sp" in qs
and check_write
# "c" allows to create new objects
and "c" not in qs["sp"][0]
# "w" allows to write to objects. Either permission is valid to
# write to an objstorage
and "w" not in qs["sp"][0]
):
# We have neither "c" or "w" permissions, this is read-only
return False
return True
@timed
def __contains__(self, obj_id: HashDict) -> bool:
"""Does the storage contains the obj_id."""
return self._async_loop.run_until_complete(self._contains_async(obj_id))
async def _contains_async(self, obj_id: HashDict) -> bool:
"""Coroutine implementing ``__contains__(obj_id)`` using azure-storage-blob's
asynchronous implementation.
"""
hex_obj_id = self._internal_id(obj_id)
client = self.get_async_blob_client(hex_obj_id)
try:
await client.get_blob_properties()
except ResourceNotFoundError:
return False
else:
return True
@timed
def add(
self, content: bytes, obj_id: HashDict, check_presence: bool = True
) -> None:
"""Add an obj in storage if it's not there already."""
self._async_loop.run_until_complete(
self._add_async(content, obj_id, check_presence)
)
async def _add_async(
self,
content: bytes,
obj_id: HashDict,
check_presence: bool = True,
) -> Tuple[int, int]:
"""Coroutine implementing ``add(obj_id)`` using azure-storage-blob's
asynchronous implementation.
While ``add(obj_id)`` does not need asynchronicity, this is useful to
``add_batch(obj_ids)``, as it can run multiple ``_add_async`` tasks
concurrently.
Returns the number of inserted contents, and their size
(ie. `(0, 0)` if the content already existed, `(1, len(content))` otherwise)."""
if check_presence and await self._contains_async(obj_id):
return (0, 0)
hex_obj_id = self._internal_id(obj_id)
# Send the compressed content
data = self.compress(content)
client = self.get_async_blob_client(hex_obj_id)
try:
await client.upload_blob(data=data, length=len(data))
except ResourceExistsError:
# There's a race condition between check_presence and upload_blob,
# that we can't get rid of as the azure api doesn't allow atomic
# replaces or renaming a blob. As the restore operation explicitly
# removes the blob, it should be safe to just ignore the error.
return (0, 0)
else:
return (1, len(content))
async def _add_batch_async(
self, contents: Iterable[Tuple[HashDict, bytes]], check_presence: bool = True
) -> Dict:
try:
results = await asyncio.gather(
*[
self._add_async(content, obj_id, check_presence)
for (obj_id, content) in contents
]
)
summary = {"object:add": 0, "object:add:bytes": 0}
for added, length in results:
summary["object:add"] += added
summary["object:add:bytes"] += length
except Exception:
import traceback
traceback.print_exc()
raise
return summary
@timed
def add_batch(
self, contents: Iterable[Tuple[HashDict, bytes]], check_presence: bool = True
) -> Dict:
return self._async_loop.run_until_complete(
self._add_batch_async(contents, check_presence)
)
def restore(self, content: bytes, obj_id: HashDict) -> None:
"""Restore a content."""
if obj_id in self:
self.delete(obj_id)
return self.add(content, obj_id, check_presence=False)
@timed
def get(self, obj_id: HashDict) -> bytes:
"""retrieve blob's content if found."""
return self._async_loop.run_until_complete(self._get_async(obj_id))
async def _get_async(self, obj_id, container_clients=None):
"""Coroutine implementing ``get(obj_id)`` using azure-storage-blob's
asynchronous implementation.
While ``get(obj_id)`` does not need asynchronicity, this is useful to
``get_batch(obj_ids)``, as it can run multiple ``_get_async`` tasks
concurrently."""
hex_obj_id = self._internal_id(obj_id)
client = self.get_async_blob_client(hex_obj_id)
try:
download = await client.download_blob()
except ResourceNotFoundError:
raise ObjNotFoundError(obj_id) from None
else:
data = await download.content_as_bytes()
return self.decompress(data, hex_obj_id)
async def _get_async_or_none(self, obj_id):
"""Like ``get_async(obj_id)``, but returns None instead of raising
ResourceNotFoundError. Used by ``get_batch`` so other blobs can be returned
even if one is missing."""
try:
return await self._get_async(obj_id)
except ObjNotFoundError:
return None
async def _get_batch_async(self, obj_ids):
return await asyncio.gather(
*[self._get_async_or_none(obj_id) for obj_id in obj_ids]
)
def get_batch(self, obj_ids: Iterable[HashDict]) -> Iterator[Optional[bytes]]:
"""Retrieve objects' raw content in bulk from storage, concurrently."""
return self._async_loop.run_until_complete(self._get_batch_async(obj_ids))
def delete(self, obj_id: HashDict):
"""Delete an object."""
super().delete(obj_id) # Check delete permission
return self._async_loop.run_until_complete(self._delete_async(obj_id))
async def _delete_async(self, obj_id: HashDict) -> bool:
hex_obj_id = self._internal_id(obj_id)
client = self.get_async_blob_client(hex_obj_id)
try:
await client.delete_blob()
except ResourceNotFoundError:
raise ObjNotFoundError(obj_id) from None
return True
def download_url(
self,
obj_id: HashDict,
content_disposition: Optional[str] = None,
expiry: Optional[timedelta] = None,
) -> Optional[str]:
return self._async_loop.run_until_complete(
self._download_url_async(obj_id, content_disposition, expiry)
)
async def _download_url_async(
self,
obj_id: HashDict,
content_disposition: Optional[str] = None,
expiry: Optional[timedelta] = None,
) -> Optional[str]:
hex_obj_id = self._internal_id(obj_id)
client = self.get_async_blob_client(hex_obj_id)
try:
await client.get_blob_properties()
except ResourceNotFoundError:
raise ObjNotFoundError(obj_id)
else:
assert client.account_name is not None
signature = generate_blob_sas(
client.account_name,
client.container_name,
hex_obj_id,
account_key=client.credential.account_key,
permission=BlobSasPermissions(read=True),
expiry=datetime.now() + (expiry or timedelta(hours=24)),
content_disposition=content_disposition,
)
if self.use_secondary:
return f"{client.secondary_endpoint}?{signature}"
else:
return f"{client.primary_endpoint}?{signature}"
[docs]
class AzureCloudObjStorage(_BaseAzureCloudObjStorage):
"""ObjStorage backend for Azure blob storage accounts.
Args:
container_url: the URL of the container in which the objects are stored.
account_name: (deprecated) the name of the storage account under which objects are
stored
api_secret_key: (deprecated) the shared account key
container_name: (deprecated) the name of the container under which objects are
stored
compression: the compression algorithm used to compress objects in storage
connection_limit: maximum number of HTTP connections to Azure
use_secondary_endpoint_for_downloads: if True, use the secondary endpoint
url to generate download URLs. To configure the secondary endpoint, use
the BlobSecondaryEndpoint entry of the connection string.
Notes:
The container url should contain the credentials via a "Shared Access
Signature". The :func:`get_container_url` helper can be used to generate
such a URL from the account's access keys. The ``account_name``,
``api_secret_key`` and ``container_name`` arguments are deprecated.
"""
name: str = "azure"
def __init__(
self,
*,
container_url: Optional[str] = None,
account_name: Optional[str] = None,
api_secret_key: Optional[str] = None,
container_name: Optional[str] = None,
connection_string: Optional[str] = None,
compression: CompressionFormat | None = None,
connection_limit: int = 30, # avg batch size + 20%
**kwargs,
):
if container_url is None and connection_string is None:
if account_name is None or api_secret_key is None or container_name is None:
raise ValueError(
"AzureCloudObjStorage must have a container_url, a connection_string,"
"or all three account_name, api_secret_key and container_name"
)
else:
warnings.warn(
"The Azure objstorage account secret key parameters are "
"deprecated, please use container URLs instead.",
DeprecationWarning,
)
container_url = get_container_url(
account_name=account_name,
account_key=api_secret_key,
container_name=container_name,
access_policy="full",
)
elif connection_string:
if container_name is None:
raise ValueError(
"container_name is required when using connection_string."
)
self.container_name = container_name
super().__init__(**kwargs, compression=compression)
self.container_url = container_url
self.connection_string = connection_string
async def create_connector() -> aiohttp.TCPConnector:
# aiohttp.TCPConnector is a sync function, but it needs to be initialized
# while an event loop is running
return aiohttp.TCPConnector(limit=connection_limit)
connector = self._async_loop.run_until_complete(create_connector())
session = aiohttp.ClientSession(connector=connector)
self._transport = AioHttpTransport(session=session)
self._container_client: Optional[ContainerClient] = None
if self.connection_string:
self._async_container_client = AsyncContainerClient.from_connection_string(
self.connection_string, self.container_name, transport=self._transport
)
else:
assert self.container_url is not None
self._async_container_client = AsyncContainerClient.from_container_url(
self.container_url, transport=self._transport
)
self._async_loop.run_until_complete(
self._async_exit_stack.enter_async_context(self._async_container_client)
)
[docs]
def get_async_container_clients(self) -> Iterable[AsyncContainerClient]:
yield self._async_container_client
[docs]
def get_async_blob_client(self, hex_obj_id) -> AsyncBlobClient:
return self._async_container_client.get_blob_client(blob=hex_obj_id)
[docs]
class PrefixedAzureCloudObjStorage(_BaseAzureCloudObjStorage):
"""ObjStorage with azure capabilities, striped by prefix.
Args:
connection_limit_per_container: maximum number of HTTP connections
to each Azure container.
accounts is a dict containing entries of the form:
<prefix>: <container_url_for_prefix>
"""
def __init__(
self,
accounts: Mapping[str, Union[str, Dict[str, str]]],
name: str = "azure-prefixed",
compression: CompressionFormat | None = None,
# rationale: avg batch size is 25, and we have 16 containers in prod
# so it is unlikely for a batch to insert more than 6 objects in
# the same container.
connection_limit_per_container: int = 6,
**kwargs,
):
super().__init__(**kwargs, compression=compression)
self.name = name
# Definition sanity check
prefix_lengths = set(len(prefix) for prefix in accounts)
if not len(prefix_lengths) == 1:
raise ValueError(
"Inconsistent prefixes, found lengths %s"
% ", ".join(str(lst) for lst in sorted(prefix_lengths))
)
self.prefix_len = prefix_lengths.pop()
expected_prefixes = set(
"".join(letters)
for letters in product(
set(string.hexdigits.lower()), repeat=self.prefix_len
)
)
missing_prefixes = expected_prefixes - set(accounts)
if missing_prefixes:
raise ValueError(
"Missing prefixes %s" % ", ".join(sorted(missing_prefixes))
)
do_warning = False
self.container_urls: Dict[str, str] = {}
for prefix, container_url in accounts.items():
if isinstance(container_url, dict):
do_warning = True
container_url = get_container_url(
account_name=container_url["account_name"],
account_key=container_url["api_secret_key"],
container_name=container_url["container_name"],
access_policy="full",
)
self.container_urls[prefix] = container_url
if do_warning:
warnings.warn(
"The Azure objstorage account secret key parameters are "
"deprecated, please use container URLs instead.",
DeprecationWarning,
)
self._async_container_clients: dict[str, AsyncContainerClient] = {}
async def create_container_client(prefix, url) -> None:
connector = aiohttp.TCPConnector(limit=connection_limit_per_container)
session = aiohttp.ClientSession(connector=connector)
transport = AioHttpTransport(session=session)
container_client = AsyncContainerClient.from_container_url(
url, transport=transport
)
await self._async_exit_stack.enter_async_context(container_client)
self._async_container_clients[prefix] = container_client
async def create_container_clients() -> None:
await asyncio.gather(
*[
create_container_client(prefix, url)
for (prefix, url) in self.container_urls.items()
],
)
self._async_loop.run_until_complete(create_container_clients())
[docs]
def get_async_container_clients(self) -> Iterable[AsyncContainerClient]:
yield from self._async_container_clients.values()
[docs]
def get_async_blob_client(self, hex_obj_id) -> AsyncBlobClient:
prefix = hex_obj_id[: self.prefix_len]
return self._async_container_clients[prefix].get_blob_client(blob=hex_obj_id)
