Migrate to Astro Private Cloud 1.0 from 0.37

# This Pod mutation hook runs for all Pods dynamically created by Airflow in Kubernetes (K8s).
This includes K8s executor Airflow workers, which is an alternative executor to Celery, and K8s Pod Operator (KPO) Pods, which is an Airflow operator that launches a task in a K8s cluster.
This function is responsible for two things. It assigns labels to disambiguate between KPO and K8s executor Pods and it handles writing the entrypoint for K8s executor Pods.
# This function assigns labels to disambiguate between KPO and K8s executor Pods.
import warnings
from airflow.configuration import conf
from airflow.version import version
xcom_sidecar_image = "quay.io/astronomer/ap-init:3.21.3-5"
try:
    from airflow.providers.cncf.kubernetes.utils.xcom_sidecar import PodDefaults
    PodDefaults.SIDECAR_CONTAINER.image = xcom_sidecar_image
except ImportError:
    warnings.warn(
        """
        Airflow CNCF Provider not installed by default.
        This will bypass custom images for xcom and defaults to alpine.
        Refer https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/stable/index.html
        to install the provider.
        """
    )
def pod_mutation_hook_deprecated(pod):
    from airflow.contrib.kubernetes.pod import Pod
    extra_labels = {
        "kubernetes-pod-operator": "False"
    }
    if 'airflow-worker' in pod.labels.keys() and \
            conf.get('core', 'EXECUTOR') == "KubernetesExecutor":
        # extra_labels["kubernetes-executor"] = "True"
        # By default, Airflow overwrites the entrypoint
        # of KPO and K8s executor Airflow-worker pods.
        # K8s worker pods are Airflow containers, and we can assume these
        # Airflow containers are vendored by Astronomer. Astronomer provides
        # an entrypoint in these containers to handle waiting for the network
        # and the database to be ready, so we do not want this entrypoint
        # to be overwritten. This helps with the stability of the K8s executor.
        if "1.10" in version:
            if not pod.args:
                pod.args = []
            pod.args = pod.cmds + pod.args
        pod.cmds = ["tini", "--", "/entrypoint"]
    else:
        # In the case of KPO, we allow Airflow to overwrite the entrypoint
        # because we do not know what the container will be (and it's probably
        # not a container vendored by Astronomer, and it's definitely not
        # an Airflow container).
        extra_labels["kubernetes-pod-operator"] = "True"
        extra_labels["kubernetes-executor"] = "False"
    pod.labels.update(extra_labels)
def pod_mutation_hook_new(pod):
    extra_labels = {
        "kubernetes-pod-operator": "False"
    }
    if 'airflow-worker' in pod.metadata.labels.keys() and \
            conf.get('core', 'EXECUTOR') == "KubernetesExecutor":
        # extra_labels["kubernetes-executor"] = "True"
        # By default, Airflow overwrites the entrypoint
        # of KPO and K8s executor Airflow-worker pods.
        # K8s worker pods are Airflow containers, and we can assume these
        # Airflow containers are vendored by Astronomer. Astronomer provides
        # an entrypoint in these containers to handle waiting for the network
        # and the database to be ready, so we do not want this entrypoint
        # to be overwritten. This helps with the stability of the K8s executor.
        container = pod.spec.containers[0]
        if "1.10" in version:
            if not container.args:
                container.args = []
            container.args = container.command + container.args
        container.command = ["tini", "--", "/entrypoint"]
        pod.spec.containers[0] = container
    else:
        # In the case of KPO, we allow Airflow to overwrite the entrypoint
        # because we do not know what the container will be (and it's probably
        # not a container vendored by Astronomer, and it's definitely not
        # an Airflow container).
        extra_labels["kubernetes-pod-operator"] = "True"
        extra_labels["kubernetes-executor"] = "False"
    pod.metadata.labels.update(extra_labels)
def pod_mutation_hook(pod):
    try:
        pod_mutation_hook_new(pod)
    except Exception:
        pod_mutation_hook_deprecated(pod)