apache-airflow-providers-amazon 9.6.1rc1__py3-none-any.whl → 9.7.1a1__py3-none-any.whl

airflow/providers/amazon/__init__.py

@@ -29,11 +29,11 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "9.6.1"
+__version__ = "9.7.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
-    "2.9.0"
+    "2.10.0"
 ):
     raise RuntimeError(
-        f"The package `apache-airflow-providers-amazon:{__version__}` needs Apache Airflow 2.9.0+"
+        f"The package `apache-airflow-providers-amazon:{__version__}` needs Apache Airflow 2.10.0+"
     )

airflow/providers/amazon/aws/auth_manager/avp/entities.py

@@ -35,7 +35,7 @@ class AvpEntities(Enum):
     # Resource types
     ASSET = "Asset"
     ASSET_ALIAS = "AssetAlias"
-    BACKFILL = "Backfills"
+    BACKFILL = "Backfill"
     CONFIGURATION = "Configuration"
     CONNECTION = "Connection"
     CUSTOM = "Custom"

airflow/providers/amazon/aws/auth_manager/avp/schema.json

@@ -1,6 +1,36 @@
 {
     "Airflow": {
         "actions": {
+            "Asset.GET": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Asset"]
+                }
+            },
+            "AssetAlias.GET": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["AssetAlias"]
+                }
+            },
+            "Backfill.GET": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Backfill"]
+                }
+            },
+            "Backfill.POST": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Backfill"]
+                }
+            },
+            "Backfill.PUT": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Backfill"]
+                }
+            },
             "Connection.DELETE": {
                 "appliesTo": {
                     "principalTypes": ["User"],
@@ -115,12 +145,6 @@
                     }
                 }
             },
