[Synapse] Add customer-managed key related cmdlets

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

@@ -54,6 +54,14 @@ def cf_synapse_client_ipfirewallrules_factory(cli_ctx, *_):
     return cf_synapse(cli_ctx).ip_firewall_rules
 
 
+def cf_synapse_client_cmk_factory(cli_ctx, *_):
+    return cf_synapse(cli_ctx).keys
+
+
+def cf_synapse_client_managed_identity_sqlcontrol_factory(cli_ctx, *_):
+    return cf_synapse(cli_ctx).workspace_managed_identity_sql_control_settings
+
+
 def cf_synapse_client_integrationruntimes_factory(cli_ctx, *_):
     return cf_synapse(cli_ctx).integration_runtimes
 

src/azure-cli/azure/cli/command_modules/synapse/_help.py

@@ -32,6 +32,12 @@
         az synapse workspace create --name fromcli4 --resource-group rg \\
           --storage-account /subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/rg/providers/Microsoft.Storage/storageAccounts/testadlsgen2 --file-system testfilesystem \\
           --sql-admin-login-user cliuser1 --sql-admin-login-password Password123! --location "East US"
+  - name: Create a Synapse workspace using customer-managed key
+    text: |-
+        az synapse workspace create --name fromcli4 --resource-group rg \\
+          --storage-account testadlsgen2 --file-system testfilesystem \\
+          --sql-admin-login-user cliuser1 --sql-admin-login-password Password123! --location "East US" \\
+          --key-identifier https://{keyvaultname}.vault.azure.net/keys/{keyname} --key-name testcmk
 """
 
 helps['synapse workspace list'] = """
@@ -544,6 +550,105 @@
         --state Disabled
 """
 
+helps['synapse workspace key'] = """
+type: group
+short-summary:  Manage workspace's keys.
+"""
+
+helps['synapse workspace key create'] = """
+type: command
+short-summary: Create a workspace's key.
+examples:
+  - name: Create a workspace's key.
+    text: |-
+        az synapse workspace key create --name newkey --workspace-name testsynapseworkspace \\
+        --resource-group rg --key-identifier https://{keyvaultname}.vault.azure.net/keys/{keyname}
+"""
+
+helps['synapse workspace key update'] = """
+type: command
+short-summary: Update a workspace's key or update the state of key to change the workspace state from pending to success state when the workspace is first being provisioned.
+examples:
+  - name: Update a workspace's key.
+    text: |-
+        az synapse workspace key update --name newkey --workspace-name testsynapseworkspace \\
+        --resource-group rg --key-identifier https://{keyvaultname}.vault.azure.net/keys/{keyname}
+"""
+
+helps['synapse workspace key delete'] = """
+type: command
+short-summary: Delete a workspace's key. The key at active status can't be deleted.
+examples:
+  - name: Delete a workspace's key.
+    text: |-
+        az synapse workspace key delete --name newkey --workspace-name testsynapseworkspace \\
+        --resource-group rg
+"""
+
+helps['synapse workspace key show'] = """
+type: command
+short-summary: Show a workspace's key by name.
+examples:
+  - name: Show a workspace's key.
+    text: |-
+        az synapse workspace key show --name newkey --workspace-name testsynapseworkspace \\
+        --resource-group rg
+"""
+
+helps['synapse workspace key list'] = """
+type: command
+short-summary: List keys under specified workspace.
+examples:
+  - name: List keys under specified workspace.
+    text: |-
+        az synapse workspace key list --workspace-name testsynapseworkspace --resource-group rg
+"""
+
+helps['synapse workspace key wait'] = """
+type: command
+short-summary: Place the CLI in a waiting state until a condition of a workspace key is met.
+"""
+
+helps['synapse workspace managed-identity'] = """
+type: group
+short-summary:  Manage workspace's managed-identity.
+"""
+
+helps['synapse workspace managed-identity show-sql-access'] = """
+type: command
+short-summary: Show workspace's sql-access state to managed-identity.
+examples:
+  - name: Show workspace's sql-access state to managed-identity.
+    text: |-
+        az synapse workspace managed-identity show-sql-access --workspace-name testsynapseworkspace \\
+        --resource-group rg
+"""
+
+helps['synapse workspace managed-identity revoke-sql-access'] = """
+type: command
+short-summary: Revoke workspace's sql-access to managed-identity.
+examples:
+  - name: Revoke workspace's sql-access to managed-identity.
+    text: |-
+        az synapse workspace managed-identity revoke-sql-access --workspace-name testsynapseworkspace \\
+        --resource-group rg
+"""
+
+helps['synapse workspace managed-identity grant-sql-access'] = """
+type: command
+short-summary: Grant workspace's sql-access to managed-identity.
+examples:
+  - name: Grant workspace's sql-access to managed-identity.
+    text: |-
+        az synapse workspace managed-identity grant-sql-access --workspace-name testsynapseworkspace \\
+        --resource-group rg
+"""
+
+helps['synapse workspace managed-identity wait'] = """
+type: command
+short-summary: Place the CLI in a waiting state until a condition of sql-access state to managed-identity is met.
+"""
+
 helps['synapse workspace firewall-rule'] = """
 type: group
 short-summary:  Manage a workspace's firewall rules.

