fix: Allow the admin arn to be set for the KMS key

main.tf

@@ -6,6 +6,8 @@ module "kms" {
   key_alias  = var.kms_key_alias == null ? "${var.namespace}-kms-alias" : var.kms_key_alias
   key_policy = var.kms_key_policy
 
+  policy_administrator_arn = var.kms_key_policy_administrator_arn
+
   create_clickhouse_key = var.enable_clickhouse
   clickhouse_key_alias  = var.kms_clickhouse_key_alias == null ? "${var.namespace}-kms-clickhouse-alias" : var.kms_clickhouse_key_alias
   clickhouse_key_policy = var.kms_clickhouse_key_policy

modules/kms/main.tf

@@ -1,5 +1,9 @@
 data "aws_caller_identity" "current" {}
 
+locals {
+  policy_administrator_arn = var.policy_administrator_arn != "" ? var.policy_administrator_arn : data.aws_caller_identity.current.arn
+}
+
 resource "aws_kms_key" "key" {
   deletion_window_in_days = var.key_deletion_window
   description             = "AWS KMS Customer-managed key to encrypt Weights & Biases resources"
@@ -11,7 +15,7 @@ resource "aws_kms_key" "key" {
       {
         "Sid" : "Allow administration of the key",
         "Effect" : "Allow",
-        "Principal" : { "AWS" : "${data.aws_caller_identity.current.arn}" },
+        "Principal" : { "AWS" : "${local.policy_administrator_arn}" },
         "Action" : "kms:*",
         "Resource" : "*"
       },

modules/kms/variables.tf

@@ -14,6 +14,12 @@ variable "iam_principal_arn" {
   default     = ""
 }
 
+variable "policy_administrator_arn" {
+  description = "The IAM principal (role or user) ARN that will be authorized to manage the key."
+  type        = string
+  default     = ""
+}
+
 variable "key_policy" {
   description = "The policy that will define the permissions for the kms key."
   type        = string

variables.tf

@@ -199,6 +199,12 @@ variable "kms_key_policy" {
   default     = ""
 }
 
+variable "kms_key_policy_administrator_arn" {
+  type        = string
+  description = "The principal that will be allowed to manage the kms key."
+  default     = ""
+}
+
 variable "kms_clickhouse_key_alias" {
   type        = string
   description = "KMS key alias for AWS KMS Customer managed key used by Clickhouse CMEK."