apache-airflow-providers-amazon 8.19.0rc1__py3-none-any.whl → 8.20.0__py3-none-any.whl

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

@@ -0,0 +1,160 @@
+# 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 datetime
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Dict, List
+
+from airflow.providers.amazon.aws.executors.utils.base_config_keys import BaseConfigKeys
+from airflow.utils.state import State
+
+if TYPE_CHECKING:
+    from airflow.models.taskinstance import TaskInstanceKey
+
+CommandType = List[str]
+ExecutorConfigType = Dict[str, Any]
+
+CONFIG_GROUP_NAME = "aws_batch_executor"
+
+CONFIG_DEFAULTS = {
+    "conn_id": "aws_default",
+    "max_submit_job_attempts": "3",
+    "check_health_on_startup": "True",
+}
+
+
+@dataclass
+class BatchQueuedJob:
+    """Represents a Batch job that is queued. The job will be run in the next heartbeat."""
+
+    key: TaskInstanceKey
+    command: CommandType
+    queue: str
+    executor_config: ExecutorConfigType
+    attempt_number: int
+    next_attempt_time: datetime.datetime
+
+
+@dataclass
+class BatchJobInfo:
+    """Contains information about a currently running Batch job."""
+
+    cmd: CommandType
+    queue: str
+    config: ExecutorConfigType
+
+
+class BatchJob:
+    """Data Transfer Object for an AWS Batch Job."""
+
+    STATE_MAPPINGS = {
+        "SUBMITTED": State.QUEUED,
+        "PENDING": State.QUEUED,
+        "RUNNABLE": State.QUEUED,
+        "STARTING": State.QUEUED,
+        "RUNNING": State.RUNNING,
+        "SUCCEEDED": State.SUCCESS,
+        "FAILED": State.FAILED,
+    }
+
+    def __init__(self, job_id: str, status: str, status_reason: str | None = None):
+        self.job_id = job_id
+        self.status = status
+        self.status_reason = status_reason
+
+    def get_job_state(self) -> str:
+        """Return the state of the job."""
+        return self.STATE_MAPPINGS.get(self.status, State.QUEUED)
+
+    def __repr__(self):
+        """Return a visual representation of the Job status."""
+        return f"({self.job_id} -> {self.status}, {self.get_job_state()})"
+
+
+class BatchJobCollection:
+    """A collection to manage running Batch Jobs."""
+
+    def __init__(self):
+        self.key_to_id: dict[TaskInstanceKey, str] = {}
+        self.id_to_key: dict[str, TaskInstanceKey] = {}
+        self.id_to_failure_counts: dict[str, int] = defaultdict(int)
+        self.id_to_job_info: dict[str, BatchJobInfo] = {}
+
+    def add_job(
+        self,
+        job_id: str,
+        airflow_task_key: TaskInstanceKey,
+        airflow_cmd: CommandType,
+        queue: str,
+        exec_config: ExecutorConfigType,
+        attempt_number: int,
+    ):
+        """Add a job to the collection."""
+        self.key_to_id[airflow_task_key] = job_id
+        self.id_to_key[job_id] = airflow_task_key
+        self.id_to_failure_counts[job_id] = attempt_number
+        self.id_to_job_info[job_id] = BatchJobInfo(cmd=airflow_cmd, queue=queue, config=exec_config)
+
+    def pop_by_id(self, job_id: str) -> TaskInstanceKey:
+        """Delete job from collection based off of Batch Job ID."""
+        task_key = self.id_to_key[job_id]
+        del self.key_to_id[task_key]
+        del self.id_to_key[job_id]
+        del self.id_to_failure_counts[job_id]
+        return task_key
+
+    def failure_count_by_id(self, job_id: str) -> int:
+        """Get the number of times a job has failed given a Batch Job Id."""
+        return self.id_to_failure_counts[job_id]
+
+    def increment_failure_count(self, job_id: str):
+        """Increment the failure counter given a Batch Job Id."""
+        self.id_to_failure_counts[job_id] += 1
+
+    def get_all_jobs(self) -> list[str]:
+        """Get all AWS ARNs in collection."""
+        return list(self.id_to_key.keys())
+
+    def __len__(self):
+        """Return the number of jobs in collection."""
+        return len(self.key_to_id)
+
+
+class BatchSubmitJobKwargsConfigKeys(BaseConfigKeys):
+    """Keys loaded into the config which are valid Batch submit_job kwargs."""
+
+    JOB_NAME = "job_name"
+    JOB_QUEUE = "job_queue"
+    JOB_DEFINITION = "job_definition"
+    EKS_PROPERTIES_OVERRIDE = "eks_properties_override"
+    NODE_OVERRIDE = "node_override"
+
+
+class AllBatchConfigKeys(BatchSubmitJobKwargsConfigKeys):
+    """All keys loaded into the config which are related to the Batch Executor."""
+
+    MAX_SUBMIT_JOB_ATTEMPTS = "max_submit_job_attempts"
+    AWS_CONN_ID = "conn_id"
+    SUBMIT_JOB_KWARGS = "submit_job_kwargs"
+    REGION_NAME = "region_name"
+    CHECK_HEALTH_ON_STARTUP = "check_health_on_startup"
+
+
+class BatchExecutorException(Exception):
+    """Thrown when something unexpected has occurred within the AWS Batch ecosystem."""

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

