[Quantum]Add Azure quantum api key support

src/quantum/HISTORY.rst

@@ -3,6 +3,11 @@
 Release History
 ===============
 
+1.0.1
+++++++
+* [2024-02-14] Version intended to work with QDK version 0.29.0
+* Add azure quantum connection string/ api key functionalities.
+
 1.0.0b1
 ++++++
 * [2024-02-08] Version intended to work with QDK version 0.29.0

src/quantum/azext_quantum/__init__.py

@@ -13,7 +13,7 @@
 # This is the version reported by the CLI to the service when submitting requests.
 # This should be in sync with the extension version in 'setup.py', unless we need to
 # submit using a different version.
-CLI_REPORTED_VERSION = "1.0.0b1"
+CLI_REPORTED_VERSION = "1.0.1"
 
 
 class QuantumCommandsLoader(AzCommandsLoader):

src/quantum/azext_quantum/_client_factory.py

@@ -42,7 +42,7 @@ def get_appid():
 def cf_quantum_mgmt(cli_ctx, *_):
     from azure.cli.core.commands.client_factory import get_mgmt_service_client
     from .vendored_sdks.azure_mgmt_quantum import AzureQuantumManagementClient
-    client = get_mgmt_service_client(cli_ctx, AzureQuantumManagementClient)
+    client = get_mgmt_service_client(cli_ctx, AzureQuantumManagementClient, base_url_bound=False)
     # Add user agent on the management client to include extension information
     client._config.user_agent_policy.add_user_agent(get_appid())
     return client
@@ -52,6 +52,10 @@ def cf_workspaces(cli_ctx, *_):
     return cf_quantum_mgmt(cli_ctx).workspaces
 
 
+def cf_workspace(cli_ctx, *_):
+    return cf_quantum_mgmt(cli_ctx).workspace
+
+
 def cf_offerings(cli_ctx, *_):
     return cf_quantum_mgmt(cli_ctx).offerings
 

src/quantum/azext_quantum/_help.py

