com.codacy.scoobydoo.kubernetes.KubernetesWrapper Maven / Gradle / Ivy

Go to download
package com.codacy.scoobydoo.kubernetes;

import com.codacy.scoobydoo.LoggingHelper;
import com.codacy.scoobydoo.kubernetes.handlers.FindPodByLabelHandler;
import com.codacy.scoobydoo.kubernetes.handlers.FindPodByNameHandler;
import io.fabric8.kubernetes.api.model.LabelSelector;
import io.fabric8.kubernetes.api.model.LabelSelectorBuilder;
import io.fabric8.kubernetes.api.model.Node;
import io.fabric8.kubernetes.api.model.Pod;
import io.fabric8.kubernetes.client.KubernetesClient;
import io.fabric8.kubernetes.client.LocalPortForward;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.ResponseBody;

import java.io.IOException;
import java.util.Collections;
import java.util.Hashtable;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.concurrent.Callable;
import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import static org.awaitility.Awaitility.await;

public class KubernetesWrapper {

    public static final String PORT_FORWARD_BASE_URL = "http://127.0.0.1:";
    public static final String POD_FOUND_MESSAGE = "Pod found: ";
    public static final String POD_NOT_FOUND_MESSAGE = "Pod not found: ";
    public static final String POD_WITH_LABEL_NOT_FOUND_MESSAGE = "No pod found for label: ";
    public static final String NODE_WITH_LABEL_NOT_FOUND_MESSAGE = "No node found for label: ";
    public static final String POD_NOT_FOUND_WITH_REGEX_MESSAGE = "No Pod found for regex: ";
    public static final String TOO_MANY_PODS_FOUND_FOR_REGEX_MESSAGE = "Too many pods found for regex: ";
    public static final String PORT_FORWARDED_MESSAGE = "Port forwarded for ";
    public static final String NEW_PORT_FORWARD_LIST_MESSAGE = "New Port-Forward list for ";
    public static final String KEY_NOT_FOUND_IN_CONFIGMAP = "Configmap doesn't contain key: ";
    public static final int KUBERNETES_API_TIMEOUT_IN_SECONDS = 60;
    public static final int KUBERNETES_API_POLLING_IN_SECONDS = 5;

    protected KubernetesClient kubernetesClient;
    protected Map> openPortForwards;

    public KubernetesWrapper(KubernetesClient kubernetesClient) {
        this.openPortForwards = new Hashtable<>();
        this.kubernetesClient = kubernetesClient;
    }

    public Pod openPortForward(Pattern podNameRegex, String namespace, int localPort, int targetPort, boolean testPortForward) {
        Pod pod = findPod(podNameRegex, namespace);
        openPortForward(localPort, targetPort, namespace, pod, testPortForward);
        return pod;
    }

    public Pod openPortForward(String podName, String namespace, int localPort, int targetPort, boolean testPortForward) {
        Pod pod = findPod(podName, namespace);
        openPortForward(localPort, targetPort, namespace, pod, testPortForward);
        return pod;
    }

    public Pod openPortForward(String labelKey, String labelValue, String namespace, int localPort, int targetPort, boolean testPortForward) {
        Pod pod = findPod(labelKey, labelValue, namespace);
        openPortForward(localPort, targetPort, namespace, pod, testPortForward);
        return pod;
    }

    public Pod findPod(String podName, String namespace) {
        Callable handler = new FindPodByNameHandler(kubernetesClient, podName, namespace);
        syncWaitForHandler(handler, String.format("Could not find pod by name [%s].",podName));

        Pod pod = kubernetesClient
                .pods()
                .inNamespace(namespace)
                .withName(podName).get();
        if (pod == null) {
            handleNoSuchElement(POD_NOT_FOUND_MESSAGE, podName);
        } else {
            LoggingHelper.info(POD_FOUND_MESSAGE + podName);
        }

        return pod;
    }

    private void syncWaitForHandler(Callable handler, String errorMessage) {
        try {
            await()
                    .atMost(KUBERNETES_API_TIMEOUT_IN_SECONDS, TimeUnit.SECONDS)
                    .with()
                    .pollInterval(KUBERNETES_API_POLLING_IN_SECONDS, TimeUnit.SECONDS)
                    .until(handler);
        } catch (AssertionError e) {
            LoggingHelper.error(errorMessage, e);
            throw e;
        }
    }

    private void handleNoSuchElement(String message, String element) {
        final String errorMessage = String.format("%s %s", message, element);
        NoSuchElementException exception = new NoSuchElementException(errorMessage);
        LoggingHelper.error(errorMessage, exception);
        throw exception;
    }

