Merge pull request #8 from questdb/feat/kafka_native_timestamps

feat: optionally use Kafka timestamps as designated timestamps

Author: jerrinot

connector/src/main/java/io/questdb/kafka/QuestDBSinkConnector.java

@@ -1,5 +1,6 @@
 package io.questdb.kafka;
 
+import org.apache.kafka.common.config.Config;
 import org.apache.kafka.common.config.ConfigDef;
 import org.apache.kafka.connect.connector.Task;
 import org.apache.kafka.connect.sink.SinkConnector;
@@ -43,4 +44,14 @@ public void stop() {
     public ConfigDef config() {
         return QuestDBSinkConnectorConfig.conf();
     }
+
+    @Override
+    public Config validate(Map<String, String> connectorConfigs) {
+        String s = connectorConfigs.get(QuestDBSinkConnectorConfig.DESIGNATED_TIMESTAMP_KAFKA_NATIVE_CONFIG);
+        if (Boolean.parseBoolean(s) && connectorConfigs.get(QuestDBSinkConnectorConfig.DESIGNATED_TIMESTAMP_COLUMN_NAME_CONFIG) != null) {
+            throw new IllegalArgumentException("Cannot use " + QuestDBSinkConnectorConfig.DESIGNATED_TIMESTAMP_COLUMN_NAME_CONFIG
+                    + " with " + QuestDBSinkConnectorConfig.DESIGNATED_TIMESTAMP_KAFKA_NATIVE_CONFIG +". These options are mutually exclusive.");
+        }
+        return super.validate(connectorConfigs);
+    }
 }

connector/src/main/java/io/questdb/kafka/QuestDBSinkConnectorConfig.java

