[Databricks]az databricks workspace create: Add new parameters --publ…

src/databricks/azext_databricks/_client_factory.py

@@ -6,8 +6,8 @@
 
 def cf_databricks(cli_ctx, *_):
     from azure.cli.core.commands.client_factory import get_mgmt_service_client
-    from .vendored_sdks.databricks import DatabricksClient
-    return get_mgmt_service_client(cli_ctx, DatabricksClient)
+    from .vendored_sdks.databricks import AzureDatabricksManagementClient
+    return get_mgmt_service_client(cli_ctx, AzureDatabricksManagementClient)
 
 
 def cf_workspaces(cli_ctx, *_):

src/databricks/azext_databricks/_params.py

@@ -13,6 +13,7 @@
     get_three_state_flag
 )
 
+from .vendored_sdks.databricks.models import PublicNetworkAccess, RequiredNsgRules
 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
 
@@ -31,6 +32,9 @@ def load_arguments(self, _):
         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.')

src/databricks/azext_databricks/custom.py

@@ -25,12 +25,16 @@ def create_databricks_workspace(cmd, client,
                                 prepare_encryption=None,
                                 require_infrastructure_encryption=None,
                                 enable_no_public_ip=None,
+                                public_network_access=None,
+                                required_nsg_rules=None,
                                 no_wait=False):
     body = {}
     body['tags'] = tags  # dictionary
     body['location'] = location  # str
     body['managed_resource_group_id'] = managed_resource_group  # str
     body.setdefault('sku', {})['name'] = sku_name  # str
+    body['public_network_access'] = public_network_access
+    body['required_nsg_rules'] = required_nsg_rules
 
     parameters = {}
     _set_parameter_value(parameters, 'custom_virtual_network_id', custom_virtual_network_id)  # str

src/databricks/azext_databricks/tests/latest/test_databricks_scenario.py

@@ -126,6 +126,30 @@ def test_databricks(self, resource_group, key_vault):
                  '-y',
                  checks=[])
 
+    @ResourceGroupPreparer(name_prefix='cli_test_databricks_v1', location="westus")
+    def test_databricks_v1(self, resource_group):
+        self.kwargs.update({
+            'workspace_name': 'my-test-workspace'
+        })
+
+        self.cmd('az databricks workspace create '
+                 '--resource-group {rg} '
+                 '--name {workspace_name} '
+                 '--location westus '
+                 '--sku premium '
+                 '--public-network-access Enabled '
+                 '--required-nsg-rules AllRules',
+                 checks=[self.check('name', '{workspace_name}'),
+                         self.check('sku.name', 'premium'),
+                         self.check('publicNetworkAccess', 'Enabled'),
+                         self.check('requiredNsgRules', 'AllRules')])
+
+        self.cmd('az databricks workspace delete '
+                 '--resource-group {rg} '
+                 '--name {workspace_name} '
+                 '-y',
+                 checks=[])
+
 
 class DatabricksVNetPeeringScenarioTest(ScenarioTest):
 

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

@@ -6,11 +6,11 @@
 # Changes may cause incorrect behavior and will be lost if the code is regenerated.
 # --------------------------------------------------------------------------
 
-from ._databricks_client import DatabricksClient
+from ._azure_databricks_management_client import AzureDatabricksManagementClient
 from ._version import VERSION
 
 __version__ = VERSION
-__all__ = ['DatabricksClient']
+__all__ = ['AzureDatabricksManagementClient']
 
 try:
     from ._patch import patch_sdk  # type: ignore

src/databricks/azext_databricks/vendored_sdks/databricks/_databricks_client.py → src/databricks/azext_databricks/vendored_sdks/databricks/_azure_databricks_management_client.py

@@ -16,23 +16,30 @@
     from typing import Any, Optional
 
     from azure.core.credentials import TokenCredential
+    from azure.core.pipeline.transport import HttpRequest, HttpResponse
 
-from ._configuration import DatabricksClientConfiguration
+from ._configuration import AzureDatabricksManagementClientConfiguration
 from .operations import WorkspacesOperations
-from .operations import VNetPeeringOperations
 from .operations import Operations