-            "Dataset.GET": {
-                "appliesTo": {
-                    "principalTypes": ["User"],
-                    "resourceTypes": ["Dataset"]
-                }
-            },
             "Menu.MENU": {
                 "appliesTo": {
                     "principalTypes": ["User"],
@@ -183,11 +207,13 @@
             }
         },
         "entityTypes": {
+            "Asset":  {},
+            "AssetAlias":  {},
+            "Backfill":  {},
             "Configuration": {},
             "Connection": {},
             "Custom": {},
             "Dag": {},
-            "Dataset": {},
             "Menu": {},
             "Pool": {},
             "Group": {},

airflow/providers/amazon/aws/auth_manager/aws_auth_manager.py

@@ -72,13 +72,11 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
     authentication and authorization in Airflow.
     """
 
-    def __init__(self) -> None:
+    def init(self) -> None:
         if not AIRFLOW_V_3_0_PLUS:
             raise AirflowOptionalProviderFeatureException(
                 "AWS auth manager is only compatible with Airflow versions >= 3.0.0"
             )
-
-        super().__init__()
         self._check_avp_schema_version()
 
     @cached_property
@@ -90,7 +88,12 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
         return conf.get("api", "base_url", fallback="/")
 
     def deserialize_user(self, token: dict[str, Any]) -> AwsAuthManagerUser:
-        return AwsAuthManagerUser(user_id=token.pop("sub"), **token)
+        return AwsAuthManagerUser(
+            user_id=token.pop("sub"),
+            groups=token.get("groups", []),
+            username=token.get("username"),
+            email=token.get("email"),
+        )
 
     def serialize_user(self, user: AwsAuthManagerUser) -> dict[str, Any]:
         return {
@@ -364,7 +367,7 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
         ]
 
     def get_fastapi_app(self) -> FastAPI | None:
-        from airflow.providers.amazon.aws.auth_manager.router.login import login_router
+        from airflow.providers.amazon.aws.auth_manager.routes.login import login_router
 
         app = FastAPI(
             title="AWS auth manager sub application",

airflow/providers/amazon/aws/auth_manager/cli/avp_commands.py

@@ -14,8 +14,6 @@
 # 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
@@ -71,6 +69,8 @@ def init_avp(args):
 @providers_configuration_loaded
 def update_schema(args):
     """Update Amazon Verified Permissions policy store schema."""
+    if not args.policy_store_id:
+        raise ValueError("Parameter '--policy-store-id' is required.")
     client = _get_client()
     _set_schema(client, args.policy_store_id, args)
 
@@ -98,9 +98,8 @@ def _create_policy_store(client: BaseClient, args) -> tuple[str | None, bool]:
         if policy_store.get("description") == args.policy_store_description
     ]
 
-    if args.verbose:
-        log.debug("Policy stores found: %s", policy_stores)
-        log.debug("Existing policy stores found: %s", existing_policy_stores)
+    log.debug("Policy stores found: %s", policy_stores)
+    log.debug("Existing policy stores found: %s", existing_policy_stores)
 
     if len(existing_policy_stores) > 0:
         print(
@@ -118,8 +117,7 @@ def _create_policy_store(client: BaseClient, args) -> tuple[str | None, bool]:
         },
         description=args.policy_store_description,
     )
-    if args.verbose:
-        log.debug("Response from create_policy_store: %s", response)
+    log.debug("Response from create_policy_store: %s", response)
 
     print(f"Policy store created: '{response['policyStoreId']}'")
 
@@ -141,7 +139,6 @@ def _set_schema(client: BaseClient, policy_store_id: str, args) -> None:
             },
         )
 
-        if args.verbose:
-            log.debug("Response from put_schema: %s", response)
+        log.debug("Response from put_schema: %s", response)
 
     print("Policy store schema updated.")

airflow/providers/amazon/aws/auth_manager/cli/definition.py

@@ -27,22 +27,12 @@ from airflow.cli.cli_config import (
 # # ARGS # #
 ############
 
-ARG_VERBOSE = Arg(("-v", "--verbose"), help="Make logging output more verbose", action="store_true")
-
 ARG_DRY_RUN = Arg(
     ("--dry-run",),
     help="Perform a dry run",
     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"
@@ -59,12 +49,12 @@ AWS_AUTH_MANAGER_COMMANDS = (
         name="init-avp",
         help="Initialize Amazon Verified resources to be used by AWS manager",
         func=lazy_load_command("airflow.providers.amazon.aws.auth_manager.cli.avp_commands.init_avp"),
-        args=(ARG_POLICY_STORE_DESCRIPTION, ARG_DRY_RUN, ARG_VERBOSE),
+        args=(ARG_POLICY_STORE_DESCRIPTION, ARG_DRY_RUN),
     ),
     ActionCommand(
         name="update-avp-schema",
         help="Update Amazon Verified permissions policy store schema to the latest version in 'airflow/providers/amazon/aws/auth_manager/avp/schema.json'",
         func=lazy_load_command("airflow.providers.amazon.aws.auth_manager.cli.avp_commands.update_schema"),
-        args=(ARG_POLICY_STORE_ID, ARG_DRY_RUN, ARG_VERBOSE),
+        args=(ARG_POLICY_STORE_ID, ARG_DRY_RUN),
     ),
 )

airflow/providers/amazon/aws/auth_manager/datamodels/login.py

@@ -0,0 +1,26 @@
+#
+# 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
+
+from airflow.api_fastapi.core_api.base import BaseModel
+
+
+class LoginResponse(BaseModel):
+    """Login serializer for responses."""
+
+    access_token: str

airflow/providers/amazon/aws/auth_manager/routes/__init__.py

@@ -0,0 +1,16 @@
+# 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.

airflow/providers/amazon/aws/auth_manager/{router → routes}/login.py

@@ -25,11 +25,15 @@ from fastapi import HTTPException, Request
 from starlette import status
 from starlette.responses import RedirectResponse
 
-from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX, get_auth_manager
+from airflow.api_fastapi.app import (
+    AUTH_MANAGER_FASTAPI_APP_PREFIX,
+    get_auth_manager,
+)
 from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
 from airflow.api_fastapi.common.router import AirflowRouter
 from airflow.configuration import conf
 from airflow.providers.amazon.aws.auth_manager.constants import CONF_SAML_METADATA_URL_KEY, CONF_SECTION_NAME
+from airflow.providers.amazon.aws.auth_manager.datamodels.login import LoginResponse
 from airflow.providers.amazon.aws.auth_manager.user import AwsAuthManagerUser
 
 try:
@@ -49,9 +53,17 @@ login_router = AirflowRouter(tags=["AWSAuthManagerLogin"])
 
 @login_router.get("/login")
 def login(request: Request):
-    """Authenticate the user."""
+    """Initiate the authentication."""
     saml_auth = _init_saml_auth(request)
-    callback_url = saml_auth.login()
+    callback_url = saml_auth.login("login-redirect")
+    return RedirectResponse(url=callback_url)
+
+
+@login_router.get("/login/token")
+def login_token(request: Request) -> RedirectResponse:
+    """Initiate the authentication to create a token."""
+    saml_auth = _init_saml_auth(request)
+    callback_url = saml_auth.login("login-token")
     return RedirectResponse(url=callback_url)
 
 
@@ -76,22 +88,29 @@ def login_callback(request: Request):
     attributes = saml_auth.get_attributes()
     user = AwsAuthManagerUser(
         user_id=attributes["id"][0],
-        groups=attributes["groups"],
+        groups=attributes["groups"] or [],
         username=saml_auth.get_nameid(),
         email=attributes["email"][0] if "email" in attributes else None,
     )
     url = conf.get("api", "base_url", fallback="/")
     token = get_auth_manager().generate_jwt(user)
-    response = RedirectResponse(url=url, status_code=303)
 
-    secure = conf.has_option("api", "ssl_cert")
-    response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=secure)
-    return response
+    form_data = anyio.from_thread.run(request.form)
+    relay_state = form_data["RelayState"]
+
+    if relay_state == "login-redirect":
+        response = RedirectResponse(url=url, status_code=303)
+        secure = bool(conf.get("api", "ssl_cert", fallback=""))
+        response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=secure)
+        return response
+    if relay_state == "login-token":
+        return LoginResponse(access_token=token)
+    raise HTTPException(status.HTTP_500_INTERNAL_SERVER_ERROR, f"Invalid relay state: {relay_state}")
 
 
 def _init_saml_auth(request: Request) -> OneLogin_Saml2_Auth:
     request_data = _prepare_request(request)
-    base_url = conf.get(section="api", key="base_url", fallback="/")
+    base_url = conf.get(section="api", key="base_url")
     settings = {
         # We want to keep this flag on in case of errors.
         # It provides an error reasons, if turned off, it does not
@@ -99,7 +118,7 @@ def _init_saml_auth(request: Request) -> OneLogin_Saml2_Auth:
         "sp": {
             "entityId": "aws-auth-manager-saml-client",
             "assertionConsumerService": {
-                "url": f"{base_url}{AUTH_MANAGER_FASTAPI_APP_PREFIX}/login_callback",
+                "url": f"{base_url.rstrip('/')}{AUTH_MANAGER_FASTAPI_APP_PREFIX}/login_callback",
                 "binding": "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST",
             },
         },

airflow/providers/amazon/aws/executors/batch/batch_executor.py

@@ -324,11 +324,7 @@ class AwsBatchExecutor(BaseExecutor):
                     exec_config=exec_config,
                     attempt_number=attempt_number,
                 )
-                with suppress(AttributeError):
-                    # TODO: Remove this when min_airflow_version is 2.10.0 or higher in Amazon provider.
-                    # running_state is added in Airflow 2.10 and only needed to support task adoption
-                    # (an optional executor feature).
-                    self.running_state(key, job_id)
+                self.running_state(key, job_id)
 
     def _describe_jobs(self, job_ids) -> list[BatchJob]:
         all_jobs = []

airflow/providers/amazon/aws/executors/ecs/ecs_executor.py

@@ -23,7 +23,6 @@ Each Airflow task gets delegated out to an Amazon ECS Task.
 
 from __future__ import annotations
 
-import contextlib
 import time
 from collections import defaultdict, deque
 from collections.abc import Sequence
@@ -450,11 +449,7 @@ class AwsEcsExecutor(BaseExecutor):
             else:
                 task = run_task_response["tasks"][0]
                 self.active_workers.add_task(task, task_key, queue, cmd, exec_config, attempt_number)
-                with contextlib.suppress(AttributeError):
-                    # running_state is newly added, and only needed to support task adoption (an optional
-                    # executor feature).
-                    # TODO: remove when min airflow version >= 2.9.2
-                    self.running_state(task_key, task.task_arn)
+                self.running_state(task_key, task.task_arn)
 
     def _run_task(
         self, task_id: TaskInstanceKey, cmd: CommandType, queue: str, exec_config: ExecutorConfigType

airflow/providers/amazon/aws/hooks/redshift_sql.py

@@ -26,7 +26,6 @@ from sqlalchemy.engine.url import URL
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
-from airflow.providers.amazon.version_compat import AIRFLOW_V_2_10_PLUS
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 
 if TYPE_CHECKING:
@@ -260,6 +259,4 @@ class RedshiftSQLHook(DbApiHook):
 
     def get_openlineage_default_schema(self) -> str | None:
         """Return current schema. This is usually changed with ``SEARCH_PATH`` parameter."""
-        if AIRFLOW_V_2_10_PLUS:
-            return self.get_first("SELECT CURRENT_SCHEMA();")[0]
-        return super().get_openlineage_default_schema()
+        return self.get_first("SELECT CURRENT_SCHEMA();")[0]