fix: Azure storage and Vault simplified

README.md

@@ -78,15 +78,13 @@ resources that lack official modules.
 | <a name="input_cluster_sku_tier"></a> [cluster\_sku\_tier](#input\_cluster\_sku\_tier) | The Azure AKS SKU Tier to use for this cluster (https://learn.microsoft.com/en-us/azure/aks/free-standard-pricing-tiers) | `string` | `"Free"` | no |
 | <a name="input_create_private_link"></a> [create\_private\_link](#input\_create\_private\_link) | Use for the azure private link. | `bool` | `false` | no |
 | <a name="input_create_redis"></a> [create\_redis](#input\_create\_redis) | Boolean indicating whether to provision an redis instance (true) or not (false). | `bool` | `false` | no |
-| <a name="input_customer_storage_vault_key_id"></a> [customer\_storage\_vault\_key\_id](#input\_customer\_storage\_vault\_key\_id) | The Azure Key Vault key ID for customer-provided storage encryption keys. Must match the pattern 'https://<vault-name>.vault.azure.net/keys/<key-name>/<key-version>', or be null. | `string` | `null` | no |
 | <a name="input_database_availability_mode"></a> [database\_availability\_mode](#input\_database\_availability\_mode) | n/a | `string` | `"SameZone"` | no |
 | <a name="input_database_sku_name"></a> [database\_sku\_name](#input\_database\_sku\_name) | Specifies the SKU Name for this MySQL Server | `string` | `"GP_Standard_D4ds_v4"` | no |
 | <a name="input_database_version"></a> [database\_version](#input\_database\_version) | Version for MySQL | `string` | `"5.7"` | no |
 | <a name="input_deletion_protection"></a> [deletion\_protection](#input\_deletion\_protection) | If the instance should have deletion protection enabled. The database / Bucket can't be deleted when this value is set to `true`. | `bool` | `true` | no |
 | <a name="input_disable_storage_vault_key_id"></a> [disable\_storage\_vault\_key\_id](#input\_disable\_storage\_vault\_key\_id) | Flag to disable the `customer_managed_key` block, the properties 'encryption.identity, encryption.keyvaultproperties' cannot be updated in a single operation. | `bool` | `false` | no |
 | <a name="input_domain_name"></a> [domain\_name](#input\_domain\_name) | Domain for accessing the Weights & Biases UI. | `string` | `null` | no |
 | <a name="input_enable_database_vault_key"></a> [enable\_database\_vault\_key](#input\_enable\_database\_vault\_key) | Flag to enable managed key encryption for the database. Once enabled, cannot be disabled. | `bool` | `false` | no |
-| <a name="input_enable_purge_protection"></a> [enable\_purge\_protection](#input\_enable\_purge\_protection) | Flag to enable purge protection for the Azure Key Vault. Once enabled, cannot be disabled. | `bool` | `false` | no |
 | <a name="input_enable_storage_vault_key"></a> [enable\_storage\_vault\_key](#input\_enable\_storage\_vault\_key) | Flag to enable managed key encryption for the storage account. | `bool` | `false` | no |
 | <a name="input_external_bucket"></a> [external\_bucket](#input\_external\_bucket) | config an external bucket | `any` | `null` | no |
 | <a name="input_kubernetes_instance_type"></a> [kubernetes\_instance\_type](#input\_kubernetes\_instance\_type) | Use for the Kubernetes cluster. | `string` | `"Standard_D4a_v4"` | no |
@@ -132,8 +130,6 @@ resources that lack official modules.
 | <a name="output_private_link_resource_id"></a> [private\_link\_resource\_id](#output\_private\_link\_resource\_id) | n/a |
 | <a name="output_private_link_sub_resource_name"></a> [private\_link\_sub\_resource\_name](#output\_private\_link\_sub\_resource\_name) | n/a |
 | <a name="output_standardized_size"></a> [standardized\_size](#output\_standardized\_size) | n/a |
-| <a name="output_storage_account"></a> [storage\_account](#output\_storage\_account) | n/a |
-| <a name="output_storage_container"></a> [storage\_container](#output\_storage\_container) | n/a |
 | <a name="output_tenant_id"></a> [tenant\_id](#output\_tenant\_id) | n/a |
 | <a name="output_url"></a> [url](#output\_url) | The URL to the W&B application |
 <!-- END_TF_DOCS -->

examples/byob/README.md

@@ -38,18 +38,16 @@ To configure Customer Managed Key encryption, ensure you are using the latest ve
 
 - create_cmk
 - disable_storage_vault_key_id
-- enable_purge_protection
 - tenant_id
 - client_id
 
-You need to obtain the `tenant_id` and `client_id` from a Solutions Architect at W&B for an already instantiated instance of a Weights & Biases managed deployment.
+You need to obtain the `tenant_id` and `client_id` from the `https://${WANDB_BASE_URL}/console/settings/advanced/spec/active` at W&B for an already instantiated instance of a Weights & Biases managed deployment.
 
 Set the follwoing new variabels to enable the CMK:
 
 ```ini terraform.tfvars
 create_cmk = true
 
-enable_purge_protection       = true
 disable_storage_vault_key_id  = false
 
 tenant_id = "<tenant_id>"
@@ -63,7 +61,7 @@ terraform init -upgrade
 terraform apply -var-file=terraform.tfvars
 ```
 
-Upon successful execution, you will receive the following output which needs to the shared with the SA at Weights & Biases.
+Upon successful execution, you will receive the following output which needs to be set in the system connection settings `https://${WANDB_BASE_URL}/console/settings/system`
 
 ```bash
 blob_container = "<storage_account_name>/wandb"
@@ -72,4 +70,4 @@ storage_key = <sensitive>
 storage_vault_key_id = "https://<key_vault_name>.vault.azure.net/keys/<key_name>/<key_version>"
 ```
 
-Retrieve the storage key as shown above and also share that with the SA at Weights & Biases.
+Retrieve the storage key as shown above.

examples/byob/main.tf

@@ -9,7 +9,6 @@ module "byob" {
   prefix                       = var.prefix
   deletion_protection          = var.deletion_protection
   create_cmk                   = var.create_cmk
-  enable_purge_protection      = var.enable_purge_protection
   client_id                    = var.client_id
   tenant_id                    = var.tenant_id
   tags                         = var.tags

examples/byob/variables.tf

@@ -31,20 +31,14 @@ variable "disable_storage_vault_key_id" {
   description = "Flag to disable the `customer_managed_key` block, the properties 'encryption.identity, encryption.keyvaultproperties' cannot be updated in a single operation."
 }
 
-variable "enable_purge_protection" {
-  type        = bool
-  default     = false
-  description = "Flag to enable purge protection for the Azure Key Vault. Once enabled, cannot be disabled."
-}
-
 variable "tenant_id" {
   type        = string
-  description = "The tenant ID for the Key Vault Access Policy. Get from W&B SA"
+  description = "The tenant ID for the Key Vault Access Policy. Get from `https://<WANDB_BASE_URL>/console/settings/advanced/spec/active`"
 }
 
 variable "client_id" {
   type        = string
-  description = "The client ID (object id) for the Key Vault Access Policy. Get from W&B SA"
+  description = "The client ID (object id) for the Key Vault Access Policy. Get from `https://<WANDB_BASE_URL>/console/settings/advanced/spec/active`"
 }
 
 variable "tags" {

main.tf

@@ -2,8 +2,6 @@ locals {
   fqdn       = var.subdomain == null ? var.domain_name : "${var.subdomain}.${var.domain_name}"
   url_prefix = var.ssl ? "https" : "http"
   url        = "${local.url_prefix}://${local.fqdn}"
-
-  enable_internal_storage = var.blob_container == "" && var.external_bucket == null
 }
 
 resource "azurerm_resource_group" "default" {
@@ -67,30 +65,29 @@ module "redis" {
 module "vault" {
   source = "./modules/vault"
 
-  identity_object_id      = module.identity.identity.principal_id
-  location                = azurerm_resource_group.default.location
-  namespace               = var.namespace
-  resource_group          = azurerm_resource_group.default
-  enable_purge_protection = var.enable_purge_protection
+  identity_object_id = module.identity.identity.principal_id
+  location           = azurerm_resource_group.default.location
+  namespace          = var.namespace
+  resource_group     = azurerm_resource_group.default
 
   enable_database_vault_key = var.enable_database_vault_key
-  enable_storage_vault_key  = var.enable_storage_vault_key && local.enable_internal_storage
+  enable_storage_vault_key  = var.enable_storage_vault_key
 
   tags = var.tags
 }
 
 module "storage" {
   source = "./modules/storage"
 
-  count               = local.enable_internal_storage ? 1 : 0
+  count               = 1
   namespace           = var.namespace
   resource_group_name = azurerm_resource_group.default.name
   location            = azurerm_resource_group.default.location
   create_queue        = !var.use_internal_queue
   identity_ids        = module.identity.identity.id
   deletion_protection = var.deletion_protection
 
-  storage_key_id               = var.customer_storage_vault_key_id == null ? try(module.vault.vault_internal_keys[module.vault.vault_key_map.storage].id, null) : var.customer_storage_vault_key_id
+  storage_key_id               = try(module.vault.vault_internal_keys[module.vault.vault_key_map.storage].id, null)
   disable_storage_vault_key_id = var.disable_storage_vault_key_id
 
   tags = var.tags
@@ -129,21 +126,6 @@ module "app_aks" {
   max_pods              = var.node_max_pods
   tags                  = var.tags
 }
-
-locals {
-  container_name  = try(module.storage[0].container.name, "")
-  account_name    = try(module.storage[0].account.name, "")
-  access_key      = try(module.storage[0].account.primary_access_key, "")
-  queue_name      = try(module.storage[0].queue.name, "")
-  blob_container  = var.external_bucket == null ? coalesce(var.blob_container, local.container_name) : ""
-  storage_account = var.external_bucket == null ? coalesce(var.storage_account, local.account_name) : ""
-  storage_key     = var.external_bucket == null ? coalesce(var.storage_key, local.access_key) : ""
-  bucket          = "az://${local.storage_account}/${local.blob_container}"
-  queue           = (var.use_internal_queue || var.blob_container == "" || var.external_bucket == null) ? "internal://" : "az://${local.account_name}/${local.queue_name}"
-
-  redis_connection_string = "redis://:${module.redis.instance.primary_access_key}@${module.redis.instance.hostname}:${module.redis.instance.port}"
-}
-
 locals {
   service_account_name         = "wandb-app"
   private_endpoint_approval_sa = "private-endpoint-sa"
@@ -231,6 +213,21 @@ module "cert_manager" {
   depends_on = [module.app_aks]
 }
 
+locals {
+  use_customer_bucket = (
+    var.storage_account != "" &&
+    var.blob_container != "" &&
+    var.storage_key != ""
+  )
+  default_bucket_config = {
+    provider  = "az"
+    name      = var.storage_account
+    path      = var.blob_container
+    accessKey = var.storage_key
+  }
+  bucket_config = var.external_bucket != null ? var.external_bucket : (local.use_customer_bucket ? local.default_bucket_config : null)
+}
+
 module "wandb" {
   source  = "wandb/wandb/helm"
   version = "1.2.0"
@@ -250,12 +247,22 @@ module "wandb" {
         host          = local.url
         license       = var.license
         cloudProvider = "azure"
-        bucket = var.external_bucket == null ? {
+        bucket        = local.bucket_config == null ?  {
           provider  = "az"
-          name      = local.storage_account
-          path      = local.blob_container
-          accessKey = local.storage_key
-        } : var.external_bucket
+          name      = module.storage[0].account.name
+          path      = module.storage[0].container.name
+          accessKey = module.storage[0].account.primary_access_key
+        } : local.bucket_config
+        defaultBucket = {
+          provider  = "az"
+          name      = module.storage[0].account.name
+          path      = module.storage[0].container.name
+          accessKey = module.storage[0].account.primary_access_key
+        }
+        azureIdentityForBucket = {
+          clientID = module.identity.identity.client_id
+          tenantID = module.identity.identity.tenant_id
+        }
 
         mysql = {
           host     = module.database.address

modules/byob/main.tf

@@ -17,7 +17,6 @@ module "vault" {
   identity_object_id       = module.identity[0].identity.principal_id
   depends_on               = [module.identity]
   tags                     = var.tags
-  enable_purge_protection  = var.enable_purge_protection
   enable_storage_vault_key = var.create_cmk
 }
 

modules/byob/variables.tf

@@ -30,12 +30,6 @@ variable "tags" {
   default     = {}
 }
 
-variable "enable_purge_protection" {
-  type        = bool
-  default     = false
-  description = "Flag to enable purge protection for the Azure Key Vault. Once enabled, cannot be disabled."
-}
-
 variable "disable_storage_vault_key_id" {
   type        = bool
   default     = false
@@ -44,10 +38,10 @@ variable "disable_storage_vault_key_id" {
 
 variable "tenant_id" {
   type        = string
-  description = "The tenant ID for the Key Vault Access Policy. Get from W&B SA"
+  description = "The tenant ID for the Key Vault Access Policy. Get from `https://<WANDB_BASE_URL>/console/settings/advanced/spec/active`"
 }
 
 variable "client_id" {
   type        = string
-  description = "The client ID (object id) for the Key Vault Access Policy. Get from W&B SA"
+  description = "The client ID (object id) for the Key Vault Access Policy. Get from `https://<WANDB_BASE_URL>/console/settings/advanced/spec/active`"
 }

modules/vault/main.tf

@@ -16,7 +16,7 @@ resource "azurerm_key_vault" "default" {
   location                 = var.location
   resource_group_name      = var.resource_group.name
   tenant_id                = data.azurerm_client_config.current.tenant_id
-  purge_protection_enabled = var.enable_purge_protection
+  purge_protection_enabled = true
   # https://learn.microsoft.com/en-us/azure/mysql/flexible-server/concepts-customer-managed-key#requirements-for-configuring-data-encryption-for-azure-database-for-mysql-flexible-server
   soft_delete_retention_days  = 90 # This must be 90 for azure msyql flex server encryption. 
   enabled_for_disk_encryption = true

modules/vault/variables.tf

@@ -33,9 +33,3 @@ variable "enable_database_vault_key" {
   default     = false
   description = "Flag to enable managed key encryption for the database. Once enabled, cannot be disabled or you will loose access to the database."
 }
-
-variable "enable_purge_protection" {
-  type        = bool
-  default     = false
-  description = "Flag to enable purge protection for the Azure Key Vault. Once enabled, cannot be disabled."
-}

outputs.tf

@@ -33,14 +33,6 @@ output "oidc_issuer_url" {
   value = module.app_aks.oidc_issuer_url
 }
 
-output "storage_account" {
-  value = var.external_bucket != null ? "" : coalesce(var.storage_account, local.account_name, "")
-}
-
-output "storage_container" {
-  value = var.external_bucket != null ? "" : coalesce(var.blob_container, local.container_name)
-}
-
 output "private_link_resource_id" {
   value = module.app_lb.gateway_id
 }

variables.tf

@@ -271,29 +271,12 @@ variable "disable_storage_vault_key_id" {
   description = "Flag to disable the `customer_managed_key` block, the properties 'encryption.identity, encryption.keyvaultproperties' cannot be updated in a single operation."
 }
 
-variable "customer_storage_vault_key_id" {
-  type        = string
-  default     = null
-  description = "The Azure Key Vault key ID for customer-provided storage encryption keys. Must match the pattern 'https://<vault-name>.vault.azure.net/keys/<key-name>/<key-version>', or be null."
-
-  validation {
-    condition     = var.customer_storage_vault_key_id == null || can(regex("^https://[a-zA-Z0-9-]+.vault.azure.net/keys/[a-zA-Z0-9-]+/[a-f0-9]+$", var.customer_storage_vault_key_id))
-    error_message = "The customer_storage_vault_key_id must be null or a valid Azure Key Vault key ID in the format 'https://<vault-name>.vault.azure.net/keys/<key-name>/<key-version>'."
-  }
-}
-
 variable "enable_database_vault_key" {
   type        = bool
   default     = false
   description = "Flag to enable managed key encryption for the database. Once enabled, cannot be disabled."
 }
 
-variable "enable_purge_protection" {
-  type        = bool
-  default     = false
-  description = "Flag to enable purge protection for the Azure Key Vault. Once enabled, cannot be disabled."
-}
-
 ## To support otel azure monitor sql and redis metrics need operator-wandb chart minimum version 0.14.0 
 variable "azuremonitor" {
   type    = bool