Add WriteExceptionHandler

kafka-connect/src/main/java/io/tabular/iceberg/connect/IcebergSinkConfig.java

@@ -96,6 +96,9 @@ public class IcebergSinkConfig extends AbstractConfig {
   private static final String NAME_PROP = "name";
   private static final String BOOTSTRAP_SERVERS_PROP = "bootstrap.servers";
 
+  private static final String WRITE_EXCEPTION_HANDLER_CLASS_PROP =
+      "experimental.write-exception-handler.class";
+
   private static final String DEFAULT_CATALOG_NAME = "iceberg";
   private static final String DEFAULT_CONTROL_TOPIC = "control-iceberg";
   public static final String DEFAULT_CONTROL_GROUP_PREFIX = "cg-control-";
@@ -237,6 +240,12 @@ private static ConfigDef newConfigDef() {
         null,
         Importance.MEDIUM,
         "Coordinator threads to use for table commits, default is (cores * 2)");
+    configDef.define(
+        WRITE_EXCEPTION_HANDLER_CLASS_PROP,
+        Type.STRING,
+        null,
+        Importance.MEDIUM,
+        "The WriteExceptionHandler implementation to use to handle exceptions when writing records to files");
     return configDef;
   }
 
@@ -451,6 +460,10 @@ public JsonConverter jsonConverter() {
     return jsonConverter;
   }
 
