This is an automated email from the ASF dual-hosted git repository.
vincbeck 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 56c27f8f5a Create CLI commands for AWS auth manager to create AWS
Identity Center related resources (#37407)
56c27f8f5a is described below
commit 56c27f8f5a3c547147389253ea9653a374ad29f1
Author: Vincent <[email protected]>
AuthorDate: Wed Feb 14 11:43:16 2024 -0800
Create CLI commands for AWS auth manager to create AWS Identity Center
related resources (#37407)
---
.../amazon/aws/auth_manager/cli/avp_commands.py | 4 +-
.../amazon/aws/auth_manager/cli/definition.py | 14 ++
.../amazon/aws/auth_manager/cli/idc_commands.py | 148 +++++++++++++++++++++
.../amazon/aws/auth_manager/cli/test_definition.py | 2 +-
.../aws/auth_manager/cli/test_idc_commands.py | 134 +++++++++++++++++++
5 files changed, 299 insertions(+), 3 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 24f7907104..fd49cbcb5a 100644
--- a/airflow/providers/amazon/aws/auth_manager/cli/avp_commands.py
+++ b/airflow/providers/amazon/aws/auth_manager/cli/avp_commands.py
@@ -55,7 +55,7 @@ def init_avp(args):
if not is_new_policy_store:
print(
f"Since an existing policy store with description
'{args.policy_store_description}' has been found in Amazon Verified
Permissions, "
- "the CLI nade no changes to this policy store for security
reasons. "
+ "the CLI made no changes to this policy store for security
reasons. "
"Any modification to this policy store must be done manually.",
)
else:
@@ -115,7 +115,7 @@ def _create_policy_store(client: BaseClient, args) ->
tuple[str | None, bool]:
print(f"No policy store with description
'{args.policy_store_description}' found, creating one.")
if args.dry_run:
print(
- "Dry run, not creating the policy store with description
'{args.policy_store_description}'."
+ f"Dry run, not creating the policy store with description
'{args.policy_store_description}'."
)
return None, True
diff --git a/airflow/providers/amazon/aws/auth_manager/cli/definition.py
b/airflow/providers/amazon/aws/auth_manager/cli/definition.py
index 4355dcce9b..846dc56171 100644
--- a/airflow/providers/amazon/aws/auth_manager/cli/definition.py
+++ b/airflow/providers/amazon/aws/auth_manager/cli/definition.py
@@ -35,6 +35,14 @@ ARG_DRY_RUN = Arg(
action="store_true",
)
+# AWS IAM Identity Center
+ARG_INSTANCE_NAME = Arg(("--instance-name",), help="Instance name in Identity
Center", default="Airflow")
+
+ARG_APPLICATION_NAME = Arg(
+ ("--application-name",), help="Application name in Identity Center",
default="Airflow"
+)
+
+
# Amazon Verified Permissions
ARG_POLICY_STORE_DESCRIPTION = Arg(
("--policy-store-description",), help="Policy store description",
default="Airflow"
@@ -47,6 +55,12 @@ ARG_POLICY_STORE_ID = Arg(("--policy-store-id",),
help="Policy store ID")
################
AWS_AUTH_MANAGER_COMMANDS = (
+ ActionCommand(
+ name="init-identity-center",
+ help="Initialize AWS IAM identity Center resources to be used by AWS
manager",
+
func=lazy_load_command("airflow.providers.amazon.aws.auth_manager.cli.idc_commands.init_idc"),
+ args=(ARG_INSTANCE_NAME, ARG_APPLICATION_NAME, ARG_DRY_RUN,
ARG_VERBOSE),
+ ),
ActionCommand(
name="init-avp",
help="Initialize Amazon Verified resources to be used by AWS manager",
diff --git a/airflow/providers/amazon/aws/auth_manager/cli/idc_commands.py
b/airflow/providers/amazon/aws/auth_manager/cli/idc_commands.py
new file mode 100644
index 0000000000..3639dd4b91
--- /dev/null
+++ b/airflow/providers/amazon/aws/auth_manager/cli/idc_commands.py
@@ -0,0 +1,148 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""User sub-commands."""
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import boto3
+from botocore.exceptions import ClientError
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowOptionalProviderFeatureException
+from airflow.providers.amazon.aws.auth_manager.constants import
CONF_REGION_NAME_KEY, CONF_SECTION_NAME
+from airflow.utils import cli as cli_utils
+
+try:
+ from airflow.utils.providers_configuration_loader import
providers_configuration_loaded
+except ImportError:
+ raise AirflowOptionalProviderFeatureException(
+ "Failed to import avp_commands. This feature is only available in
Airflow "
+ "version >= 2.8.0 where Auth Managers are introduced."
+ )
+
+if TYPE_CHECKING:
+ from botocore.client import BaseClient
+
+log = logging.getLogger(__name__)
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def init_idc(args):
+ """Initialize AWS IAM Identity Center resources."""
+ client = _get_client()
+
+ # Create the instance if needed
+ instance_arn = _create_instance(client, args)
+
+ # Create the application if needed
+ _create_application(client, instance_arn, args)
+
+ if not args.dry_run:
+ print("AWS IAM Identity Center resources created successfully.")
+
+
+def _get_client():
+ """Return AWS IAM Identity Center client."""
+ region_name = conf.get(CONF_SECTION_NAME, CONF_REGION_NAME_KEY)
+ return boto3.client("sso-admin", region_name=region_name)
+
+
+def _create_instance(client: BaseClient, args) -> str | None:
+ """Create if needed AWS IAM Identity Center instance."""
+ instances = client.list_instances()
+
+ if args.verbose:
+ log.debug("Instances found: %s", instances)
+
+ if len(instances["Instances"]) > 0:
+ print(
+ f"There is already an instance configured in AWS IAM Identity
Center: '{instances['Instances'][0]['InstanceArn']}'. "
+ "No need to create a new one."
+ )
+ return instances["Instances"][0]["InstanceArn"]
+ else:
+ print("No instance configured in AWS IAM Identity Center, creating
one.")
+ if args.dry_run:
+ print("Dry run, not creating the instance.")
+ return None
+
+ response = client.create_instance(Name=args.instance_name)
+ if args.verbose:
+ log.debug("Response from create_instance: %s", response)
+
+ print(f"Instance created: '{response['InstanceArn']}'")
+
+ return response["InstanceArn"]
+
+
+def _create_application(client: BaseClient, instance_arn: str | None, args) ->
str | None:
+ """Create if needed AWS IAM identity Center application."""
+ paginator = client.get_paginator("list_applications")
+ pages = paginator.paginate(InstanceArn=instance_arn or "")
+ applications = [application for page in pages for application in
page["Applications"]]
+ existing_applications = [
+ application for application in applications if application["Name"] ==
args.application_name
+ ]
+
+ if args.verbose:
+ log.debug("Applications found: %s", applications)
+ log.debug("Existing applications found: %s", existing_applications)
+
+ if len(existing_applications) > 0:
+ print(
+ f"There is already an application named '{args.application_name}'
in AWS IAM Identity Center: '{existing_applications[0]['ApplicationArn']}'. "
+ "Using this application."
+ )
+ return existing_applications[0]["ApplicationArn"]
+ else:
+ print(f"No application named {args.application_name} found, creating
one.")
+ if args.dry_run:
+ print("Dry run, not creating the application.")
+ return None
+
+ try:
+ response = client.create_application(
+
ApplicationProviderArn="arn:aws:sso::aws:applicationProvider/custom-saml",
+ Description="Application automatically created through the
Airflow CLI. This application is used to access Airflow environment.",
+ InstanceArn=instance_arn,
+ Name=args.application_name,
+ PortalOptions={
+ "SignInOptions": {
+ "Origin": "IDENTITY_CENTER",
+ },
+ "Visibility": "ENABLED",
+ },
+ Status="ENABLED",
+ )
+ if args.verbose:
+ log.debug("Response from create_application: %s", response)
+ except ClientError as e:
+ # This is needed because as of today, the create_application in
AWS Identity Center does not support SAML application
+ # Remove this part when it is supported
+ if "is not supported for this action" in
e.response["Error"]["Message"]:
+ print(
+ "Creation of SAML applications is only supported in AWS
console today. "
+ "Please create the application through the console."
+ )
+ raise
+
+ print(f"Application created: '{response['ApplicationArn']}'")
+
+ return response["ApplicationArn"]
diff --git a/tests/providers/amazon/aws/auth_manager/cli/test_definition.py
b/tests/providers/amazon/aws/auth_manager/cli/test_definition.py
index 079df886f6..5866aa594f 100644
--- a/tests/providers/amazon/aws/auth_manager/cli/test_definition.py
+++ b/tests/providers/amazon/aws/auth_manager/cli/test_definition.py
@@ -21,4 +21,4 @@ from airflow.providers.amazon.aws.auth_manager.cli.definition
import AWS_AUTH_MA
class TestAwsCliDefinition:
def test_aws_auth_manager_cli_commands(self):
- assert len(AWS_AUTH_MANAGER_COMMANDS) == 2
+ assert len(AWS_AUTH_MANAGER_COMMANDS) == 3
diff --git a/tests/providers/amazon/aws/auth_manager/cli/test_idc_commands.py
b/tests/providers/amazon/aws/auth_manager/cli/test_idc_commands.py
new file mode 100644
index 0000000000..9913e1ae7c
--- /dev/null
+++ b/tests/providers/amazon/aws/auth_manager/cli/test_idc_commands.py
@@ -0,0 +1,134 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import importlib
+from unittest.mock import Mock, patch
+
+import pytest
+
+from airflow.cli import cli_parser
+from airflow.providers.amazon.aws.auth_manager.cli.idc_commands import init_idc
+from tests.test_utils.config import conf_vars
+
+mock_boto3 = Mock()
+
+
[email protected]_test
+class TestIdcCommands:
+ def setup_method(self):
+ mock_boto3.reset_mock()
+
+ @classmethod
+ def setup_class(cls):
+ with conf_vars(
+ {
+ (
+ "core",
+ "auth_manager",
+ ):
"airflow.providers.amazon.aws.auth_manager.aws_auth_manager.AwsAuthManager"
+ }
+ ):
+ importlib.reload(cli_parser)
+ cls.arg_parser = cli_parser.get_parser()
+
+ @pytest.mark.parametrize(
+ "dry_run, verbose",
+ [
+ (False, False),
+ (True, True),
+ ],
+ )
+
@patch("airflow.providers.amazon.aws.auth_manager.cli.idc_commands._get_client")
+ def test_init_idc_with_no_existing_resources(self, mock_get_client,
dry_run, verbose):
+ mock_get_client.return_value = mock_boto3
+
+ instance_name = "test-instance"
+ instance_arn = "test-instance-arn"
+ application_name = "test-application"
+ application_arn = "test-application-arn"
+
+ paginator = Mock()
+ paginator.paginate.return_value = []
+
+ mock_boto3.list_instances.return_value = {"Instances": []}
+ mock_boto3.create_instance.return_value = {"InstanceArn": instance_arn}
+ mock_boto3.get_paginator.return_value = paginator
+ mock_boto3.create_application.return_value = {"ApplicationArn":
application_arn}
+
+ with conf_vars({("database", "check_migrations"): "False"}):
+ params = [
+ "aws-auth-manager",
+ "init-identity-center",
+ "--instance-name",
+ instance_name,
+ "--application-name",
+ application_name,
+ ]
+ if dry_run:
+ params.append("--dry-run")
+ if verbose:
+ params.append("--verbose")
+ init_idc(self.arg_parser.parse_args(params))
+
+ mock_boto3.list_instances.assert_called_once_with()
+ if not dry_run:
+
mock_boto3.create_instance.assert_called_once_with(Name=instance_name)
+ mock_boto3.create_application.assert_called_once()
+
+ @pytest.mark.parametrize(
+ "dry_run, verbose",
+ [
+ (False, False),
+ (True, True),
+ ],
+ )
+
@patch("airflow.providers.amazon.aws.auth_manager.cli.idc_commands._get_client")
+ def test_init_idc_with_existing_resources(self, mock_get_client, dry_run,
verbose):
+ mock_get_client.return_value = mock_boto3
+
+ instance_name = "test-instance"
+ instance_arn = "test-instance-arn"
+ application_name = "test-application"
+ application_arn = "test-application-arn"
+
+ paginator = Mock()
+ paginator.paginate.return_value = [
+ {"Applications": [{"Name": application_name, "ApplicationArn":
application_arn}]}
+ ]
+
+ mock_boto3.list_instances.return_value = {"Instances":
[{"InstanceArn": instance_arn}]}
+ mock_boto3.get_paginator.return_value = paginator
+
+ with conf_vars({("database", "check_migrations"): "False"}):
+ params = [
+ "aws-auth-manager",
+ "init-identity-center",
+ "--instance-name",
+ instance_name,
+ "--application-name",
+ application_name,
+ ]
+ if dry_run:
+ params.append("--dry-run")
+ if verbose:
+ params.append("--verbose")
+ init_idc(self.arg_parser.parse_args(params))
+
+ mock_boto3.list_instances.assert_called_once_with()
+ mock_boto3.create_instance.assert_not_called()
+ mock_boto3.create_application.assert_not_called()
