From 4c1bbd0cb7f4c7bbf7639250419bc9fd22adaa28 Mon Sep 17 00:00:00 2001
From: James Stott <158563996+jamesstottmoj@users.noreply.github.com>
Date: Thu, 5 Sep 2024 12:46:15 +0100
Subject: [PATCH] Fix/database opt in (#1333)

* Fixed potential issue with granting opt in on db with multiple tables.

* Fixed failing tests
---
 controlpanel/api/aws.py                | 43 +++++++++++++++++++++-----
 tests/frontend/views/test_databases.py | 12 +++++--
 2 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/controlpanel/api/aws.py b/controlpanel/api/aws.py
index 3d8158cd5..f7b62324c 100644
--- a/controlpanel/api/aws.py
+++ b/controlpanel/api/aws.py
@@ -1250,10 +1250,13 @@ def create_lf_opt_in(self, database_name, table_name, principal_arn, catalog_id=
             },
         }
 
-        self.client.create_lake_formation_opt_in(
-            Resource=database_resource,
-            Principal={"DataLakePrincipalIdentifier": principal_arn},
-        )
+        opt_ins = self.list_lf_opt_ins(database_name, principal_arn, catalog_id)
+
+        if opt_ins["database"] is None:
+            self.client.create_lake_formation_opt_in(
+                Resource=database_resource,
+                Principal={"DataLakePrincipalIdentifier": principal_arn},
+            )
 
         self.client.create_lake_formation_opt_in(
             Resource=table_resource,
@@ -1277,16 +1280,40 @@ def delete_lf_opt_in(self, database_name, table_name, principal_arn, catalog_id=
             },
         }
 
-        self.client.delete_lake_formation_opt_in(
-            Resource=database_resource,
-            Principal={"DataLakePrincipalIdentifier": principal_arn},
-        )
+        opt_ins = self.list_lf_opt_ins(database_name, principal_arn, catalog_id)
+
+        if len(opt_ins["tables"]) == 1:
+            self.client.delete_lake_formation_opt_in(
+                Resource=database_resource,
+                Principal={"DataLakePrincipalIdentifier": principal_arn},
+            )
 
         self.client.delete_lake_formation_opt_in(
             Resource=table_resource,
             Principal={"DataLakePrincipalIdentifier": principal_arn},
         )
 
+    def list_lf_opt_ins(self, database_name, principal_arn, catalog_id=None):
+        catalog_id = catalog_id or settings.AWS_DATA_ACCOUNT_ID
+        result = {
+            "database": None,
+            "tables": [],
+        }
+
+        response = self.client.list_lake_formation_opt_ins(
+            Principal={"DataLakePrincipalIdentifier": principal_arn}
+        )["LakeFormationOptInsInfoList"]
+
+        for resource_object in response:
+            resource = resource_object["Resource"]
+            if "Database" in resource and resource["Database"]["Name"] == database_name:
+                result["database"] = resource
+            elif "Table" in resource:
+                if resource["Table"]["DatabaseName"] == database_name:
+                    result["tables"].append(resource)
+
+        return result
+
 
 class AWSGlue(AWSService):
 
diff --git a/tests/frontend/views/test_databases.py b/tests/frontend/views/test_databases.py
index 8c166d19c..ce49c3adf 100644
--- a/tests/frontend/views/test_databases.py
+++ b/tests/frontend/views/test_databases.py
@@ -14,8 +14,16 @@
 
 # Mocked botocore _make_api_call function
 def mock_make_api_call(self, operation_name, kwarg):
-    if operation_name == "CreateLakeFormationOptIn":
-        return {}
+    op_names = [
+        {"CreateLakeFormationOptIn": {}},
+        {"DeleteLakeFormationOptIn": {}},
+        {"ListLakeFormationOptIns": {"LakeFormationOptInsInfoList": []}},
+    ]
+
+    for operation in op_names:
+        if operation_name in operation:
+            return operation[operation_name]
+
     # If we don't want to patch the API call
     return orig(self, operation_name, kwarg)