    public Pod findPod(String labelKey, String labelValue, String namespace) {
        String labelAsString = labelKey + "=" + labelValue;
        Callable handler = new FindPodByLabelHandler(kubernetesClient, labelKey, labelValue, namespace);

        syncWaitForHandler(handler, String.format("Could not find pod by label [%s].",labelAsString));

        List podList = kubernetesClient
                .pods()
                .inNamespace(namespace)
                .withLabel(labelKey, labelValue)
                .list()
                .getItems();

        if (podList.isEmpty()) {
            handleNoSuchElement(POD_WITH_LABEL_NOT_FOUND_MESSAGE, labelAsString);
        }

        Pod pod = podList.get(0);
        String parsedPodName = pod.getMetadata().getName();
        LoggingHelper.info(POD_FOUND_MESSAGE + parsedPodName);
        return pod;
    }

    public Pod findPod(Pattern podNamePattern, String namespace) {
        Pod foundPod = null;
        for (Pod pod : kubernetesClient.pods().inNamespace(namespace).list().getItems()) {
            String podName = pod.getMetadata().getName();
            Matcher matcher = podNamePattern.matcher(podName);
            if (matcher.find()) {
                LoggingHelper.info("Found pod [" + podName +"].");
                if (foundPod != null) {
                    handleNoSuchElement(TOO_MANY_PODS_FOUND_FOR_REGEX_MESSAGE, podNamePattern.toString());
                }

                foundPod = pod;
                foundPod.getSpec().getNodeName();
            }
        }
        if (foundPod == null) {
            handleNoSuchElement(POD_NOT_FOUND_WITH_REGEX_MESSAGE, podNamePattern.toString());
        }
        return foundPod;
    }

    public List findNode(String nodeLabelKey, String nodeLabelValue) {
        String labelAsString = nodeLabelKey + "=" + nodeLabelValue;
        LabelSelector selector = new LabelSelectorBuilder().withMatchLabels(Collections.singletonMap(nodeLabelKey, nodeLabelValue)).build();
        List nodeList = kubernetesClient.nodes().withLabelSelector(selector).list().getItems();

        if (nodeList.isEmpty()) {
            handleNoSuchElement(NODE_WITH_LABEL_NOT_FOUND_MESSAGE, labelAsString);
        }

        return nodeList;
    }

    private LocalPortForward openPortForward(int localPort, int targetPort, String namespace, Pod pod, Boolean testPortForward) {

        if (pod == null) {
            return null;
        } else {

            String parsedPodName = pod.getMetadata().getName();
            if (openPortForwards.get(pod) == null) {
                openPortForwards.put(pod, new Hashtable<>());
                LoggingHelper.info(NEW_PORT_FORWARD_LIST_MESSAGE + parsedPodName);
            }

            LocalPortForward portForward = kubernetesClient
                    .pods()
                    .inNamespace(namespace)
                    .withName(parsedPodName)
                    .portForward(targetPort, localPort);

            openPortForwards.get(pod).put(localPort, portForward);
            LoggingHelper.info(PORT_FORWARDED_MESSAGE + parsedPodName + " " + PORT_FORWARD_BASE_URL + portForward.getLocalPort());

            // NOTE: this should always be the last step of the method and should not be inside any block
            // (except for the pod != null one) to ensure it works as expected.
            if (testPortForward) {
                testPortForward(portForward);
            }
            return portForward;
        }
    }

    private void testPortForward(LocalPortForward portForward) {
        try {
            await()
                    .atMost(KUBERNETES_API_TIMEOUT_IN_SECONDS, TimeUnit.SECONDS)
                    .with()
                    .pollInterval(KUBERNETES_API_POLLING_IN_SECONDS, TimeUnit.SECONDS)
                    .until(() -> testPortForwardCallback(portForward));
        } catch (AssertionError e) {
            LoggingHelper.error("Failed portForward is not alive yet.", e);
            throw e;
        }
    }

    protected boolean testPortForwardCallback(LocalPortForward portForward) {
        try {
            LoggingHelper.info("Check forwarded port.");
            final ResponseBody responseBody = new OkHttpClient()
                    .newCall(new Request.Builder().get().url(PORT_FORWARD_BASE_URL + portForward.getLocalPort()).build()).execute()
                    .body();
            LoggingHelper.info(String.format("Response: %s", responseBody != null ? responseBody.string() : "[Empty Body]"));
            return true;
        } catch (Exception e) {
            LoggingHelper.error("Failed to check forwarded port.", e);
        }
        return false;
    }

    public void closePortForward(Pod pod, int localPort) throws Throwable {
        LocalPortForward localPortForward = this.openPortForwards.get(pod).get(localPort);
        localPortForward.close();
        this.openPortForwards.get(pod).remove(localPort);
        LoggingHelper.info(String.format("Close forwarded port [%s] on [%s].", localPort, pod.getMetadata().getName()));
    }

    public String getConfigMapValue(String key, String configMapName, String namespace) {
        Map configMapData = kubernetesClient.configMaps().inNamespace(namespace).withName(configMapName).get().getData();
        String value = configMapData.get(key);
        if (value == null) {
            handleNoSuchElement(KEY_NOT_FOUND_IN_CONFIGMAP, key);
        }
        return value;
    }
}