+from .operations import PrivateLinkResourcesOperations
+from .operations import PrivateEndpointConnectionsOperations
+from .operations import VNetPeeringOperations
 from . import models
 
 
-class DatabricksClient(object):
-    """ARM Databricks.
+class AzureDatabricksManagementClient(object):
+    """The Microsoft Azure management APIs allow end users to operate on Azure Databricks Workspace resources.
 
     :ivar workspaces: WorkspacesOperations operations
-    :vartype workspaces: azure.mgmt.databricks.operations.WorkspacesOperations
-    :ivar vnet_peering: VNetPeeringOperations operations
-    :vartype vnet_peering: azure.mgmt.databricks.operations.VNetPeeringOperations
+    :vartype workspaces: azure_databricks_management_client.operations.WorkspacesOperations
     :ivar operations: Operations operations
-    :vartype operations: azure.mgmt.databricks.operations.Operations
+    :vartype operations: azure_databricks_management_client.operations.Operations
+    :ivar private_link_resources: PrivateLinkResourcesOperations operations
+    :vartype private_link_resources: azure_databricks_management_client.operations.PrivateLinkResourcesOperations
+    :ivar private_endpoint_connections: PrivateEndpointConnectionsOperations operations
+    :vartype private_endpoint_connections: azure_databricks_management_client.operations.PrivateEndpointConnectionsOperations
+    :ivar vnet_peering: VNetPeeringOperations operations
+    :vartype vnet_peering: azure_databricks_management_client.operations.VNetPeeringOperations
     :param credential: Credential needed for the client to connect to Azure.
     :type credential: ~azure.core.credentials.TokenCredential
     :param subscription_id: The ID of the target subscription.
