<commit_before># 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.
import tempfile
from typing import Generator, Optional, Tuple, Union

import yaml
from kubernetes import client, config, watch

from airflow.exceptions import AirflowException
from airflow.hooks.base_hook import BaseHook


def _load_body_to_dict(body):
    try:
        body_dict = yaml.safe_load(body)
    except yaml.YAMLError as e:
        raise AirflowException("Exception when loading resource definition: %s\n" % e)
    return body_dict


class KubernetesHook(BaseHook):
    """
    Creates Kubernetes API connection.

    :param conn_id: the connection to Kubernetes cluster
    :type conn_id: str
    """

    def __init__(self, conn_id: str = "kubernetes_default"):
        super().__init__()
        self.conn_id = conn_id

    def get_conn(self):
        """
        Returns kubernetes api session for use with requests
        """
        connection = self.get_connection(self.conn_id)
        extras = connection.extra_dejson
        if extras.get("extra__kubernetes__in_cluster"):
            self.log.debug("loading kube_config from: in_cluster configuration")
            config.load_incluster_config()
        elif extras.get("extra__kubernetes__kube_config") is None:
            self.log.debug("loading kube_config from: default file")
            config.load_kube_config()
        else:
            with tempfile.NamedTemporaryFile() as temp_config:
                self.log.debug("loading kube_config from: connection kube_config")
                temp_config.write(extras.get("extra__kubernetes__kube_config").encode())
                temp_config.flush()
                config.load_kube_config(temp_config.name)
        return client.ApiClient()

    def create_custom_resource_definition(
        self, group: str, version: str, plural: str, body: Union[str, dict], namespace: Optional[str] = None
    ):
        """
        Creates custom resource definition object in Kubernetes

        :param group: api group
        :type group: str
        :param version: api version
        :type version: str
        :param plural: api plural
        :type plural: str
        :param body: crd object definition
        :type body: Union[str, dict]
        :param namespace: kubernetes namespace
        :type namespace: str
        """
        api = client.CustomObjectsApi(self.get_conn())
        if namespace is None:
            namespace = self.get_namespace()
        if isinstance(body, str):
            body = _load_body_to_dict(body)
        try:
            response = api.create_namespaced_custom_object(
                group=group, version=version, namespace=namespace, plural=plural, body=body
            )
            self.log.debug("Response: %s", response)
            return response
        except client.rest.ApiException as e:
            raise AirflowException("Exception when calling -> create_custom_resource_definition: %s\n" % e)

    def get_custom_resource_definition(
        self, group: str, version: str, plural: str, name: str, namespace: Optional[str] = None
    ):
        """
        Get custom resource definition object from Kubernetes

        :param group: api group
        :type group: str
        :param version: api version
        :type version: str
        :param plural: api plural
        :type plural: str
        :param name: crd object name
        :type name: str
        :param namespace: kubernetes namespace
        :type namespace: str
        """
        custom_resource_definition_api = client.CustomObjectsApi(self.get_conn())
        if namespace is None:
            namespace = self.get_namespace()
        try:
            response = custom_resource_definition_api.get_namespaced_custom_object(
                group=group, version=version, namespace=namespace, plural=plural, name=name
            )
            return response
        except client.rest.ApiException as e:
            raise AirflowException("Exception when calling -> get_custom_resource_definition: %s\n" % e)

    def get_namespace(self):
        """
        Returns the namespace that defined in the connection
        """
        connection = self.get_connection(self.conn_id)
        extras = connection.extra_dejson
        namespace = extras.get("extra__kubernetes__namespace", "default")
        return namespace

    def get_pod_log_stream(
        self,
        pod_name: str,
        container: Optional[str] = "",
        namespace: Optional[str] = None,
    ) -> Tuple[watch.Watch, Generator[str, None, None]]:
        """
        Retrieves a log stream for a container in a kubernetes pod.

        :param pod_name: pod name
        :type pod_name: str
        :param container: container name
        :type version: str
        :param namespace: kubernetes namespace
        :type namespace: str
        """

        api = client.CoreV1Api(self.get_conn())
        watcher = watch.Watch()
        return (
            watcher,
            watcher.stream(
                api.read_namespaced_pod_log,
                name=pod_name,
                container=container,
                namespace=namespace if namespace else self.get_namespace(),
            ),
        )

    def get_pod_logs(
        self,
        pod_name: str,
        container: Optional[str] = "",
        namespace: Optional[str] = None,
    ):
        """
        Retrieves a container's log from the specified pod.

        :param pod_name: pod name
        :type pod_name: str
        :param container: container name
        :type version: str
        :param namespace: kubernetes namespace
        :type namespace: str
        """
        api = client.CoreV1Api(self.get_conn())
        return api.read_namespaced_pod_log(
            name=pod_name,
            container=container,
            _preload_content=False,
            namespace=namespace if namespace else self.get_namespace(),
        )<commit_msg>Add connection caching to KubernetesHook (#10447)<commit_after># 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.
import tempfile
from typing import Generator, Optional, Tuple, Union

import yaml
from cached_property import cached_property

from kubernetes import client, config, watch

from airflow.exceptions import AirflowException
from airflow.hooks.base_hook import BaseHook


def _load_body_to_dict(body):
    try:
        body_dict = yaml.safe_load(body)
    except yaml.YAMLError as e:
        raise AirflowException("Exception when loading resource definition: %s\n" % e)
    return body_dict


class KubernetesHook(BaseHook):
    """
    Creates Kubernetes API connection.

    :param conn_id: the connection to Kubernetes cluster
    :type conn_id: str
    """

    def __init__(self, conn_id: str = "kubernetes_default"):
        super().__init__()
        self.conn_id = conn_id

    def get_conn(self):
        """
        Returns kubernetes api session for use with requests
        """
        connection = self.get_connection(self.conn_id)
        extras = connection.extra_dejson
        if extras.get("extra__kubernetes__in_cluster"):
            self.log.debug("loading kube_config from: in_cluster configuration")
            config.load_incluster_config()
        elif extras.get("extra__kubernetes__kube_config") is None:
            self.log.debug("loading kube_config from: default file")
            config.load_kube_config()
        else:
            with tempfile.NamedTemporaryFile() as temp_config:
                self.log.debug("loading kube_config from: connection kube_config")
                temp_config.write(extras.get("extra__kubernetes__kube_config").encode())
                temp_config.flush()
                config.load_kube_config(temp_config.name)
        return client.ApiClient()

    @cached_property
    def api_client(self):
        """Cached Kubernetes API client"""
        return self.get_conn()

    def create_custom_resource_definition(
        self, group: str, version: str, plural: str, body: Union[str, dict], namespace: Optional[str] = None
    ):
        """
        Creates custom resource definition object in Kubernetes

        :param group: api group
        :type group: str
        :param version: api version
        :type version: str
        :param plural: api plural
        :type plural: str
        :param body: crd object definition
        :type body: Union[str, dict]
        :param namespace: kubernetes namespace
        :type namespace: str
        """
        api = client.CustomObjectsApi(self.api_client)
        if namespace is None:
            namespace = self.get_namespace()
        if isinstance(body, str):
            body = _load_body_to_dict(body)
        try:
            response = api.create_namespaced_custom_object(
                group=group, version=version, namespace=namespace, plural=plural, body=body
            )
            self.log.debug("Response: %s", response)
            return response
        except client.rest.ApiException as e:
            raise AirflowException("Exception when calling -> create_custom_resource_definition: %s\n" % e)

    def get_custom_resource_definition(
        self, group: str, version: str, plural: str, name: str, namespace: Optional[str] = None
    ):
        """
        Get custom resource definition object from Kubernetes

        :param group: api group
        :type group: str
        :param version: api version
        :type version: str
        :param plural: api plural
        :type plural: str
        :param name: crd object name
        :type name: str
        :param namespace: kubernetes namespace
        :type namespace: str
        """
        custom_resource_definition_api = client.CustomObjectsApi(self.api_client)
        if namespace is None:
            namespace = self.get_namespace()
        try:
            response = custom_resource_definition_api.get_namespaced_custom_object(
                group=group, version=version, namespace=namespace, plural=plural, name=name
            )
            return response
        except client.rest.ApiException as e:
            raise AirflowException("Exception when calling -> get_custom_resource_definition: %s\n" % e)

    def get_namespace(self):
        """
        Returns the namespace that defined in the connection
        """
        connection = self.get_connection(self.conn_id)
        extras = connection.extra_dejson
        namespace = extras.get("extra__kubernetes__namespace", "default")
        return namespace

    def get_pod_log_stream(
        self,
        pod_name: str,
        container: Optional[str] = "",
        namespace: Optional[str] = None,
    ) -> Tuple[watch.Watch, Generator[str, None, None]]:
        """
        Retrieves a log stream for a container in a kubernetes pod.

        :param pod_name: pod name
        :type pod_name: str
        :param container: container name
        :param namespace: kubernetes namespace
        :type namespace: str
        """

        api = client.CoreV1Api(self.api_client)
        watcher = watch.Watch()
        return (
            watcher,
            watcher.stream(
                api.read_namespaced_pod_log,
                name=pod_name,
                container=container,
                namespace=namespace if namespace else self.get_namespace(),
            ),
        )

    def get_pod_logs(
        self,
        pod_name: str,
        container: Optional[str] = "",
        namespace: Optional[str] = None,
    ):
        """
        Retrieves a container's log from the specified pod.

        :param pod_name: pod name
        :type pod_name: str
        :param container: container name
        :param namespace: kubernetes namespace
        :type namespace: str
        """
        api = client.CoreV1Api(self.api_client)
        return api.read_namespaced_pod_log(
            name=pod_name,
            container=container,
            _preload_content=False,
            namespace=namespace if namespace else self.get_namespace(),
        )
{'fraction_non_alphanumeric': 0.07859742961147018, 'fraction_numerical': 0.0012629076591635093, 'mean_word_length': 2.413286004056795, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5446449', 'n_tokens_mistral': 3545, 'n_tokens_neox': 3448, 'n_words': 1094}