Tables aad (#19299)

closes #19244

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

@@ -3,6 +3,7 @@
 ## 12.0.1 (Unreleased)
 
 ### Features Added
+* Storage Accounts only: `TableClient` and `TableServiceClient`s can now use `azure-identity` credentials for authentication. Note: A `TableClient` authenticated with a `TokenCredential` cannot use the `get_table_access_policy` or `set_table_access_policy` methods.
 
 ### Breaking Changes
 

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

@@ -394,7 +394,7 @@ def format_query_string(sas_token, credential):
             "You cannot use AzureSasCredential when the resource URI also contains a Shared Access Signature.")
     if sas_token and not credential:
         query_str += sas_token
-    elif isinstance(credential, (AzureSasCredential, AzureNamedKeyCredential)):
+    elif credential:
         return "", credential
     return query_str.rstrip("?&"), None
 

sdk/tables/azure-data-tables/azure/data/tables/_table_client.py

@@ -72,11 +72,12 @@ def __init__(  # pylint: disable=missing-client-constructor-parameter-credential
         :param str table_name: The table name.
         :keyword credential:
             The credentials with which to authenticate. This is optional if the
-            account URL already has a SAS token. The value can be one of AzureNamedKeyCredential
-            or AzureSasCredential from azure-core.
+            account URL already has a SAS token. The value can be one of AzureNamedKeyCredential (azure-core),
+            AzureSasCredential (azure-core), or TokenCredentials from azure-identity.
         :paramtype credential:
             :class:`~azure.core.credentials.AzureNamedKeyCredential` or
-            :class:`~azure.core.credentials.AzureSasCredential`
+            :class:`~azure.core.credentials.AzureSasCredential` or
+            :class:`~azure.core.credentials.TokenCredential`
         :returns: None
         """
         if not table_name:

sdk/tables/azure-data-tables/azure/data/tables/_table_service_client.py

@@ -44,11 +44,12 @@ class TableServiceClient(TablesBaseClient):
         authenticated with a SAS token.
     :keyword credential:
         The credentials with which to authenticate. This is optional if the
-        account URL already has a SAS token. The value can be one of AzureNamedKeyCredential
-        or AzureSasCredential from azure-core.
+        account URL already has a SAS token. The value can be one of AzureNamedKeyCredential (azure-core),
+        AzureSasCredential (azure-core), or TokenCredentials from azure-identity.
     :paramtype credential:
         :class:`~azure.core.credentials.AzureNamedKeyCredential` or
-        :class:`~azure.core.credentials.AzureSasCredential`
+        :class:`~azure.core.credentials.AzureSasCredential` or
+        :class:`~azure.core.credentials.TokenCredential`
     :keyword str api_version:
         The Storage API version to use for requests. Default value is '2019-02-02'.
         Setting to an older version may result in reduced feature compatibility.

sdk/tables/azure-data-tables/azure/data/tables/aio/_table_client_async.py

@@ -61,11 +61,12 @@ def __init__(  # pylint: disable=missing-client-constructor-parameter-credential
         :param str table_name: The table name.
         :keyword credential:
             The credentials with which to authenticate. This is optional if the
-            account URL already has a SAS token. The value can be one of AzureNamedKeyCredential
-            or AzureSasCredential from azure-core.
+            account URL already has a SAS token. The value can be one of AzureNamedKeyCredential (azure-core),
+            AzureSasCredential (azure-core), or TokenCredentials from azure-identity.
         :paramtype credential:
             :class:`~azure.core.credentials.AzureNamedKeyCredential` or
-            :class:`~azure.core.credentials.AzureSasCredential`
+            :class:`~azure.core.credentials.AzureSasCredential` or
+            :class:`~azure.core.credentials.TokenCredential`
 
         :returns: None
         """

sdk/tables/azure-data-tables/azure/data/tables/aio/_table_service_client_async.py

@@ -47,11 +47,12 @@ class TableServiceClient(AsyncTablesBaseClient):
         authenticated with a SAS token.
     :keyword credential:
         The credentials with which to authenticate. This is optional if the
-        account URL already has a SAS token. The value can be one of AzureNamedKeyCredential
-        or AzureSasCredential from azure-core.
+        account URL already has a SAS token. The value can be one of AzureNamedKeyCredential (azure-core),
+        AzureSasCredential (azure-core), or TokenCredentials from azure-identity.
     :paramtype credential:
         :class:`~azure.core.credentials.AzureNamedKeyCredential` or
-        :class:`~azure.core.credentials.AzureSasCredential`
+        :class:`~azure.core.credentials.AzureSasCredential` or
+        :class:`~azure.core.credentials.TokenCredential`
     :keyword str api_version:
         The Storage API version to use for requests. Default value is '2019-02-02'.
         Setting to an older version may result in reduced feature compatibility.

sdk/tables/azure-data-tables/tests/_shared/asynctestcase.py

@@ -16,6 +16,7 @@
     EdmType,
 )
 from azure.data.tables.aio import TableServiceClient
