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 = [