{Databricks} Migrate databricks with CodeGen v2

src/databricks/HISTORY.rst

@@ -2,6 +2,14 @@
 
 Release History
 ===============
+0.9.0
++++++
+* Upgrade API version from 2021-04-01-preview to 2022-04-01-preview
+* az databricks access-connector: Support create/update/list/show/delete access connector.
+* az databricks workspace private-endpoint-connection: Support create/update/list/show/delete private endpoint connection.
+* az databricks workspace private-link-resource: Support list/show private link resource.
+* az databricks workspace outbound-endpoint: Support list outbound endpoint.
+
 0.8.0
 +++++
 * az databricks workspace create: Add --public-network-access to allow creating workspace with network access from public internet

src/databricks/azext_databricks/__init__.py

@@ -12,15 +12,24 @@ class DatabricksClientCommandsLoader(AzCommandsLoader):
 
     def __init__(self, cli_ctx=None):
         from azure.cli.core.commands import CliCommandType
-        from azext_databricks._client_factory import cf_databricks
         databricks_custom = CliCommandType(
-            operations_tmpl='azext_databricks.custom#{}',
-            client_factory=cf_databricks)
+            operations_tmpl='azext_databricks.custom#{}')
         super().__init__(cli_ctx=cli_ctx,
                          custom_command_type=databricks_custom)
 
     def load_command_table(self, args):
         from azext_databricks.commands import load_command_table
+        from azure.cli.core.aaz import load_aaz_command_table
+        try:
+            from . import aaz
+        except ImportError:
+            aaz = None
+        if aaz:
+            load_aaz_command_table(
+                loader=self,
+                aaz_pkg_name=aaz.__name__,
+                args=args
+            )
         load_command_table(self, args)
         return self.command_table
 

src/databricks/azext_databricks/_help.py

@@ -7,148 +7,3 @@
 # pylint: disable=too-many-lines
 # pylint: disable=line-too-long
 from knack.help_files import helps  # pylint: disable=unused-import
-
-
-helps['databricks'] = """
-    type: group
-    short-summary: Manage databricks workspaces.
-"""
-
-helps['databricks workspace'] = """
-    type: group
-    short-summary: Commands to manage databricks workspace.
-"""
-
-helps['databricks workspace create'] = """
-    type: command
-    short-summary: Create a new workspace.
-    examples:
-      - name: Create a workspace
-        text: |-
-               az databricks workspace create --resource-group MyResourceGroup --name MyWorkspace --location westus --sku standard
-      - name: Create a workspace with managed identity for storage account
-        text: |-
-               az databricks workspace create --resource-group MyResourceGroup --name MyWorkspace --location eastus2euap --sku premium --prepare-encryption
-"""
-
-helps['databricks workspace update'] = """
-    type: command
-    short-summary: Update the workspace.
-    examples:
-      - name: Update the workspace's tags.
-        text: |-
-               az databricks workspace update --resource-group MyResourceGroup --name MyWorkspace --tags key1=value1 key2=value2
-      - name: Clean the workspace's tags.
-        text: |-
-               az databricks workspace update --resource-group MyResourceGroup --name MyWorkspace --tags ""
-      - name: Prepare for CMK encryption by assigning identity for storage account.
-        text: |-
-               az databricks workspace update --resource-group MyResourceGroup --name MyWorkspace --prepare-encryption
-      - name: Configure CMK encryption
-        text: |-
-               az databricks workspace update --resource-group MyResourceGroup --name MyWorkspace --key-source Microsoft.KeyVault \
---key-name MyKey --key-vault https://myKeyVault.vault.azure.net/ --key-version 00000000000000000000000000000000
-      - name: Revert encryption to Microsoft Managed Keys
-        text: |-
-               az databricks workspace update --resource-group MyResourceGroup --name MyWorkspace --key-source Default
-"""
-
-helps['databricks workspace delete'] = """
-    type: command
-    short-summary: Delete the workspace.
-    examples:
-      - name: Delete the workspace
-        text: |-
-               az databricks workspace delete --resource-group MyResourceGroup --name MyWorkspace
-"""
-
-helps['databricks workspace show'] = """
-    type: command
-    short-summary: Show the workspace.
-    examples:
-      - name: Show the workspace
-        text: |-
-               az databricks workspace show --resource-group MyResourceGroup --name MyWorkspace
-"""
-
-helps['databricks workspace list'] = """
-    type: command
-    short-summary: Get all the workspaces.
-    examples:
-      - name: List workspaces within a resource group
-        text: |-
-               az databricks workspace list --resource-group MyResourceGroup
-      - name: List workspaces within the default subscription
-        text: |-
-               az databricks workspace list
-"""
-
-helps['databricks workspace wait'] = """
-    type: command
-    short-summary: Place the CLI in a waiting state until a condition of the Databricks workspace is met.
-    examples:
-        - name: Pause executing next line of CLI script until the Databricks workspace is successfully provisioned.
-          text: az databricks workspace wait --resource-group MyResourceGroup --name MyWorkspace \\
-                --created
-"""
-
-helps['databricks workspace vnet-peering'] = """
-    type: group
-    short-summary: Commands to manage databricks workspace vnet peering.
-"""
-
-helps['databricks workspace vnet-peering create'] = """
-    type: command
-    short-summary: Create a vnet peering for a workspace.
-    examples:
-      - name: Create a vnet peering for a workspace
-        text: |-
-               az databricks workspace vnet-peering create --resource-group MyResourceGroup --workspace-name MyWorkspace \\
-               -n MyPeering --remote-vnet /subscriptions/000000-0000-0000/resourceGroups/MyRG/providers/Microsoft.Network/virtualNetworks/MyVNet
-"""
-
-helps['databricks workspace vnet-peering update'] = """
-    type: command
-    short-summary: Update the vnet peering.
-    examples:
-      - name: Update the vnet peering (enable gateway transit and disable virtual network access).
-        text: |-
-               az databricks workspace vnet-peering update --resource-group MyResourceGroup --workspace-name MyWorkspace \\
-               -n MyPeering --allow-gateway-transit --allow-virtual-network-access false
-"""
-
-helps['databricks workspace vnet-peering list'] = """
-    type: command
-    short-summary: List vnet peerings under a workspace.
-    examples:
-      - name: List vnet peerings under a workspace.
-        text: |-
-               az databricks workspace vnet-peering list --resource-group MyResourceGroup --workspace-name MyWorkspace
-"""
-
-helps['databricks workspace vnet-peering delete'] = """
-    type: command
-    short-summary: Delete the vnet peering.
-    examples:
-      - name: Delete the vnet peering.
-        text: |-
-               az databricks workspace vnet-peering delete --resource-group MyResourceGroup --workspace-name MyWorkspace -n MyPeering
-"""
-
-helps['databricks workspace vnet-peering show'] = """
-    type: command
-    short-summary: Show the vnet peering.
-    examples:
-      - name: Show the vnet peering
-        text: |-
-               az databricks workspace vnet-peering show --resource-group MyResourceGroup --workspace-name MyWorkspace -n MyPeering
-"""
-
-helps['databricks workspace vnet-peering wait'] = """
-    type: command
-    short-summary: Place the CLI in a waiting state until a condition of the Databricks workspace vnet peering is met.
-    examples:
-        - name: Pause executing next line of CLI script until the Databricks workspace vnet peering is successfully provisioned.
-          text: az databricks workspace vnet-peering wait --resource-group MyResourceGroup --workspace-name MyWorkspace \\
-                --name MyPeering --created
-"""

