Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 2 commits into from
Jan 24, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/databricks/HISTORY.rst
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,10 @@

Release History
===============
0.8.0
+++++
* az databricks workspace create: Add --public-network-access to allow creating workspace with network access from public internet
* az databricks workspace create: Add --required-nsg-rules to allow creating workspace with nsg rule for internal

0.7.3
+++++
Expand Down
4 changes: 2 additions & 2 deletions src/databricks/azext_databricks/_client_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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, *_):
Expand Down
4 changes: 4 additions & 0 deletions src/databricks/azext_databricks/_params.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@


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.')
Expand All @@ -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.')
Expand Down
4 changes: 4 additions & 0 deletions src/databricks/azext_databricks/custom.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1,429 changes: 389 additions & 1,040 deletions src/databricks/azext_databricks/tests/latest/recordings/test_databricks.yaml

Large diffs are not rendered by default.

Large diffs are not rendered by default.

Large diffs are not rendered by default.

Original file line number Diff line number Diff line change
Expand Up @@ -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):

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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)}
Expand All @@ -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

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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)
Expand Down
Original file line number Diff line number Diff line change
@@ -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"
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Original file line number Diff line number Diff line change
Expand Up @@ -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']
Loading