Place offset manager in commons

commons/src/main/java/io/aiven/kafka/connect/common/ClosableIterator.java

@@ -0,0 +1,84 @@
+/*
+ * Copyright 2024 Aiven Oy
+ *
+ * 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
+ *
+ * http://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 io.aiven.kafka.connect.common;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Iterator;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * An iterator that returns closable items and will close the last item returned in one of two cases:
+ * <ul>
+ * <li>when {@link #next} is called.</li>
+ * <li>when {@link #hasNext} returns false.</li>
+ * </ul>
+ *
+ * @param <T>
+ *            The type of Closeable object to return.
+ */
+public final class ClosableIterator<T extends Closeable> implements Iterator<T>, Closeable {
+    private static final Logger LOGGER = LoggerFactory.getLogger(ClosableIterator.class);
+    private final Iterator<T> delegate;
+    private T nextItem;
+    private T lastItem;
+
+    public static <T extends Closeable> ClosableIterator<T> wrap(final Iterator<T> iterator) {
+        return iterator instanceof ClosableIterator ? (ClosableIterator<T>) iterator : new ClosableIterator<>(iterator);
+    }
+    private ClosableIterator(final Iterator<T> iterator) {
+        delegate = iterator;
+    }
+
+    @Override
+    public boolean hasNext() {
+        if (nextItem == null) {
+            if (delegate.hasNext()) {
+                nextItem = delegate.next();
+            } else {
+                closeItem(lastItem);
+            }
+        }
+        return nextItem != null;
+    }
+
+    private void closeItem(final T item) {
+        if (item != null) {
+            try {
+                item.close();
+            } catch (IOException e) {
+                LOGGER.error(String.format("Error closing %s:", item), e);
+            }
+        }
+    }
+
+    @Override
+    public T next() {
+        closeItem(lastItem);
+        lastItem = nextItem;
+        nextItem = null; // NOPMD assign null
+        return lastItem;
+    }
+
+    @Override
+    public void close() {
+        closeItem(lastItem);
+        closeItem(nextItem);
+    }
+}

commons/src/main/java/io/aiven/kafka/connect/common/OffsetManager.java