src/azure-cli/azure/cli/command_modules/synapse/_params.py

@@ -48,7 +48,7 @@ def _configure_security_or_audit_policy_storage_params(arg_ctx):
 
 def load_arguments(self, _):
     # synapse workspace
-    for scope in ['show', 'create', 'update', 'delete']:
+    for scope in ['show', 'create', 'update', 'delete', 'activate']:
         with self.argument_context('synapse workspace ' + scope) as c:
             c.argument('workspace_name', arg_type=name_type, id_part='name', help='The workspace name.')
 
@@ -57,6 +57,8 @@ def load_arguments(self, _):
             c.argument('sql_admin_login_password', options_list=['--sql-admin-login-password', '-p'],
                        help='The sql administrator login password.')
             c.argument('tags', arg_type=tags_type)
+            c.argument('allowed_aad_tenant_ids', options_list=['--allowed-tenant-ids'], nargs='+', help="The approved Azure AD tenants which outbound data traffic allowed to. The Azure AD tenant of the current user will be included by default. Use ""(\'""\' in PowerShell) to disable all allowed tenant ids.")
+            c.argument('key_name', help='The workspace customer-managed key display name. All existing keys can be found using "az synapse workspace key list" cmdlet.')
 
     with self.argument_context('synapse workspace create') as c:
         c.argument('location', get_location_type(self.cli_ctx), validator=get_default_location_from_resource_group)
@@ -69,6 +71,9 @@ def load_arguments(self, _):
                                                                    '--enable-managed-virtual-network'],
                    arg_type=get_three_state_flag(),
                    help='The flag indicates whether enable managed virtual network.')
+        c.argument('prevent_data_exfiltration', arg_type=get_three_state_flag(),
+                   help='The flag indicates whether enable data exfiltration.', options_list=['--prevent-exfiltration', '--prevent-data-exfiltration'])
+        c.argument('key_identifier', help='The customer-managed key used to encrypt all data at rest in the workspace. Key identifier should be in the format of: https://{keyvaultname}.vault.azure.net/keys/{keyname}.', options_list=['--key-identifier', '--cmk'])
 
     with self.argument_context('synapse workspace check-name') as c:
         c.argument('name', arg_type=name_type, help='The name you wanted to check.')
@@ -311,6 +316,32 @@ def load_arguments(self, _):
             c.argument('end_ip_address', help='The end IP address of the firewall rule. Must be IPv4 format. '
                                               'Must be greater than or equal to startIpAddress.')
 
