Add KMS Key for S3

src/nebari_plugin_mlflow_aws/__about__.py

@@ -1 +1 @@
-__version__ = "0.0.10"
+__version__ = "0.0.11"

src/nebari_plugin_mlflow_aws/__init__.py

@@ -13,12 +13,10 @@
 
 CLIENT_NAME = "mlflow"
 
-# TODO this only works for AWS.  How to check
-
-
 class MlflowConfig(Base):
     name: Optional[str] = "mlflow"
     namespace: Optional[str] = None
+    enable_s3_encryption: Optional[bool] = True
     values: Optional[Dict[str, Any]] = {}
 
 
@@ -171,6 +169,7 @@ def input_vars(self, stage_outputs: Dict[str, Dict[str, Any]]):
                 "namespace": self.config.namespace,
             },
             "create_namespace": create_ns,
+            "enable_s3_encryption": self.config.mlflow.enable_s3_encryption,
             "namespace": chart_ns,
             "ingress_host": domain,
             "cluster_oidc_issuer_url": cluster_oidc_issuer_url,

src/nebari_plugin_mlflow_aws/terraform/main.tf

@@ -21,6 +21,24 @@ resource "aws_s3_bucket" "artifact_storage" {
   }
 }
 
+# If enable_s3_encryption is true, create a key and apply Server Side Encryption to S3 bucket
+resource "aws_kms_key" "mlflow_kms_key" {
+  count                   = var.enable_s3_encryption ? 1 : 0
+  description             = "This key is used to encrypt bucket objects for the AWS MLflow extension"
+}
+
+resource "aws_s3_bucket_server_side_encryption_configuration" "mlflow_s3_encryption" {
+  count  = var.enable_s3_encryption ? 1 : 0
+  bucket = aws_s3_bucket.artifact_storage.id
+
+  rule {
+    apply_server_side_encryption_by_default {
+      kms_master_key_id = aws_kms_key.mlflow_kms_key[0].arn
+      sse_algorithm     = "aws:kms"
+    }
+  }
+}
+
 # --------------------------------------------------------------------------
 # Create IAM Resources for IRSA
 # --------------------------------------------------------------------------

src/nebari_plugin_mlflow_aws/terraform/variables.tf

@@ -60,6 +60,11 @@ variable "namespace" {
   type = string
 }
 
+variable "enable_s3_encryption" {
+  type = bool
+  default = true
+}
+
 variable "overrides" {
   type    = any
   default = {}