Track 2 Azure CLI Support Extension (#3738)

* Track 2

* Apply suggestions from code review

Co-authored-by: Yu Chen <ychenu@microsoft.com>

src/support/HISTORY.md

@@ -1,6 +1,10 @@
 Release History
 ===============
 
+1.0.3
+-----
+* Migrate to track 2 SDK
+
 1.0.2
 -----
 * Removed custom error message in lieu of error message coming from backend.

src/support/azext_support/__init__.py

@@ -16,8 +16,8 @@ def __init__(self, cli_ctx=None):
         support_custom = CliCommandType(
             operations_tmpl='azext_support.custom#{}',
             client_factory=cf_support)
-        super(SupportCommandsLoader, self).__init__(cli_ctx=cli_ctx,
-                                                    custom_command_type=support_custom)
+        super().__init__(cli_ctx=cli_ctx,
+                         custom_command_type=support_custom)
 
     def load_command_table(self, args):
         from azext_support.commands import load_command_table

src/support/azext_support/_utils.py

@@ -6,8 +6,9 @@
 import re
 
 from azure.cli.core._profile import Profile
-from knack.log import get_logger
+from azure.cli.core.azclierror import UnauthorizedError
 from knack.util import CLIError
+from knack.log import get_logger
 
 logger = get_logger(__name__)
 
@@ -48,8 +49,9 @@ def get_bearer_token(cmd, tenant_id):
     try:
         logger.debug("Retrieving access token for tenant %s", tenant_id)
         creds, _, _ = client.get_raw_token(tenant=tenant_id)
-    except CLIError:
-        raise CLIError("Can't find authorization for {0}. ".format(tenant_id) +
-                       "Run \'az login -t <tenant_name> --allow-no-subscriptions\' and try again.")
+    except CLIError as unauthorized_error:
+        raise UnauthorizedError("Can't find authorization for {0}. ".format(tenant_id) +
+                                "Run \'az login -t <tenant_name> --allow-no-subscriptions\' and try again.") from \
+            unauthorized_error
 
     return "Bearer " + creds[1]

src/support/azext_support/_validators.py

@@ -41,15 +41,17 @@ def validate_communication_create(cmd, namespace):
 
 def _validate_communication_name(cmd, ticket_name, communication_name):
     client = cf_communications(cmd.cli_ctx)
-    rsp = client.check_name_availability(support_ticket_name=ticket_name, name=communication_name,
-                                         type="Microsoft.Support/communications")
+    check_name_availability_input = {"name": communication_name, "type": "Microsoft.Support/communications"}
+    rsp = client.check_name_availability(support_ticket_name=ticket_name,
+                                         check_name_availability_input=check_name_availability_input)
     if not rsp.name_available:
         raise CLIError(rsp.message)
 
 
 def _validate_ticket_name(cmd, ticket_name):
     client = cf_support_tickets(cmd.cli_ctx)
-    rsp = client.check_name_availability(name=ticket_name, type="Microsoft.Support/supportTickets")
+    check_name_availability_input = {"name": ticket_name, "type": "Microsoft.Support/supportTickets"}
+    rsp = client.check_name_availability(check_name_availability_input=check_name_availability_input)
     if not rsp.name_available:
         raise CLIError(rsp.message)
 

src/support/azext_support/custom.py

@@ -135,12 +135,12 @@ def create_support_tickets(cmd, client,
     logger.debug("Sending create request with below payload: ")
     logger.debug(json.dumps(body, indent=4))
 
-    custom_headers = {}
     if partner_tenant_id is not None:
-        custom_headers["x-ms-authorization-auxiliary"] = get_bearer_token(cmd, partner_tenant_id)
+        external_bearer_token = get_bearer_token(cmd, partner_tenant_id)
+        return client.begin_create(support_ticket_name=ticket_name, create_support_ticket_parameters=body,
+                                   headers={'x-ms-authorization-auxiliary': external_bearer_token})
 
-    return client.create(support_ticket_name=ticket_name, create_support_ticket_parameters=body,
-                         custom_headers=custom_headers)
+    return client.begin_create(support_ticket_name=ticket_name, create_support_ticket_parameters=body)
 
 
 def create_support_tickets_communications(cmd, client,
@@ -154,5 +154,5 @@ def create_support_tickets_communications(cmd, client,
     body["subject"] = communication_subject
     body["body"] = communication_body
 
-    return client.create(support_ticket_name=ticket_name, communication_name=communication_name,
-                         create_communication_parameters=body)
+    return client.begin_create(support_ticket_name=ticket_name, communication_name=communication_name,
+                               create_communication_parameters=body)