content: /*
   * Copyright 2022 The FeatHub Authors
   *
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *     https://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  package org.apache.flink.streaming.connectors.kafka.table;
  
  import org.apache.flink.api.common.eventtime.WatermarkStrategy;
  import org.apache.flink.api.common.serialization.DeserializationSchema;
  import org.apache.flink.api.common.typeinfo.TypeInformation;
  import org.apache.flink.api.connector.source.Boundedness;
  import org.apache.flink.connector.kafka.source.KafkaSource;
  import org.apache.flink.connector.kafka.source.KafkaSourceBuilder;
  import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
  import org.apache.flink.connector.kafka.source.reader.deserializer.KafkaRecordDeserializationSchema;
  import org.apache.flink.streaming.api.datastream.DataStream;
  import org.apache.flink.streaming.api.datastream.DataStreamSource;
  import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
  import org.apache.flink.streaming.connectors.kafka.KafkaDeserializationSchema;
  import org.apache.flink.streaming.connectors.kafka.config.StartupMode;
  import org.apache.flink.streaming.connectors.kafka.internals.KafkaTopicPartition;
  import org.apache.flink.table.connector.Projection;
  import org.apache.flink.table.connector.ProviderContext;
  import org.apache.flink.table.connector.format.DecodingFormat;
  import org.apache.flink.table.connector.source.DataStreamScanProvider;
  import org.apache.flink.table.connector.source.DynamicTableSource;
  import org.apache.flink.table.connector.source.ScanTableSource;
  import org.apache.flink.table.data.RowData;
  import org.apache.flink.table.types.DataType;
  import org.apache.flink.table.types.utils.DataTypeUtils;
  
  import org.apache.kafka.clients.consumer.ConsumerConfig;
  import org.apache.kafka.clients.consumer.OffsetResetStrategy;
  import org.apache.kafka.common.TopicPartition;
  
  import javax.annotation.Nullable;
  
  import java.util.Arrays;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Locale;
  import java.util.Map;
  import java.util.Properties;
  import java.util.regex.Pattern;
  import java.util.stream.Collectors;
  import java.util.stream.IntStream;
  import java.util.stream.Stream;
  
  // TODO: This class will be removed after FLINK-24456 resolved.
  /**
   * {@link BoundedKafkaDynamicSource} extends {@link KafkaDynamicSource} to support bounded Kafka
   * source. The bounded Kafka source stops at the latest offsets of the partitions when the
   * KafkaSource starts to run.
   */
  public class BoundedKafkaDynamicSource extends KafkaDynamicSource {
      private static final String KAFKA_TRANSFORMATION = "bounded-kafka";
  
      public BoundedKafkaDynamicSource(
              DataType physicalDataType,
              @Nullable DecodingFormat<DeserializationSchema<RowData>> keyDecodingFormat,
              DecodingFormat<DeserializationSchema<RowData>> valueDecodingFormat,
              int[] keyProjection,
              int[] valueProjection,
              @Nullable String keyPrefix,
              @Nullable List<String> topics,
              @Nullable Pattern topicPattern,
              Properties properties,
              StartupMode startupMode,
              Map<KafkaTopicPartition, Long> specificStartupOffsets,
              long startupTimestampMillis,
              boolean upsertMode,
              String tableIdentifier) {
          super(
                  physicalDataType,
                  keyDecodingFormat,
                  valueDecodingFormat,
                  keyProjection,
                  valueProjection,
                  keyPrefix,
                  topics,
                  topicPattern,
                  properties,
                  startupMode,
                  specificStartupOffsets,
                  startupTimestampMillis,
                  upsertMode,
                  tableIdentifier);
      }
  
      @Override
      public ScanTableSource.ScanRuntimeProvider getScanRuntimeProvider(
              ScanTableSource.ScanContext context) {
          final DeserializationSchema<RowData> keyDeserialization =
                  createDeserialization(context, keyDecodingFormat, keyProjection, keyPrefix);
  
          final DeserializationSchema<RowData> valueDeserialization =
                  createDeserialization(context, valueDecodingFormat, valueProjection, null);
  
          final TypeInformation<RowData> producedTypeInfo =
                  context.createTypeInformation(producedDataType);
  
          final KafkaSource<RowData> kafkaSource =
                  createKafkaSource(keyDeserialization, valueDeserialization, producedTypeInfo);
  
          return new DataStreamScanProvider() {
              @Override
              public DataStream<RowData> produceDataStream(
                      ProviderContext providerContext, StreamExecutionEnvironment execEnv) {
                  if (watermarkStrategy == null) {
                      watermarkStrategy = WatermarkStrategy.noWatermarks();
                  }
                  DataStreamSource<RowData> sourceStream =
                          execEnv.fromSource(
                                  kafkaSource,
                                  watermarkStrategy,
                                  "BoundedKafkaSource-" + tableIdentifier);
                  providerContext.generateUid(KAFKA_TRANSFORMATION).ifPresent(sourceStream::uid);
                  return sourceStream;
              }
  
              @Override
              public boolean isBounded() {
                  return kafkaSource.getBoundedness() == Boundedness.BOUNDED;
              }
          };
      }
  
      private @Nullable DeserializationSchema<RowData> createDeserialization(
              DynamicTableSource.Context context,
              @Nullable DecodingFormat<DeserializationSchema<RowData>> format,
              int[] projection,
              @Nullable String prefix) {
          if (format == null) {
              return null;
          }
          DataType physicalFormatDataType = Projection.of(projection).project(this.physicalDataType);
          if (prefix != null) {
              physicalFormatDataType = DataTypeUtils.stripRowPrefix(physicalFormatDataType, prefix);
          }
          return format.createRuntimeDecoder(context, physicalFormatDataType);
      }
  
      @Override
      protected KafkaSource<RowData> createKafkaSource(
              DeserializationSchema<RowData> keyDeserialization,
              DeserializationSchema<RowData> valueDeserialization,
              TypeInformation<RowData> producedTypeInfo) {
  
          final KafkaDeserializationSchema<RowData> kafkaDeserializer =
                  createKafkaDeserializationSchema(
                          keyDeserialization, valueDeserialization, producedTypeInfo);
  
          final KafkaSourceBuilder<RowData> kafkaSourceBuilder = KafkaSource.builder();
  
          if (topics != null) {
              kafkaSourceBuilder.setTopics(topics);
          } else {
              kafkaSourceBuilder.setTopicPattern(topicPattern);
          }
  
          switch (startupMode) {
              case EARLIEST:
                  kafkaSourceBuilder.setStartingOffsets(OffsetsInitializer.earliest());
                  break;
              case LATEST:
                  kafkaSourceBuilder.setStartingOffsets(OffsetsInitializer.latest());
                  break;
              case GROUP_OFFSETS:
                  String offsetResetConfig =
                          properties.getProperty(
                                  ConsumerConfig.AUTO_OFFSET_RESET_CONFIG,
                                  OffsetResetStrategy.NONE.name());
                  OffsetResetStrategy offsetResetStrategy = getResetStrategy(offsetResetConfig);
                  kafkaSourceBuilder.setStartingOffsets(
                          OffsetsInitializer.committedOffsets(offsetResetStrategy));
                  break;
              case SPECIFIC_OFFSETS:
                  Map<TopicPartition, Long> offsets = new HashMap<>();
                  specificStartupOffsets.forEach(
                          (tp, offset) ->
                                  offsets.put(
                                          new TopicPartition(tp.getTopic(), tp.getPartition()),
                                          offset));
                  kafkaSourceBuilder.setStartingOffsets(OffsetsInitializer.offsets(offsets));
                  break;
              case TIMESTAMP:
                  kafkaSourceBuilder.setStartingOffsets(
                          OffsetsInitializer.timestamp(startupTimestampMillis));
                  break;
          }
  
          kafkaSourceBuilder.setBounded(OffsetsInitializer.latest());
  
          kafkaSourceBuilder
                  .setProperties(properties)
                  .setDeserializer(KafkaRecordDeserializationSchema.of(kafkaDeserializer));
  
          return kafkaSourceBuilder.build();
      }
  
      private OffsetResetStrategy getResetStrategy(String offsetResetConfig) {
          return Arrays.stream(OffsetResetStrategy.values())
                  .filter(ors -> ors.name().equals(offsetResetConfig.toUpperCase(Locale.ROOT)))
                  .findAny()
                  .orElseThrow(
                          () ->
                                  new IllegalArgumentException(
                                          String.format(
                                                  "%s can not be set to %s. Valid values: [%s]",
                                                  ConsumerConfig.AUTO_OFFSET_RESET_CONFIG,
                                                  offsetResetConfig,
                                                  Arrays.stream(OffsetResetStrategy.values())
                                                          .map(Enum::name)
                                                          .map(String::toLowerCase)
                                                          .collect(Collectors.joining(",")))));
      }
  
      private KafkaDeserializationSchema<RowData> createKafkaDeserializationSchema(
              DeserializationSchema<RowData> keyDeserialization,
              DeserializationSchema<RowData> valueDeserialization,
              TypeInformation<RowData> producedTypeInfo) {
          final DynamicKafkaDeserializationSchema.MetadataConverter[] metadataConverters =
                  metadataKeys.stream()
                          .map(
                                  k ->
                                          Stream.of(ReadableMetadata.values())
                                                  .filter(rm -> rm.key.equals(k))
                                                  .findFirst()
                                                  .orElseThrow(IllegalStateException::new))
                          .map(m -> m.converter)
                          .toArray(DynamicKafkaDeserializationSchema.MetadataConverter[]::new);
  
          // check if connector metadata is used at all
          final boolean hasMetadata = metadataKeys.size() > 0;
  
          // adjust physical arity with value format's metadata
          final int adjustedPhysicalArity =
                  DataType.getFieldDataTypes(producedDataType).size() - metadataKeys.size();
  
          // adjust value format projection to include value format's metadata columns at the end
          final int[] adjustedValueProjection =
                  IntStream.concat(
                                  IntStream.of(valueProjection),
                                  IntStream.range(
                                          keyProjection.length + valueProjection.length,
                                          adjustedPhysicalArity))
                          .toArray();
  
          return new DynamicKafkaDeserializationSchema(
                  adjustedPhysicalArity,
                  keyDeserialization,
                  keyProjection,
                  valueDeserialization,
                  adjustedValueProjection,
                  hasMetadata,
                  metadataConverters,
                  producedTypeInfo,
                  upsertMode);
      }
  
      @Override
      public DynamicTableSource copy() {
          final BoundedKafkaDynamicSource copy =
                  new BoundedKafkaDynamicSource(
                          physicalDataType,
                          keyDecodingFormat,
                          valueDecodingFormat,
                          keyProjection,
                          valueProjection,
                          keyPrefix,
                          topics,
                          topicPattern,
                          properties,
                          startupMode,
                          specificStartupOffsets,
                          startupTimestampMillis,
                          upsertMode,
                          tableIdentifier);
          copy.producedDataType = producedDataType;
          copy.metadataKeys = metadataKeys;
          copy.watermarkStrategy = watermarkStrategy;
          return copy;
      }
  }
  