@@ -0,0 +1,197 @@
+/*
+ * Copyright 2024 Aiven Oy
+ *
+ * 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
+ *
+ * http://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 io.aiven.kafka.connect.common;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
+
+import org.apache.kafka.connect.source.SourceTaskContext;
+
+public class OffsetManager<E extends OffsetManager.OffsetManagerEntry<E>> {
+
+    /**
+     * The local manager data.
+     */
+    private final Map<Map<String, Object>, Map<String, Object>> offsets;
+
+    /**
+     * The context in which this is running.
+     */
+    private final SourceTaskContext context;
+
+    /**
+     * Constructor
+     *
+     * @param context
+     *            the context for this instance to use.
+     */
+    public OffsetManager(final SourceTaskContext context) {
+        this(context, new ConcurrentHashMap<>());
+    }
+
+    /**
+     * Package private for testing.
+     *
+     * @param context
+     *            the context for this instance to use.
+     * @param offsets
+     *            the offsets
+     */
+    protected OffsetManager(final SourceTaskContext context,
+            final Map<Map<String, Object>, Map<String, Object>> offsets) {
+        this.context = context;
+        this.offsets = offsets;
+    }
+
+    /**
+     * Get an entry from the offset manager. This method will return the local copy if it has been created otherwise
+     * will get the data from Kafka.
+     *
+     * @param key
+     *            the key for the entry.
+     * @param creator
+     *            a function to create the connector defined offset entry from a Map of string to object.
+     * @return the entry.
+     */
+    public E getEntry(final OffsetManagerKey key, final Function<Map<String, Object>, E> creator) {
+        final Map<String, Object> data = offsets.compute(key.getPartitionMap(), (k, v) -> {
+            if (v == null) {
+                final Map<String, Object> kafkaData = context.offsetStorageReader().offset(key.getPartitionMap());
+                return kafkaData == null || kafkaData.isEmpty() ? new HashMap<>(key.getPartitionMap()) : kafkaData;
+            } else {
+                return v;
+            }
+        });
+        return creator.apply(data);
+    }
+
+    /**
+     * Copies the entry into the offset manager data.
+     *
+     * @param entry
+     *            the entry to update.
+     */
+    public void updateCurrentOffsets(final E entry) {
+        offsets.compute(entry.getManagerKey().getPartitionMap(), (k, v) -> {
+            if (v == null) {
+                return new HashMap<>(entry.getProperties());
+            } else {
+                v.putAll(entry.getProperties());
+                return v;
+            }
+        });
+    }
+
+    /**
+     * The definition of an entry in the OffsetManager.
+     */
+    public interface OffsetManagerEntry<T extends OffsetManagerEntry<T>> extends Comparable<T> {
+
+        /**
+         * Creates a new OffsetManagerEntry by wrapping the properties with the current implementation. This method may
+         * throw a RuntimeException if requried properties are not defined in the map.
+         *
+         * @param properties
+         *            the properties to wrap. May be {@code null}.
+         * @return an OffsetManagerProperty
+         */
+        T fromProperties(Map<String, Object> properties);
+
+        /**
+         * Extracts the data from the entry in the correct format to return to Kafka.
+         *
+         * @return the properties in a format to return to Kafka.
+         */
+        Map<String, Object> getProperties();
+
+        /**
+         * Gets the value of the named property. The value returned from a {@code null} key is implementation dependant.
+         *
+         * @param key
+         *            the property to retrieve.
+         * @return the value associated with the property or @{code null} if not set.
+         * @throws NullPointerException
+         *             if a {@code null} key is not supported.
+         */
+        Object getProperty(String key);
+
+        /**
+         * Sets a key/value pair. Will overwrite any existing value. Implementations of OffsetManagerEntry may declare
+         * specific keys as restricted. These are generally keys that are managed internally by the OffsetManagerEntry
+         * and may not be set except through provided setter methods or the constructor.
+         *
+         * @param key
+         *            the key to set.
+         * @param value
+         *            the value to set.
+         * @throws IllegalArgumentException
+         *             if the key is restricted.
+         */
+        void setProperty(String key, Object value);
+
+        /**
+         * ManagerKey getManagerKey
+         *
+         * @return The offset manager key for this entry.
+         */
+        OffsetManagerKey getManagerKey();
+
+        /**
+         * Gets the Kafka topic for this entry.
+         *
+         * @return The Kafka topic for this entry.
+         */
+        String getTopic();
+
+        /**
+         * Gets the Kafka partition for this entry.
+         *
+         * @return The Kafka partition for this entry.
+         */
+        Integer getPartition();
+
+        /**
+         * Gets the number of records to skip to get to this record. This is the same as the zero-based index of this
+         * record if all records were in an array.
+         *
+         * @return The number of records to skip to get to this record.
+         */
+        default long skipRecords() {
+            return 0;
+        }
+
+        /**
+         * Increments the record count.
+         */
+        void incrementRecordCount();
+    }
+
+    /**
+     * The OffsetManager Key. Must override hashCode() and equals().
+     */
+    @FunctionalInterface
+    public interface OffsetManagerKey {
+        /**
+         * gets the partition map used by Kafka to identify this Offset entry.
+         *
+         * @return The partition map used by Kafka to identify this Offset entry.
+         */
+        Map<String, Object> getPartitionMap();
+    }
+}

commons/src/main/java/io/aiven/kafka/connect/common/config/SourceCommonConfig.java

@@ -22,6 +22,8 @@
 
 import io.aiven.kafka.connect.common.config.enums.ErrorsTolerance;
 import io.aiven.kafka.connect.common.source.input.InputFormat;
