diff --git a/src/aks-preview/HISTORY.rst b/src/aks-preview/HISTORY.rst
index faeea199722..9cb7c9be5c1 100644
--- a/src/aks-preview/HISTORY.rst
+++ b/src/aks-preview/HISTORY.rst
@@ -12,6 +12,10 @@ To release a new version, please select a new version number (usually plus 1 to
 Pending
 +++++++
 
+18.0.0b31
++++++++
+* Remove DATA_STORAGE references
+
 18.0.0b30
 +++++++
 * Remove interactive AI-powered debugging tool `az aks agent`. To use the debugging tool, please enable it in the aks-agent extension by running `az extension add --name aks-agent`.
diff --git a/src/aks-preview/azext_aks_preview/aks_diagnostics.py b/src/aks-preview/azext_aks_preview/aks_diagnostics.py
index 33bcc728c1c..2bac3d71d8b 100644
--- a/src/aks-preview/azext_aks_preview/aks_diagnostics.py
+++ b/src/aks-preview/azext_aks_preview/aks_diagnostics.py
@@ -28,6 +28,7 @@
     get_mgmt_service_client,
     get_subscription_id,
 )
+from azure.cli.core.profiles import ResourceType, get_sdk
 from azure.core.exceptions import HttpResponseError
 from azure.mgmt.core.tools import is_valid_resource_id, parse_resource_id, resource_id
 from knack.log import get_logger
@@ -98,24 +99,20 @@ def aks_kollect_cmd(cmd,    # pylint: disable=too-many-statements,too-many-local
             cmd.cli_ctx, parsed_storage_account['subscription'])
         storage_account_keys = storage_client.storage_accounts.list_keys(parsed_storage_account['resource_group'],
                                                                          storage_account_name)
-        kwargs = {
-            'account_name': storage_account_name,
-            'account_key': storage_account_keys.keys[0].value
-        }
-        cloud_storage_client = _cloud_storage_account_service_factory(
-            cmd.cli_ctx, kwargs)
-
-        sas_token = cloud_storage_client.generate_shared_access_signature(
-            'b',
-            'sco',
-            'rwdlacup',
-            datetime.datetime.utcnow() + datetime.timedelta(days=1))
-
-        readonly_sas_token = cloud_storage_client.generate_shared_access_signature(
-            'b',
-            'sco',
-            'rl',
-            datetime.datetime.utcnow() + datetime.timedelta(days=1))
+
+        t_generate_blob_service_sas = get_sdk(cmd.cli_ctx, ResourceType.DATA_STORAGE_BLOB, '#generate_account_sas')
+
+        sas_token = t_generate_blob_service_sas(storage_account_name,
+                                                storage_account_keys.keys[0].value,
+                                                resource_types='sco',
+                                                permission='rwdlacup',
+                                                expiry=datetime.datetime.utcnow() + datetime.timedelta(days=1))
+
+        readonly_sas_token = t_generate_blob_service_sas(storage_account_name,
+                                                         storage_account_keys.keys[0].value,
+                                                         resource_types='sco',
+                                                         permission='rl',
+                                                         expiry=datetime.datetime.utcnow() + datetime.timedelta(days=1))
 
         readonly_sas_token = readonly_sas_token.strip('?')
 
@@ -483,17 +480,6 @@ def _display_diagnostics_report(temp_kubeconfig_path):   # pylint: disable=too-m
                        "Please run 'az aks kanalyze' command later to get the analysis results.")
 
 
-def _cloud_storage_account_service_factory(cli_ctx, kwargs):
-    from azure.cli.core.profiles import ResourceType, get_sdk
-    t_cloud_storage_account = get_sdk(
-        cli_ctx, ResourceType.DATA_STORAGE, 'common#CloudStorageAccount')
-    account_name = kwargs.pop('account_name', None)
-    account_key = kwargs.pop('account_key', None)
-    sas_token = kwargs.pop('sas_token', None)
-    kwargs.pop('connection_string', None)
-    return t_cloud_storage_account(account_name, account_key, sas_token)
-
-
 def _format_hyperlink(the_link):
     return f'\033[1m{colorama.Style.BRIGHT}{colorama.Fore.BLUE}{the_link}{colorama.Style.RESET_ALL}'
 
diff --git a/src/aks-preview/setup.py b/src/aks-preview/setup.py
index 8e4cfa67575..3fab14b6467 100644
--- a/src/aks-preview/setup.py
+++ b/src/aks-preview/setup.py
@@ -9,7 +9,7 @@
 
 from setuptools import find_packages, setup
 
-VERSION = "18.0.0b30"
+VERSION = "18.0.0b31"
 
 CLASSIFIERS = [
     "Development Status :: 4 - Beta",