refactor(feature-store): Replace input() with acknowledge_risk param

Add acknowledge_risk: Optional[bool] = None to enable_lake_formation()
and LakeFormationConfig. None triggers interactive input() prompt, True
proceeds without prompting, False aborts with RuntimeError.

Removes all builtins.input mocking from unit and integration tests.
Tests now pass acknowledge_risk=True or False directly. Removes one
duplicate test that became identical after the refactor.
---
X-AI-Prompt: add y/n confirmation for disable_hybrid_access_mode=True, then refactor to use acknowledge_risk param instead of input()
X-AI-Tool: kiro-cli

Author: BassemHalim

sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_group_manager.py

@@ -45,12 +45,17 @@ class LakeFormationConfig(Base):
             may break existing jobs that access the table via IAM-based permissions. After
             this change, all principals must be granted access through Lake Formation.
             If False, IAM-based access remains alongside Lake Formation permissions.
+        acknowledge_risk: Controls confirmation behavior for risky Lake Formation operations.
+            If True, skips interactive confirmation prompts and proceeds. If False, raises
+            RuntimeError without proceeding. If None (default), prompts the user interactively
+            via input().
     """
 
     enabled: bool = False
     use_service_linked_role: bool = True
     registration_role_arn: Optional[str] = None
     disable_hybrid_access_mode: bool
+    acknowledge_risk: Optional[bool] = None
 
 
 class FeatureGroupManager(FeatureGroup):
@@ -375,6 +380,7 @@ def _generate_s3_deny_statements(
     def enable_lake_formation(
         self,
         disable_hybrid_access_mode: bool,
+        acknowledge_risk: Optional[bool] = None,
         session: Optional[Session] = None,
         region: Optional[str] = None,
         use_service_linked_role: bool = True,
@@ -399,6 +405,9 @@ def enable_lake_formation(
                 the table via IAM-based permissions. After this change, all principals must
                 be granted access through Lake Formation. If False, prompts the user for
                 confirmation before proceeding with hybrid access.
+            acknowledge_risk: Controls confirmation behavior for risky operations.
+                If True, skips interactive prompts and proceeds. If False, raises
+                RuntimeError without proceeding. If None (default), prompts interactively.
             session: Boto3 session.
             region: Region name.
             use_service_linked_role: Whether to use the Lake Formation service-linked role
@@ -496,12 +505,15 @@ def enable_lake_formation(
                 f"to the table is still allowed alongside Lake Formation permissions. "
                 f"For more info: https://docs.aws.amazon.com/lake-formation/latest/dg/hybrid-access-mode.html"
             )
-            proceed = input(
-                "Hybrid access mode is not disabled. IAM-based access to the Glue table will "
-                "still be allowed. Do you want to proceed without revoking IAMAllowedPrincipal "
-                "permissions? (y/n): "
-            ).strip().lower()
-            if proceed != "y":
+            if acknowledge_risk is None:
+                proceed = input(
+                    "Hybrid access mode is not disabled. IAM-based access to the Glue table will "
+                    "still be allowed. Do you want to proceed without revoking IAMAllowedPrincipal "
+                    "permissions? (y/n): "
+                ).strip().lower() == "y"
+            else:
+                proceed = acknowledge_risk
+            if not proceed:
                 raise RuntimeError(
                     "User chose not to proceed without disabling hybrid access mode. "
                     "Re-run with disable_hybrid_access_mode=True to revoke IAMAllowedPrincipal permissions."
@@ -514,6 +526,18 @@ def enable_lake_formation(
                 f"After this change, all principals must be granted access through Lake Formation. "
                 f"For more info: https://docs.aws.amazon.com/lake-formation/latest/dg/hybrid-access-mode.html"
             )
+            if acknowledge_risk is None:
+                proceed = input(
+                    "This will revoke IAMAllowedPrincipal permissions and may break existing jobs "
+                    "that rely on IAM-based access. Do you want to proceed? (y/n): "
+                ).strip().lower() == "y"
+            else:
+                proceed = acknowledge_risk
+            if not proceed:
+                raise RuntimeError(
+                    "User chose not to proceed with disabling hybrid access mode. "
+                    "Re-run with disable_hybrid_access_mode=False to keep IAMAllowedPrincipal permissions."
+                )
 
 
         results = {
@@ -805,5 +829,6 @@ def create(
                 use_service_linked_role=lake_formation_config.use_service_linked_role,
                 registration_role_arn=lake_formation_config.registration_role_arn,
                 disable_hybrid_access_mode=lake_formation_config.disable_hybrid_access_mode,
+                acknowledge_risk=lake_formation_config.acknowledge_risk,
             )
         return feature_group

sagemaker-mlops/tests/integ/test_feature_store_lakeformation.py

@@ -11,7 +11,6 @@
 
 import logging
 import uuid
-
 import boto3
 import pytest
 from botocore.exceptions import ClientError
@@ -161,7 +160,7 @@ def test_create_feature_group_and_enable_lake_formation(s3_uri, role, region):
         assert fg.feature_group_status == "Created"
 
         # Enable Lake Formation governance
-        result = fg.enable_lake_formation(disable_hybrid_access_mode=True)
+        result = fg.enable_lake_formation(disable_hybrid_access_mode=True, acknowledge_risk=True)
 
         # Verify all phases completed successfully
         assert result["s3_location_registered"] is True
@@ -197,7 +196,8 @@ def test_create_feature_group_with_lake_formation_enabled(s3_uri, role, region):
         offline_store_config = OfflineStoreConfig(s3_storage_config=S3StorageConfig(s3_uri=s3_uri))
         lake_formation_config = LakeFormationConfig(
             enabled=True,
-            disable_hybrid_access_mode = True
+            disable_hybrid_access_mode = True,
+            acknowledge_risk=True,
         )
 
         fg = FeatureGroupManager.create(
@@ -460,6 +460,7 @@ def test_enable_lake_formation_fails_with_nonexistent_role(
             use_service_linked_role=False,
             registration_role_arn=nonexistent_role,
             disable_hybrid_access_mode=True,
+            acknowledge_risk=True,
         )
 
     # Verify we got an appropriate error
@@ -501,7 +502,7 @@ def test_enable_lake_formation_full_flow_with_policy_output(s3_uri, role, region
 
         # Enable Lake Formation governance
         with caplog.at_level(logging.WARNING, logger="sagemaker.mlops.feature_store.feature_group_manager"):
-            result = fg.enable_lake_formation(disable_hybrid_access_mode=True)
+            result = fg.enable_lake_formation(disable_hybrid_access_mode=True, acknowledge_risk=True)
 
         # Verify all phases completed successfully
         assert result["s3_location_registered"] is True
@@ -544,7 +545,7 @@ def test_enable_lake_formation_default_logs_recommended_policy(s3_uri, role, reg
 
         # Enable Lake Formation governance with disable_hybrid_access_mode=True
         with caplog.at_level(logging.WARNING, logger="sagemaker.mlops.feature_store.feature_group_manager"):
-            result = fg.enable_lake_formation(disable_hybrid_access_mode=True)
+            result = fg.enable_lake_formation(disable_hybrid_access_mode=True, acknowledge_risk=True)
 
         # Verify phases completed successfully
         assert result["s3_location_registered"] is True
@@ -588,6 +589,7 @@ def test_enable_lake_formation_with_custom_role_logs_policy(s3_uri, role, region
                 use_service_linked_role=False,
                 registration_role_arn=role,
                 disable_hybrid_access_mode=True,
+                acknowledge_risk=True,
             )
 
         # Verify all phases completed successfully

sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_group_manager.py

@@ -440,7 +440,7 @@ def test_wait_for_active_calls_wait_for_status(
         mock_revoke.return_value = True
 
         # Call with wait_for_active=True
-        fg.enable_lake_formation(wait_for_active=True, disable_hybrid_access_mode=True)
+        fg.enable_lake_formation(wait_for_active=True, disable_hybrid_access_mode=True, acknowledge_risk=True)
 
         # Verify wait_for_status was called with "Created"
         mock_wait.assert_called_once_with(target_status="Created")
@@ -478,7 +478,7 @@ def test_wait_for_active_false_does_not_call_wait(
         mock_revoke.return_value = True
 
         # Call with wait_for_active=False (default)
-        fg.enable_lake_formation(wait_for_active=False, disable_hybrid_access_mode=True)
+        fg.enable_lake_formation(wait_for_active=False, disable_hybrid_access_mode=True, acknowledge_risk=True)
 
         # Verify wait_for_status was NOT called
         mock_wait.assert_not_called()
@@ -564,7 +564,7 @@ def test_fail_fast_phase_execution(
         with pytest.raises(
             RuntimeError, match="Failed to register S3 location with Lake Formation"
         ):
-            fg.enable_lake_formation(disable_hybrid_access_mode=True)
+            fg.enable_lake_formation(disable_hybrid_access_mode=True, acknowledge_risk=True)
 
         # Verify Phase 1 was called but Phase 2 and 3 were not
         mock_register.assert_called_once()
@@ -583,7 +583,7 @@ def test_fail_fast_phase_execution(
         mock_revoke.return_value = True
 
         with pytest.raises(RuntimeError, match="Failed to grant Lake Formation permissions"):
-            fg.enable_lake_formation(disable_hybrid_access_mode=True)
+            fg.enable_lake_formation(disable_hybrid_access_mode=True, acknowledge_risk=True)
 
         # Verify Phase 1 and 2 were called but Phase 3 was not
         mock_register.assert_called_once()
@@ -603,7 +603,7 @@ def test_fail_fast_phase_execution(
         mock_revoke.side_effect = Exception("Phase 3 failed")
 
         with pytest.raises(RuntimeError, match="Failed to revoke IAMAllowedPrincipal permissions"):
-            fg.enable_lake_formation(disable_hybrid_access_mode=True)
+            fg.enable_lake_formation(disable_hybrid_access_mode=True, acknowledge_risk=True)
 
         # Verify all phases were called
         mock_register.assert_called_once()
@@ -967,6 +967,7 @@ def test_enable_lake_formation_called_when_enabled(
             use_service_linked_role=True,
             registration_role_arn=None,
             disable_hybrid_access_mode=False,
+            acknowledge_risk=None,
         )
         # Verify the feature group was returned
         assert result == mock_fg
@@ -1175,6 +1176,7 @@ def test_use_service_linked_role_extraction_from_config(
             use_service_linked_role=use_slr,
             registration_role_arn=expected_registration_role,
             disable_hybrid_access_mode=False,
+            acknowledge_risk=None,
         )
         # Verify the feature group was returned
         assert result == mock_fg
@@ -1213,45 +1215,34 @@ def test_revoke_called_when_disable_hybrid_access_mode_true(
         mock_grant.return_value = True
         mock_revoke.return_value = True
 
-        result = self.fg.enable_lake_formation(disable_hybrid_access_mode=True)
+        result = self.fg.enable_lake_formation(disable_hybrid_access_mode=True, acknowledge_risk=True)
 
         mock_revoke.assert_called_once()
         assert result["hybrid_access_mode_disabled"] is True
 
-    @patch("builtins.input", return_value="y")
     @patch.object(FeatureGroupManager, "refresh")
     @patch.object(FeatureGroupManager, "_register_s3_with_lake_formation")
     @patch.object(FeatureGroupManager, "_grant_lake_formation_permissions")
     @patch.object(FeatureGroupManager, "_revoke_iam_allowed_principal")
     def test_revoke_not_called_when_disable_hybrid_access_mode_false(
-        self, mock_revoke, mock_grant, mock_register, mock_refresh, mock_input
+        self, mock_revoke, mock_grant, mock_register, mock_refresh
     ):
         """Test that IAMAllowedPrincipal is NOT revoked when disable_hybrid_access_mode=False."""
         mock_register.return_value = True
         mock_grant.return_value = True
 
-        result = self.fg.enable_lake_formation(disable_hybrid_access_mode=False)
+        result = self.fg.enable_lake_formation(disable_hybrid_access_mode=False, acknowledge_risk=True)
 
         mock_revoke.assert_not_called()
         assert result["hybrid_access_mode_disabled"] is False
 
-    @patch("builtins.input", return_value="n")
     @patch.object(FeatureGroupManager, "refresh")
     def test_raises_error_when_user_declines_hybrid_access_prompt(
-        self, mock_refresh, mock_input
+        self, mock_refresh
     ):
         """Test that RuntimeError is raised when user declines the hybrid access prompt."""
         with pytest.raises(RuntimeError, match="User chose not to proceed"):
-            self.fg.enable_lake_formation(disable_hybrid_access_mode=False)
-
-    @patch("builtins.input", return_value="")
-    @patch.object(FeatureGroupManager, "refresh")
-    def test_raises_error_when_user_enters_empty_at_hybrid_access_prompt(
-        self, mock_refresh, mock_input
-    ):
-        """Test that RuntimeError is raised when user enters empty string at prompt."""
-        with pytest.raises(RuntimeError, match="User chose not to proceed"):
-            self.fg.enable_lake_formation(disable_hybrid_access_mode=False)
+            self.fg.enable_lake_formation(disable_hybrid_access_mode=False, acknowledge_risk=False)
 
 
 class TestCreateWithLakeFormationDisableHybridAccessMode:
@@ -1305,6 +1296,7 @@ def test_enable_lake_formation_called_with_disable_hybrid_access_mode(
             use_service_linked_role=True,
             registration_role_arn=None,
             disable_hybrid_access_mode=True,
+            acknowledge_risk=None,
         )
 
 
@@ -1623,7 +1615,7 @@ def test_uses_service_linked_role_arn_when_use_service_linked_role_true(
         mock_revoke.return_value = True
         mock_generate.return_value = []
 
-        fg.enable_lake_formation(use_service_linked_role=True, disable_hybrid_access_mode=True)
+        fg.enable_lake_formation(use_service_linked_role=True, disable_hybrid_access_mode=True, acknowledge_risk=True)
 
         expected_slr_arn = "arn:aws:iam::123456789012:role/aws-service-role/lakeformation.amazonaws.com/AWSServiceRoleForLakeFormationDataAccess"
         mock_generate.assert_called_once()
@@ -1666,7 +1658,7 @@ def test_uses_service_linked_role_arn_by_default(
         mock_revoke.return_value = True
         mock_generate.return_value = []
 
-        fg.enable_lake_formation(disable_hybrid_access_mode=True)
+        fg.enable_lake_formation(disable_hybrid_access_mode=True, acknowledge_risk=True)
 
         expected_slr_arn = "arn:aws:iam::987654321098:role/aws-service-role/lakeformation.amazonaws.com/AWSServiceRoleForLakeFormationDataAccess"
         mock_generate.assert_called_once()
@@ -1709,7 +1701,7 @@ def test_service_linked_role_arn_uses_correct_account_id(
         mock_revoke.return_value = True
         mock_generate.return_value = []
 
-        fg.enable_lake_formation(use_service_linked_role=True, disable_hybrid_access_mode=True)
+        fg.enable_lake_formation(use_service_linked_role=True, disable_hybrid_access_mode=True, acknowledge_risk=True)
 
         expected_slr_arn = f"arn:aws:iam::{account_id}:role/aws-service-role/lakeformation.amazonaws.com/AWSServiceRoleForLakeFormationDataAccess"
         mock_generate.assert_called_once()
@@ -1763,6 +1755,7 @@ def test_uses_registration_role_arn_when_use_service_linked_role_false(
             use_service_linked_role=False,
             registration_role_arn=custom_registration_role,
             disable_hybrid_access_mode=True,
+            acknowledge_risk=True,
         )
 
         mock_generate.assert_called_once()
@@ -1813,6 +1806,7 @@ def test_registration_role_arn_passed_to_s3_registration(
             use_service_linked_role=False,
             registration_role_arn=custom_registration_role,
             disable_hybrid_access_mode=True,
+            acknowledge_risk=True,
         )
 
         mock_register.assert_called_once()
@@ -1860,6 +1854,7 @@ def test_different_registration_role_arns_produce_different_policies(
             use_service_linked_role=False,
             registration_role_arn=first_role,
             disable_hybrid_access_mode=True,
+            acknowledge_risk=True,
         )
         first_call_kwargs = mock_generate.call_args[1]
         first_lf_role = first_call_kwargs["lake_formation_role_arn"]
@@ -1874,6 +1869,7 @@ def test_different_registration_role_arns_produce_different_policies(
             use_service_linked_role=False,
             registration_role_arn=second_role,
             disable_hybrid_access_mode=True,
+            acknowledge_risk=True,
         )
         second_call_kwargs = mock_generate.call_args[1]
         second_lf_role = second_call_kwargs["lake_formation_role_arn"]
@@ -1972,7 +1968,7 @@ def test_iceberg_strips_data_suffix_for_s3_registration(
         mock_grant.return_value = True
         mock_revoke.return_value = True
 
-        self.fg.enable_lake_formation(disable_hybrid_access_mode=True)
+        self.fg.enable_lake_formation(disable_hybrid_access_mode=True, acknowledge_risk=True)
 
         # The registered S3 location should NOT end with /data
         call_args = mock_register.call_args
@@ -1997,7 +1993,7 @@ def test_non_iceberg_keeps_full_s3_path(
         mock_grant.return_value = True
         mock_revoke.return_value = True
 
-        self.fg.enable_lake_formation(disable_hybrid_access_mode=True)
+        self.fg.enable_lake_formation(disable_hybrid_access_mode=True, acknowledge_risk=True)
 
         call_args = mock_register.call_args
         registered_location = call_args[0][0]
@@ -2036,7 +2032,7 @@ def test_raises_error_when_feature_group_arn_is_none(
         mock_revoke.return_value = True
 
         with pytest.raises(ValueError, match="Feature Group ARN is required"):
-            fg.enable_lake_formation(disable_hybrid_access_mode=True)
+            fg.enable_lake_formation(disable_hybrid_access_mode=True, acknowledge_risk=True)
 
 
 class TestEnableLakeFormationHappyPath:
@@ -2071,7 +2067,7 @@ def test_returns_all_true_on_success(
         mock_grant.return_value = True
         mock_revoke.return_value = True
 
-        result = self.fg.enable_lake_formation(disable_hybrid_access_mode=True)
+        result = self.fg.enable_lake_formation(disable_hybrid_access_mode=True, acknowledge_risk=True)
 
         assert result == {
             "s3_location_registered": True,
@@ -2128,6 +2124,7 @@ def test_session_and_region_passed_to_enable_lake_formation(
             use_service_linked_role=True,
             registration_role_arn=None,
             disable_hybrid_access_mode=False,
+            acknowledge_risk=None,
         )