+import io.aiven.kafka.connect.common.source.input.Transformer;
+import io.aiven.kafka.connect.common.source.input.TransformerFactory;
 
 public class SourceCommonConfig extends CommonConfig {
 
@@ -71,4 +73,8 @@ public int getMaxPollRecords() {
         return sourceConfigFragment.getMaxPollRecords();
     }
 
+    public Transformer getTransformer() {
+        return TransformerFactory.getTransformer(schemaRegistryFragment.getInputFormat());
+    }
+
 }

commons/src/main/java/io/aiven/kafka/connect/common/source/input/AvroTransformer.java

@@ -16,18 +16,16 @@
 
 package io.aiven.kafka.connect.common.source.input;
 
-import static io.aiven.kafka.connect.common.config.SchemaRegistryFragment.SCHEMA_REGISTRY_URL;
-
 import java.io.IOException;
 import java.io.InputStream;
-import java.nio.charset.StandardCharsets;
-import java.util.Map;
 import java.util.function.Consumer;
 
 import org.apache.kafka.common.config.AbstractConfig;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaAndValue;
 
+import io.aiven.kafka.connect.common.OffsetManager;
+
 import io.confluent.connect.avro.AvroData;
 import org.apache.avro.file.DataFileStream;
 import org.apache.avro.generic.GenericDatumReader;
@@ -37,26 +35,35 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-public class AvroTransformer extends Transformer<GenericRecord> {
-
+/**
+ * A transformer that reads the inputstream as Avro data.
+ */
+public class AvroTransformer extends Transformer {
+    /** The AvroData instance to read with */
     private final AvroData avroData;
-
+    /** The logger for this transformer */
     private static final Logger LOGGER = LoggerFactory.getLogger(AvroTransformer.class);
 
+    /**
+     * The constructor.
+     *
+     * @param avroData
+     *            the AvroData object to read with.
+     */
     AvroTransformer(final AvroData avroData) {
         super();
         this.avroData = avroData;
     }
 
     @Override
-    public void configureValueConverter(final Map<String, String> config, final AbstractConfig sourceConfig) {
-        config.put(SCHEMA_REGISTRY_URL, sourceConfig.getString(SCHEMA_REGISTRY_URL));
+    public Schema getKeySchema() {
+        return Schema.OPTIONAL_BYTES_SCHEMA;
     }
 
     @Override
-    public StreamSpliterator<GenericRecord> createSpliterator(final IOSupplier<InputStream> inputStreamIOSupplier,
-            final String topic, final int topicPartition, final AbstractConfig sourceConfig) {
-        return new StreamSpliterator<>(LOGGER, inputStreamIOSupplier) {
+    public StreamSpliterator createSpliterator(final IOSupplier<InputStream> inputStreamIOSupplier,
+            final OffsetManager.OffsetManagerEntry<?> offsetManagerEntry, final AbstractConfig sourceConfig) {
+        return new StreamSpliterator(LOGGER, inputStreamIOSupplier, offsetManagerEntry) {
             private DataFileStream<GenericRecord> dataFileStream;
             private final DatumReader<GenericRecord> datumReader = new GenericDatumReader<>();
 
@@ -78,26 +85,14 @@ public void doClose() {
             }
 
             @Override
-            protected boolean doAdvance(final Consumer<? super GenericRecord> action) {
+            protected boolean doAdvance(final Consumer<? super SchemaAndValue> action) {
                 if (dataFileStream.hasNext()) {
-                    action.accept(dataFileStream.next());
+                    final GenericRecord record = dataFileStream.next();
+                    action.accept(avroData.toConnectData(record.getSchema(), record));
                     return true;
                 }
                 return false;
             }
         };
     }
-
-    @Override
-    public SchemaAndValue getValueData(final GenericRecord record, final String topic,
-            final AbstractConfig sourceConfig) {
-        return avroData.toConnectData(record.getSchema(), record);
-    }
-
-    @Override
-    public SchemaAndValue getKeyData(final Object cloudStorageKey, final String topic,
-            final AbstractConfig sourceConfig) {
-        return new SchemaAndValue(Schema.OPTIONAL_BYTES_SCHEMA,
-                ((String) cloudStorageKey).getBytes(StandardCharsets.UTF_8));
-    }
 }