Merge branch 'apache:main' into teradata_release_1.0.0

Author: satish-chinthanippu

airflow/providers/amazon/aws/log/cloudwatch_task_handler.py

@@ -17,7 +17,7 @@
 # under the License.
 from __future__ import annotations
 
-from datetime import date, datetime, timedelta
+from datetime import date, datetime, timedelta, timezone
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
@@ -163,7 +163,7 @@ def get_cloudwatch_logs(self, stream_name: str, task_instance: TaskInstance) ->
         return "\n".join(self._event_to_str(event) for event in events)
 
     def _event_to_str(self, event: dict) -> str:
-        event_dt = datetime.utcfromtimestamp(event["timestamp"] / 1000.0)
+        event_dt = datetime.fromtimestamp(event["timestamp"] / 1000.0, tz=timezone.utc)
         formatted_event_dt = event_dt.strftime("%Y-%m-%d %H:%M:%S,%f")[:-3]
         message = event["message"]
         return f"[{formatted_event_dt}] {message}"

airflow/providers/amazon/aws/utils/task_log_fetcher.py

@@ -18,7 +18,7 @@
 from __future__ import annotations
 
 import time
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from threading import Event, Thread
 from typing import TYPE_CHECKING, Generator
 
@@ -87,7 +87,7 @@ def _get_log_events(self, skip_token: AwsLogsHook.ContinuationToken | None = Non
 
     @staticmethod
     def event_to_str(event: dict) -> str:
-        event_dt = datetime.utcfromtimestamp(event["timestamp"] / 1000.0)
+        event_dt = datetime.fromtimestamp(event["timestamp"] / 1000.0, tz=timezone.utc)
         formatted_event_dt = event_dt.strftime("%Y-%m-%d %H:%M:%S,%f")[:-3]
         message = event["message"]
         return f"[{formatted_event_dt}] {message}"

airflow/providers/cncf/kubernetes/hooks/kubernetes.py

@@ -37,7 +37,7 @@
 from airflow.utils import yaml
 
 if TYPE_CHECKING:
-    from kubernetes.client.models import V1Deployment, V1Pod
+    from kubernetes.client.models import V1Deployment, V1Job, V1Pod
 
 LOADING_KUBE_CONFIG_FILE_RESOURCE = "Loading Kubernetes configuration file kube_config from {}..."
 
@@ -290,6 +290,10 @@ def apps_v1_client(self) -> client.AppsV1Api:
     def custom_object_client(self) -> client.CustomObjectsApi:
         return client.CustomObjectsApi(api_client=self.api_client)
 
+    @cached_property
+    def batch_v1_client(self) -> client.BatchV1Api:
+        return client.BatchV1Api(api_client=self.api_client)
+
     def create_custom_object(
         self, group: str, version: str, plural: str, body: str | dict, namespace: str | None = None
     ):
@@ -472,6 +476,32 @@ def get_deployment_status(
         except Exception as exc:
             raise exc
 
+    def create_job(
+        self,
+        job: V1Job,
+        **kwargs,
+    ) -> V1Job:
+        """
+        Run Job.
+
+        :param job: A kubernetes Job object
+        """
+        sanitized_job = self.batch_v1_client.api_client.sanitize_for_serialization(job)
+        json_job = json.dumps(sanitized_job, indent=2)
+
+        self.log.debug("Job Creation Request: \n%s", json_job)
+        try:
+            resp = self.batch_v1_client.create_namespaced_job(
+                body=sanitized_job, namespace=job.metadata.namespace, **kwargs
+            )
+            self.log.debug("Job Creation Response: %s", resp)
+        except Exception as e:
+            self.log.exception(
+                "Exception when attempting to create Namespaced Job: %s", str(json_job).replace("\n", " ")
+            )
+            raise e
+        return resp
+
 
 def _get_bool(val) -> bool | None:
     """Convert val to bool if can be done with certainty; if we cannot infer intention we return None."""

airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py

@@ -19,13 +19,15 @@
 import logging
 import secrets
 import string
+import warnings
 from typing import TYPE_CHECKING
 
 import pendulum
 from slugify import slugify
 
 from airflow.compat.functools import cache
 from airflow.configuration import conf
+from airflow.exceptions import AirflowProviderDeprecationWarning
 
 if TYPE_CHECKING:
     from airflow.models.taskinstancekey import TaskInstanceKey
@@ -45,6 +47,18 @@ def rand_str(num):
     return "".join(secrets.choice(alphanum_lower) for _ in range(num))
 
 
+def add_unique_suffix(*, name: str, rand_len: int = 8, max_len: int = POD_NAME_MAX_LENGTH) -> str:
+    """Add random string to pod or job name while staying under max length.
+
+    :param name: name of the pod or job
+    :param rand_len: length of the random string to append
+    :param max_len: maximum length of the pod name
+    :meta private:
+    """
+    suffix = "-" + rand_str(rand_len)
+    return name[: max_len - len(suffix)].strip("-.") + suffix
+
+
 def add_pod_suffix(*, pod_name: str, rand_len: int = 8, max_len: int = POD_NAME_MAX_LENGTH) -> str:
     """Add random string to pod name while staying under max length.
 
@@ -53,10 +67,48 @@ def add_pod_suffix(*, pod_name: str, rand_len: int = 8, max_len: int = POD_NAME_
     :param max_len: maximum length of the pod name
     :meta private:
     """
+    warnings.warn(
+        "This function is deprecated. Please use `add_unique_suffix`.",
+        AirflowProviderDeprecationWarning,
+        stacklevel=2,
+    )
+
     suffix = "-" + rand_str(rand_len)
     return pod_name[: max_len - len(suffix)].strip("-.") + suffix
 
 
+def create_unique_id(
+    dag_id: str | None = None,
+    task_id: str | None = None,
+    *,
+    max_length: int = POD_NAME_MAX_LENGTH,
+    unique: bool = True,
+) -> str:
+    """
+    Generate unique pod or job ID given a dag_id and / or task_id.
+
+    :param dag_id: DAG ID
+    :param task_id: Task ID
+    :param max_length: max number of characters
+    :param unique: whether a random string suffix should be added
+    :return: A valid identifier for a kubernetes pod name
+    """
+    if not (dag_id or task_id):
+        raise ValueError("Must supply either dag_id or task_id.")
+    name = ""
+    if dag_id:
+        name += dag_id
+    if task_id:
+        if name:
+            name += "-"
+        name += task_id
+    base_name = slugify(name, lowercase=True)[:max_length].strip(".-")
+    if unique:
+        return add_pod_suffix(pod_name=base_name, rand_len=8, max_len=max_length)
+    else:
+        return base_name
+
+
 def create_pod_id(
     dag_id: str | None = None,
     task_id: str | None = None,
@@ -73,6 +125,12 @@ def create_pod_id(
     :param unique: whether a random string suffix should be added
     :return: A valid identifier for a kubernetes pod name
     """
+    warnings.warn(
+        "This function is deprecated. Please use `create_unique_id`.",
+        AirflowProviderDeprecationWarning,
+        stacklevel=2,
+    )
+
     if not (dag_id or task_id):
         raise ValueError("Must supply either dag_id or task_id.")
     name = ""