@@ -283,3 +283,39 @@
         text: |-
             az quantum workspace show -g MyResourceGroup -w MyWorkspace
 """
+
+helps['quantum workspace update'] = """
+    type: command
+    short-summary: Enable or disable Azure Quantum workspace api keys.
+    examples:
+      - name: Enable a provided Azure Quantum workspace api keys.
+        text: |-
+            az quantum workspace update --enable-api-key True
+      - name: Disable a provided Azure Quantum workspace api keys.
+        text: |-
+            az quantum workspace update --enable-api-key False
+"""
+
+helps['quantum workspace keys list'] = """
+    type: command
+    short-summary: Display Azure Quantum workspace api keys.
+    examples:
+      - name: Show the currently selected default Azure Quantum workspace api keys.
+        text: |-
+            az quantum workspace keys list
+"""
+
+helps['quantum workspace keys regenerate'] = """
+    type: command
+    short-summary: Regenerate Azure Quantum workspace api key.
+    examples:
+      - name: Regenerate the currently selected default Azure Quantum workspace primary api key.
+        text: |-
+            az quantum workspace keys regenerate --key-type Primary
+      - name: Regenerate the currently selected default Azure Quantum workspace secondary api key.
+        text: |-
+            az quantum workspace keys regenerate --key-type Secondary
+      - name: Regenerate the currently selected default Azure Quantum workspace secondary api key.
+        text: |-
+            az quantum workspace keys regenerate --key-type Primary,Secondary
+"""

src/quantum/azext_quantum/_params.py

@@ -58,6 +58,8 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals
     entry_point_type = CLIArgumentType(help='The entry point for the QIR program or circuit. Required for QIR. Ignored on Q# jobs.')
     item_type = CLIArgumentType(help='The item index in a batching job.')
     skip_autoadd_type = CLIArgumentType(help='If specified, the plans that offer free credits will not automatically be added.')
+    key_type = CLIArgumentType(options_list=['--key-type'], help='The api keys to be regenerated, should be Primary and/or Secondary.')
+    enable_key = CLIArgumentType(options_list=['--enable-api-key'], help='Enable or disable api key authentication.')
 
     with self.argument_context('quantum workspace') as c:
         c.argument('workspace_name', workspace_name_type)
@@ -135,3 +137,14 @@ def load_arguments(self, _):  # pylint: disable=too-many-locals
 
     with self.argument_context('quantum offerings list') as c:
         c.argument('autoadd_only', autoadd_only_type)
+
+    with self.argument_context('quantum workspace keys list') as c:
+        c.argument('workspace_name', workspace_name_type)
+
+    with self.argument_context('quantum workspace keys regenerate') as c:
+        c.argument('workspace_name', workspace_name_type)
+        c.argument('key_type', key_type)
+
+    with self.argument_context('quantum workspace update') as c:
+        c.argument('workspace_name', workspace_name_type)
+        c.argument('enable_key', enable_key)

src/quantum/azext_quantum/commands.py

@@ -182,6 +182,9 @@ def load_command_table(self, _):
         w.command('set', 'set', validator=validate_workspace_info)
         w.command('clear', 'clear')
         w.command('quotas', 'quotas', validator=validate_workspace_info)
+        w.command('keys list', 'list_keys')
+        w.command('keys regenerate', 'regenerate_keys')
+        w.command('update', 'enable_keys')
 
     with self.command_group('quantum target', target_ops) as t:
         t.command('list', 'list', validator=validate_workspace_info, table_transformer=transform_targets)

src/quantum/azext_quantum/operations/job.py

@@ -382,9 +382,9 @@ def _submit_directly_to_service(cmd, resource_group_name, workspace_name, locati
     if storage is None:
         from .workspace import get as ws_get
         ws = ws_get(cmd)
-        if ws.storage_account is None:
+        if ws.properties.storage_account is None:
             raise RequiredArgumentMissingError("No storage account specified or linked with workspace.")
-        storage = ws.storage_account.split('/')[-1]
+        storage = ws.properties.storage_account.split('/')[-1]
     job_id = str(uuid.uuid4())
     container_name = "quantum-job-" + job_id
     connection_string_dict = show_storage_account_connection_string(cmd, resource_group_name, storage)

src/quantum/azext_quantum/operations/workspace.py

@@ -19,10 +19,11 @@
                                        RequiredArgumentMissingError, ResourceNotFoundError)
 
 from msrestazure.azure_exceptions import CloudError
-from .._client_factory import cf_workspaces, cf_quotas, cf_offerings, _get_data_credentials
+from .._client_factory import cf_workspaces, cf_quotas, cf_workspace, cf_offerings, _get_data_credentials
 from ..vendored_sdks.azure_mgmt_quantum.models import QuantumWorkspace
 from ..vendored_sdks.azure_mgmt_quantum.models import QuantumWorkspaceIdentity
-from ..vendored_sdks.azure_mgmt_quantum.models import Provider
+from ..vendored_sdks.azure_mgmt_quantum.models import Provider, APIKeys, WorkspaceResourceProperties
+from ..vendored_sdks.azure_mgmt_quantum.models._enums import KeyType
 from .offerings import accept_terms, _get_publisher_and_offer_from_provider_id, _get_terms_from_marketplace, OFFER_NOT_AVAILABLE, PUBLISHER_NOT_AVAILABLE
 
 DEFAULT_WORKSPACE_LOCATION = 'westus'
@@ -228,6 +229,10 @@ def create(cmd, resource_group_name, workspace_name, location, storage_account,
     # Until the "--skip-role-assignment" parameter is deprecated, use the old non-ARM code to create a workspace without doing a role assignment
     if skip_role_assignment:
         _add_quantum_providers(cmd, quantum_workspace, provider_sku_list, auto_accept, skip_autoadd)
+        properties = WorkspaceResourceProperties()
+        properties.providers = quantum_workspace.providers
+        properties.api_key_enabled = True
+        quantum_workspace.properties = properties
         poller = client.begin_create_or_update(info.resource_group, info.name, quantum_workspace, polling=False)
         while not poller.done():
             time.sleep(POLLING_TIME_DURATION)
@@ -380,3 +385,62 @@ def clear(cmd):
     info = WorkspaceInfo(cmd)
     info.clear()
     info.save(cmd)
+
+
+def list_keys(cmd, resource_group_name=None, workspace_name=None):
+    """
+    List Azure Quantum workspace api keys.
+    """
+    client = cf_workspace(cmd.cli_ctx)
+    info = WorkspaceInfo(cmd, resource_group_name, workspace_name, None)
+    if (not info.resource_group) or (not info.name):
+        raise ResourceNotFoundError("Please run 'az quantum workspace set' first to select a default Quantum Workspace.")
+
+    keys = client.list_keys(resource_group_name=info.resource_group, workspace_name=info.name)
+    return keys
+
+
+def regenerate_keys(cmd, resource_group_name=None, workspace_name=None, key_type=None):
+    """
+    Regenerate Azure Quantum workspace api keys.
+    """
+    client = cf_workspace(cmd.cli_ctx)
+    info = WorkspaceInfo(cmd, resource_group_name, workspace_name, None)
+    if (not info.resource_group) or (not info.name):
+        raise ResourceNotFoundError("Please run 'az quantum workspace set' first to select a default Quantum Workspace.")
+
+    if (not key_type):
+        raise RequiredArgumentMissingError("Please select the api key to regenerate.")
+
+    keys = []
+    if key_type is not None:
+        for key in key_type.split(','):
+            keys.append(KeyType[key])
+
+    key_specification = APIKeys(keys=keys)
+    response = client.regenerate_keys(resource_group_name=info.resource_group, workspace_name=info.name, key_specification=key_specification)
+    return response
+
+
+def enable_keys(cmd, resource_group_name=None, workspace_name=None, enable_key=None):
+    """
+    Update the default Azure Quantum workspace.
+    """
+    client = cf_workspaces(cmd.cli_ctx)
+    info = WorkspaceInfo(cmd, resource_group_name, workspace_name, None)
+    if (not info.resource_group) or (not info.name):
+        raise ResourceNotFoundError("Please run 'az quantum workspace set' first to select a default Quantum Workspace.")
+
+    if (enable_key != "True") and (enable_key != "true") and (enable_key != "False") and (enable_key != "false"):
+        raise InvalidArgumentValueError("Please set –-enable-api-key to be True/true or False/false.")
+
+    ws = client.get(info.resource_group, info.name)
+
+    if (enable_key == "True" or enable_key == "true"):
+        ws.properties.api_key_enabled = True
+    elif (enable_key == "False" or enable_key == "false"):
+        ws.properties.api_key_enabled = False
+    ws = client.begin_create_or_update(info.resource_group, info.name, ws)
+    if ws:
+        info.save(cmd)
+    return ws