+    # synapse workspace key
+    with self.argument_context('synapse workspace key') as c:
+        c.argument('workspace_name', id_part='name', help='The workspace name.')
+
+    with self.argument_context('synapse workspace key list') as c:
+        c.argument('workspace_name', id_part=None, help='The workspace name.')
+
+    for scope in ['show', 'create', 'delete', 'update']:
+        with self.argument_context('synapse workspace key ' + scope) as c:
+            c.argument('key_name', arg_type=name_type, id_part='child_name_1', help='The workspace customer-managed key display name. All existing keys can be found using /"az synapse workspace key list/" cmdlet.')
+
+    with self.argument_context('synapse workspace key create') as c:
+        c.argument('key_identifier', help='The Key Vault Url of the workspace encryption key. should be in the format of: https://{keyvaultname}.vault.azure.net/keys/{keyname}.')
+
+    with self.argument_context('synapse workspace key update') as c:
+        c.argument('key_identifier', help='The Key Vault Url of the workspace encryption key. should be in the format of: https://{keyvaultname}.vault.azure.net/keys/{keyname}.')
+        c.argument('is_active', arg_type=get_three_state_flag(), help='Set True to change the workspace state from pending to success state.')
+
+    # synapse workspace managed-identity
+    with self.argument_context('synapse workspace managed-identity') as c:
+        c.argument('workspace_name', id_part='name', help='The workspace name.')
+
+    for scope in ['grant-sql-access', 'revoke-sql-access', ' show-sql-access']:
+        with self.argument_context('synapse workspace managed-identity ' + scope) as c:
+            c.argument('workspace_name', id_part='name', help='The workspace name.')
+
     # synapse spark job
     for scope in ['job', 'session', 'statement']:
         with self.argument_context('synapse spark ' + scope) as c:

src/azure-cli/azure/cli/command_modules/synapse/commands.py

@@ -13,11 +13,13 @@ def load_command_table(self, _):
     from ._client_factory import cf_synapse_client_bigdatapool_factory
     from ._client_factory import cf_synapse_client_sqlpool_factory
     from ._client_factory import cf_synapse_client_ipfirewallrules_factory
+    from ._client_factory import cf_synapse_client_cmk_factory
     from ._client_factory import cf_synapse_client_sqlpool_sensitivity_labels_factory
     from ._client_factory import cf_synapse_client_restorable_dropped_sqlpools_factory
     from ._client_factory import cf_synapse_client_sqlpool_transparent_data_encryptions_factory
     from ._client_factory import cf_synapse_client_sqlpool_security_alert_policies_factory
     from ._client_factory import cf_synapse_client_sqlpool_blob_auditing_policies_factory
+    from ._client_factory import cf_synapse_client_managed_identity_sqlcontrol_factory
     from ._client_factory import cf_synapse_client_workspace_aad_admins_factory
     from ._client_factory import cf_synapse_client_sqlserver_blob_auditing_policies_factory
     from ._client_factory import cf_synapse_client_integrationruntimes_factory
@@ -89,6 +91,14 @@ def get_custom_sdk(custom_module, client_factory):
         operations_tmpl='azure.mgmt.synapse.operations#IpFirewallRulesOperations.{}',
         client_factory=cf_synapse_client_ipfirewallrules_factory)
 
+    synapse_cmk_sdk = CliCommandType(
+        operations_tmpl='azure.mgmt.synapse.operations#KeysOperations.{}',
+        client_factory=cf_synapse_client_cmk_factory)
+
+    synapse_managedidentitysqlcontrol_sdk = CliCommandType(
+        operations_tmpl='azure.mgmt.synapse.operations#WorkspaceManagedIdentitySqlControlSettingsOperations.{}',
+        client_factory=cf_synapse_client_managed_identity_sqlcontrol_factory)
+
     synapse_integrationruntimes_sdk = CliCommandType(
         operations_tmpl='azure.mgmt.synapse.operations#IntegrationRuntimesOperations.{}',
         client_factory=cf_synapse_client_integrationruntimes_factory)
@@ -311,6 +321,26 @@ def get_custom_sdk(custom_module, client_factory):
                   client_factory=cf_synapse_client_integrationruntimestatus_factory)
         g.wait_command('wait')
 