@@ -51,7 +58,7 @@ def __init__(
         # type: (...) -> None
         if not base_url:
             base_url = 'https://management.azure.com'
-        self._config = DatabricksClientConfiguration(credential, subscription_id, **kwargs)
+        self._config = AzureDatabricksManagementClientConfiguration(credential, subscription_id, **kwargs)
         self._client = ARMPipelineClient(base_url=base_url, config=self._config, **kwargs)
 
         client_models = {k: v for k, v in models.__dict__.items() if isinstance(v, type)}
@@ -61,17 +68,39 @@ def __init__(
 
         self.workspaces = WorkspacesOperations(
             self._client, self._config, self._serialize, self._deserialize)
-        self.vnet_peering = VNetPeeringOperations(
-            self._client, self._config, self._serialize, self._deserialize)
         self.operations = Operations(
             self._client, self._config, self._serialize, self._deserialize)
+        self.private_link_resources = PrivateLinkResourcesOperations(
+            self._client, self._config, self._serialize, self._deserialize)
+        self.private_endpoint_connections = PrivateEndpointConnectionsOperations(
+            self._client, self._config, self._serialize, self._deserialize)
+        self.vnet_peering = VNetPeeringOperations(
+            self._client, self._config, self._serialize, self._deserialize)
+
+    def _send_request(self, http_request, **kwargs):
+        # type: (HttpRequest, Any) -> HttpResponse
+        """Runs the network request through the client's chained policies.
+
+        :param http_request: The network request you want to make. Required.
+        :type http_request: ~azure.core.pipeline.transport.HttpRequest
+        :keyword bool stream: Whether the response payload will be streamed. Defaults to True.
+        :return: The response of your network call. Does not do error handling on your response.
+        :rtype: ~azure.core.pipeline.transport.HttpResponse
+        """
+        path_format_arguments = {
+            'subscriptionId': self._serialize.url("self._config.subscription_id", self._config.subscription_id, 'str'),
+        }
+        http_request.url = self._client.format_url(http_request.url, **path_format_arguments)
+        stream = kwargs.pop("stream", True)
+        pipeline_response = self._client._pipeline.run(http_request, stream=stream, **kwargs)
+        return pipeline_response.http_response
 
     def close(self):
         # type: () -> None
         self._client.close()
 
     def __enter__(self):
-        # type: () -> DatabricksClient
+        # type: () -> AzureDatabricksManagementClient
         self._client.__enter__()
         return self
 

src/databricks/azext_databricks/vendored_sdks/databricks/_configuration.py

@@ -21,8 +21,8 @@
     from azure.core.credentials import TokenCredential
 
 
-class DatabricksClientConfiguration(Configuration):
-    """Configuration for DatabricksClient.
+class AzureDatabricksManagementClientConfiguration(Configuration):
+    """Configuration for AzureDatabricksManagementClient.
 
     Note that all parameters used to create this instance are saved as instance
     attributes.
@@ -44,11 +44,10 @@ def __init__(
             raise ValueError("Parameter 'credential' must not be None.")
         if subscription_id is None:
             raise ValueError("Parameter 'subscription_id' must not be None.")
-        super(DatabricksClientConfiguration, self).__init__(**kwargs)
+        super(AzureDatabricksManagementClientConfiguration, self).__init__(**kwargs)
 
         self.credential = credential
         self.subscription_id = subscription_id
-        self.api_version = "2018-04-01"
         self.credential_scopes = kwargs.pop('credential_scopes', ['https://management.azure.com/.default'])
         kwargs.setdefault('sdk_moniker', 'mgmt-databricks/{}'.format(VERSION))
         self._configure(**kwargs)

src/databricks/azext_databricks/vendored_sdks/databricks/_metadata.json

@@ -0,0 +1,107 @@
+{
+    "chosen_version": "",
+    "total_api_version_list": ["2018-04-01", "2021-04-01-preview"],
+    "client": {
+        "name": "AzureDatabricksManagementClient",
+        "filename": "_azure_databricks_management_client",
+        "description": "The Microsoft Azure management APIs allow end users to operate on Azure Databricks Workspace resources.",
+        "base_url": "\u0027https://management.azure.com\u0027",
+        "custom_base_url": null,
+        "azure_arm": true,
+        "has_lro_operations": true,
+        "client_side_validation": false,
+        "sync_imports": "{\"typing\": {\"azurecore\": {\"azure.core.credentials\": [\"TokenCredential\"]}}, \"regular\": {\"azurecore\": {\"azure.profiles\": [\"KnownProfiles\", \"ProfileDefinition\"], \"azure.profiles.multiapiclient\": [\"MultiApiClientMixin\"], \"msrest\": [\"Deserializer\", \"Serializer\"], \"azure.mgmt.core\": [\"ARMPipelineClient\"]}, \"local\": {\"._configuration\": [\"AzureDatabricksManagementClientConfiguration\"]}}, \"conditional\": {\"stdlib\": {\"typing\": [\"Any\", \"Optional\"]}, \"azurecore\": {\"azure.core.pipeline.transport\": [\"HttpRequest\", \"HttpResponse\"]}}}",
+        "async_imports": "{\"typing\": {\"azurecore\": {\"azure.core.credentials_async\": [\"AsyncTokenCredential\"]}}, \"regular\": {\"azurecore\": {\"azure.profiles\": [\"KnownProfiles\", \"ProfileDefinition\"], \"azure.profiles.multiapiclient\": [\"MultiApiClientMixin\"], \"msrest\": [\"Deserializer\", \"Serializer\"], \"azure.mgmt.core\": [\"AsyncARMPipelineClient\"]}, \"local\": {\"._configuration\": [\"AzureDatabricksManagementClientConfiguration\"]}}, \"conditional\": {\"stdlib\": {\"typing\": [\"Any\", \"Optional\"]}, \"azurecore\": {\"azure.core.pipeline.transport\": [\"AsyncHttpResponse\", \"HttpRequest\"]}}}"
+    },
+    "global_parameters": {
+        "sync": {
+            "credential": {
+                "signature": "credential,  # type: \"TokenCredential\"",
+                "description": "Credential needed for the client to connect to Azure.",
+                "docstring_type": "~azure.core.credentials.TokenCredential",
+                "required": true
+            },
+            "subscription_id": {
+                "signature": "subscription_id,  # type: str",
+                "description": "The ID of the target subscription.",
+                "docstring_type": "str",
+                "required": true
+            }
+        },
+        "async": {
+            "credential": {
+                "signature": "credential: \"AsyncTokenCredential\",",
+                "description": "Credential needed for the client to connect to Azure.",
+                "docstring_type": "~azure.core.credentials_async.AsyncTokenCredential",
+                "required": true
+            },
+            "subscription_id": {
+                "signature": "subscription_id: str,",
+                "description": "The ID of the target subscription.",
+                "docstring_type": "str",
+                "required": true
+            }
+        },
+        "constant": {
+        },
+        "call": "credential, subscription_id",
+        "service_client_specific": {
+            "sync": {
+                "api_version": {
+                    "signature": "api_version=None, # type: Optional[str]",
+                    "description": "API version to use if no profile is provided, or if missing in profile.",
+                    "docstring_type": "str",
+                    "required": false
+                },
+                "base_url": {
+                    "signature": "base_url=None,  # type: Optional[str]",
+                    "description": "Service URL",
+                    "docstring_type": "str",
+                    "required": false
+                },
+                "profile": {
+                    "signature": "profile=KnownProfiles.default, # type: KnownProfiles",
+                    "description": "A profile definition, from KnownProfiles to dict.",
+                    "docstring_type": "azure.profiles.KnownProfiles",
+                    "required": false
+                }
+            },
+            "async": {
+                "api_version": {
+                    "signature": "api_version: Optional[str] = None,",
+                    "description": "API version to use if no profile is provided, or if missing in profile.",
+                    "docstring_type": "str",
+                    "required": false
+                },
+                "base_url": {
+                    "signature": "base_url: Optional[str] = None,",
+                    "description": "Service URL",
+                    "docstring_type": "str",
+                    "required": false
+                },
+                "profile": {
+                    "signature": "profile: KnownProfiles = KnownProfiles.default,",
+                    "description": "A profile definition, from KnownProfiles to dict.",
+                    "docstring_type": "azure.profiles.KnownProfiles",
+                    "required": false
+                }
+            }
+        }
+    },
+    "config": {
+        "credential": true,
+        "credential_scopes": ["https://management.azure.com/.default"],
+        "credential_default_policy_type": "BearerTokenCredentialPolicy",
+        "credential_default_policy_type_has_async_version": true,
+        "credential_key_header_name": null,
+        "sync_imports": "{\"regular\": {\"azurecore\": {\"azure.core.configuration\": [\"Configuration\"], \"azure.core.pipeline\": [\"policies\"], \"azure.mgmt.core.policies\": [\"ARMHttpLoggingPolicy\"]}, \"local\": {\"._version\": [\"VERSION\"]}}, \"conditional\": {\"stdlib\": {\"typing\": [\"Any\"]}}, \"typing\": {\"azurecore\": {\"azure.core.credentials\": [\"TokenCredential\"]}}}",
+        "async_imports": "{\"regular\": {\"azurecore\": {\"azure.core.configuration\": [\"Configuration\"], \"azure.core.pipeline\": [\"policies\"], \"azure.mgmt.core.policies\": [\"ARMHttpLoggingPolicy\"]}, \"local\": {\".._version\": [\"VERSION\"]}}, \"conditional\": {\"stdlib\": {\"typing\": [\"Any\"]}}, \"typing\": {\"azurecore\": {\"azure.core.credentials_async\": [\"AsyncTokenCredential\"]}}}"
+    },
+    "operation_groups": {
+        "workspaces": "WorkspacesOperations",
+        "operations": "Operations",
+        "private_link_resources": "PrivateLinkResourcesOperations",
+        "private_endpoint_connections": "PrivateEndpointConnectionsOperations",
+        "vnet_peering": "VNetPeeringOperations"
+    }
+}

src/databricks/azext_databricks/vendored_sdks/databricks/_version.py

@@ -6,4 +6,4 @@
 # Changes may cause incorrect behavior and will be lost if the code is regenerated.
 # --------------------------------------------------------------------------
 
-VERSION = "1.0.0"
+VERSION = "1.1.0b1"

src/databricks/azext_databricks/vendored_sdks/databricks/aio/__init__.py

@@ -6,5 +6,5 @@
 # Changes may cause incorrect behavior and will be lost if the code is regenerated.
 # --------------------------------------------------------------------------
 
-from ._databricks_client import DatabricksClient
-__all__ = ['DatabricksClient']
+from ._azure_databricks_management_client import AzureDatabricksManagementClient
+__all__ = ['AzureDatabricksManagementClient']