Skip to content

Commit 88d3875

Browse files
seesharprunseesharprun
authored
Add support for Azure Cosmos DB for Table Entra auth (#38121)
1 parent b1585d0 commit 88d3875

11 files changed

+162
-76
lines changed

‎sdk/tables/azure-data-tables/CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
* Added to support custom encoder in entity CRUD operations.
77
* Added to support custom Entity type.
88
* Added to support Entity property in Tuple and Enum types.
9+
* Added support for Microsoft Entra auth with Azure Cosmos DB for Table's OAuth scope (`https://cosmos.azure.com/.default`).
910

1011
### Bugs Fixed
1112
* Fixed a bug in encoder when Entity property has "@odata.type" provided.

‎sdk/tables/azure-data-tables/assets.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,5 +2,5 @@
22
"AssetsRepo": "Azure/azure-sdk-assets",
33
"AssetsRepoPrefixPath": "python",
44
"TagPrefix": "python/tables/azure-data-tables",
5-
"Tag": "python/tables/azure-data-tables_1fb1a4af1a"
5+
"Tag": "python/tables/azure-data-tables_032477adf3"
66
}

‎sdk/tables/azure-data-tables/azure/data/tables/_authentication.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919

2020
from ._common_conversion import _sign_string
2121
from ._error import _wrap_exception
22-
from ._constants import STORAGE_OAUTH_SCOPE
22+
from ._constants import STORAGE_OAUTH_SCOPE, COSMOS_OAUTH_SCOPE
2323

2424

2525
class AzureSigningError(ClientAuthenticationError):
@@ -236,11 +236,15 @@ def _configure_credential(credential: None) -> None: ...
236236

237237

238238
def _configure_credential(
239-
credential: Optional[Union[AzureNamedKeyCredential, AzureSasCredential, TokenCredential, SharedKeyCredentialPolicy]]
239+
credential: Optional[
240+
Union[AzureNamedKeyCredential, AzureSasCredential, TokenCredential, SharedKeyCredentialPolicy]
241+
],
242+
cosmos_endpoint: bool = False,
240243
) -> Optional[Union[BearerTokenChallengePolicy, AzureSasCredentialPolicy, SharedKeyCredentialPolicy]]:
241244
if hasattr(credential, "get_token"):
242245
credential = cast(TokenCredential, credential)
243-
return BearerTokenChallengePolicy(credential, STORAGE_OAUTH_SCOPE)
246+
scope = COSMOS_OAUTH_SCOPE if cosmos_endpoint else STORAGE_OAUTH_SCOPE
247+
return BearerTokenChallengePolicy(credential, scope)
244248
if isinstance(credential, SharedKeyCredentialPolicy):
245249
return credential
246250
if isinstance(credential, AzureSasCredential):

‎sdk/tables/azure-data-tables/azure/data/tables/_base_client.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -223,7 +223,7 @@ def _format_url(self, hostname):
223223
return f"{self.scheme}://{hostname}{self._query_str}"
224224

225225
def _configure_policies(self, **kwargs):
226-
credential_policy = _configure_credential(self.credential)
226+
credential_policy = _configure_credential(self.credential, self._cosmos_endpoint)
227227
return [
228228
RequestIdPolicy(**kwargs),
229229
StorageHeadersPolicy(**kwargs),

‎sdk/tables/azure-data-tables/azure/data/tables/_constants.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
DEFAULT_COSMOS_ENDPOINT_SUFFIX = "cosmos.azure.com"
1414

1515
STORAGE_OAUTH_SCOPE = "https://storage.azure.com/.default"
16+
COSMOS_OAUTH_SCOPE = "https://cosmos.azure.com/.default"
1617

1718
NEXT_TABLE_NAME = "x-ms-continuation-NextTableName"
1819
NEXT_PARTITION_KEY = "x-ms-continuation-NextPartitionKey"

‎sdk/tables/azure-data-tables/azure/data/tables/aio/_authentication_async.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
from azure.core.pipeline import PipelineResponse, PipelineRequest
1111
from azure.core.pipeline.policies import AsyncBearerTokenCredentialPolicy
1212

13-
from .._constants import STORAGE_OAUTH_SCOPE
13+
from .._constants import STORAGE_OAUTH_SCOPE, COSMOS_OAUTH_SCOPE
1414
from .._authentication import _HttpChallenge, AzureSasCredentialPolicy, SharedKeyCredentialPolicy
1515

1616

@@ -94,11 +94,13 @@ def _configure_credential(credential: None) -> None: ...
9494
def _configure_credential(
9595
credential: Optional[
9696
Union[AzureNamedKeyCredential, AzureSasCredential, AsyncTokenCredential, SharedKeyCredentialPolicy]
97-
]
97+
],
98+
cosmos_endpoint: bool = False,
9899
) -> Optional[Union[AsyncBearerTokenChallengePolicy, AzureSasCredentialPolicy, SharedKeyCredentialPolicy]]:
99100
if hasattr(credential, "get_token"):
100101
credential = cast(AsyncTokenCredential, credential)
101-
return AsyncBearerTokenChallengePolicy(credential, STORAGE_OAUTH_SCOPE)
102+
scope = COSMOS_OAUTH_SCOPE if cosmos_endpoint else STORAGE_OAUTH_SCOPE
103+
return AsyncBearerTokenChallengePolicy(credential, scope)
102104
if isinstance(credential, SharedKeyCredentialPolicy):
103105
return credential
104106
if isinstance(credential, AzureSasCredential):

