def create_client(
stage_info: Dict[str, Any],
use_accelerate_endpoint: bool = False) -> BlobServiceClient:
"""Creates a client object with a stage credential.
Args:
stage_info: Information about the stage.
use_accelerate_endpoint: Not used for Azure client.
Returns:
The client to communicate with GCS.
"""
stage_credentials = stage_info["creds"]
sas_token = stage_credentials["AZURE_SAS_TOKEN"]
if sas_token and sas_token.startswith("?"):
sas_token = sas_token[1:]
end_point = stage_info["endPoint"]
if end_point.startswith("blob."):
end_point = end_point[len("blob."):]
client = BlobServiceClient(
account_url=
f"https://{stage_info['storageAccount']}.blob.{end_point}",
credential=sas_token,
)
client._config.retry_policy = ExponentialRetry(initial_backoff=1,
increment_base=2,
max_attempts=60,
random_jitter_range=2)
return client
def test_location_lock(self):
# Arrange
retry = ExponentialRetry(retry_to_secondary=True, initial_backoff=1, increment_base=2)
service = self._create_storage_service(
BlobServiceClient, self.settings, retry_policy=retry)
# Act
# Fail the first request and set the retry policy to retry to secondary
response_callback = ResponseCallback(status=200, new_status=408).override_first_status
#context = _OperationContext(location_lock=True)
# Assert
# Confirm that the first request gets retried to secondary
# The given test account must be GRS
def retry_callback(retry_count=None, location_mode=None, **kwargs):
self.assertEqual(LocationMode.SECONDARY, location_mode)
# Confirm that the second list request done with the same context sticks
# to the final location of the first list request (aka secondary) despite
# the client normally trying primary first
requests = []
def request_callback(request):
if not requests:
requests.append(request)
else:
self.assertNotEqual(-1, request.http_request.url.find('-secondary'))
containers = service.list_containers(
results_per_page=1, retry_hook=retry_callback)
next(containers)
next(containers)
def test_retry_to_secondary_with_put(self):
# Arrange
container_name = self.get_resource_name()
retry = ExponentialRetry(retry_to_secondary=True,
initial_backoff=1,
increment_base=2)
service = self._create_storage_service(BlobServiceClient,
self.settings,
retry_policy=retry)
# Act
try:
# Fail the first create attempt
response_callback = ResponseCallback(
status=201, new_status=408).override_first_status
# Assert
# Confirm that the create request does *not* get retried to secondary
# This should actually throw InvalidPermissions if sent to secondary,
# but validate the location_mode anyways.
def retry_callback(location_mode=None, **kwargs):
self.assertEqual(LocationMode.PRIMARY, location_mode)
with self.assertRaises(ResourceExistsError):
service.create_container(container_name,
raw_response_hook=response_callback,
retry_hook=retry_callback)
finally:
service.delete_container(container_name)
def test_retry_to_secondary_with_get(self):
# Arrange
container_name = self.get_resource_name()
retry = ExponentialRetry(retry_to_secondary=True,
initial_backoff=1,
increment_base=2)
service = self._create_storage_service(BlobServiceClient,
self.settings,
retry_policy=retry)
# Act
try:
container = service.create_container(container_name)
response_callback = ResponseCallback(
status=200, new_status=408).override_first_status
# Assert
# Confirm that the get request gets retried to secondary
def retry_callback(retry_count=None, location_mode=None, **kwargs):
# Only check this every other time, sometimes the secondary location fails due to delay
if retry_count % 2 == 0:
self.assertEqual(LocationMode.SECONDARY, location_mode)
container.get_container_properties(
raw_response_hook=response_callback, retry_hook=retry_callback)
finally:
service.delete_container(container_name)
def create_client(stage_info, use_accelerate_endpoint: bool = False):
"""Creates a client object with a stage credential.
Args:
stage_info: Information about the stage.
use_accelerate_endpoint: Not used for Azure client.
Returns:
The client to communicate with GCS.
"""
stage_credentials = stage_info['creds']
sas_token = stage_credentials['AZURE_SAS_TOKEN']
if sas_token and sas_token.startswith('?'):
sas_token = sas_token[1:]
end_point = stage_info['endPoint']
if end_point.startswith('blob.'):
end_point = end_point[len('blob.'):]
client = BlobServiceClient(
account_url="https://{}.blob.{}".format(
stage_info['storageAccount'],
end_point
),
credential=sas_token)
client._config.retry_policy = ExponentialRetry(
initial_backoff=1,
increment_base=2,
max_attempts=60,
random_jitter_range=2
)
return client
def test_exponential_retry(self, resource_group, location, storage_account,
storage_account_key):
# Arrange
container_name = self.get_resource_name('utcontainer')
retry = ExponentialRetry(initial_backoff=1,
increment_base=3,
retry_total=3)
service = self._create_storage_service(BlobServiceClient,
storage_account,
storage_account_key,
retry_policy=retry)
try:
container = service.create_container(container_name)
# Force the create call to 'timeout' with a 408
callback = ResponseCallback(status=200, new_status=408)
# Act
with self.assertRaises(HttpResponseError):
container.get_container_properties(
raw_response_hook=callback.override_status)
# Assert the response was called the right number of times (1 initial request + 3 retries)
self.assertEqual(callback.count, 1 + 3)
finally:
# Clean up
service.delete_container(container_name)
def test_secondary_location_mode(self):
# Arrange
container_name = self.get_resource_name()
retry = ExponentialRetry(initial_backoff=1, increment_base=2)
service = self._create_storage_service(BlobServiceClient,
self.settings,
retry_policy=retry)
# Act
try:
container = service.create_container(container_name)
container.location_mode = LocationMode.SECONDARY
# Override the response from secondary if it's 404 as that simply means
# the container hasn't replicated. We're just testing we try secondary,
# so that's fine.
response_callback = ResponseCallback(
status=404, new_status=200).override_first_status
# Assert
def request_callback(request):
self.assertNotEqual(
-1, request.http_request.url.find('-secondary'))
request_callback = request_callback
container.get_container_properties(
raw_request_hook=request_callback,
raw_response_hook=response_callback)
finally:
# Delete will go to primary, so disable the request validation
service.delete_container(container_name)
def test_invalid_account_key(self, resource_group, location,
storage_account, storage_account_key):
# Arrange
container_name = self.get_resource_name('utcontainer')
retry = ExponentialRetry(initial_backoff=1,
increment_base=3,
retry_total=3)
service = self._create_storage_service(BlobServiceClient,
storage_account,
storage_account_key,
retry_policy=retry)
service.credential.account_name = "dummy_account_name"
service.credential.account_key = "dummy_account_key"
# Shorten retries and add counter
retry_counter = RetryCounter()
retry_callback = retry_counter.simple_count
# Act
with self.assertRaises(ClientAuthenticationError):
service.create_container(container_name,
retry_callback=retry_callback)
# Assert
# No retry should be performed since the signing error is fatal
self.assertEqual(retry_counter.count, 0)
def test_exponential_retry_interval(self):
# Arrange
retry_policy = ExponentialRetry(initial_backoff=1,
increment_base=3,
random_jitter_range=3)
context_stub = {}
for i in range(10):
# Act
context_stub['count'] = 0
backoff = retry_policy.get_backoff_time(context_stub)
# Assert backoff interval is within +/- 3 of 1
self.assertTrue(0 <= backoff <= 4)
# Act
context_stub['count'] = 1
backoff = retry_policy.get_backoff_time(context_stub)
# Assert backoff interval is within +/- 3 of 4(1+3^1)
self.assertTrue(1 <= backoff <= 7)
# Act
context_stub['count'] = 2
backoff = retry_policy.get_backoff_time(context_stub)
# Assert backoff interval is within +/- 3 of 10(1+3^2)
self.assertTrue(7 <= backoff <= 13)
# Act
context_stub['count'] = 3
backoff = retry_policy.get_backoff_time(context_stub)
# Assert backoff interval is within +/- 3 of 28(1+3^3)
self.assertTrue(25 <= backoff <= 31)
def create_client(stage_info, use_accelerate_endpoint: bool = False):
"""Creates a client object with a stage credential.
Args:
stage_info: Information about the stage.
use_accelerate_endpoint: Not used for Azure client.
Returns:
The client to communicate with GCS.
"""
stage_credentials = stage_info['creds']
sas_token = stage_credentials['AZURE_SAS_TOKEN']
if sas_token and sas_token.startswith('?'):
sas_token = sas_token[1:]
end_point = stage_info['endPoint']
if end_point.startswith('blob.'):
end_point = end_point[len('blob.'):]
if use_new_azure_api:
client = BlobServiceClient(account_url="https://{}.blob.{}".format(
stage_info['storageAccount'], end_point),
credential=sas_token)
client._config.retry_policy = ExponentialRetry(
initial_backoff=1,
increment_base=2,
max_attempts=60,
random_jitter_range=2)
else:
client = BlockBlobService(
account_name=stage_info['storageAccount'],
sas_token=sas_token,
endpoint_suffix=end_point)
client._httpclient = RawBodyReadingClient(
session=requests.session(), protocol="https", timeout=2000)
client.retry = ExponentialRetry(initial_backoff=1,
increment_base=2,
max_attempts=60,
random_jitter_range=2).retry
return client
def list_blobs():
# 接続先BlobアカウントのURL作る
blob_url = "https://{}.blob.core.windows.net".format(
os.getenv("AZURE_STORAGE_ACCOUNT_NAME"))
# DefaultAzureCredentialを使い、Blobに接続するためのCredentialを自動で取得する。
# DefaultAzureCredentialを使うと、次の順番でCredentialの取得を試みる。
# なので、Azure上ではManaged IDの資格情報、ローカル開発環境上ではVSCodeの資格情報が使われるといったことが自動的に行われる。
# 1. EnvironmentCredential
# 環境変数に設定されてるCredentialを使う
# https://docs.microsoft.com/en-us/python/api/azure-identity/azure.identity.environmentcredential?view=azure-python
# 2. ManagedIdentityCredential
# AzureのManaged Identityを使う
# https://docs.microsoft.com/en-us/python/api/azure-identity/azure.identity.managedidentitycredential?view=azure-python
# 3. SharedTokenCacheCredential
# WindowsのVisual Studio等でログインした際のCredentialを使う
# https://docs.microsoft.com/en-us/python/api/azure-identity/azure.identity.sharedtokencachecredential?view=azure-python
# 4. VisualStudioCodeCredential
# Visual Studio CodeのAzure Account拡張機能でログインした際のCredentialを使う。
# Windows、macOS、Linux対応。
# https://marketplace.visualstudio.com/items?itemName=ms-vscode.azure-account
# https://docs.microsoft.com/en-us/python/api/azure-identity/azure.identity.visualstudiocodecredential?view=azure-python
# 5. AzureCliCredential
# AzureのCLIでログインした際のCredentialを使う。
# https://docs.microsoft.com/en-us/python/api/azure-identity/azure.identity.azureclicredential?view=azure-python
cred = DefaultAzureCredential()
# Blobに接続する際、パラメータを明示したExponentialRetryを使う。
# デフォルトだとExponentialRetryが使われるがその際のデフォルトパラメータは
# initial_backoff=15, increment_base=3, retry_total=3, random_jitter_range=3
#
# なので、リトライ分含め合計4回接続を試み、リトライの間隔は
# (15+3^1) = 18±3秒、(15+3^2) = 24±3秒、(15+3^3) = 42±3秒
# になるので、Flaskに接続してくるclientのHTTP Connectionを長時間保持したままになってしまう。
#
# それがイヤだったら明示的にパラメータを設定して早めにBlobに対してリトライをかける。
# このコードの例だと、
# (0.5+1.2^1) = 1.7±0.2秒、(0.5+1.2^2) = 1.94±0.2秒、(0.5+1.2^3) = 2.228±0.2秒
# の間隔でのリトライになる。
retry = ExponentialRetry(initial_backoff=0.5,
increment_base=1.2,
random_jitter_range=0.2)
client = BlobServiceClient(blob_url, cred, retry_policy=retry)
containers = client.list_containers()
container_names = [
container.get("name", "unknown") for container in containers
]
return ", ".join(container_names)
def test_retry_on_timeout(self):
# Arrange
container_name = self.get_resource_name()
retry = ExponentialRetry(initial_backoff=1, increment_base=2)
service = self._create_storage_service(
BlobServiceClient, self.settings, retry_policy=retry)
callback = ResponseCallback(status=201, new_status=408).override_status
# Act
try:
# The initial create will return 201, but we overwrite it and retry.
# The retry will then get a 409 and return false.
with self.assertRaises(ResourceExistsError):
service.create_container(container_name, raw_response_hook=callback)
finally:
service.delete_container(container_name)
def test_invalid_retry(self):
# Arrange
container_name = self.get_resource_name()
retry = ExponentialRetry(initial_backoff=1, increment_base=2)
service = self._create_storage_service(
BlobServiceClient, self.settings, retry_policy=retry)
# Force the create call to fail by pretending it's a teapot
callback = ResponseCallback(status=201, new_status=418).override_status
# Act
try:
with self.assertRaises(HttpResponseError) as error:
service.create_container(container_name, raw_response_hook=callback)
self.assertEqual(error.exception.response.status_code, 418)
self.assertEqual(error.exception.reason, 'Created')
finally:
service.delete_container(container_name)
def test_retry_with_deserialization(self):
# Arrange
container_name = self.get_resource_name(prefix='retry')
retry = ExponentialRetry(initial_backoff=1, increment_base=2)
service = self._create_storage_service(
BlobServiceClient, self.settings, retry_policy=retry)
try:
created = service.create_container(container_name)
# Act
callback = ResponseCallback(status=200, new_status=408).override_first_status
containers = service.list_containers(name_starts_with='retry', raw_response_hook=callback)
# Assert
containers = list(containers)
self.assertTrue(len(containers) >= 1)
finally:
service.delete_container(container_name)
def setUp(self):
super(StorageBlobRetryTest, self).setUp()
url = self._get_account_url()
credential = self._get_shared_key_credential()
retry = ExponentialRetry(initial_backoff=1,
increment_base=2,
retry_total=3)
self.bs = BlobServiceClient(url,
credential=credential,
retry_policy=retry)
self.container_name = self.get_resource_name('utcontainer')
if not self.is_playback():
try:
self.bs.create_container(self.container_name)
except ResourceExistsError:
pass
def test_retry_secondary(self, resource_group, location, storage_account,
storage_account_key):
"""Secondary location test.
This test is special, since in pratical term, we don't have time to wait
for the georeplication to be done (can take a loooooong time).
So for the purpose of this test, we fake a 408 on the primary request,
and then we check we do a 408. AND DONE.
It's not really perfect, since we didn't tested it would work on
a real geo-location.
Might be changed to live only as loooooong test with a polling on
the current geo-replication status.
"""
# Arrange
# Fail the first request and set the retry policy to retry to secondary
# The given test account must be GRS
class MockTransport(RequestsTransport):
CALL_NUMBER = 1
ENABLE = False
def send(self, request, **kwargs):
if MockTransport.ENABLE:
if MockTransport.CALL_NUMBER == 2:
if request.method != 'PUT':
assert '-secondary' in request.url
# Here's our hack
# Replace with primary so the test works even
# if secondary is not ready
request.url = request.url.replace('-secondary', '')
response = super(MockTransport, self).send(request, **kwargs)
if MockTransport.ENABLE:
assert response.status_code in [200, 201, 409]
if MockTransport.CALL_NUMBER == 1:
response.status_code = 408
elif MockTransport.CALL_NUMBER == 2:
if response.status_code == 409: # We can't really retry on PUT
response.status_code = 201
else:
pytest.fail(
"This test is not supposed to do more calls")
MockTransport.CALL_NUMBER += 1
return response
retry = ExponentialRetry(retry_to_secondary=True,
initial_backoff=1,
increment_base=2)
service = self._create_storage_service(BlobServiceClient,
storage_account,
storage_account_key,
retry_policy=retry,
transport=MockTransport())
# Act
MockTransport.ENABLE = True
# Assert
# Try put
def put_retry_callback(retry_count=None, location_mode=None, **kwargs):
# This call should be called once, with the decision to try secondary
put_retry_callback.called = True
if MockTransport.CALL_NUMBER == 1:
self.assertEqual(LocationMode.PRIMARY, location_mode)
elif MockTransport.CALL_NUMBER == 2:
self.assertEqual(LocationMode.PRIMARY, location_mode)
else:
pytest.fail(
"This test is not supposed to retry more than once")
put_retry_callback.called = False
container = service.get_container_client('containername')
created = container.create_container(retry_hook=put_retry_callback)
assert put_retry_callback.called
def retry_callback(retry_count=None, location_mode=None, **kwargs):
# This call should be called once, with the decision to try secondary
retry_callback.called = True
if MockTransport.CALL_NUMBER == 1:
self.assertEqual(LocationMode.SECONDARY, location_mode)
elif MockTransport.CALL_NUMBER == 2:
self.assertEqual(LocationMode.SECONDARY, location_mode)
else:
pytest.fail(
"This test is not supposed to retry more than once")
retry_callback.called = False
# Try list
MockTransport.CALL_NUMBER = 1
retry_callback.called = False
containers = service.list_containers(results_per_page=1,
retry_hook=retry_callback)
next(containers)
assert retry_callback.called
# Try get
MockTransport.CALL_NUMBER = 1
retry_callback.called = False
container.get_container_properties(retry_hook=retry_callback)
assert retry_callback.called
def setUp(self):
self.retry = ExponentialRetry(initial_backoff=1,
increment_base=2,
retry_total=3)