@@ -26,7 +26,7 @@ from __future__ import annotations
 import time
 from collections import defaultdict, deque
 from copy import deepcopy
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Sequence
 
 from botocore.exceptions import ClientError, NoCredentialsError
 
@@ -47,12 +47,13 @@ from airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry impo
     exponential_backoff_retry,
 )
 from airflow.providers.amazon.aws.hooks.ecs import EcsHook
+from airflow.stats import Stats
 from airflow.utils import timezone
 from airflow.utils.helpers import merge_dicts
 from airflow.utils.state import State
 
 if TYPE_CHECKING:
-    from airflow.models.taskinstance import TaskInstanceKey
+    from airflow.models.taskinstance import TaskInstance, TaskInstanceKey
     from airflow.providers.amazon.aws.executors.ecs.utils import (
         CommandType,
         ExecutorConfigType,
@@ -182,7 +183,7 @@ class AwsEcsExecutor(BaseExecutor):
             AllEcsConfigKeys.AWS_CONN_ID,
             fallback=CONFIG_DEFAULTS[AllEcsConfigKeys.AWS_CONN_ID],
         )
-        region_name = conf.get(CONFIG_GROUP_NAME, AllEcsConfigKeys.REGION_NAME)
+        region_name = conf.get(CONFIG_GROUP_NAME, AllEcsConfigKeys.REGION_NAME, fallback=None)
         self.ecs = EcsHook(aws_conn_id=aws_conn_id, region_name=region_name).conn
         self.attempts_since_last_successful_connection += 1
         self.last_connection_reload = timezone.utcnow()
@@ -208,7 +209,7 @@ class AwsEcsExecutor(BaseExecutor):
             if error_code in INVALID_CREDENTIALS_EXCEPTIONS:
                 self.IS_BOTO_CONNECTION_HEALTHY = False
                 self.log.warning(
-                    f"AWS credentials are either missing or expired: {error}.\nRetrying connection"
+                    "AWS credentials are either missing or expired: %s.\nRetrying connection", error
                 )
 
         except Exception:
@@ -240,21 +241,19 @@ class AwsEcsExecutor(BaseExecutor):
         # Get state of current task.
         task_state = task.get_task_state()
         task_key = self.active_workers.arn_to_key[task.task_arn]
+
         # Mark finished tasks as either a success/failure.
-        if task_state == State.FAILED:
-            self.fail(task_key)
+        if task_state == State.FAILED or task_state == State.REMOVED:
             self.__log_container_failures(task_arn=task.task_arn)