‎sdk/tables/azure-data-tables/azure/data/tables/aio/_base_client_async.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -216,7 +216,7 @@ def _format_url(self, hostname):
216216
return f"{self.scheme}://{hostname}{self._query_str}"
217217

218218
def _configure_policies(self, **kwargs):
219-
credential_policy = _configure_credential(self.credential)
219+
credential_policy = _configure_credential(self.credential, self._cosmos_endpoint)
220220
return [
221221
RequestIdPolicy(**kwargs),
222222
StorageHeadersPolicy(**kwargs),

‎sdk/tables/azure-data-tables/tests/test_table_client_async.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -518,7 +518,6 @@ async def test_table_service_client_with_token_credential(
518518

519519
async with TableServiceClient(base_url, credential=default_azure_credential) as client:
520520
await client.create_table(table_name)
521-
name_filter = "TableName eq '{}'".format(table_name)
522521
count = 0
523522
result = client.query_tables(name_filter)
524523
async for table in result:

‎sdk/tables/azure-data-tables/tests/test_table_client_cosmos.py

Lines changed: 48 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -362,11 +362,10 @@ def test_table_service_client_with_sas_token_credential(
362362
@cosmos_decorator
363363
@recorded_by_proxy
364364
def test_table_client_with_token_credential(self, tables_cosmos_account_name, tables_primary_cosmos_account_key):
365-
# DefaultAzureCredential doesn't work on Cosmos
366365
base_url = self.account_url(tables_cosmos_account_name, "cosmos")
367366
table_name = self.get_resource_name("mytable")
368367
default_azure_credential = self.get_token_credential()
369-
sas_token = self.generate_sas(
368+
self.sas_token = self.generate_sas(
370369
generate_account_sas,
371370
tables_primary_cosmos_account_key,
372371
resource_types=ResourceTypes.from_string("sco"),
@@ -375,41 +374,50 @@ def test_table_client_with_token_credential(self, tables_cosmos_account_name, ta
375374
)
376375

377376
with TableClient(base_url, table_name, credential=default_azure_credential) as client:
378-
with pytest.raises(HttpResponseError) as ex:
379-
client.create_table()
380-
ex_msg = "Authorization header doesn't confirm to the required format. Please verify and try again."
381-
assert ex_msg in str(ex.value)
377+
table = client.create_table()
378+
assert table.name == table_name
382379

383380
with TableClient.from_table_url(f"{base_url}/{table_name}", credential=default_azure_credential) as client:
384-
with pytest.raises(HttpResponseError) as ex:
385-
client.create_table()
386-
ex_msg = "Authorization header doesn't confirm to the required format. Please verify and try again."
387-
assert ex_msg in str(ex.value)
381+
entities = client.query_entities(
382+
query_filter="PartitionKey eq @pk",
383+
parameters={"pk": "dummy-pk"},
384+
)
385+
for e in entities:
386+
pass
388387

389-
with TableClient(f"{base_url}/?{sas_token}", table_name, credential=default_azure_credential) as client:
390-
with pytest.raises(HttpResponseError) as ex:
391-
client.create_table()
392-
ex_msg = "Authorization header doesn't confirm to the required format. Please verify and try again."
393-
assert ex_msg in str(ex.value)
388+
# DefaultAzureCredential is actually in use
389+
with TableClient(f"{base_url}/?{self.sas_token}", table_name, credential=default_azure_credential) as client:
390+
entities = client.query_entities(
391+
query_filter="PartitionKey eq @pk",
392+
parameters={"pk": "dummy-pk"},
393+
raw_request_hook=self.check_request_auth,
394+
)
395+
for e in entities:
396+
pass
394397

398+
# DefaultAzureCredential is actually in use
395399
with TableClient.from_table_url(
396-
f"{base_url}/{table_name}?{sas_token}", credential=default_azure_credential
400+
f"{base_url}/{table_name}?{self.sas_token}", credential=default_azure_credential
397401
) as client:
398-
with pytest.raises(HttpResponseError) as ex:
399-
client.create_table()
400-
ex_msg = "Authorization header doesn't confirm to the required format. Please verify and try again."
401-
assert ex_msg in str(ex.value)
402+
entities = client.query_entities(
403+
query_filter="PartitionKey eq @pk",
404+
parameters={"pk": "dummy-pk"},
405+
raw_request_hook=self.check_request_auth,
406+
)
407+
for e in entities:
408+
pass
409+
client.delete_table()
402410

403411
@cosmos_decorator
404412
@recorded_by_proxy
405413
def test_table_service_client_with_token_credential(
406414
self, tables_cosmos_account_name, tables_primary_cosmos_account_key
407415
):
408-
# DefaultAzureCredential doesn't work on Cosmos
409416
base_url = self.account_url(tables_cosmos_account_name, "cosmos")
410417
table_name = self.get_resource_name("mytable")
411418
default_azure_credential = self.get_token_credential()
412-
sas_token = self.generate_sas(
419+
name_filter = "TableName eq '{}'".format(table_name)
420+
self.sas_token = self.generate_sas(
413421
generate_account_sas,
414422
tables_primary_cosmos_account_key,
415423
resource_types=ResourceTypes.from_string("sco"),
@@ -418,16 +426,24 @@ def test_table_service_client_with_token_credential(
418426
)
419427

420428
with TableServiceClient(base_url, credential=default_azure_credential) as client:
421-
with pytest.raises(HttpResponseError) as ex:
422-
client.create_table(table_name)
423-
ex_msg = "Authorization header doesn't confirm to the required format. Please verify and try again."
424-
assert ex_msg in str(ex.value)
425-
426-
with TableServiceClient(f"{base_url}/?{sas_token}", credential=default_azure_credential) as client:
427-
with pytest.raises(HttpResponseError) as ex:
428-
client.create_table(table_name)
429-
ex_msg = "Authorization header doesn't confirm to the required format. Please verify and try again."
430-
assert ex_msg in str(ex.value)
429+
client.create_table(table_name)
430+
result = client.query_tables(name_filter)
431+
assert len(list(result)) == 1
432+
433+
# DefaultAzureCredential is actually in use
434+
with TableServiceClient(f"{base_url}/?{self.sas_token}", credential=default_azure_credential) as client:
435+
entities = client.get_table_client(table_name).query_entities(
436+
query_filter="PartitionKey eq @pk",
437+
parameters={"pk": "dummy-pk"},
438+
raw_request_hook=self.check_request_auth,
439+
)
440+
for e in entities:
441+
pass
442+
client.delete_table(table_name)
443+
444+
def check_request_auth(self, pipeline_request):
445+
assert self.sas_token not in pipeline_request.http_request.url
446+
assert pipeline_request.http_request.headers.get("Authorization") is not None
431447

432448
@cosmos_decorator
433449
@recorded_by_proxy

‎sdk/tables/azure-data-tables/tests/test_table_client_cosmos_async.py

Lines changed: 52 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -376,11 +376,10 @@ async def test_table_service_client_with_sas_token_credential(
376376
async def test_table_client_with_token_credential(
377377
self, tables_cosmos_account_name, tables_primary_cosmos_account_key
378378
):
379-
# DefaultAzureCredential doesn't work on Cosmos
380379
base_url = self.account_url(tables_cosmos_account_name, "cosmos")
381380
table_name = self.get_resource_name("mytable")
382381
default_azure_credential = self.get_token_credential()
383-
sas_token = self.generate_sas(
382+
self.sas_token = self.generate_sas(
384383
generate_account_sas,
385384
tables_primary_cosmos_account_key,
386385
resource_types=ResourceTypes.from_string("sco"),
@@ -389,43 +388,52 @@ async def test_table_client_with_token_credential(
389388
)
390389

391390
async with TableClient(base_url, table_name, credential=default_azure_credential) as client:
392-
with pytest.raises(HttpResponseError) as ex:
393-
await client.create_table()
394-
ex_msg = "Authorization header doesn't confirm to the required format. Please verify and try again."
395-
assert ex_msg in str(ex.value)
391+
table = await client.create_table()
392+
assert table.name == table_name
396393

397394
async with TableClient.from_table_url(
398395
f"{base_url}/{table_name}", credential=default_azure_credential
399396
) as client:
400-
with pytest.raises(HttpResponseError) as ex:
401-
await client.create_table()
402-
ex_msg = "Authorization header doesn't confirm to the required format. Please verify and try again."
403-
assert ex_msg in str(ex.value)
397+
entities = client.query_entities(
398+
query_filter="PartitionKey eq @pk",
399+
parameters={"pk": "dummy-pk"},
400+
)
401+
async for e in entities:
402+
pass
404403

405-
async with TableClient(f"{base_url}/?{sas_token}", table_name, credential=default_azure_credential) as client:
406-
with pytest.raises(HttpResponseError) as ex:
407-
await client.create_table()
408-
ex_msg = "Authorization header doesn't confirm to the required format. Please verify and try again."
409-
assert ex_msg in str(ex.value)
404+
async with TableClient(
405+
f"{base_url}/?{self.sas_token}", table_name, credential=default_azure_credential
406+
) as client:
407+
entities = client.query_entities(
408+
query_filter="PartitionKey eq @pk",
409+
parameters={"pk": "dummy-pk"},
410+
raw_request_hook=self.check_request_auth,
411+
)
412+
async for e in entities:
413+
pass
410414

411415
async with TableClient.from_table_url(
412-
f"{base_url}/{table_name}?{sas_token}", credential=default_azure_credential
416+
f"{base_url}/{table_name}?{self.sas_token}", credential=default_azure_credential
413417
) as client:
414-
with pytest.raises(HttpResponseError) as ex:
415-
await client.create_table()
416-
ex_msg = "Authorization header doesn't confirm to the required format. Please verify and try again."
417-
assert ex_msg in str(ex.value)
418+
entities = client.query_entities(
419+
query_filter="PartitionKey eq @pk",
420+
parameters={"pk": "dummy-pk"},
421+
raw_request_hook=self.check_request_auth,
422+
)
423+
async for e in entities:
424+
pass
425+
await client.delete_table()
418426

419427
@cosmos_decorator_async
420428
@recorded_by_proxy_async
421429
async def test_table_service_client_with_token_credential(
422430
self, tables_cosmos_account_name, tables_primary_cosmos_account_key
423431
):
424-
# DefaultAzureCredential doesn't work on Cosmos
425432
base_url = self.account_url(tables_cosmos_account_name, "cosmos")
426433
table_name = self.get_resource_name("mytable")
427434
default_azure_credential = self.get_token_credential()
428-
sas_token = self.generate_sas(
435+
name_filter = "TableName eq '{}'".format(table_name)
436+
self.sas_token = self.generate_sas(
429437
generate_account_sas,
430438
tables_primary_cosmos_account_key,
431439
resource_types=ResourceTypes.from_string("sco"),
@@ -434,16 +442,28 @@ async def test_table_service_client_with_token_credential(
434442
)
435443

436444
async with TableServiceClient(base_url, credential=default_azure_credential) as client:
437-
with pytest.raises(HttpResponseError) as ex:
438-
await client.create_table(table_name)
439-
ex_msg = "Authorization header doesn't confirm to the required format. Please verify and try again."
440-
assert ex_msg in str(ex.value)
441-
442-
async with TableServiceClient(f"{base_url}/?{sas_token}", credential=default_azure_credential) as client:
443-
with pytest.raises(HttpResponseError) as ex:
444-
await client.create_table(table_name)
445-
ex_msg = "Authorization header doesn't confirm to the required format. Please verify and try again."
446-
assert ex_msg in str(ex.value)
445+
await client.create_table(table_name)
446+
name_filter = "TableName eq '{}'".format(table_name)
447+
count = 0
448+
result = client.query_tables(name_filter)
449+
async for table in result:
450+
count += 1
451+
assert count == 1
452+
453+
# DefaultAzureCredential is actually in use
454+
async with TableServiceClient(f"{base_url}/?{self.sas_token}", credential=default_azure_credential) as client:
455+
entities = client.get_table_client(table_name).query_entities(
456+
query_filter="PartitionKey eq @pk",
457+
parameters={"pk": "dummy-pk"},
458+
raw_request_hook=self.check_request_auth,
459+
)
460+
async for e in entities:
461+
pass
462+
client.delete_table(table_name)
463+
464+
def check_request_auth(self, pipeline_request):
465+
assert self.sas_token not in pipeline_request.http_request.url
466+
assert pipeline_request.http_request.headers.get("Authorization") is not None
447467

448468
@cosmos_decorator_async
449469
@recorded_by_proxy_async

0 commit comments

Comments
 (0)