[DebeziumIO] Implement offsetRetainer for Kafka (#28248)

Author: tkaymak

sdks/go/pkg/beam/io/xlang/debeziumio/debezium.go

@@ -77,6 +77,8 @@ type readFromDebeziumSchema struct {
 	MaxNumberOfRecords   *int64
 	MaxTimeToRun         *int64
 	ConnectionProperties []string
+	StartOffset          []string
+	OffsetStoragePath    *string
 }
 
 type debeziumConfig struct {
@@ -149,6 +151,44 @@ func ConnectionProperties(cp []string) readOption {
 	}
 }
 
+// StartOffset specifies the offset from which the connector should resume consuming
+// changes. Each entry must be a "key=value" string, where numeric values are encoded
+// as their decimal string representation.
+//
+// Example for PostgreSQL:
+//
+//	debeziumio.StartOffset([]string{"lsn=28160840"})
+//
+// Example for MySQL:
+//
+//	debeziumio.StartOffset([]string{"file=binlog.000001", "pos=156"})
+//
+// Obtain the offset from the output of a previous pipeline run. Numeric values such
+// as LSN or binlog position are automatically parsed to Long on the Java side.
+func StartOffset(offset []string) readOption {
+	return func(cfg *debeziumConfig) {
+		cfg.readSchema.StartOffset = offset
+	}
+}
+
+// OffsetStoragePath sets a path where the connector offset is automatically saved after each
+// checkpoint and loaded on pipeline startup, allowing the pipeline to resume from where it
+// left off without any manual offset management.
+//
+// The path can be on any filesystem supported by the active Beam runner
+// (local disk, GCS, S3, etc.).
+//
+// Example:
+//
+//	debeziumio.OffsetStoragePath("gs://my-bucket/debezium/orders-offset.json")
+//
+// When set, takes precedence over StartOffset.
+func OffsetStoragePath(path string) readOption {
+	return func(cfg *debeziumConfig) {
+		cfg.readSchema.OffsetStoragePath = &path
+	}
+}
+
 // ExpansionAddr sets the expansion service address to use for DebeziumIO cross-langauage transform.
 func ExpansionAddr(expansionAddr string) readOption {
 	return func(cfg *debeziumConfig) {

sdks/java/io/debezium/src/main/java/org/apache/beam/io/debezium/DebeziumIO.java

@@ -144,6 +144,10 @@ public abstract static class Read<T> extends PTransform<PBegin, PCollection<T>>
 
     abstract @Nullable Long getPollingTimeout();
 
+    abstract @Nullable Map<String, Object> getStartOffset();
+
+    abstract @Nullable OffsetRetainer getOffsetRetainer();
+
     abstract @Nullable Coder<T> getCoder();
 
     abstract Builder<T> toBuilder();
@@ -162,6 +166,10 @@ abstract static class Builder<T> {
 
       abstract Builder<T> setPollingTimeout(Long miliseconds);
 
+      abstract Builder<T> setStartOffset(Map<String, Object> startOffset);
+
+      abstract Builder<T> setOffsetRetainer(OffsetRetainer retainer);
+
       abstract Read<T> build();
     }
 
@@ -230,6 +238,74 @@ public Read<T> withPollingTimeout(Long miliseconds) {
       return toBuilder().setPollingTimeout(miliseconds).build();
     }
 
+    /**
+     * Sets a starting offset so the connector resumes consuming changes from a previously seen
+     * position rather than from the beginning of the change stream.
+     *
+     * <p>The offset format is connector-specific. You can capture the current offset for each
+     * processed record inside your {@link SourceRecordMapper} via {@link
+     * org.apache.kafka.connect.source.SourceRecord#sourceOffset()} and persist it externally (for
+     * example in Cloud Storage, a database, or a local file). On the next pipeline run, pass the
+     * last saved offset here.
+     *
+     * <p>Example (PostgreSQL):
+     *
+     * <pre>{@code
+     * // Capture the offset inside the SourceRecordMapper:
+     * Map<String, Object> offset = sourceRecord.sourceOffset();
+     * // Persist 'offset' externally, then on restart:
+     * DebeziumIO.read()
+     *     .withConnectorConfiguration(config)
+     *     .withStartOffset(savedOffset)
+     *     .withFormatFunction(myMapper);
+     * }</pre>
+     *
+     * @param startOffset A map representing the resumption point, as returned by {@code
+     *     SourceRecord#sourceOffset()}.
+     * @return PTransform {@link #read}
+     */
+    public Read<T> withStartOffset(Map<String, Object> startOffset) {
+      checkArgument(startOffset != null, "startOffset can not be null");
+      return toBuilder().setStartOffset(startOffset).build();
+    }
+
+    /**
+     * Sets an {@link OffsetRetainer} that automatically saves and restores the connector offset,
+     * allowing the pipeline to resume from where it left off after a restart without any manual
+     * offset management.
+     *
+     * <p>When a retainer is configured:
+     *
+     * <ol>
+     *   <li>At pipeline startup, {@link OffsetRetainer#loadOffset()} is called. If a saved offset
+     *       is found, the connector resumes from that position; otherwise it starts from the
+     *       beginning of the change stream.
+     *   <li>After each successful checkpoint ({@code task.commit()}), {@link
+     *       OffsetRetainer#saveOffset(Map)} is called with the latest committed offset.
+     * </ol>
+     *
+     * <p>The built-in {@link FileSystemOffsetRetainer} persists the offset as a JSON file on any
+     * Beam-compatible filesystem (local, GCS, S3, etc.):
+     *
+     * <pre>{@code
+     * DebeziumIO.read()
+     *     .withConnectorConfiguration(config)
+     *     .withOffsetRetainer(
+     *         new FileSystemOffsetRetainer("gs://my-bucket/debezium/orders-offset.json"))
+     *     .withFormatFunction(myMapper);
+     * }</pre>
+     *
+     * <p>When both a retainer and {@link #withStartOffset(Map)} are set, the retainer takes
+     * precedence. Use {@link #withStartOffset(Map)} alone for a one-time manual override.
+     *
+     * @param retainer The {@link OffsetRetainer} to use for loading and saving offsets.
+     * @return PTransform {@link #read}
+     */
+    public Read<T> withOffsetRetainer(OffsetRetainer retainer) {
+      checkArgument(retainer != null, "retainer can not be null");
+      return toBuilder().setOffsetRetainer(retainer).build();
+    }
+
     protected Schema getRecordSchema() {
       KafkaSourceConsumerFn<T> fn =
           new KafkaSourceConsumerFn<>(getConnectorConfiguration().getConnectorClass().get(), this);

sdks/java/io/debezium/src/main/java/org/apache/beam/io/debezium/DebeziumTransformRegistrar.java

@@ -18,6 +18,7 @@
 package org.apache.beam.io.debezium;
 
 import com.google.auto.service.AutoService;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import org.apache.beam.sdk.expansion.ExternalTransformRegistrar;
@@ -78,6 +79,8 @@ public static class Configuration extends CrossLanguageConfiguration {
       private @Nullable List<String> connectionProperties;
       private @Nullable Long maxNumberOfRecords;
       private @Nullable Long maxTimeToRun;
+      private @Nullable List<String> startOffset;
+      private @Nullable String offsetStoragePath;
 
       public void setConnectionProperties(@Nullable List<String> connectionProperties) {
         this.connectionProperties = connectionProperties;
@@ -90,6 +93,14 @@ public void setMaxNumberOfRecords(@Nullable Long maxNumberOfRecords) {
       public void setMaxTimeToRun(@Nullable Long maxTimeToRun) {
         this.maxTimeToRun = maxTimeToRun;
       }
+
+      public void setStartOffset(@Nullable List<String> startOffset) {
+        this.startOffset = startOffset;
+      }
+
+      public void setOffsetStoragePath(@Nullable String offsetStoragePath) {
+        this.offsetStoragePath = offsetStoragePath;
+      }
     }
 
     @Override
@@ -123,6 +134,27 @@ public PTransform<PBegin, PCollection<String>> buildExternal(Configuration confi
         readTransform = readTransform.withMaxTimeToRun(configuration.maxTimeToRun);
       }
 
+      if (configuration.startOffset != null) {
+        Map<String, Object> startOffsetMap = new HashMap<>();
+        for (String property : configuration.startOffset) {
+          String[] parts = property.split("=", 2);
+          String key = parts[0];
+          String value = parts[1];
+          try {
+            startOffsetMap.put(key, Long.parseLong(value));
+          } catch (NumberFormatException e) {
+            startOffsetMap.put(key, value);
+          }
+        }
+        readTransform = readTransform.withStartOffset(startOffsetMap);
+      }
+
+      if (configuration.offsetStoragePath != null) {
+        readTransform =
+            readTransform.withOffsetRetainer(
+                new FileSystemOffsetRetainer(configuration.offsetStoragePath));
+      }
+
       return readTransform;
     }
   }

sdks/java/io/debezium/src/main/java/org/apache/beam/io/debezium/FileSystemOffsetRetainer.java

@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.beam.io.debezium;
+
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import java.io.FileNotFoundException;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.nio.channels.Channels;
+import java.nio.channels.ReadableByteChannel;
+import java.nio.channels.WritableByteChannel;
+import java.util.Map;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * An {@link OffsetRetainer} that persists the Debezium connector offset as a JSON file using Beam's
+ * {@link FileSystems} abstraction.
+ *
+ * <p>The {@code path} argument can point to any filesystem supported by the active Beam runner,
+ * including local disk, Google Cloud Storage, Amazon S3, and others
+ *
+ * <p>On every {@code task.commit()}, the latest offset is serialised to JSON and written to the
+ * given path (overwriting the previous file). On pipeline startup the file is read back and the
+ * connector resumes from the stored position. If the file does not yet exist the connector starts
+ * from the beginning of the change stream.
+ *
+ * <p>Example — resume from GCS:
+ *
+ * <pre>{@code
+ * DebeziumIO.read()
+ *     .withConnectorConfiguration(config)
+ *     .withOffsetRetainer(
+ *         new FileSystemOffsetRetainer("gs://my-bucket/debezium/orders-offset.json"))
+ *     .withFormatFunction(myMapper);
+ * }</pre>
+ *
+ * <p>Example — local filesystem (useful for testing):
+ *
+ * <pre>{@code
+ * DebeziumIO.read()
+ *     .withConnectorConfiguration(config)
+ *     .withOffsetRetainer(new FileSystemOffsetRetainer("/tmp/debezium-offset.json"))
+ *     .withFormatFunction(myMapper);
+ * }</pre>
+ *
+ * <p><b>Note:</b> writes are not atomic. If the pipeline is killed mid-write, the offset file may
+ * be corrupt. In that case, delete the file and the connector will restart from the beginning.
+ */
+public class FileSystemOffsetRetainer implements OffsetRetainer {
+
+  private static final Logger LOG = LoggerFactory.getLogger(FileSystemOffsetRetainer.class);
+  private static final TypeReference<Map<String, Object>> MAP_TYPE = new TypeReference<>() {};
+
+  private final String path;
+
+  // ObjectMapper is thread-safe after configuration and does not need to be serialised.
+  private transient @Nullable ObjectMapper objectMapper;
+
+  public FileSystemOffsetRetainer(String path) {
+    this.path = path;
+  }
+
+  private ObjectMapper mapper() {
+    if (objectMapper == null) {
+      objectMapper = new ObjectMapper();
+    }
+    return objectMapper;
+  }
+
+  /**
+   * Reads the offset JSON file and returns its contents, or {@code null} if the file does not exist
+   * or cannot be read.
+   */
+  @Override
+  public @Nullable Map<String, Object> loadOffset() {
+    try {
+      ResourceId resourceId = FileSystems.matchNewResource(path, /* isDirectory= */ false);
+      try (ReadableByteChannel channel = FileSystems.open(resourceId);
+          InputStream stream = Channels.newInputStream(channel)) {
+        Map<String, Object> offset = mapper().readValue(stream, MAP_TYPE);
+        LOG.info("OffsetRetainer: loaded offset from {}: {}", path, offset);
+        return offset;
+      }
+    } catch (FileNotFoundException e) {
+      LOG.info("OffsetRetainer: no offset file found at {}; starting from the beginning.", path);
+      return null;
+    } catch (IOException e) {
+      LOG.warn(
+          "OffsetRetainer: failed to load offset from {}; starting from the beginning.", path, e);
+      return null;
+    }
+  }
+
+  /**
+   * Serialises {@code offset} to JSON and writes it to the configured path, overwriting any
+   * existing file. Errors are logged as warnings and swallowed so the pipeline continues.
+   */
+  @Override
+  public void saveOffset(Map<String, Object> offset) {
+    try {
+      ResourceId resourceId = FileSystems.matchNewResource(path, /* isDirectory= */ false);
+      try (WritableByteChannel channel = FileSystems.create(resourceId, "application/json");
+          OutputStream stream = Channels.newOutputStream(channel)) {
+        mapper().writeValue(stream, offset);
+      }
+      LOG.debug("OffsetRetainer: saved offset to {}: {}", path, offset);
+    } catch (IOException e) {
+      LOG.warn(
+          "OffsetRetainer: failed to save offset to {}."
+              + " The offset will be lost if the pipeline restarts.",
+          path,
+          e);
+    }
+  }
+}

sdks/java/io/debezium/src/main/java/org/apache/beam/io/debezium/KafkaSourceConsumerFn.java

@@ -122,7 +122,21 @@ public class KafkaSourceConsumerFn<T> extends DoFn<Map<String, String>, T> {
   @GetInitialRestriction
   public OffsetHolder getInitialRestriction(@Element Map<String, String> unused)
       throws IOException {
-    return new OffsetHolder(null, null, null, spec.getMaxNumberOfRecords(), spec.getMaxTimeToRun());
+    Map<String, Object> initialOffset = null;
+
+    // Retainer takes precedence: it reflects the most recently committed position.
+    OffsetRetainer retainer = spec.getOffsetRetainer();
+    if (retainer != null) {
+      initialOffset = retainer.loadOffset();
+    }
+
+    // Fall back to the explicit one-time override when the retainer has no saved offset.
+    if (initialOffset == null) {
+      initialOffset = spec.getStartOffset();
+    }
+
+    return new OffsetHolder(
+        initialOffset, null, null, spec.getMaxNumberOfRecords(), spec.getMaxTimeToRun());
   }
 
   @NewTracker
@@ -284,6 +298,16 @@ public ProcessContinuation process(
           receiver.outputWithTimestamp(json, recordInstant);
         }
         task.commit();
+
+        // Persist the offset after every successful commit so the pipeline can resume
+        // from this position on restart.
+        OffsetRetainer retainer = spec.getOffsetRetainer();
+        @SuppressWarnings("unchecked")
+        Map<String, Object> committedOffset =
+            (Map<String, Object>) tracker.currentRestriction().offset;
+        if (retainer != null && committedOffset != null) {
+          retainer.saveOffset(committedOffset);
+        }
       }
     } catch (Exception ex) {
       throw new RuntimeException("Error occurred when consuming changes from Database. ", ex);