Network Security for Azure Databricks in a Virtual Network

Question

Pulumi · Accepted Answer

When deploying Azure Databricks in a more secure and isolated environment, it's common to set it up within an Azure Virtual Network (VNet). This allows you to control the inbound and outbound network traffic to and from Azure Databricks and other services with which it interacts.

To secure Azure Databricks within a VNet, you often have to:

1. Create a VNet or use an existing one.
2. Create a Databricks workspace configured to deploy in your VNet.
3. Set up network security groups (NSGs) and rules to control traffic.
4. Optionally, set up private link connections if you need to access Azure Databricks over a private endpoint.

Here's a Pulumi program in Python that demonstrates how to set up a secure Databricks workspace within a Virtual Network in Azure:

```python
import pulumi
import pulumi_azure_native as azure_native

# Create an Azure Resource Group
resource_group = azure_native.resources.ResourceGroup("resource_group")

# Create an Azure Virtual Network
vnet = azure_native.network.VirtualNetwork(
    "virtual_network",
    resource_group_name=resource_group.name,
    location=resource_group.location,
    address_space=azure_native.network.AddressSpaceArgs(
        address_prefixes=["10.0.0.0/16"],
    ),
    subnets=[azure_native.network.SubnetArgs(
        name="default",
        address_prefix="10.0.0.0/24",
    )]
)

# Create an Azure Databricks workspace in the Virtual Network
databricks_workspace = azure_native.databricks.Workspace(
    "databricks_workspace",
    resource_group_name=resource_group.name,
    location=resource_group.location,
    managed_resource_group_id=resource_group.id,
    workspace_name="my-databricks-workspace",
    sku=azure_native.databricks.SkuArgs(
        name="standard"  # Choose between 'standard', 'premium', and 'trial'
    ),
    parameters=azure_native.databricks.WorkspaceCustomParametersArgs(
        custom_virtual_network_id=azure_native.databricks.WorkspaceCustomStringParameterArgs(
            value=vnet.id,
        ),
        custom_public_subnet_name=azure_native.databricks.WorkspaceCustomStringParameterArgs(
            value="default",
        ),
        custom_private_subnet_name=azure_native.databricks.WorkspaceCustomStringParameterArgs(
            value="default",
        ),
        enable_no_public_ip=azure_native.databricks.WorkspaceCustomBooleanParameterArgs(
            value=True,  # Restricts public IPs for a more secure environment
        ),
    )
)

# Create a Network Security Group (NSG) to secure the Virtual Network
nsg = azure_native.network.NetworkSecurityGroup(
    "network_security_group",
    resource_group_name=resource_group.name,
    location=resource_group.location
)

# Associate the NSG to the Databricks Subnet within the VNet
subnet_association = azure_native.network.Subnet(
    "databricks_subnet_association",
    subnet_name="default",
    virtual_network_name=vnet.name,
    resource_group_name=resource_group.name,
    address_prefix="10.0.1.0/24",
    network_security_group=nsg
)

# Export the URLs needed to access resources
pulumi.export("resource_group_name", resource_group.name)
pulumi.export("databricks_workspace_url", databricks_workspace.workspace_url)
pulumi.export("vnet_name", vnet.name)
```

Let's break down the code:

- We start by creating a Resource Group, which is a logical container for grouping related Azure resources.
- We then create a Virtual Network (VNet) with a subnet. This defines a private network space where your resources live.
- Next, we set up the Azure Databricks workspace within this VNet. The `parameters` for the workspace specify the VNet's details and enable the "No Public IP" feature for added security.
- We then create a Network Security Group (NSG) that will be used to define security rules for the subnet where the Databricks workspace is located.
- We associate the Network Security Group with the subnet by creating a new `Subnet` resource with the `network_security_group` property set to the NSG we created.
- Lastly, we export some of the outputs that might be useful, like the names of the resource group, Databricks workspace, and VNet.

By completing the above steps, we've established a secure network environment for Azure Databricks that can be controlled and limited as per the organization's security policies. Remember to add NSG rules according to your security requirements to control the access to the Databricks workspace.