+    # Management Plane Commands --Keys
+    with self.command_group('synapse workspace key', command_type=synapse_cmk_sdk,
+                            custom_command_type=get_custom_sdk('workspace', cf_synapse_client_cmk_factory),
+                            client_factory=cf_synapse_client_cmk_factory) as g:
+        g.command('list', 'list_by_workspace')
+        g.show_command('show', 'get')
+        g.custom_command('create', 'create_workspace_key', supports_no_wait=True)
+        g.custom_command('update', 'update_workspace_key', supports_no_wait=True)
+        g.command('delete', 'delete', confirmation=True, supports_no_wait=True)
+        g.wait_command('wait')
+
+    # Management Plane Commands --Managed-Identity
+    with self.command_group('synapse workspace managed-identity', command_type=synapse_managedidentitysqlcontrol_sdk,
+                            custom_command_type=get_custom_sdk('workspace', cf_synapse_client_managed_identity_sqlcontrol_factory),
+                            client_factory=cf_synapse_client_managed_identity_sqlcontrol_factory) as g:
+        g.show_command('show-sql-access', 'get')
+        g.custom_command('grant-sql-access', 'grant_sql_access_to_managed_identity', supports_no_wait=True)
+        g.custom_command('revoke-sql-access', 'revoke_sql_access_to_managed_identity', supports_no_wait=True)
+        g.wait_command('wait')
+
     with self.command_group('synapse integration-runtime-node', command_type=synapse_integrationruntimenodes_sdk,
                             client_factory=cf_synapse_client_integrationruntimenodes_factory) as g:
         g.show_command('show', 'get')

src/azure-cli/azure/cli/command_modules/synapse/operations/workspace.py

@@ -2,10 +2,11 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
-# pylint: disable=unused-argument
+# pylint: disable=unused-argument, line-too-long
 from azure.cli.core.util import sdk_no_wait, CLIError
 from azure.mgmt.synapse.models import Workspace, WorkspacePatchInfo, ManagedIdentity, \
-    DataLakeStorageAccountDetails
+    DataLakeStorageAccountDetails, WorkspaceKeyDetails, CustomerManagedKeyDetails, EncryptionDetails, ManagedVirtualNetworkSettings, \
+    ManagedIdentitySqlControlSettingsModelPropertiesGrantSqlControlToManagedIdentity
 
 
 # Synapse workspace
@@ -15,12 +16,29 @@ def list_workspaces(cmd, client, resource_group_name=None):
 
 
 def create_workspace(cmd, client, resource_group_name, workspace_name, storage_account, file_system,
-                     sql_admin_login_user, sql_admin_login_password, location=None, enable_managed_virtual_network=None,
-                     tags=None, no_wait=False):
+                     sql_admin_login_user, sql_admin_login_password, location=None, key_name="default", key_identifier=None, enable_managed_virtual_network=None,
+                     allowed_aad_tenant_ids=None, prevent_data_exfiltration=None, tags=None, no_wait=False):
     identity_type = "SystemAssigned"
     identity = ManagedIdentity(type=identity_type)
     account_url = "https://{}.dfs.{}".format(storage_account, cmd.cli_ctx.cloud.suffixes.storage_endpoint)
     default_data_lake_storage = DataLakeStorageAccountDetails(account_url=account_url, filesystem=file_system)
