Use AwsLogsHook when fetching Glue job logs (#38010)

Author: VincentChantreau

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

@@ -19,12 +19,13 @@
 
 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 @@ async def async_get_job_state(self, job_name: str, run_id: str) -> str:
             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 @@ def print_job_logs(
         :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:
@@ -264,7 +272,6 @@ def display_logs_from(log_group: str, continuation_token: str | None) -> str | N
         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(

tests/providers/amazon/aws/hooks/test_glue.py

@@ -29,6 +29,7 @@
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.hooks.glue import GlueJobHook
+from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
 
 if TYPE_CHECKING:
     from unittest.mock import MagicMock
@@ -376,12 +377,12 @@ def test_initialize_job(self, mock_conn, mock_get_job_state):
         glue_job_run_state = glue_job_hook.get_job_state(glue_job_run["JobName"], glue_job_run["JobRunId"])
         assert glue_job_run_state == mock_job_run_state, "Mocks but be equal"
 
-    @mock.patch("airflow.providers.amazon.aws.hooks.glue.boto3.client")
+    @mock.patch.object(AwsLogsHook, "get_conn")
     @mock.patch.object(GlueJobHook, "conn")
-    def test_print_job_logs_returns_token(self, conn_mock: MagicMock, client_mock: MagicMock, caplog):
-        hook = GlueJobHook()
+    def test_print_job_logs_returns_token(self, conn_mock: MagicMock, log_client_mock: MagicMock, caplog):
+        hook = GlueJobHook(job_name="test")
         conn_mock().get_job_run.return_value = {"JobRun": {"LogGroupName": "my_log_group"}}
-        client_mock().get_paginator().paginate.return_value = [
+        log_client_mock().get_paginator().paginate.return_value = [
             # first response : 2 log lines
             {
                 "events": [
@@ -399,12 +400,11 @@ def test_print_job_logs_returns_token(self, conn_mock: MagicMock, client_mock: M
         tokens = GlueJobHook.LogContinuationTokens()
         with caplog.at_level("INFO"):
             hook.print_job_logs("name", "run", tokens)
-
         assert "\thello\n\tworld\n" in caplog.text
         assert tokens.output_stream_continuation == "my_continuation_token"
         assert tokens.error_stream_continuation == "my_continuation_token"
 
-    @mock.patch("airflow.providers.amazon.aws.hooks.glue.boto3.client")
+    @mock.patch.object(AwsLogsHook, "get_conn")
     @mock.patch.object(GlueJobHook, "conn")
     def test_print_job_logs_no_stream_yet(self, conn_mock: MagicMock, client_mock: MagicMock):
         hook = GlueJobHook()