code: /*
   * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
   */
  
  package io.airbyte.integrations.destination.cassandra;
  
  import io.airbyte.commons.json.Jsons;
  import io.airbyte.integrations.base.AirbyteStreamNameNamespacePair;
  import io.airbyte.integrations.base.FailureTrackingAirbyteMessageConsumer;
  import io.airbyte.protocol.models.AirbyteMessage;
  import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
  import java.util.Map;
  import java.util.function.Consumer;
  import java.util.stream.Collectors;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  class CassandraMessageConsumer extends FailureTrackingAirbyteMessageConsumer {
  
    private static final Logger LOGGER = LoggerFactory.getLogger(CassandraMessageConsumer.class);
  
    private final CassandraConfig cassandraConfig;
  
    private final Consumer<AirbyteMessage> outputRecordCollector;
  
    private final Map<AirbyteStreamNameNamespacePair, CassandraStreamConfig> cassandraStreams;
  
    private final CassandraCqlProvider cassandraCqlProvider;
  
    public CassandraMessageConsumer(final CassandraConfig cassandraConfig,
                                    final ConfiguredAirbyteCatalog configuredCatalog,
                                    final CassandraCqlProvider provider,
                                    final Consumer<AirbyteMessage> outputRecordCollector) {
      this.cassandraConfig = cassandraConfig;
      this.outputRecordCollector = outputRecordCollector;
      this.cassandraCqlProvider = provider;
      var nameTransformer = new CassandraNameTransformer(cassandraConfig);
      this.cassandraStreams = configuredCatalog.getStreams().stream()
          .collect(Collectors.toUnmodifiableMap(
              AirbyteStreamNameNamespacePair::fromConfiguredAirbyteSteam,
              k -> new CassandraStreamConfig(
                  nameTransformer.outputKeyspace(k.getStream().getNamespace()),
                  nameTransformer.outputTable(k.getStream().getName()),
                  nameTransformer.outputTmpTable(k.getStream().getName()),
                  k.getDestinationSyncMode())));
    }
  
    @Override
    protected void startTracked() {
      cassandraStreams.forEach((k, v) -> {
        cassandraCqlProvider.createKeySpaceIfNotExists(v.getKeyspace(), cassandraConfig.getReplication());
        cassandraCqlProvider.createTableIfNotExists(v.getKeyspace(), v.getTempTableName());
      });
    }
  
    @Override
    protected void acceptTracked(final AirbyteMessage message) {
      if (message.getType() == AirbyteMessage.Type.RECORD) {
        var messageRecord = message.getRecord();
        var streamConfig =
            cassandraStreams.get(AirbyteStreamNameNamespacePair.fromRecordMessage(messageRecord));
        if (streamConfig == null) {
          throw new IllegalArgumentException("Unrecognized destination stream");
        }
        var data = Jsons.serialize(messageRecord.getData());
        cassandraCqlProvider.insert(streamConfig.getKeyspace(), streamConfig.getTempTableName(), data);
      } else if (message.getType() == AirbyteMessage.Type.STATE) {
        outputRecordCollector.accept(message);
      } else {
        LOGGER.warn("Unsupported airbyte message type: {}", message.getType());
      }
    }
  
    @Override
    protected void close(final boolean hasFailed) {
      if (!hasFailed) {
        cassandraStreams.forEach((k, v) -> {
          try {
            cassandraCqlProvider.createTableIfNotExists(v.getKeyspace(), v.getTableName());
            switch (v.getDestinationSyncMode()) {
              case APPEND -> {
                cassandraCqlProvider.copy(v.getKeyspace(), v.getTempTableName(), v.getTableName());
              }
              case OVERWRITE -> {
                cassandraCqlProvider.truncate(v.getKeyspace(), v.getTableName());
                cassandraCqlProvider.copy(v.getKeyspace(), v.getTempTableName(), v.getTableName());
              }
              default -> throw new UnsupportedOperationException();
            }
          } catch (final Exception e) {
            LOGGER.error("Error while copying data to table {}: : ", v.getTableName(), e);
          }
        });
      }
  
      cassandraStreams.forEach((k, v) -> {
        try {
          cassandraCqlProvider.dropTableIfExists(v.getKeyspace(), v.getTempTableName());
        } catch (final Exception e) {
          LOGGER.error("Error while deleting temp table {} with reason: ", v.getTempTableName(), e);
        }
      });
      cassandraCqlProvider.close();
  
    }
  
  }
  