+from azure.identity.aio import DefaultAzureCredential
 
 from devtools_testutils import is_live
 
@@ -38,6 +39,11 @@ async def get_token(self, *args):
 
 
 class AsyncTableTestCase(TableTestCase):
+    def get_token_credential(self):
+        if is_live():
+            return DefaultAzureCredential()
+        return self.generate_fake_token()
+
     def generate_fake_token(self):
         return AsyncFakeTokenCredential()
 
@@ -116,9 +122,9 @@ async def _insert_random_entity(self, pk=None, rk=None):
         metadata = await self.table.create_entity(entity=entity)
         return entity, metadata["etag"]
 
-    async def _set_up(self, account_name, account_key, url="table"):
+    async def _set_up(self, account_name, credential, url="table"):
         account_url = self.account_url(account_name, url)
-        self.ts = TableServiceClient(account_url, credential=account_key)
+        self.ts = TableServiceClient(account_url, credential=credential)
         self.table_name = self.get_resource_name("uttable")
         self.table = self.ts.get_table_client(self.table_name)
         if self.is_live:

sdk/tables/azure-data-tables/tests/_shared/testcase.py

@@ -23,13 +23,22 @@
     TableMetrics,
     TableServiceClient,
 )
+from azure.identity import DefaultAzureCredential
 
 from devtools_testutils import is_live
 
 SLEEP_DELAY = 30
 
 TEST_TABLE_PREFIX = "pytablesync"
 
