{Batch} Upgrade api to track2 v2024-02-01

src/azure-cli-core/azure/cli/core/profiles/_shared.py

@@ -122,6 +122,7 @@ class ResourceType(Enum):  # pylint: disable=too-few-public-methods
     DATA_STORAGE_FILESHARE = ('azure.multiapi.storagev2.fileshare', None)
     DATA_STORAGE_QUEUE = ('azure.multiapi.storagev2.queue', None)
     DATA_STORAGE_TABLE = ('azure.data.tables', None)
+    DATA_BATCH = ('azure.batch', None)
 
     def __init__(self, import_prefix, client_name):
         """Constructor.

src/azure-cli/azure/cli/command_modules/batch/_client_factory.py

@@ -5,97 +5,63 @@
 
 
 def mgmt_batch_account_client_factory(cli_ctx, _):
-    return batch_client_factory(cli_ctx).batch_account
+    return batch_mgmt_client_factory(cli_ctx).batch_account
 
 
 def mgmt_pool_client_factory(cli_ctx, _):
-    return batch_client_factory(cli_ctx).pool
+    return batch_mgmt_client_factory(cli_ctx).pool
 
 
 def mgmt_private_link_resource_client_factory(cli_ctx, _):
-    return batch_client_factory(cli_ctx).private_link_resource
+    return batch_mgmt_client_factory(cli_ctx).private_link_resource
 
 
 def mgmt_private_endpoint_connection_client_factory(cli_ctx, _):
-    return batch_client_factory(cli_ctx).private_endpoint_connection
+    return batch_mgmt_client_factory(cli_ctx).private_endpoint_connection
 
 
 def mgmt_application_client_factory(cli_ctx, _):
-    return batch_client_factory(cli_ctx).application
+    return batch_mgmt_client_factory(cli_ctx).application
 
 
 def mgmt_application_package_client_factory(cli_ctx, _):
-    return batch_client_factory(cli_ctx).application_package
+    return batch_mgmt_client_factory(cli_ctx).application_package
 
 
 def mgmt_location_client_factory(cli_ctx, _):
-    return batch_client_factory(cli_ctx).location
+    return batch_mgmt_client_factory(cli_ctx).location
 
 
-def application_client_factory(cli_ctx, kwargs):
-    return batch_data_service_factory(cli_ctx, kwargs).application
-
-
-def account_client_factory(cli_ctx, kwargs):
-    return batch_data_service_factory(cli_ctx, kwargs).account
-
-
-def certificate_client_factory(cli_ctx, kwargs):
-    return batch_data_service_factory(cli_ctx, kwargs).certificate
-
-
-def pool_client_factory(cli_ctx, kwargs):
-    return batch_data_service_factory(cli_ctx, kwargs).pool
-
-
-def job_client_factory(cli_ctx, kwargs):
-    return batch_data_service_factory(cli_ctx, kwargs).job
-
-
-def job_schedule_client_factory(cli_ctx, kwargs):
-    return batch_data_service_factory(cli_ctx, kwargs).job_schedule
-
-
-def task_client_factory(cli_ctx, kwargs):
-    return batch_data_service_factory(cli_ctx, kwargs).task
-
-
-def file_client_factory(cli_ctx, kwargs):
-    return batch_data_service_factory(cli_ctx, kwargs).file
-
-
-def compute_node_client_factory(cli_ctx, kwargs):
-    return batch_data_service_factory(cli_ctx, kwargs).compute_node
-
-
-def batch_client_factory(cli_ctx, **_):
+def batch_mgmt_client_factory(cli_ctx, **_):
     from azure.mgmt.batch import BatchManagementClient
     from azure.cli.core.commands.client_factory import get_mgmt_service_client
     return get_mgmt_service_client(cli_ctx, BatchManagementClient)
 
 
-def batch_data_service_factory(cli_ctx, kwargs):
-    import azure.batch._batch_service_client as batch
-    import azure.batch.batch_auth as batchauth
+def batch_data_client_factory(cli_ctx, kwargs):
+    from azure.batch import BatchClient
 
     account_name = kwargs.pop('account_name', None)
     account_key = kwargs.pop('account_key', None)
     account_endpoint = kwargs.pop('account_endpoint', None)
-    kwargs.pop('yes', None)
+    token_credential = kwargs.pop('token_credential', None)
+
+    credential = None
 
-    credentials = None
-    if not account_key:
+    if not token_credential and not account_key:
         from azure.cli.core._profile import Profile
         profile = Profile(cli_ctx=cli_ctx)
         resource = cli_ctx.cloud.endpoints.batch_resource_id
-        credentials, _, _ = profile.get_login_credentials(resource=resource)
+        token_credential, _, _ = profile.get_login_credentials(resource=resource)
+
+    if account_key:
+        from azure.core.credentials import AzureNamedKeyCredential
+        credential = AzureNamedKeyCredential(name=account_name, key=account_key)
     else:
-        # Verify all values are populated and display readable error
-        if not all([account_name, account_key, account_endpoint]):
-            raise ValueError(
-                'usage error: --account-name NAME --account-key KEY --account-endpoint ENDPOINT')
-        credentials = batchauth.SharedKeyCredentials(account_name, account_key)
+        credential = token_credential
+
     if not (account_endpoint.startswith('https://') or
             account_endpoint.startswith('http://')):
         account_endpoint = 'https://' + account_endpoint
-    return batch.BatchServiceClient(credentials, batch_url=account_endpoint.rstrip('/'))
+
+    return BatchClient(credential=credential, endpoint=account_endpoint.rstrip('/'))