address review comments

Author: ferenc-csaky

connectors/kafka-safe-connector/src/main/java/com/datasqrl/flinkrunner/connector/kafka/KafkaRecordTimestampWatermarkStrategy.java

@@ -15,6 +15,7 @@
  */
 package com.datasqrl.flinkrunner.connector.kafka;
 
+import com.datasqrl.flinkrunner.connector.kafka.SourceWatermarkOptions.SourceWatermarkConfig;
 import java.io.Serial;
 import java.time.Duration;
 import java.util.Arrays;
@@ -25,6 +26,7 @@
 import org.apache.flink.api.common.eventtime.WatermarkOutput;
 import org.apache.flink.api.common.eventtime.WatermarkStrategy;
 import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.watermark.WatermarkEmitStrategy;
 
 /** Adaptive watermark strategy based on Kafka record timestamps. */
 @Internal
@@ -37,18 +39,36 @@ public final class KafkaRecordTimestampWatermarkStrategy implements WatermarkStr
   static final long MIN_OUT_OF_ORDERNESS_MILLIS = 50L;
   static final long MAX_OUT_OF_ORDERNESS_MILLIS = Duration.ofDays(1).toMillis();
 
-  private static final double OUT_OF_ORDERNESS_QUANTILE = 0.95D;
+  static final double OUT_OF_ORDERNESS_QUANTILE = 0.95D;
 
-  public static final KafkaRecordTimestampWatermarkStrategy INSTANCE =
-      new KafkaRecordTimestampWatermarkStrategy();
+  private final WatermarkEmitStrategy emitStrategy;
+  private final SourceWatermarkConfig configuration;
 
-  private KafkaRecordTimestampWatermarkStrategy() {}
+  public KafkaRecordTimestampWatermarkStrategy() {
+    this(WatermarkEmitStrategy.ON_PERIODIC);
+  }
+
+  public KafkaRecordTimestampWatermarkStrategy(WatermarkEmitStrategy emitStrategy) {
+    this(
+        emitStrategy,
+        new SourceWatermarkConfig(
+            MIN_RECORDS,
+            MIN_OUT_OF_ORDERNESS_MILLIS,
+            MAX_OUT_OF_ORDERNESS_MILLIS,
+            OUT_OF_ORDERNESS_QUANTILE));
+  }
+
+  public KafkaRecordTimestampWatermarkStrategy(
+      WatermarkEmitStrategy emitStrategy, SourceWatermarkConfig configuration) {
+    this.emitStrategy = emitStrategy;
+    this.configuration = configuration;
+  }
 
   @Override
   public WatermarkGenerator<RowData> createWatermarkGenerator(
       WatermarkGeneratorSupplier.Context context) {
 
-    return new AdaptiveKafkaRecordTimestampWatermarkGenerator();
+    return new AdaptiveKafkaRecordTimestampWatermarkGenerator(emitStrategy, configuration);
   }
 
   /**
@@ -75,6 +95,9 @@ private static final class AdaptiveKafkaRecordTimestampWatermarkGenerator
      */
     private final long[] latenessSamples = new long[SAMPLE_SIZE];
 
+    private final WatermarkEmitStrategy emitStrategy;
+    private final SourceWatermarkConfig configuration;
+
     /** Highest Kafka record timestamp observed by this generator. */
     private long maxTimestamp = Long.MIN_VALUE;
 
@@ -89,6 +112,12 @@ private static final class AdaptiveKafkaRecordTimestampWatermarkGenerator
     /** Next slot in the circular lateness sample buffer. */
     private int nextSampleIndex;
 