+SERVICE_UNAVAILABLE_RESP_BODY = '<?xml version="1.0" encoding="utf-8"?><StorageServiceStats><GeoReplication><Status' \
+                                '>unavailable</Status><LastSyncTime></LastSyncTime></GeoReplication' \
+                                '></StorageServiceStats> '
+
+SERVICE_LIVE_RESP_BODY = '<?xml version="1.0" encoding="utf-8"?><StorageServiceStats><GeoReplication><Status' \
+                         '>live</Status><LastSyncTime>Wed, 19 Jan 2021 22:28:43 GMT</LastSyncTime></GeoReplication' \
+                         '></StorageServiceStats> '
+
 
 class FakeTokenCredential(object):
     """Protocol for classes able to provide OAuth tokens.
@@ -81,6 +90,11 @@ def generate_sas_token(self):
             expiry=datetime.now() + timedelta(days=8),
         )
 
+    def get_token_credential(self):
+        if is_live():
+            return DefaultAzureCredential()
+        return self.generate_fake_token()
+
     def generate_fake_token(self):
         return FakeTokenCredential()
 
@@ -421,10 +435,10 @@ def _insert_random_entity(self, pk=None, rk=None):
         metadata = self.table.create_entity(entity)
         return entity, metadata["etag"]
 
-    def _set_up(self, account_name, account_key, url="table"):
+    def _set_up(self, account_name, credential, url="table"):
         self.table_name = self.get_resource_name("uttable")
         self.ts = TableServiceClient(
-            self.account_url(account_name, url), credential=account_key, table_name=self.table_name
+            self.account_url(account_name, url), credential=credential, table_name=self.table_name
         )
         self.table = self.ts.get_table_client(self.table_name)
         if self.is_live:
@@ -449,6 +463,13 @@ def _assert_stats_unavailable(self, stats):
         assert stats["geo_replication"]["status"] == "unavailable"
         assert stats["geo_replication"]["last_sync_time"] is None
 
+    @staticmethod
+    def override_response_body_with_unavailable_status(response):
+        response.http_response.text = lambda _: SERVICE_UNAVAILABLE_RESP_BODY
+
+    @staticmethod
+    def override_response_body_with_live_status(response):
+        response.http_response.text = lambda _: SERVICE_LIVE_RESP_BODY
 
 class ResponseCallback(object):
     def __init__(self, status=None, new_status=None):

sdk/tables/azure-data-tables/tests/recordings/test_table_aad.test_aad_access_policy_error.yaml

@@ -0,0 +1,92 @@
+interactions:
+- request:
+    body: null
+    headers:
+      Accept:
+      - application/xml
+      Accept-Encoding:
+      - gzip, deflate
+      Connection:
+      - keep-alive
+      Date:
+      - Thu, 17 Jun 2021 16:53:03 GMT
+      User-Agent:
+      - azsdk-python-data-tables/12.0.1 Python/3.9.0rc1 (Windows-10-10.0.19041-SP0)
+      x-ms-date:
+      - Thu, 17 Jun 2021 16:53:03 GMT
+      x-ms-version:
+      - '2019-02-02'
+    method: GET
+    uri: https://fake_table_account.table.core.windows.net/pytablesync7b551169?comp=acl
+  response:
+    body:
+      string: '<?xml version="1.0" encoding="utf-8"?><m:error xmlns:m="http://schemas.microsoft.com/ado/2007/08/dataservices/metadata"><m:code>ResourceNotFound</m:code><m:message
+        xml:lang="en-US">The specified resource does not exist.
+
+        RequestId:7c3b01c2-0002-006a-1999-63305d000000
+
+        Time:2021-06-17T16:53:05.6341722Z</m:message></m:error>'
+    headers:
+      content-length:
+      - '322'
+      content-type:
+      - application/xml
+      date:
+      - Thu, 17 Jun 2021 16:53:04 GMT
+      server:
+      - Windows-Azure-Table/1.0 Microsoft-HTTPAPI/2.0
+      x-ms-error-code:
+      - ResourceNotFound
+      x-ms-version:
+      - '2019-02-02'
+    status:
+      code: 404
+      message: The specified resource does not exist.
+- request:
+    body: null
+    headers:
+      Accept:
+      - application/xml
+      Accept-Encoding:
+      - gzip, deflate
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '0'
+      Content-Type:
+      - application/xml
+      Date:
+      - Thu, 17 Jun 2021 16:53:04 GMT
+      User-Agent:
+      - azsdk-python-data-tables/12.0.1 Python/3.9.0rc1 (Windows-10-10.0.19041-SP0)
+      x-ms-date:
+      - Thu, 17 Jun 2021 16:53:04 GMT
+      x-ms-version:
+      - '2019-02-02'
+    method: PUT
+    uri: https://fake_table_account.table.core.windows.net/pytablesync7b551169?comp=acl
+  response:
+    body:
+      string: '<?xml version="1.0" encoding="utf-8"?><m:error xmlns:m="http://schemas.microsoft.com/ado/2007/08/dataservices/metadata"><m:code>ResourceNotFound</m:code><m:message
+        xml:lang="en-US">The specified resource does not exist.
+
+        RequestId:7c3b0201-0002-006a-5499-63305d000000
+
+        Time:2021-06-17T16:53:05.7602607Z</m:message></m:error>'
+    headers:
+      content-length:
+      - '322'
+      content-type:
+      - application/xml
+      date:
+      - Thu, 17 Jun 2021 16:53:04 GMT
+      server:
+      - Windows-Azure-Table/1.0 Microsoft-HTTPAPI/2.0
+      x-ms-error-code:
+      - ResourceNotFound
+      x-ms-version:
+      - '2019-02-02'
+    status:
+      code: 404
+      message: The specified resource does not exist.
+version: 1