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.util.*;
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 static final LoggingHelper loggingHelper = new LoggingHelper();


    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(String.format("%s%s", 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.getMessage());
            throw e;
        }
    }

    private void handleNoSuchElement(String message, String element) {
        NoSuchElementException noSuchElementException = new NoSuchElementException(String.format("%s %s", message, element));
        loggingHelper.error(String.format("%s%s", message, element), noSuchElementException.getMessage());
        throw noSuchElementException;
    }

    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(String.format("%s%s", 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 m = podNamePattern.matcher(podName);
            if (m.find()) {
                loggingHelper.info("Found " + 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) {
            String parsedPodName = pod.getMetadata().getName();
            if (openPortForwards.get(pod) == null) {
                openPortForwards.put(pod, new Hashtable<>());
                loggingHelper.info(String.format("%s%s", 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(String.format("%s%s %s:%s", PORT_FORWARDED_MESSAGE, parsedPodName, PORT_FORWARD_BASE_URL, portForward.getLocalPort()));

            if (testPortForward) {
                testPortForward(portForward);
            }
            return portForward;
        }
        return null;
    }

    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.getMessage());
            throw e;
        }
    }

    protected boolean testPortForwardCallback(LocalPortForward portForward) {
        try {
            loggingHelper.info("Checking 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: \n%s", responseBody != null ? responseBody.string() : "[Empty Body]"));
            return true;
        } catch (Exception e) {
            loggingHelper.error("Exception occurred: " + e.getMessage(), e.getMessage());
        }
        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("Closed 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;
    }
}