-        elif task_state == State.SUCCESS:
-            self.success(task_key)
-        elif task_state == State.REMOVED:
             self.__handle_failed_task(task.task_arn, task.stopped_reason)
-        if task_state in (State.FAILED, State.SUCCESS):
+        elif task_state == State.SUCCESS:
             self.log.debug(
                 "Airflow task %s marked as %s after running on ECS Task (arn) %s",
                 task_key,
                 task_state,
                 task.task_arn,
             )
+            self.success(task_key)
             self.active_workers.pop_by_key(task_key)
 
     def __describe_tasks(self, task_arns):
@@ -287,7 +286,14 @@ class AwsEcsExecutor(BaseExecutor):
             )
 
     def __handle_failed_task(self, task_arn: str, reason: str):
-        """If an API failure occurs, the task is rescheduled."""
+        """
+        If an API failure occurs, the task is rescheduled.
+
+        This function will determine whether the task has been attempted the appropriate number
+        of times, and determine whether the task should be marked failed or not. The task will
+        be removed active_workers, and marked as FAILED, or set into pending_tasks depending on
+        how many times it has been retried.
+        """
         task_key = self.active_workers.arn_to_key[task_arn]
         task_info = self.active_workers.info_by_key(task_key)
         task_cmd = task_info.cmd
@@ -303,8 +309,7 @@ class AwsEcsExecutor(BaseExecutor):
                 self.__class__.MAX_RUN_TASK_ATTEMPTS,
                 task_arn,
             )
