(airflow) branch main updated: AWS auth manager CLI: persist the policy store description when doing updates (#37946)

[potiuk]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new fbbe03ea40 AWS auth manager CLI: persist the policy store description 
when doing updates (#37946)
fbbe03ea40 is described below

commit fbbe03ea40727d7a737c8cce9500fcbe9d9965ef
Author: Vincent <97131062+vincb...@users.noreply.github.com>
AuthorDate: Wed Mar 6 15:33:56 2024 -0500

    AWS auth manager CLI: persist the policy store description when doing 
updates (#37946)
---
 .../amazon/aws/auth_manager/cli/avp_commands.py    | 22 +++++++++++++++++-----
 .../aws/auth_manager/cli/test_avp_commands.py      |  1 +
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/airflow/providers/amazon/aws/auth_manager/cli/avp_commands.py 
b/airflow/providers/amazon/aws/auth_manager/cli/avp_commands.py
index fd49cbcb5a..4d94e723ef 100644
--- a/airflow/providers/amazon/aws/auth_manager/cli/avp_commands.py
+++ b/airflow/providers/amazon/aws/auth_manager/cli/avp_commands.py
@@ -63,8 +63,9 @@ def init_avp(args):
         _set_schema(client, policy_store_id, args)
 
     if not args.dry_run:
-        print("Amazon Verified Permissions resources created successfully.")
-        print("Please set them in Airflow configuration under 
AIRFLOW__AWS_AUTH_MANAGER__<config name>.")
+        print(
+            "Please set configs below in Airflow configuration under 
AIRFLOW__AWS_AUTH_MANAGER__<config name>."
+        )
         print(json.dumps({"avp_policy_store_id": policy_store_id}, indent=4))
 
 
@@ -75,9 +76,6 @@ def update_schema(args):
     client = _get_client()
     _set_schema(client, args.policy_store_id, args)
 
-    if not args.dry_run:
-        print("Amazon Verified Permissions policy store schema updated 
successfully.")
-
 
 def _get_client():
     """Return Amazon Verified Permissions client."""
@@ -139,6 +137,16 @@ def _set_schema(client: BaseClient, policy_store_id: str, 
args) -> None:
         print(f"Dry run, not updating the schema of the policy store with ID 
'{policy_store_id}'.")
         return
 
+    if args.verbose:
+        log.debug("Getting the policy store details")
+
+    details = client.get_policy_store(
+        policyStoreId=policy_store_id,
+    )
+
+    if args.verbose:
+        log.debug("Response from get_policy_store: %s", details)
+
     if args.verbose:
         log.debug("Disabling schema validation before updating schema")
 
@@ -147,6 +155,7 @@ def _set_schema(client: BaseClient, policy_store_id: str, 
args) -> None:
         validationSettings={
             "mode": "OFF",
         },
+        description=details["description"],
     )
 
     if args.verbose:
@@ -164,6 +173,8 @@ def _set_schema(client: BaseClient, policy_store_id: str, 
args) -> None:
         if args.verbose:
             log.debug("Response from put_schema: %s", response)
 
+    print("Policy store schema updated.")
+
     if args.verbose:
         log.debug("Enabling schema validation after updating schema")
 
@@ -172,6 +183,7 @@ def _set_schema(client: BaseClient, policy_store_id: str, 
args) -> None:
         validationSettings={
             "mode": "STRICT",
         },
+        description=details["description"],
     )
 
     if args.verbose:
diff --git a/tests/providers/amazon/aws/auth_manager/cli/test_avp_commands.py 
b/tests/providers/amazon/aws/auth_manager/cli/test_avp_commands.py
index c65dae87e6..8787c22fa8 100644
--- a/tests/providers/amazon/aws/auth_manager/cli/test_avp_commands.py
+++ b/tests/providers/amazon/aws/auth_manager/cli/test_avp_commands.py
@@ -65,6 +65,7 @@ class TestAvpCommands:
 
         mock_boto3.get_paginator.return_value = paginator
         mock_boto3.create_policy_store.return_value = {"policyStoreId": 
policy_store_id}
+        mock_boto3.get_policy_store.return_value = {"description": 
policy_store_description}
 
         with conf_vars({("database", "check_migrations"): "False"}):
             params = [