io.contek.tusk.yaml.TuskLoader Maven / Gradle / Ivy

Go to download
package io.contek.tusk.yaml;

import io.contek.tusk.MetricClient;
import io.contek.tusk.Table;
import io.contek.tusk.Tusk;
import io.contek.tusk.heartbeat.TuskHeartbeat;
import io.contek.tusk.heartbeat.TuskHeartbeatConfig;
import io.contek.tusk.log4j.TuskLogging;
import io.contek.tusk.log4j.TuskLoggingConfig;
import org.yaml.snakeyaml.Yaml;

import javax.annotation.Nullable;
import javax.annotation.concurrent.Immutable;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.time.Duration;
import java.util.logging.Logger;

import static java.lang.String.format;
import static java.util.Objects.requireNonNull;
import static java.util.logging.Level.SEVERE;

@Immutable
public final class TuskLoader {

  private static final Logger LOGGER = Logger.getLogger(TuskLoader.class.getName());

  public static void loadNullable(@Nullable Path tuskYmlFile) {
    if (tuskYmlFile == null) {
      return;
    }

    try (InputStream input = Files.newInputStream(tuskYmlFile)) {
      TuskYml tuskYml = new Yaml().loadAs(input, TuskYml.class);
      loadNullable(tuskYml);
    } catch (IOException e) {
      LOGGER.log(SEVERE, format("Failed to load %s.", tuskYmlFile));
    }
  }

  public static void loadNullable(@Nullable TuskYml yml) {
    if (yml != null) {
      load(yml);
    }
  }

  public static void load(TuskYml yml) {
    if (yml.enabled != null && !yml.enabled) {
      return;
    }

    MetricClient.Builder builder =
        MetricClient.newBuilder()
            .setHost(requireNonNull(yml.host))
            .setUser(requireNonNull(yml.user));
    if (yml.password != null) {
      builder.setPassword(yml.password);
    }
    if (yml.default_database != null) {
      builder.setDatabase(yml.default_database);
    }
    if (yml.connect_timeout_seconds != null) {
      builder.setConnectTimeout(Duration.ofSeconds(yml.connect_timeout_seconds));
    }
    if (yml.socket_timeout_seconds != null) {
      builder.setSocketTimeout(Duration.ofSeconds(yml.socket_timeout_seconds));
    }
    Tusk.setClient(builder.build());

    if (yml.heartbeat != null) {
      setUpHeartbeat(yml.heartbeat);
    }

    if (yml.logging != null) {
      setUpLogging(yml.logging);
    }
  }

  private static void setUpHeartbeat(TuskYml.Heartbeat heartbeat) {
    Table table =
        Table.newBuilder()
            .setDatabase(requireNonNull(heartbeat.database))
            .setName(requireNonNull(heartbeat.table))
            .build();

    TuskHeartbeatConfig.Builder config = TuskHeartbeatConfig.newBuilder().setTable(table);
    if (heartbeat.app_column != null) {
      config.setAppColumn(heartbeat.app_column);
    }
    if (heartbeat.task_column != null) {
      config.setTaskColumn(heartbeat.task_column);
    }
    if (heartbeat.sequence_column != null) {
      config.setSequenceColumn(heartbeat.sequence_column);
    }
    if (heartbeat.expiry_column != null) {
      config.setExpiryColumn(heartbeat.expiry_column);
    }
    TuskHeartbeat.init(config);
  }

  private static void setUpLogging(TuskYml.Logging logging) {
    Table table =
        Table.newBuilder()
            .setDatabase(requireNonNull(logging.database))
            .setName(requireNonNull(logging.table))
            .build();

    TuskLoggingConfig.Builder config = TuskLoggingConfig.newBuilder().setTable(table);
    if (logging.app_column != null) {
      config.setAppColumn(logging.app_column);
    }
    if (logging.level_column != null) {
      config.setLevelColumn(logging.level_column);
    }
    if (logging.logger_column != null) {
      config.setLoggerColumn(logging.logger_column);
    }
    if (logging.error_column != null) {
      config.setErrorColumn(logging.error_column);
    }
    if (logging.message_column != null) {
      config.setMessageColumn(logging.message_column);
    }
    if (logging.stacktrace_column != null) {
      config.setStacktraceColumn(logging.stacktrace_column);
    }
    if (logging.context_columns != null) {
      config.setContextColumns(logging.context_columns);
    }
    TuskLogging.init(config);
  }
}