Source code for airflow.providers.amazon.aws.hooks.ecs

#
# 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 typing import TYPE_CHECKING

from airflow.providers.amazon.aws.exceptions import EcsOperatorError, EcsTaskFailToStart
from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
from airflow.providers.amazon.aws.utils import _StringCompareEnum
from airflow.typing_compat import Protocol, runtime_checkable

if TYPE_CHECKING:
    from botocore.waiter import Waiter


[docs]def should_retry(exception: Exception):
    """Check if exception is related to ECS resource quota (CPU, MEM)."""
    if isinstance(exception, EcsOperatorError):
        return any(
            quota_reason in failure["reason"]
            for quota_reason in ["RESOURCE:MEMORY", "RESOURCE:CPU"]
            for failure in exception.failures
        )
    return False


[docs]def should_retry_eni(exception: Exception):
    """Check if exception is related to ENI (Elastic Network Interfaces)."""
    if isinstance(exception, EcsTaskFailToStart):
        return any(
            eni_reason in exception.message
            for eni_reason in ["network interface provisioning", "ResourceInitializationError"]
        )
    return False


[docs]class EcsClusterStates(_StringCompareEnum):
    """Contains the possible State values of an ECS Cluster."""

[docs]    ACTIVE = "ACTIVE"
[docs]    PROVISIONING = "PROVISIONING"
[docs]    DEPROVISIONING = "DEPROVISIONING"
[docs]    FAILED = "FAILED"
[docs]    INACTIVE = "INACTIVE"


[docs]class EcsTaskDefinitionStates(_StringCompareEnum):
    """Contains the possible State values of an ECS Task Definition."""

[docs]    ACTIVE = "ACTIVE"
[docs]    INACTIVE = "INACTIVE"
[docs]    DELETE_IN_PROGRESS = "DELETE_IN_PROGRESS"


[docs]class EcsTaskStates(_StringCompareEnum):
    """Contains the possible State values of an ECS Task."""

[docs]    PROVISIONING = "PROVISIONING"
[docs]    PENDING = "PENDING"
[docs]    ACTIVATING = "ACTIVATING"
[docs]    RUNNING = "RUNNING"
[docs]    DEACTIVATING = "DEACTIVATING"
[docs]    STOPPING = "STOPPING"
[docs]    DEPROVISIONING = "DEPROVISIONING"
[docs]    STOPPED = "STOPPED"
[docs]    NONE = "NONE"


[docs]class EcsHook(AwsGenericHook):
    """
    Interact with Amazon Elastic Container Service (ECS).

    Provide thin wrapper around :external+boto3:py:class:`boto3.client("ecs") <ECS.Client>`.

    Additional arguments (such as ``aws_conn_id``) may be specified and
    are passed down to the underlying AwsBaseHook.

    .. seealso::
        - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
        - `Amazon Elastic Container Service \
        <https://docs.aws.amazon.com/AmazonECS/latest/APIReference/Welcome.html>`__
    """

    def __init__(self, *args, **kwargs) -> None:
        kwargs["client_type"] = "ecs"
        super().__init__(*args, **kwargs)

[docs]    def get_cluster_state(self, cluster_name: str) -> str:
        """
        Get ECS Cluster state.

        .. seealso::
            - :external+boto3:py:meth:`ECS.Client.describe_clusters`

        :param cluster_name: ECS Cluster name or full cluster Amazon Resource Name (ARN) entry.
        """
        return self.conn.describe_clusters(clusters=[cluster_name])["clusters"][0]["status"]

[docs]    def get_task_definition_state(self, task_definition: str) -> str:
        """
        Get ECS Task Definition state.

        .. seealso::
            - :external+boto3:py:meth:`ECS.Client.describe_task_definition`

        :param task_definition: The family for the latest ACTIVE revision,
            family and revision ( family:revision ) for a specific revision in the family,
            or full Amazon Resource Name (ARN) of the task definition to describe.
        """
        return self.conn.describe_task_definition(taskDefinition=task_definition)["taskDefinition"]["status"]

[docs]    def get_task_state(self, cluster, task) -> str:
        """
        Get ECS Task state.

        .. seealso::
            - :external+boto3:py:meth:`ECS.Client.describe_tasks`

        :param cluster: The short name or full Amazon Resource Name (ARN)
            of the cluster that hosts the task or tasks to describe.
        :param task: Task ID or full ARN entry.
        """
        return self.conn.describe_tasks(cluster=cluster, tasks=[task])["tasks"][0]["lastStatus"]


@runtime_checkable
[docs]class EcsProtocol(Protocol):
    """
    A structured Protocol for ``boto3.client('ecs')``.

    This is used for type hints on :py:meth:`.EcsOperator.client`.

    .. seealso::

        - https://mypy.readthedocs.io/en/latest/protocols.html
        - https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs.html
    """

[docs]    def run_task(self, **kwargs) -> dict:
        """Run a task.

        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs.html#ECS.Client.run_task
        """
        ...

[docs]    def get_waiter(self, x: str) -> Waiter:
        """Get a waiter.

        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs.html#ECS.Client.get_waiter
        """
        ...

[docs]    def describe_tasks(self, cluster: str, tasks) -> dict:
        """Describe tasks.

        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs.html#ECS.Client.describe_tasks
        """
        ...

[docs]    def stop_task(self, cluster, task, reason: str) -> dict:
        """Stop a task.

        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs.html#ECS.Client.stop_task
        """
        ...

[docs]    def describe_task_definition(self, taskDefinition: str) -> dict:
        """Describe a task definition.

        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs.html#ECS.Client.describe_task_definition
        """
        ...

[docs]    def list_tasks(self, cluster: str, launchType: str, desiredStatus: str, family: str) -> dict:
        """List tasks.

        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs.html#ECS.Client.list_tasks
        """
        ...