+  public String writeExceptionHandlerClassName() {
+    return getString(WRITE_EXCEPTION_HANDLER_CLASS_PROP);
+  }
+
   private Map<String, String> loadWorkerProps() {
     String javaCmd = System.getProperty("sun.java.command");
     if (javaCmd != null && !javaCmd.isEmpty()) {

kafka-connect/src/main/java/io/tabular/iceberg/connect/WriteExceptionHandler.java

@@ -0,0 +1,70 @@
+/*
+ * 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 io.tabular.iceberg.connect;
+
+import org.apache.kafka.connect.sink.SinkRecord;
+import org.apache.kafka.connect.sink.SinkTaskContext;
+
+public interface WriteExceptionHandler {
+  void initialize(SinkTaskContext context, IcebergSinkConfig config);
+
+  class Result {
+    private final SinkRecord sinkRecord;
+    private final String tableName;
+
+    public Result(SinkRecord sinkRecord, String tableName) {
+      this.sinkRecord = sinkRecord;
+      this.tableName = tableName;
+    }
+
+    public SinkRecord sinkRecord() {
+      return sinkRecord;
+    }
+
+    public String tableName() {
+      return tableName;
+    }
+  }
+
+  /**
+   * This method will be invoked whenever the connector runs into an exception while trying to write
+   * SinkRecords to a table. Implementations of this method have 3 general options:
+   *
+   * <ol>
+   *   <li>Return a SinkRecord and the name of the table to write to (wrapped inside a {@link
+   *       Result})
+   *   <li>Return null to drop the SinkRecord
+   *   <li>Throw an exception to fail the task
+   * </ol>
+   *
+   * @param sinkRecord The SinkRecord that couldn't be written
+   * @param tableName The table the SinkRecord couldn't be written to
+   * @param exception The exception encountered while trying to write the SinkRecord
+   */
+  Result handle(SinkRecord sinkRecord, String tableName, Exception exception) throws Exception;
+
+  /**
+   * This method will be invoked prior to committing allowing advanced {@link WriteExceptionHandler}
+   * implementations to complete any inflight work before the connector commits.
+   *
+   * <p>Note that there is no guarantee that the connector will successfully commit after this
+   * method is called.
+   */
+  void preCommit();
+}

kafka-connect/src/main/java/io/tabular/iceberg/connect/channel/ExceptionHandlingTableWriter.java

@@ -0,0 +1,68 @@
+/*
+ * 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 io.tabular.iceberg.connect.channel;
+
+import io.tabular.iceberg.connect.WriteExceptionHandler;
+import io.tabular.iceberg.connect.data.Utilities;
+import io.tabular.iceberg.connect.data.WriterResult;
+import java.util.List;
+import org.apache.kafka.connect.sink.SinkRecord;
+
+class ExceptionHandlingTableWriter implements TableWriter, AutoCloseable {
+
+  private final TableWriter underlying;
+  private final WriteExceptionHandler exceptionHandler;
+
+  ExceptionHandlingTableWriter(TableWriter underlying, WriteExceptionHandler exceptionHandler) {
+    this.underlying = underlying;
+    this.exceptionHandler = exceptionHandler;
+  }
+
+  @Override
+  public void write(SinkRecord sinkRecord, String tableName, boolean ignoreMissingTable) {
+    try {
+      underlying.write(sinkRecord, tableName, ignoreMissingTable);
+    } catch (Exception exception) {
+      final WriteExceptionHandler.Result result;
+      try {
+        result = exceptionHandler.handle(sinkRecord, tableName, exception);
+      } catch (Exception e) {
+        throw new RuntimeException(e);
+      }
+
+      if (result != null) {
+        // TODO: ignoreMissingTables=false make sense here? I _think_ so but I could also expose
+        // it
+        write(result.sinkRecord(), result.tableName(), false);
+      }
+    }
+  }
+
+  @Override
+  public List<WriterResult> committable() {
+    exceptionHandler.preCommit();
+    return underlying.committable();
+  }
+
+  @Override
+  public void close() {
+    Utilities.close(exceptionHandler);
+    Utilities.close(underlying);
+  }
+}

kafka-connect/src/main/java/io/tabular/iceberg/connect/channel/MultiTableWriter.java

@@ -0,0 +1,74 @@
+/*
+ * 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 io.tabular.iceberg.connect.channel;
+
+import io.tabular.iceberg.connect.data.IcebergWriterFactory;
+import io.tabular.iceberg.connect.data.RecordWriter;
+import io.tabular.iceberg.connect.data.Utilities;
+import io.tabular.iceberg.connect.data.WriterResult;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+import org.apache.kafka.connect.sink.SinkRecord;
+
+class MultiTableWriter implements TableWriter, AutoCloseable {
+
+  private final IcebergWriterFactory writerFactory;
+  private final Map<String, RecordWriter> writers;
+
+  MultiTableWriter(IcebergWriterFactory writerFactory) {
+    this.writerFactory = writerFactory;
+    this.writers = Maps.newHashMap();
+  }
+
+  private RecordWriter writerForTable(
+      String tableName, SinkRecord sample, boolean ignoreMissingTable) {
+    return writers.computeIfAbsent(
+        tableName, notUsed -> writerFactory.createWriter(tableName, sample, ignoreMissingTable));
+  }
+
+  @Override
+  public void write(SinkRecord sinkRecord, String tableName, boolean ignoreMissingTable) {
+    writerForTable(tableName, sinkRecord, ignoreMissingTable).write(sinkRecord);
+  }
+
+  private void closeWriters() {
+    writers.values().forEach(Utilities::close);
+    writers.clear();
+  }
+
+  @Override
+  public List<WriterResult> committable() {
+    List<WriterResult> writerResults =
+        writers.values().stream()
+            .flatMap(writer -> writer.complete().stream())
+            .collect(Collectors.toList());
+
+    closeWriters();
+
+    return writerResults;
+  }
+
+  @Override
+  public void close() {
+    closeWriters();
+    Utilities.close(writerFactory);
+  }
+}

kafka-connect/src/main/java/io/tabular/iceberg/connect/channel/TableWriter.java

@@ -0,0 +1,29 @@
+/*
+ * 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 io.tabular.iceberg.connect.channel;
+
+import io.tabular.iceberg.connect.data.WriterResult;
+import java.util.List;
+import org.apache.kafka.connect.sink.SinkRecord;
+
+interface TableWriter {
+  void write(SinkRecord sinkRecord, String tableName, boolean ignoreMissingTable);
+
+  List<WriterResult> committable();
+}

kafka-connect/src/main/java/io/tabular/iceberg/connect/channel/TableWriterImpl.java

@@ -0,0 +1,75 @@
+/*
+ * 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 io.tabular.iceberg.connect.channel;
+
+import io.tabular.iceberg.connect.IcebergSinkConfig;
+import io.tabular.iceberg.connect.WriteExceptionHandler;
+import io.tabular.iceberg.connect.data.IcebergWriterFactory;
+import io.tabular.iceberg.connect.data.Utilities;
+import io.tabular.iceberg.connect.data.WriterResult;
+import java.util.List;
+import org.apache.kafka.connect.sink.SinkRecord;
+import org.apache.kafka.connect.sink.SinkTaskContext;
+
+class TableWriterImpl implements TableWriter, AutoCloseable {
+
+  private final TableWriter underlying;
+
+  private WriteExceptionHandler loadHandler(String name) {
+    ClassLoader loader = this.getClass().getClassLoader();
+    final Object obj;
+    try {
+      Class<?> clazz = Class.forName(name, true, loader);
+      obj = clazz.getDeclaredConstructor().newInstance();
+    } catch (Exception e) {
+      throw new RuntimeException(String.format("Could not initialize class %s", name), e);
+    }
+
+    final WriteExceptionHandler writeExceptionHandler = (WriteExceptionHandler) obj;
+
+    return writeExceptionHandler;
+  }
+
+  TableWriterImpl(
+      SinkTaskContext context, IcebergSinkConfig config, IcebergWriterFactory writerFactory) {
+    MultiTableWriter baseTableWriter = new MultiTableWriter(writerFactory);
+    if (config.writeExceptionHandlerClassName() == null) {
+      this.underlying = baseTableWriter;
+    } else {
+      WriteExceptionHandler handler = loadHandler(config.writeExceptionHandlerClassName());
+      handler.initialize(context, config);
+      this.underlying = new ExceptionHandlingTableWriter(baseTableWriter, handler);
+    }
+  }
+
+  @Override
+  public void write(SinkRecord sinkRecord, String tableName, boolean ignoreMissingTable) {
+    underlying.write(sinkRecord, tableName, ignoreMissingTable);
+  }
+
+  @Override
+  public List<WriterResult> committable() {
+    return underlying.committable();
+  }
+
+  @Override
+  public void close() {
+    Utilities.close(underlying);
+  }
+}