+    encryption = None
+    managed_virtual_network_settings = None
+    tenant_ids_list = None
+    if key_identifier is not None:
+        workspace_key_detail = WorkspaceKeyDetails(name=key_name, key_vault_url=key_identifier)
+        encryption = EncryptionDetails(cmk=CustomerManagedKeyDetails(key=workspace_key_detail))
+
+    if [''] == allowed_aad_tenant_ids:
+        tenant_ids_list = []
+    else:
+        tenant_ids_list = allowed_aad_tenant_ids
+
+    if enable_managed_virtual_network:
+        if prevent_data_exfiltration:
+            managed_virtual_network_settings = ManagedVirtualNetworkSettings(prevent_data_exfiltration=True, allowed_aad_tenant_ids_for_linking=tenant_ids_list)
+        else:
+            managed_virtual_network_settings = ManagedVirtualNetworkSettings(prevent_data_exfiltration=False)
 
     workspace_info = Workspace(
         identity=identity,
@@ -29,14 +47,29 @@ def create_workspace(cmd, client, resource_group_name, workspace_name, storage_a
         sql_administrator_login_password=sql_admin_login_password,
         location=location,
         managed_virtual_network="default" if enable_managed_virtual_network is True else None,
+        managed_virtual_network_settings=managed_virtual_network_settings,
+        encryption=encryption,
         tags=tags
     )
     return sdk_no_wait(no_wait, client.create_or_update, resource_group_name, workspace_name, workspace_info)
 
 
 def update_workspace(cmd, client, resource_group_name, workspace_name, sql_admin_login_password=None,
-                     tags=None, no_wait=False):
-    workspace_patch_info = WorkspacePatchInfo(tags=tags, sql_admin_login_password=sql_admin_login_password)
+                     allowed_aad_tenant_ids=None, tags=None, key_name=None, no_wait=False):
+    encryption = None
+    tenant_ids_list = None
+
+    if key_name:
+        workspace_key_detail = WorkspaceKeyDetails(name=key_name)
+        encryption = EncryptionDetails(cmk=CustomerManagedKeyDetails(key=workspace_key_detail))
+
+    if allowed_aad_tenant_ids and '' in allowed_aad_tenant_ids:
+        tenant_ids_list = []
+    else:
+        tenant_ids_list = allowed_aad_tenant_ids
+
+    updated_vnet_settings = ManagedVirtualNetworkSettings(allowed_aad_tenant_ids_for_linking=tenant_ids_list) if allowed_aad_tenant_ids is not None else None
+    workspace_patch_info = WorkspacePatchInfo(tags=tags, sql_admin_login_password=sql_admin_login_password, encryption=encryption, managed_virtual_network_settings=updated_vnet_settings)
     return sdk_no_wait(no_wait, client.update, resource_group_name, workspace_name, workspace_patch_info)
 
 
@@ -67,3 +100,21 @@ def update_firewall_rule(cmd, client, resource_group_name, workspace_name, rule_
     end_ip_address = end_ip_address or firewall.end_ip_address
     return sdk_no_wait(no_wait, client.create_or_update, resource_group_name, workspace_name, rule_name,
                        start_ip_address=start_ip_address, end_ip_address=end_ip_address)
+
+
+def create_workspace_key(cmd, client, resource_group_name, workspace_name, key_name, key_identifier, no_wait=False):
+    return sdk_no_wait(no_wait, client.create_or_update, resource_group_name, workspace_name, key_name=key_name, key_vault_url=key_identifier)
+
+
+def update_workspace_key(cmd, client, resource_group_name, workspace_name, key_name, key_identifier, is_active=False, no_wait=False):
+    return sdk_no_wait(no_wait, client.create_or_update, resource_group_name, workspace_name, is_active_cmk=is_active, key_name=key_name, key_vault_url=key_identifier)
+
+
+def grant_sql_access_to_managed_identity(cmd, client, resource_group_name, workspace_name, no_wait=False):
+    grant_sql_access_setting = ManagedIdentitySqlControlSettingsModelPropertiesGrantSqlControlToManagedIdentity(desired_state="Enabled")
+    return sdk_no_wait(no_wait, client.create_or_update, resource_group_name, workspace_name, grant_sql_access_setting)
+
+
+def revoke_sql_access_to_managed_identity(cmd, client, resource_group_name, workspace_name, no_wait=False):
+    revoke_sql_access_setting = ManagedIdentitySqlControlSettingsModelPropertiesGrantSqlControlToManagedIdentity(desired_state="Disabled")
+    return sdk_no_wait(no_wait, client.create_or_update, resource_group_name, workspace_name, revoke_sql_access_setting)