src/databricks/azext_databricks/_params.py

@@ -6,68 +6,6 @@
 # pylint: disable=too-many-lines
 # pylint: disable=too-many-statements
 
-from azure.cli.core.commands.parameters import (
-    tags_type,
-    get_enum_type,
-    get_location_type,
-    get_three_state_flag
-)
 
-from azure.cli.core.commands.validators import get_default_location_from_resource_group
-from ._validators import validate_encryption_values, validate_network_id, validate_managed_resource_group
-
-
-def load_arguments(self, _):
-    from .vendored_sdks.databricks.models import PublicNetworkAccess, RequiredNsgRules
-
-    with self.argument_context('databricks workspace create') as c:
-        c.argument('workspace_name', options_list=['--name', '-n'], help='The name of the workspace.')
-        c.argument('tags', tags_type)
-        c.argument('location', arg_type=get_location_type(self.cli_ctx), validator=get_default_location_from_resource_group)
-        c.argument('managed_resource_group', validator=validate_managed_resource_group, help='The managed resource group to create. It can be either a name or a resource ID.')
-        c.argument('custom_virtual_network_id', options_list=['--vnet'], arg_group='Custom VNET', validator=validate_network_id('custom_virtual_network_id'), help='Virtual Network name or resource ID.')
-        c.argument('custom_public_subnet_name', options_list=['--public-subnet'], arg_group='Custom VNET', help='The name of a Public Subnet within the Virtual Network.')
-        c.argument('custom_private_subnet_name', options_list=['--private-subnet'], arg_group='Custom VNET', help='The name of a Private Subnet within the Virtual Network.')
-        c.argument('sku_name', options_list=['--sku'], arg_type=get_enum_type(['standard', 'premium', 'trial']), help='The SKU tier name.')
-        c.argument('prepare_encryption', action='store_true', help='Flag to enable the Managed Identity for managed storage account to prepare for CMK encryption.')
-        c.argument('require_infrastructure_encryption', action='store_true', help='Flag to enable the DBFS root file system with secondary layer of encryption with platform managed keys for data at rest.')
-        c.argument('enable_no_public_ip', action='store_true', help='Flag to enable the no public ip feature.')
-        c.argument('public_network_access', arg_type=get_enum_type(PublicNetworkAccess), help='The configuration to set whether network access from public internet to the endpoints are allowed.')
-        c.argument('required_nsg_rules', options_list='--required-nsg-rules', arg_type=get_enum_type(RequiredNsgRules),
-                   help='The type of Nsg rule for internal use only.')
-
-    with self.argument_context('databricks workspace update') as c:
-        c.argument('workspace_name', options_list=['--name', '-n'], id_part='name', help='The name of the workspace.')
-        c.argument('tags', tags_type)
-        c.argument('prepare_encryption', action='store_true', help='Flag to enable the Managed Identity for managed storage account to prepare for CMK encryption.')
-        c.argument('encryption_key_source', options_list=['--key-source'], arg_group='Encryption', arg_type=get_enum_type(['Default', 'Microsoft.Keyvault']), validator=validate_encryption_values, help='The encryption key source (provider).')
-        c.argument('encryption_key_name', options_list=['--key-name'], arg_group='Encryption', help='The name of KeyVault key.')
-        c.argument('encryption_key_version', options_list=['--key-version'], arg_group='Encryption', help='The version of KeyVault key. It is optional when updating CMK.')
-        c.argument('encryption_key_vault', options_list=['--key-vault'], arg_group='Encryption', help='The Uri of KeyVault.')
-
-    with self.argument_context('databricks workspace delete') as c:
-        c.argument('workspace_name', options_list=['--name', '-n'], id_part='name', help='The name of the workspace.')
-
-    with self.argument_context('databricks workspace show') as c:
-        c.argument('workspace_name', options_list=['--name', '-n'], id_part='name', help='The name of the workspace.')
-
-    with self.argument_context('databricks workspace list') as c:
-        pass
-
-    with self.argument_context('databricks workspace wait') as c:
-        c.argument('workspace_name', options_list=['--name', '-n'], id_part='name', help='The name of the workspace.')
-
-    for scope in ['create', 'show', 'delete', 'update', 'wait']:
-        with self.argument_context('databricks workspace vnet-peering {}'.format(scope)) as c:
-            c.argument('workspace_name', help='The name of the workspace.')
-            c.argument('peering_name', options_list=['--name', '-n'], help='The name of the vnet peering.')
-
-    with self.argument_context('databricks workspace vnet-peering create') as c:
-        c.argument('remote_virtual_network', options_list=['--remote-vnet'], validator=validate_network_id('remote_virtual_network'), help='The remote virtual network name or Resource ID.')
-
-    for scope in ['create', 'update']:
-        with self.argument_context('databricks workspace vnet-peering {}'.format(scope)) as c:
-            c.argument('allow_virtual_network_access', arg_type=get_three_state_flag(), help='Whether the VMs in the local virtual network space would be able to access the VMs in remote virtual network space.')
-            c.argument('allow_forwarded_traffic', arg_type=get_three_state_flag(), help='Whether the forwarded traffic from the VMs in the local virtual network will be allowed/disallowed in remote virtual network.')
-            c.argument('allow_gateway_transit', arg_type=get_three_state_flag(), help='If gateway links can be used in remote virtual networking to link to this virtual network.')
-            c.argument('use_remote_gateways', arg_type=get_three_state_flag(), help='If remote gateways can be used on this virtual network. If the flag is set to true, and allowGatewayTransit on remote peering is also true, virtual network will use gateways of remote virtual network for transit. Only one peering can have this flag set to true. This flag cannot be set if virtual network already has a gateway.')
+def load_arguments(self, _):  # pylint: disable=unused-argument
+    pass

src/databricks/azext_databricks/vendored_sdks/__init__.py → src/databricks/azext_databricks/aaz/__init__.py

@@ -1,6 +1,6 @@
 # --------------------------------------------------------------------------------------------
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License. See License.txt in the project root for license information.
+#
+# Code generated by aaz-dev-tools
 # --------------------------------------------------------------------------------------------
-
-__import__('pkg_resources').declare_namespace(__name__)

src/databricks/azext_databricks/aaz/latest/__init__.py

@@ -0,0 +1,6 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+#
+# Code generated by aaz-dev-tools
+# --------------------------------------------------------------------------------------------

src/databricks/azext_databricks/aaz/latest/databricks/__cmd_group.py

@@ -0,0 +1,23 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+#
+# Code generated by aaz-dev-tools
+# --------------------------------------------------------------------------------------------
+
+# pylint: skip-file
+# flake8: noqa
+
+from azure.cli.core.aaz import *
+
+
+@register_command_group(
+    "databricks",
+)
+class __CMDGroup(AAZCommandGroup):
+    """Manage databricks workspaces.
+    """
+    pass
+
+
+__all__ = ["__CMDGroup"]