-            self.active_workers.increment_failure_count(task_key)
-            self.pending_tasks.appendleft(
+            self.pending_tasks.append(
                 EcsQueuedTask(
                     task_key,
                     task_cmd,
@@ -320,8 +325,8 @@ class AwsEcsExecutor(BaseExecutor):
                 task_key,
                 failure_count,
             )
-            self.active_workers.pop_by_key(task_key)
             self.fail(task_key)
+        self.active_workers.pop_by_key(task_key)
 
     def attempt_task_runs(self):
         """
@@ -344,16 +349,17 @@ class AwsEcsExecutor(BaseExecutor):
             attempt_number = ecs_task.attempt_number
             _failure_reasons = []
             if timezone.utcnow() < ecs_task.next_attempt_time:
+                self.pending_tasks.append(ecs_task)
                 continue
             try:
                 run_task_response = self._run_task(task_key, cmd, queue, exec_config)
             except NoCredentialsError:
-                self.pending_tasks.appendleft(ecs_task)
+                self.pending_tasks.append(ecs_task)
                 raise
             except ClientError as e:
                 error_code = e.response["Error"]["Code"]
                 if error_code in INVALID_CREDENTIALS_EXCEPTIONS:
-                    self.pending_tasks.appendleft(ecs_task)
+                    self.pending_tasks.append(ecs_task)
                     raise
                 _failure_reasons.append(str(e))
             except Exception as e:
@@ -373,12 +379,12 @@ class AwsEcsExecutor(BaseExecutor):
                 for reason in _failure_reasons:
                     failure_reasons[reason] += 1
                 # Make sure the number of attempts does not exceed MAX_RUN_TASK_ATTEMPTS
-                if int(attempt_number) <= int(self.__class__.MAX_RUN_TASK_ATTEMPTS):
+                if int(attempt_number) < int(self.__class__.MAX_RUN_TASK_ATTEMPTS):
                     ecs_task.attempt_number += 1
                     ecs_task.next_attempt_time = timezone.utcnow() + calculate_next_attempt_delay(
                         attempt_number
                     )
-                    self.pending_tasks.appendleft(ecs_task)
+                    self.pending_tasks.append(ecs_task)
                 else:
                     self.log.error(
                         "ECS task %s has failed a maximum of %s times. Marking as failed",
@@ -394,6 +400,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)
+                self.queued(task_key, task.task_arn)
         if failure_reasons:
             self.log.error(
                 "Pending ECS tasks failed to launch for the following reasons: %s. Retrying later.",
@@ -494,3 +501,39 @@ class AwsEcsExecutor(BaseExecutor):
                     'container "name" must be provided in "containerOverrides" configuration'
                 )
         raise KeyError(f"No such container found by container name: {self.container_name}")
+
+    def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[TaskInstance]:
+        """
+        Adopt task instances which have an external_executor_id (the ECS task ARN).
+
+        Anything that is not adopted will be cleared by the scheduler and becomes eligible for re-scheduling.
+        """
+        with Stats.timer("ecs_executor.adopt_task_instances.duration"):
+            adopted_tis: list[TaskInstance] = []
+
+            if task_arns := [ti.external_executor_id for ti in tis if ti.external_executor_id]:
+                task_descriptions = self.__describe_tasks(task_arns).get("tasks", [])
+
+                for task in task_descriptions:
+                    ti = [ti for ti in tis if ti.external_executor_id == task.task_arn][0]
+                    self.active_workers.add_task(
+                        task,
+                        ti.key,
+                        ti.queue,
+                        ti.command_as_list(),
+                        ti.executor_config,
+                        ti.prev_attempted_tries,
+                    )
+                    adopted_tis.append(ti)
+
+            if adopted_tis:
+                tasks = [f"{task} in state {task.state}" for task in adopted_tis]
+                task_instance_str = "\n\t".join(tasks)
+                self.log.info(
+                    "Adopted the following %d tasks from a dead executor:\n\t%s",
+                    len(adopted_tis),
+                    task_instance_str,
+                )
+
+            not_adopted_tis = [ti for ti in tis if ti not in adopted_tis]
+            return not_adopted_tis

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

@@ -30,6 +30,7 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, List
 
 from inflection import camelize
 
+from airflow.providers.amazon.aws.executors.utils.base_config_keys import BaseConfigKeys
 from airflow.utils.state import State
 
 if TYPE_CHECKING:
@@ -71,36 +72,28 @@ class EcsTaskInfo:
     config: ExecutorConfigType
 
 
-class BaseConfigKeys:
-    """Base Implementation of the Config Keys class. Implements iteration for child classes to inherit."""
-
-    def __iter__(self):
-        """Return an iterator of values of non dunder attributes of Config Keys."""
-        return iter({value for (key, value) in self.__class__.__dict__.items() if not key.startswith("__")})
-
-
 class RunTaskKwargsConfigKeys(BaseConfigKeys):
     """Keys loaded into the config which are valid ECS run_task kwargs."""
 
     ASSIGN_PUBLIC_IP = "assign_public_ip"
     CAPACITY_PROVIDER_STRATEGY = "capacity_provider_strategy"
     CLUSTER = "cluster"
+    CONTAINER_NAME = "container_name"
     LAUNCH_TYPE = "launch_type"
     PLATFORM_VERSION = "platform_version"
     SECURITY_GROUPS = "security_groups"
     SUBNETS = "subnets"
     TASK_DEFINITION = "task_definition"
-    CONTAINER_NAME = "container_name"
 
 
 class AllEcsConfigKeys(RunTaskKwargsConfigKeys):
     """All keys loaded into the config which are related to the ECS Executor."""
 
-    MAX_RUN_TASK_ATTEMPTS = "max_run_task_attempts"
     AWS_CONN_ID = "conn_id"
-    RUN_TASK_KWARGS = "run_task_kwargs"
-    REGION_NAME = "region_name"
     CHECK_HEALTH_ON_STARTUP = "check_health_on_startup"
+    MAX_RUN_TASK_ATTEMPTS = "max_run_task_attempts"
+    REGION_NAME = "region_name"
+    RUN_TASK_KWARGS = "run_task_kwargs"
 
 
 class EcsExecutorException(Exception):
@@ -108,7 +101,7 @@ class EcsExecutorException(Exception):
 
 
 class EcsExecutorTask:
-    """Data Transfer Object for an ECS Fargate Task."""
+    """Data Transfer Object for an ECS Task."""
 
     def __init__(
         self,
@@ -118,6 +111,7 @@ class EcsExecutorTask:
         containers: list[dict[str, Any]],
         started_at: Any | None = None,
         stopped_reason: str | None = None,
+        external_executor_id: str | None = None,
     ):
         self.task_arn = task_arn
         self.last_status = last_status
@@ -125,6 +119,7 @@ class EcsExecutorTask:
         self.containers = containers
         self.started_at = started_at
         self.stopped_reason = stopped_reason
+        self.external_executor_id = external_executor_id
 
     def get_task_state(self) -> str:
         """

airflow/providers/amazon/aws/executors/utils/base_config_keys.py

@@ -0,0 +1,25 @@
+# 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
+
+
+class BaseConfigKeys:
+    """Base Implementation of the Config Keys class. Implements iteration for child classes to inherit."""
+
+    def __iter__(self):
+        """Return an iterator of values of non dunder attributes of Config Keys."""
+        return iter({value for (key, value) in self.__class__.__dict__.items() if not key.startswith("__")})

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

@@ -22,6 +22,7 @@ This module contains AWS Athena hook.
 
     PageIterator
 """
+
 from __future__ import annotations
 
 import warnings

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

@@ -22,6 +22,7 @@ This module contains Base AWS Hook.
     For more information on how to use this hook, take a look at the guide:
     :ref:`howto/connection:aws`
 """
+
 from __future__ import annotations
 
 import datetime

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

@@ -24,6 +24,7 @@ A client for AWS Batch services.
     - https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/batch.html
     - https://docs.aws.amazon.com/batch/latest/APIReference/Welcome.html
 """
+
 from __future__ import annotations
 
 import itertools
@@ -438,7 +439,7 @@ class BatchClientHook(AwsBaseHook):
             return None
         if len(all_info) > 1:
             self.log.warning(
-                f"AWS Batch job ({job_id}) has more than one log stream, only returning the first one."
+                "AWS Batch job (%s) has more than one log stream, only returning the first one.", job_id
             )
         return all_info[0]
 
@@ -474,7 +475,7 @@ class BatchClientHook(AwsBaseHook):
         # If the user selected another logDriver than "awslogs", then CloudWatch logging is disabled.
         if any(c.get("logDriver", "awslogs") != "awslogs" for c in log_configs):
             self.log.warning(
-                f"AWS Batch job ({job_id}) uses non-aws log drivers. AWS CloudWatch logging disabled."
+                "AWS Batch job (%s) uses non-aws log drivers. AWS CloudWatch logging disabled.", job_id
             )
             return []
 
@@ -482,7 +483,7 @@ class BatchClientHook(AwsBaseHook):
             # If this method is called very early after starting the AWS Batch job,
             # there is a possibility that the AWS CloudWatch Stream Name would not exist yet.
             # This can also happen in case of misconfiguration.
-            self.log.warning(f"AWS Batch job ({job_id}) doesn't have any AWS CloudWatch Stream.")
+            self.log.warning("AWS Batch job (%s) doesn't have any AWS CloudWatch Stream.", job_id)
             return []
 
         # Try to get user-defined log configuration options

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

@@ -23,6 +23,7 @@ AWS Batch service waiters.
     - https://boto3.amazonaws.com/v1/documentation/api/latest/guide/clients.html#waiters
     - https://github.com/boto/botocore/blob/develop/botocore/waiter.py
 """
+
 from __future__ import annotations
 
 import json

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

@@ -0,0 +1,59 @@
+# 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.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+class BedrockHook(AwsBaseHook):
+    """
+    Interact with Amazon Bedrock.
+
+    Provide thin wrapper around :external+boto3:py:class:`boto3.client("bedrock") <Bedrock.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`
+    """
+
+    client_type = "bedrock"
+
+    def __init__(self, *args, **kwargs) -> None:
+        kwargs["client_type"] = self.client_type
+        super().__init__(*args, **kwargs)
+
+
+class BedrockRuntimeHook(AwsBaseHook):
+    """
+    Interact with the Amazon Bedrock Runtime.
+
+    Provide thin wrapper around :external+boto3:py:class:`boto3.client("bedrock-runtime") <BedrockRuntime.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`
+    """
+
+    client_type = "bedrock-runtime"
+
+    def __init__(self, *args, **kwargs) -> None:
+        kwargs["client_type"] = self.client_type
+        super().__init__(*args, **kwargs)

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

@@ -17,6 +17,7 @@
 # under the License.
 
 """This module contains a web hook for Chime."""
+
 from __future__ import annotations
 
 import json

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains AWS CloudFormation Hook."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING

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

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Interact with AWS DataSync, using the AWS ``boto3`` library."""
+
 from __future__ import annotations
 
 import time

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains the Amazon DynamoDB Hook."""
+
 from __future__ import annotations
 
 from typing import Iterable

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

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Interact with Amazon EKS, using the boto3 library."""
+
 from __future__ import annotations
 
 import base64

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

@@ -19,12 +19,13 @@ from __future__ import annotations
 
 import asyncio
 import time
+from functools import cached_property
 
-import boto3
 from botocore.exceptions import ClientError
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
 
 DEFAULT_LOG_SUFFIX = "output"
 ERROR_LOG_SUFFIX = "error"
@@ -213,6 +214,13 @@ class GlueJobHook(AwsBaseHook):
             job_run = await client.get_job_run(JobName=job_name, RunId=run_id)
         return job_run["JobRun"]["JobRunState"]
 
+    @cached_property
+    def logs_hook(self):
+        """Returns an AwsLogsHook instantiated with the parameters of the GlueJobHook."""
+        return AwsLogsHook(
+            aws_conn_id=self.aws_conn_id, region_name=self.region_name, verify=self.verify, config=self.config
+        )
+
     def print_job_logs(
         self,
         job_name: str,
@@ -225,7 +233,7 @@ class GlueJobHook(AwsBaseHook):
         :param continuation_tokens: the tokens where to resume from when reading logs.
             The object gets updated with the new tokens by this method.
         """
-        log_client = boto3.client("logs")
+        log_client = self.logs_hook.get_conn()
         paginator = log_client.get_paginator("filter_log_events")
 
         def display_logs_from(log_group: str, continuation_token: str | None) -> str | None:
@@ -245,8 +253,9 @@ class GlueJobHook(AwsBaseHook):
                 if e.response["Error"]["Code"] == "ResourceNotFoundException":
                     # we land here when the log groups/streams don't exist yet
                     self.log.warning(
-                        "No new Glue driver logs so far.\nIf this persists, check the CloudWatch dashboard "
-                        f"at: https://{self.conn_region_name}.console.aws.amazon.com/cloudwatch/home"
+                        "No new Glue driver logs so far.\n"
+                        "If this persists, check the CloudWatch dashboard at: %r.",
+                        f"https://{self.conn_region_name}.console.aws.amazon.com/cloudwatch/home",
                     )
                 else:
                     raise
@@ -263,7 +272,6 @@ class GlueJobHook(AwsBaseHook):
         log_group_prefix = self.conn.get_job_run(JobName=job_name, RunId=run_id)["JobRun"]["LogGroupName"]
         log_group_default = f"{log_group_prefix}/{DEFAULT_LOG_SUFFIX}"
         log_group_error = f"{log_group_prefix}/{ERROR_LOG_SUFFIX}"
-
         # one would think that the error log group would contain only errors, but it actually contains
         # a lot of interesting logs too, so it's valuable to have both
         continuation_tokens.output_stream_continuation = display_logs_from(

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains AWS Glue Catalog Hook."""
+
 from __future__ import annotations
 
 from typing import Any