@@ -33,6 +33,9 @@ public final class QuestDBSinkConnectorConfig extends AbstractConfig {
     public static final String DESIGNATED_TIMESTAMP_COLUMN_NAME_CONFIG = "timestamp.field.name";
     private static final String DESIGNATED_TIMESTAMP_COLUMN_NAME_DOC = "Designated timestamp field name";
 
+    public static final String DESIGNATED_TIMESTAMP_KAFKA_NATIVE_CONFIG = "timestamp.kafka.native";
+    private static final String DESIGNATED_TIMESTAMP_KAFKA_NATIVE_DOC = "Use Kafka record timestamps as designated timestamp. Mutually exclusive with timestamp.field.name";
+
     public static final String TIMESTAMP_STRING_FIELDS = "timestamp.string.fields";
     private static final String TIMESTAMP_STRING_FIELDS_DOC = "Comma-separated list of string fields that should be parsed as timestamps.";
 
@@ -94,7 +97,8 @@ public static ConfigDef conf() {
                 .define(RETRY_BACKOFF_MS, Type.LONG, 3_000, Importance.LOW, RETRY_BACKOFF_MS_DOC)
                 .define(MAX_RETRIES, Type.INT, 10, Importance.LOW, MAX_RETRIES_DOC)
                 .define(TIMESTAMP_FORMAT, Type.STRING, DEFAULT_TIMESTAMP_FORMAT, TimestampFormatValidator.INSTANCE, Importance.MEDIUM, TIMESTAMP_FORMAT_DOC)
-                .define(TIMESTAMP_STRING_FIELDS, Type.STRING, null, Importance.MEDIUM, TIMESTAMP_STRING_FIELDS_DOC);
+                .define(TIMESTAMP_STRING_FIELDS, Type.STRING, null, Importance.MEDIUM, TIMESTAMP_STRING_FIELDS_DOC)
+                .define(DESIGNATED_TIMESTAMP_KAFKA_NATIVE_CONFIG, Type.BOOLEAN, false, Importance.MEDIUM, DESIGNATED_TIMESTAMP_KAFKA_NATIVE_DOC);
     }
 
     public String getHost() {
@@ -121,6 +125,10 @@ public String getValuePrefix() {
         return getString(VALUE_PREFIX_CONFIG);
     }
 
+    public boolean isDesignatedTimestampKafkaNative() {
+        return getBoolean(DESIGNATED_TIMESTAMP_KAFKA_NATIVE_CONFIG);
+    }
+
     public boolean isSkipUnsupportedTypes() {
         return getBoolean(SKIP_UNSUPPORTED_TYPES_CONFIG);
     }

connector/src/main/java/io/questdb/kafka/QuestDBSinkTask.java

@@ -46,6 +46,7 @@ public final class QuestDBSinkTask extends SinkTask {
     private int remainingRetries;
     private long batchesSinceLastError = 0;
     private DateFormat dataFormat;
+    private boolean kafkaTimestampsEnabled;
 
     @Override
     public String version() {
@@ -78,6 +79,7 @@ public void start(Map<String, String> map) {
         this.sender = createSender();
         this.remainingRetries = config.getMaxRetries();
         this.timestampColumnName = config.getDesignatedTimestampColumnName();
+        this.kafkaTimestampsEnabled = config.isDesignatedTimestampKafkaNative();
         this.timestampUnits = config.getTimestampUnitsOrNull();
     }
 
@@ -176,11 +178,17 @@ private void handleSingleRecord(SinkRecord record) {
         }
         handleObject(config.getValuePrefix(), record.valueSchema(), record.value(), PRIMITIVE_VALUE_FALLBACK_NAME);
 
+        if (kafkaTimestampsEnabled) {
+            timestampColumnValue = TimeUnit.MILLISECONDS.toNanos(record.timestamp());
+        }
         if (timestampColumnValue == Long.MIN_VALUE) {
             sender.atNow();
         } else {
-            sender.at(timestampColumnValue);
-            timestampColumnValue = Long.MIN_VALUE;
+            try {
+                sender.at(timestampColumnValue);
+            } finally {
+                timestampColumnValue = Long.MIN_VALUE;
+            }
         }
     }
 
@@ -223,6 +231,7 @@ private boolean isDesignatedColumnName(String name, String fallbackName) {
 
     private void handleObject(String name, Schema schema, Object value, String fallbackName) {
         assert !name.isEmpty() || !fallbackName.isEmpty();
+
         if (isDesignatedColumnName(name, fallbackName)) {
             assert timestampColumnValue == Long.MIN_VALUE;
             if (value == null) {

connector/src/test/java/io/questdb/kafka/QuestDBSinkConnectorEmbeddedTest.java

@@ -2,15 +2,19 @@
 
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.connect.data.Date;
 import org.apache.kafka.connect.data.Decimal;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaBuilder;
 import org.apache.kafka.connect.data.Struct;
 import org.apache.kafka.connect.data.Time;
 import org.apache.kafka.connect.data.Timestamp;
+import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.json.JsonConverter;
 import org.apache.kafka.connect.runtime.AbstractStatus;
+import org.apache.kafka.connect.runtime.Connect;
 import org.apache.kafka.connect.runtime.ConnectorConfig;
 import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
 import org.apache.kafka.connect.storage.Converter;
@@ -46,6 +50,8 @@
 import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.apache.kafka.connect.runtime.ConnectorConfig.KEY_CONVERTER_CLASS_CONFIG;
 import static org.apache.kafka.connect.runtime.ConnectorConfig.VALUE_CONVERTER_CLASS_CONFIG;
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.jupiter.api.Assertions.fail;
 
 @Testcontainers
@@ -457,6 +463,53 @@ private void testTimestampUnitResolution0(String mode) {
                 "select firstname,lastname,timestamp from " + topicName);
     }
 
+    @Test
+    public void testKafkaNativeTimestampsAndExplicitDesignatedFieldTimestampMutuallyExclusive() {
+        Map<String, String> props = baseConnectorProps(topicName);
+        props.put("value.converter.schemas.enable", "false");
+        props.put(QuestDBSinkConnectorConfig.DESIGNATED_TIMESTAMP_COLUMN_NAME_CONFIG, "born");
+        props.put(QuestDBSinkConnectorConfig.DESIGNATED_TIMESTAMP_KAFKA_NATIVE_CONFIG, "true");
+        try {
+            connect.configureConnector(CONNECTOR_NAME, props);
+            fail("Expected ConnectException");
+        } catch (ConnectException e) {
+            assertThat(e.getMessage(), containsString("timestamp.field.name with timestamp.kafka.native"));
+        }
+    }
+
+    @Test
+    public void testKafkaNativeTimestamp() {
+        connect.kafka().createTopic(topicName, 1);
+        Map<String, String> props = baseConnectorProps(topicName);
+        props.put("value.converter.schemas.enable", "false");
+        props.put(QuestDBSinkConnectorConfig.DESIGNATED_TIMESTAMP_KAFKA_NATIVE_CONFIG, "true");
+
+        connect.configureConnector(CONNECTOR_NAME, props);
+        assertConnectorTaskRunningEventually();
+
+        QuestDBUtils.assertSql(questDBContainer,
+                "{\"ddl\":\"OK\"}\n",
+                "create table " + topicName + " (firstname string, lastname string, born timestamp) timestamp(born)",
+                QuestDBUtils.Endpoint.EXEC);
+
+        java.util.Date birth = new Calendar.Builder()
+                .setTimeZone(TimeZone.getTimeZone("UTC"))
+                .setDate(2022, 9, 23) // note: month is 0-based, so it's October and not November
+                .setTimeOfDay(13, 53, 59, 123)
+                .build().getTime();
+
+        Map<String, Object> prodProps = new HashMap<>();
+        try (KafkaProducer<byte[], byte[]> producer = connect.kafka().createProducer(prodProps)) {
+            String val = "{\"firstname\":\"John\",\"lastname\":\"Doe\"}";
+            ProducerRecord<byte[], byte[]> record = new ProducerRecord<>(topicName, null, birth.getTime(), null, val.getBytes());
+            producer.send(record);
+        }
+
+        QuestDBUtils.assertSqlEventually(questDBContainer, "\"firstname\",\"lastname\",\"born\"\r\n" +
+                        "\"John\",\"Doe\",\"2022-10-23T13:53:59.123000Z\"\r\n",
+                "select * from " + topicName);
+    }
+
     @Test
     public void testTimestampSMT_parseTimestamp_schemaLess() {
         connect.kafka().createTopic(topicName, 1);