com.instaclustr.cassandra.service.kubernetes.KubernetesCqlSession Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of commons Show documentation
Common classes and utilities integrated with various projects
There is a newer version: 2.0.0
package com.instaclustr.cassandra.service.kubernetes;

import static java.util.Objects.requireNonNull;

import java.io.BufferedInputStream;
import java.io.FileInputStream;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.UnknownHostException;
import java.util.Properties;
import java.util.function.Supplier;

import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.CqlSessionBuilder;
import com.datastax.oss.driver.api.core.config.DriverConfig;
import com.datastax.oss.driver.api.core.config.DriverConfigLoader;
import com.datastax.oss.driver.internal.core.config.typesafe.DefaultDriverConfigLoader;
import com.datastax.oss.driver.internal.core.config.typesafe.TypesafeDriverConfig;
import com.google.inject.Inject;
import com.instaclustr.cassandra.service.CassandraConfigReader;
import com.instaclustr.cassandra.service.CqlSessionService;
import com.instaclustr.kubernetes.KubernetesHelper;
import com.instaclustr.operations.FunctionWithEx;
import com.typesafe.config.Config;
import com.typesafe.config.ConfigFactory;
import edu.umd.cs.findbugs.annotations.NonNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class KubernetesCqlSession implements CqlSessionService {

    private static final Logger logger = LoggerFactory.getLogger(KubernetesCqlSession.class);

    private final DriverConfigLoader loader;

    @Inject
    public KubernetesCqlSession(final @Kubernetes DriverConfigLoader loader) {
        this.loader = loader;
    }

    @Override
    public CqlSession getCqlSession() {
        return createCqlSession();
    }

    @Override
    public  T doWithCqlSession(FunctionWithEx func) throws Exception {
        try (final CqlSession session = createCqlSession()) {
            return func.apply(session);
        }
    }

    protected CqlSession createCqlSession() {

        final CqlSessionBuilder builder = CqlSession.builder().withConfigLoader(loader);

        if (KubernetesHelper.isRunningInKubernetes()) {
            try {
                builder.addContactPoint(new InetSocketAddress(InetAddress.getLocalHost().getHostName(), 9042));
                builder.withLocalDatacenter(CassandraKubernetesHelper.getLocalDataCenter());

                return builder.build();
            } catch (final UnknownHostException ex) {
                throw new IllegalStateException("Unable to resolve hostname of the local host.", ex);
            }
        } else {
            return builder.build();
        }
    }

    public static final class KubernetesDriverConfigLoader extends DefaultDriverConfigLoader {

        private final Supplier configSupplier;

        @Inject
        public KubernetesDriverConfigLoader(final @Kubernetes CassandraConfigReader cassandraConfigReader) {
            this.configSupplier = () -> {
                ConfigFactory.invalidateCaches();
                Config config =
                    ConfigFactory.defaultOverrides()
                        .withFallback(ConfigFactory.parseReader(requireNonNull(cassandraConfigReader.getSecretReader())))
                        .withFallback(ConfigFactory.parseReader(requireNonNull(cassandraConfigReader.getConfigMapReader())))
                        .withFallback(ConfigFactory.defaultReference())
                        .resolve();
                return config.getConfig(DefaultDriverConfigLoader.DEFAULT_ROOT_PATH);
            };
        }

        @NonNull
        @Override
        public DriverConfig getInitialConfig() {
            return new TypesafeDriverConfig(getConfigSupplier().get());
        }

        @NonNull
        @Override
        public Supplier getConfigSupplier() {
            return configSupplier;
        }
    }

}