[improve][pulsar-io] Added support for generic record and raw JSON string schemas to CassandraSink

pulsar-io/cassandra/pom.xml

@@ -51,6 +51,39 @@
     <dependency>
       <groupId>com.datastax.cassandra</groupId>
       <artifactId>cassandra-driver-core</artifactId>
+      <version>${cassandra.version}</version>
+      <exclusions>
+        <exclusion>
+          <groupId>io.dropwizard.metrics</groupId>
+          <artifactId>metrics-core</artifactId>
+        </exclusion>
+      </exclusions>
+    </dependency>
+
+    <dependency>
+      <groupId>org.apache.pulsar</groupId>
+      <artifactId>pulsar-functions-local-runner-original</artifactId>
+      <version>${project.version}</version>
+      <scope>test</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>org.apache.pulsar</groupId>
+      <artifactId>pulsar-io-common</artifactId>
+      <version>${project.version}</version>
+      <scope>compile</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>commons-beanutils</groupId>
+      <artifactId>commons-beanutils</artifactId>
+      <scope>compile</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>org.testcontainers</groupId>
+      <artifactId>cassandra</artifactId>
+      <scope>test</scope>
     </dependency>
 
   </dependencies>

pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/CassandraAbstractSink.java

@@ -19,58 +19,65 @@
 package org.apache.pulsar.io.cassandra;
 
 import com.datastax.driver.core.BoundStatement;
-import com.datastax.driver.core.Cluster;
 import com.datastax.driver.core.PreparedStatement;
 import com.datastax.driver.core.ResultSet;
 import com.datastax.driver.core.ResultSetFuture;
-import com.datastax.driver.core.Session;
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.MoreExecutors;
 import java.util.Map;
 import org.apache.pulsar.functions.api.Record;
-import org.apache.pulsar.io.core.KeyValue;
+import org.apache.pulsar.io.cassandra.util.BoundStatementProvider;
+import org.apache.pulsar.io.cassandra.util.CassandraConnector;
+import org.apache.pulsar.io.cassandra.util.RecordWrapper;
+import org.apache.pulsar.io.cassandra.util.TableMetadataProvider;
+import org.apache.pulsar.io.common.IOConfigUtils;
 import org.apache.pulsar.io.core.Sink;
 import org.apache.pulsar.io.core.SinkContext;
+import org.apache.pulsar.io.core.annotations.Connector;
+import org.apache.pulsar.io.core.annotations.IOType;
 