+    private AdaptiveKafkaRecordTimestampWatermarkGenerator(
+        WatermarkEmitStrategy emitStrategy, SourceWatermarkConfig configuration) {
+      this.emitStrategy = emitStrategy;
+      this.configuration = configuration;
+    }
+
     @Override
     public void onEvent(RowData event, long eventTimestamp, WatermarkOutput output) {
       if (maxTimestamp != Long.MIN_VALUE) {
@@ -100,11 +129,23 @@ public void onEvent(RowData event, long eventTimestamp, WatermarkOutput output)
 
       maxTimestamp = Math.max(maxTimestamp, eventTimestamp);
       recordCount++;
+
+      if (emitStrategy.isOnEvent()) {
+        emitIfReady(output);
+      }
     }
 
     @Override
     public void onPeriodicEmit(WatermarkOutput output) {
-      if (recordCount < MIN_RECORDS || maxTimestamp == Long.MIN_VALUE) {
+      if (!emitStrategy.isOnPeriodic()) {
+        return;
+      }
+
+      emitIfReady(output);
+    }
+
+    private void emitIfReady(WatermarkOutput output) {
+      if (recordCount < configuration.getMinRecords() || maxTimestamp == Long.MIN_VALUE) {
         // Avoid emitting during warm-up, so limited early samples won't distort the estimate.
         return;
       }
@@ -121,7 +162,7 @@ public void onPeriodicEmit(WatermarkOutput output) {
     private long calculateOutOfOrdernessMillis() {
       int sampleCount = (int) Math.min(recordCount - 1, SAMPLE_SIZE);
       if (sampleCount <= 0) {
-        return MIN_OUT_OF_ORDERNESS_MILLIS;
+        return configuration.getMinOutOfOrdernessMillis();
       }
 
       long[] samples = Arrays.copyOf(latenessSamples, sampleCount);
@@ -130,11 +171,13 @@ private long calculateOutOfOrdernessMillis() {
       // Use a high quantile rather than the maximum, so a single pathological record does not stall
       // all event-time progress until it rotates out of the sample buffer.
       int quantileIndex =
-          Math.min(sampleCount - 1, (int) Math.ceil(sampleCount * OUT_OF_ORDERNESS_QUANTILE) - 1);
+          Math.min(
+              sampleCount - 1,
+              (int) Math.ceil(sampleCount * configuration.getOutOfOrdernessQuantile()) - 1);
 
       return Math.min(
-          MAX_OUT_OF_ORDERNESS_MILLIS,
-          Math.max(MIN_OUT_OF_ORDERNESS_MILLIS, samples[quantileIndex]));
+          configuration.getMaxOutOfOrdernessMillis(),
+          Math.max(configuration.getMinOutOfOrdernessMillis(), samples[quantileIndex]));
     }
   }
 }

connectors/kafka-safe-connector/src/main/java/com/datasqrl/flinkrunner/connector/kafka/SourceWatermarkOptions.java

@@ -0,0 +1,103 @@
+/*
+ * Copyright © 2026 DataSQRL (contact@datasqrl.com)
+ *
+ * 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 com.datasqrl.flinkrunner.connector.kafka;
+
+import java.io.Serial;
+import java.io.Serializable;
+import java.time.Duration;
+import lombok.Value;
+import org.apache.flink.configuration.ConfigOption;
+import org.apache.flink.configuration.ConfigOptions;
+import org.apache.flink.configuration.ReadableConfig;
+import org.apache.flink.table.api.ValidationException;
+
+public class SourceWatermarkOptions {
+
+  public static final ConfigOption<Integer> SCAN_SOURCE_WATERMARK_MIN_RECORDS =
+      ConfigOptions.key("scan.source-watermark.min-records")
+          .intType()
+          .defaultValue(KafkaRecordTimestampWatermarkStrategy.MIN_RECORDS)
+          .withDescription(
+              "Minimum number of records to observe before emitting source watermarks.");
+
+  public static final ConfigOption<Duration> SCAN_SOURCE_WATERMARK_MIN_OUT_OF_ORDERNESS =
+      ConfigOptions.key("scan.source-watermark.min-out-of-orderness")
+          .durationType()
+          .defaultValue(
+              Duration.ofMillis(KafkaRecordTimestampWatermarkStrategy.MIN_OUT_OF_ORDERNESS_MILLIS))
+          .withDescription("Minimum adaptive out-of-orderness delay for source watermarks.");
+
+  public static final ConfigOption<Duration> SCAN_SOURCE_WATERMARK_MAX_OUT_OF_ORDERNESS =
+      ConfigOptions.key("scan.source-watermark.max-out-of-orderness")
+          .durationType()
+          .defaultValue(
+              Duration.ofMillis(KafkaRecordTimestampWatermarkStrategy.MAX_OUT_OF_ORDERNESS_MILLIS))
+          .withDescription("Maximum adaptive out-of-orderness delay for source watermarks.");
+
+  public static final ConfigOption<Double> SCAN_SOURCE_WATERMARK_OUT_OF_ORDERNESS_QUANTILE =
+      ConfigOptions.key("scan.source-watermark.out-of-orderness-quantile")
+          .doubleType()
+          .defaultValue(KafkaRecordTimestampWatermarkStrategy.OUT_OF_ORDERNESS_QUANTILE)
+          .withDescription(
+              "Quantile of observed lateness samples to use as the source watermark delay.");
+
+  public static SourceWatermarkConfig sourceWatermarkConfiguration(ReadableConfig tableOptions) {
+    var minRecords = tableOptions.get(SCAN_SOURCE_WATERMARK_MIN_RECORDS);
+    var minOutOfOrderness = tableOptions.get(SCAN_SOURCE_WATERMARK_MIN_OUT_OF_ORDERNESS);
+    var maxOutOfOrderness = tableOptions.get(SCAN_SOURCE_WATERMARK_MAX_OUT_OF_ORDERNESS);
+    var quantile = tableOptions.get(SCAN_SOURCE_WATERMARK_OUT_OF_ORDERNESS_QUANTILE);
+
+    if (minRecords <= 0) {
+      throw new ValidationException(
+          String.format("'%s' must be greater than 0.", SCAN_SOURCE_WATERMARK_MIN_RECORDS.key()));
+    }
+
+    if (minOutOfOrderness.isNegative()) {
+      throw new ValidationException(
+          String.format(
+              "'%s' must not be negative.", SCAN_SOURCE_WATERMARK_MIN_OUT_OF_ORDERNESS.key()));
+    }
+
+    if (maxOutOfOrderness.compareTo(minOutOfOrderness) < 0) {
+      throw new ValidationException(
+          String.format(
+              "'%s' must be greater than or equal to '%s'.",
+              SCAN_SOURCE_WATERMARK_MAX_OUT_OF_ORDERNESS.key(),
+              SCAN_SOURCE_WATERMARK_MIN_OUT_OF_ORDERNESS.key()));
+    }
+
+    if (quantile <= 0D || quantile > 1D) {
+      throw new ValidationException(
+          String.format(
+              "'%s' must be greater than 0 and less than or equal to 1.",
+              SCAN_SOURCE_WATERMARK_OUT_OF_ORDERNESS_QUANTILE.key()));
+    }
+
+    return new SourceWatermarkConfig(
+        minRecords, minOutOfOrderness.toMillis(), maxOutOfOrderness.toMillis(), quantile);
+  }
+
+  @Value
+  public static class SourceWatermarkConfig implements Serializable {
+
+    @Serial private static final long serialVersionUID = 1L;
+
+    int minRecords;
+    long minOutOfOrdernessMillis;
+    long maxOutOfOrdernessMillis;
+    double outOfOrdernessQuantile;
+  }
+}

connectors/kafka-safe-connector/src/main/java/org/apache/flink/streaming/connectors/kafka/table/SafeKafkaDynamicSource.java

@@ -16,6 +16,7 @@
 package org.apache.flink.streaming.connectors.kafka.table;
 
 import com.datasqrl.flinkrunner.connector.kafka.KafkaRecordTimestampWatermarkStrategy;
+import com.datasqrl.flinkrunner.connector.kafka.SourceWatermarkOptions.SourceWatermarkConfig;
 import org.apache.flink.annotation.Internal;
 import org.apache.flink.api.common.eventtime.WatermarkStrategy;
 import org.apache.flink.api.common.serialization.DeserializationSchema;
@@ -52,6 +53,7 @@
 import org.apache.flink.table.types.logical.LogicalType;
 import org.apache.flink.table.types.logical.RowType;
 import org.apache.flink.table.types.logical.utils.LogicalTypeUtils;
+import org.apache.flink.table.watermark.WatermarkEmitStrategy;
 import org.apache.flink.util.Preconditions;
 
 import com.datasqrl.flinkrunner.connector.kafka.DeserFailureHandler;
@@ -63,6 +65,7 @@
 
 import javax.annotation.Nullable;
 
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -107,6 +110,15 @@ public class SafeKafkaDynamicSource
     /** Whether to use Kafka record timestamps to generate source watermarks. */
     protected boolean sourceWatermarkEnabled;
 
+    /** Emit strategy for source watermarks. */
+    protected final WatermarkEmitStrategy sourceWatermarkEmitStrategy;
+
+    /** Idle timeout for source watermarks. */
+    protected final Optional<Duration> sourceWatermarkIdleTimeout;
+
+    /** Adaptive source watermark configuration. */
+    protected final SourceWatermarkConfig sourceWatermarkConfig;
+
     // --------------------------------------------------------------------------------------------
     // Format attributes
     // --------------------------------------------------------------------------------------------
@@ -205,7 +217,10 @@ public SafeKafkaDynamicSource(
             boolean upsertMode,
             String tableIdentifier,
             @Nullable Integer parallelism,
-            DeserFailureHandler deserFailureHandler) {
+            DeserFailureHandler deserFailureHandler,
+            WatermarkEmitStrategy sourceWatermarkEmitStrategy,
+            Optional<Duration> sourceWatermarkIdleTimeout,
+            SourceWatermarkConfig sourceWatermarkConfig) {
         // Format attributes
         this.physicalDataType =
                 Preconditions.checkNotNull(
@@ -248,6 +263,18 @@ public SafeKafkaDynamicSource(
         this.tableIdentifier = tableIdentifier;
         this.parallelism = parallelism;
         this.deserFailureHandler = deserFailureHandler;
+        this.sourceWatermarkEmitStrategy =
+                Preconditions.checkNotNull(
+                        sourceWatermarkEmitStrategy,
+                        "Source watermark emit strategy must not be null.");
+        this.sourceWatermarkIdleTimeout =
+                Preconditions.checkNotNull(
+                        sourceWatermarkIdleTimeout,
+                        "Source watermark idle timeout must not be null.");
+        this.sourceWatermarkConfig =
+                Preconditions.checkNotNull(
+                    sourceWatermarkConfig,
+                        "Source watermark configuration must not be null.");
     }
 
     @Override
@@ -374,7 +401,10 @@ public DynamicTableSource copy() {
                         upsertMode,
                         tableIdentifier,
                         parallelism,
-                        deserFailureHandler);
+                        deserFailureHandler,
+                        sourceWatermarkEmitStrategy,
+                        sourceWatermarkIdleTimeout,
+                    sourceWatermarkConfig);
         copy.producedDataType = producedDataType;
         copy.metadataKeys = metadataKeys;
         copy.watermarkStrategy = watermarkStrategy;
@@ -417,6 +447,9 @@ public boolean equals(Object o) {
                 && Objects.equals(tableIdentifier, that.tableIdentifier)
                 && Objects.equals(watermarkStrategy, that.watermarkStrategy)
                 && sourceWatermarkEnabled == that.sourceWatermarkEnabled
+                && sourceWatermarkEmitStrategy == that.sourceWatermarkEmitStrategy
+                && Objects.equals(sourceWatermarkIdleTimeout, that.sourceWatermarkIdleTimeout)
+                && Objects.equals(sourceWatermarkConfig, that.sourceWatermarkConfig)
                 && Objects.equals(parallelism, that.parallelism);
     }
 
@@ -444,6 +477,9 @@ public int hashCode() {
                 tableIdentifier,
                 watermarkStrategy,
                 sourceWatermarkEnabled,
+                sourceWatermarkEmitStrategy,
+                sourceWatermarkIdleTimeout,
+            sourceWatermarkConfig,
                 parallelism);
     }
 
@@ -530,7 +566,15 @@ private WatermarkStrategy<RowData> getWatermarkStrategy() {
         }
 
         if (sourceWatermarkEnabled) {
-            return KafkaRecordTimestampWatermarkStrategy.INSTANCE;
+            WatermarkStrategy<RowData> sourceWatermarkStrategy =
+                    new KafkaRecordTimestampWatermarkStrategy(
+                            sourceWatermarkEmitStrategy, sourceWatermarkConfig);
+            if (sourceWatermarkIdleTimeout.isPresent()) {
+                sourceWatermarkStrategy =
+                        sourceWatermarkStrategy.withIdleness(sourceWatermarkIdleTimeout.get());
+            }
+
+            return sourceWatermarkStrategy;
         }
 
         return WatermarkStrategy.noWatermarks();