From 47970b36a5220c917605ad1a872acf7f2596f554 Mon Sep 17 00:00:00 2001
From: Jason Ng <ng.j.jason@gmail.com>
Date: Mon, 19 Aug 2024 16:51:12 -0700
Subject: [PATCH] fix: Allow dbx storage credentials roles to be self-assuming
 (#620)

---
 databricks-catalog-external-location/main.tf | 26 ++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/databricks-catalog-external-location/main.tf b/databricks-catalog-external-location/main.tf
index d6c33e24..de9d0232 100644
--- a/databricks-catalog-external-location/main.tf
+++ b/databricks-catalog-external-location/main.tf
@@ -71,6 +71,21 @@ data "aws_iam_policy_document" "databricks_external_location_assume_role" {
       values = [var.databricks_external_id]
     }
   }
+
+  statement {
+    principals {
+      type        = "AWS"
+      identifiers = ["arn:aws:iam::${data.aws_caller_identity.current.account_id}:root"]
+    }
+
+    actions = ["sts:AssumeRole"]
+    condition {
+      test     = "ArnEquals"
+      variable = "aws:PrincipalArn"
+
+      values = ["arn:aws:iam::${data.aws_caller_identity.current.account_id}:role${local.path}${local.iam_role_name}"]
+    }
+  }
 }
 
 resource "aws_iam_role" "databricks_external_location_iam_role" {
@@ -98,6 +113,17 @@ data "aws_iam_policy_document" "databricks_external_location_bucket_access" {
       module.catalog_bucket.arn,
     ]
   }
+
+  statement {
+    sid    = "databricksAssumeRole"
+    effect = "Allow"
+    actions = [
+      "sts:AssumeRole"
+    ]
+    resources = [
+      "arn:aws:iam::${data.aws_caller_identity.current.account_id}:role${local.path}${local.iam_role_name}"
+    ]
+  }
 }
 
 ### Policies to access bucket