-/**
- * A Simple abstract class for Cassandra sink.
- * Users need to implement extractKeyValue function to use this sink
- */
-public abstract class CassandraAbstractSink<K, V> implements Sink<byte[]> {
+@Connector(
+        name = "cassandra",
+        type = IOType.SINK,
+        help = "The CassandraStringSink is used for moving messages from Pulsar to Cassandra.",
+        configClass = CassandraSinkConfig.class)
+public abstract class CassandraAbstractSink<T> implements Sink<T> {
 
-    // ----- Runtime fields
-    private Cluster cluster;
-    private Session session;
+    CassandraConnector connector;
     CassandraSinkConfig cassandraSinkConfig;
-    private PreparedStatement statement;
+    PreparedStatement stmt;
+    BoundStatementProvider boundStatementProvider;
 
     @Override
-    public void open(Map<String, Object> config, SinkContext sinkContext) throws Exception {
-        cassandraSinkConfig = CassandraSinkConfig.load(config);
+    public void open(Map<String, Object> config, SinkContext ctx) throws Exception {
+
+        cassandraSinkConfig = IOConfigUtils.loadWithSecrets(config, CassandraSinkConfig.class, ctx);
+
         if (cassandraSinkConfig.getRoots() == null
                 || cassandraSinkConfig.getKeyspace() == null
-                || cassandraSinkConfig.getKeyname() == null
-                || cassandraSinkConfig.getColumnFamily() == null
-                || cassandraSinkConfig.getColumnName() == null) {
+                || cassandraSinkConfig.getColumnFamily() == null) {
             throw new IllegalArgumentException("Required property not set.");
         }
-        createClient(cassandraSinkConfig.getRoots());
-        statement = session.prepare("INSERT INTO " + cassandraSinkConfig.getColumnFamily() + " ("
-                + cassandraSinkConfig.getKeyname() + ", " + cassandraSinkConfig.getColumnName() + ") VALUES (?, ?)");
-    }
 
-    @Override
-    public void close() throws Exception {
-        session.close();
-        cluster.close();
+        connector = new CassandraConnector(cassandraSinkConfig);
+        connector.connect();
+
+        boundStatementProvider = new BoundStatementProvider(
+                TableMetadataProvider.getTableDefinition(
+                        connector.getTableMetadata(),
+                        cassandraSinkConfig.getKeyspace(),
+                        cassandraSinkConfig.getColumnFamily()));
     }
 
     @Override
-    public void write(Record<byte[]> record) {
-        KeyValue<K, V> keyValue = extractKeyValue(record);
-        BoundStatement bound = statement.bind(keyValue.getKey(), keyValue.getValue());
-        ResultSetFuture future = session.executeAsync(bound);
+    public void write(Record<T> record) throws Exception {
+
+        BoundStatement bs = boundStatementProvider.bindStatement(
+                getStatement(), wrapRecord(record));
+
+        ResultSetFuture future = connector.getSession().executeAsync(bs);
+
         Futures.addCallback(future,
                 new FutureCallback<ResultSet>() {
                     @Override
@@ -85,23 +92,24 @@ public void onFailure(Throwable t) {
                 }, MoreExecutors.directExecutor());
     }
 
-    private void createClient(String roots) {
-        String[] hosts = roots.split(",");
-        if (hosts.length <= 0) {
-            throw new RuntimeException("Invalid cassandra roots");
-        }
-        Cluster.Builder b = Cluster.builder();
-        for (int i = 0; i < hosts.length; ++i) {
-            String[] hostPort = hosts[i].split(":");
-            b.addContactPoint(hostPort[0]);
-            if (hostPort.length > 1) {
-                b.withPort(Integer.parseInt(hostPort[1]));
+    @Override
+    public void close()  {
+        if (connector != null) {
+            try {
+                connector.close();
+            } catch (final Throwable t) {
+
             }
         }
-        cluster = b.withoutJMXReporting().build();
-        session = cluster.connect();
-        session.execute("USE " + cassandraSinkConfig.getKeyspace());
     }
 
-    public abstract KeyValue<K, V> extractKeyValue(Record<byte[]> record);
-}
+    abstract RecordWrapper<T> wrapRecord(Record<T> record);
+
+    PreparedStatement getStatement() {
+        if (stmt == null) {
+            stmt = connector.getPreparedStatement();
+        }
+        return stmt;
+    }
+
+}

pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/CassandraGenericRecordSink.java

@@ -0,0 +1,32 @@
+/*
+ * 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.pulsar.io.cassandra;
+
+import org.apache.pulsar.client.api.schema.GenericRecord;
+import org.apache.pulsar.functions.api.Record;
+import org.apache.pulsar.io.cassandra.util.GenericRecordWrapper;
+import org.apache.pulsar.io.cassandra.util.RecordWrapper;
+
+public class CassandraGenericRecordSink extends CassandraAbstractSink<GenericRecord> {
+
+    @Override
+    RecordWrapper<GenericRecord> wrapRecord(Record<GenericRecord> record) {
+        return new GenericRecordWrapper(record.getValue());
+    }
+}

pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/CassandraSinkConfig.java

@@ -34,31 +34,39 @@ public class CassandraSinkConfig implements Serializable {
 
     private static final long serialVersionUID = 1L;
 
+    @FieldDoc(
+            required = false,
+            defaultValue = "",
+            sensitive = true,
+            help = "Username used to connect to the database specified by `root`"
+    )
+    private String userName;
+
+    @FieldDoc(
+            required = false,
+            defaultValue = "",
+            sensitive = true,
+            help = "Password used to connect to the database specified by `root`"
+    )
+    private String password;
+
     @FieldDoc(
         required = true,
         defaultValue = "",
         help = "A comma-separated list of cassandra hosts to connect to")
     private String roots;
+
     @FieldDoc(
         required = true,
         defaultValue = "",
         help = "The key space used for writing pulsar messages to")
     private String keyspace;
-    @FieldDoc(
-        required = true,
-        defaultValue = "",
-        help = "The key name of the cassandra column family")
-    private String keyname;
+
     @FieldDoc(
         required = true,
         defaultValue = "",
         help = "The cassandra column family name")
     private String columnFamily;
-    @FieldDoc(
-        required = true,
-        defaultValue = "",
-        help = "The column name of the cassandra column family")
-    private String columnName;
 
     public static CassandraSinkConfig load(String yamlFile) throws IOException {
         ObjectMapper mapper = new ObjectMapper(new YAMLFactory());

pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/CassandraStringSink.java

@@ -19,23 +19,13 @@
 package org.apache.pulsar.io.cassandra;
 
 import org.apache.pulsar.functions.api.Record;
-import org.apache.pulsar.io.core.KeyValue;
-import org.apache.pulsar.io.core.annotations.Connector;
-import org.apache.pulsar.io.core.annotations.IOType;
+import org.apache.pulsar.io.cassandra.util.RecordWrapper;
+import org.apache.pulsar.io.cassandra.util.StringRecordWrapper;
+
+public class CassandraStringSink extends CassandraAbstractSink<String> {
 
-/**
- * Cassandra sink that treats incoming messages on the input topic as Strings
- * and write identical key/value pairs.
- */
-@Connector(
-    name = "cassandra",
-    type = IOType.SINK,
-    help = "The CassandraStringSink is used for moving messages from Pulsar to Cassandra.",
-    configClass = CassandraSinkConfig.class)
-public class CassandraStringSink extends CassandraAbstractSink<String, String> {
     @Override
-    public KeyValue<String, String> extractKeyValue(Record<byte[]> record) {
-        String key = record.getKey().orElseGet(() -> new String(record.getValue()));
-        return new KeyValue<>(key, new String(record.getValue()));
+    RecordWrapper<String> wrapRecord(Record<String> record) {
+        return new StringRecordWrapper(record.getValue());
     }
-}
+}

pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/BoundStatementProvider.java

@@ -0,0 +1,45 @@
+/*
+ * 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.pulsar.io.cassandra.util;
+
+import com.datastax.driver.core.BoundStatement;
+import com.datastax.driver.core.PreparedStatement;
+
+@SuppressWarnings("rawtypes")
+public class BoundStatementProvider {
+    final TableMetadataProvider.TableDefinition tableDefinition;
+
+    public BoundStatementProvider(TableMetadataProvider.TableDefinition tableDefinition) {
+        this.tableDefinition = tableDefinition;
+    }
+
+    public BoundStatement bindStatement(PreparedStatement stmt, RecordWrapper wrapper) {
+        Object[] boundValues = new Object[tableDefinition.getColumns().size()];
+        int idx = 0;
+
+        for (TableMetadataProvider.ColumnId column : tableDefinition.getColumns()) {
+            if (wrapper.containsKey(column.getName())) {
+                boundValues[idx] = wrapper.get(column);
+            }
+            idx++;
+        }
+        return stmt.bind(boundValues);
+    }
+
+}