diff --git a/pulsar-io/cassandra/pom.xml b/pulsar-io/cassandra/pom.xml
index 41faa38aa0a3e..2ab35a0c03649 100644
--- a/pulsar-io/cassandra/pom.xml
+++ b/pulsar-io/cassandra/pom.xml
@@ -51,6 +51,39 @@
com.datastax.cassandra
cassandra-driver-core
+ ${cassandra.version}
+
+
+ io.dropwizard.metrics
+ metrics-core
+
+
+
+
+
+ org.apache.pulsar
+ pulsar-functions-local-runner-original
+ ${project.version}
+ test
+
+
+
+ org.apache.pulsar
+ pulsar-io-common
+ ${project.version}
+ compile
+
+
+
+ commons-beanutils
+ commons-beanutils
+ compile
+
+
+
+ org.testcontainers
+ cassandra
+ test
diff --git a/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/CassandraAbstractSink.java b/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/CassandraAbstractSink.java
index 4f96df280d2c5..9615e4e9a569a 100644
--- a/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/CassandraAbstractSink.java
+++ b/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 implements Sink {
+@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 implements Sink {
- // ----- Runtime fields
- private Cluster cluster;
- private Session session;
+ CassandraConnector connector;
CassandraSinkConfig cassandraSinkConfig;
- private PreparedStatement statement;
+ PreparedStatement stmt;
+ BoundStatementProvider boundStatementProvider;
@Override
- public void open(Map config, SinkContext sinkContext) throws Exception {
- cassandraSinkConfig = CassandraSinkConfig.load(config);
+ public void open(Map 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 record) {
- KeyValue keyValue = extractKeyValue(record);
- BoundStatement bound = statement.bind(keyValue.getKey(), keyValue.getValue());
- ResultSetFuture future = session.executeAsync(bound);
+ public void write(Record record) throws Exception {
+
+ BoundStatement bs = boundStatementProvider.bindStatement(
+ getStatement(), wrapRecord(record));
+
+ ResultSetFuture future = connector.getSession().executeAsync(bs);
+
Futures.addCallback(future,
new FutureCallback() {
@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 extractKeyValue(Record record);
-}
\ No newline at end of file
+ abstract RecordWrapper wrapRecord(Record record);
+
+ PreparedStatement getStatement() {
+ if (stmt == null) {
+ stmt = connector.getPreparedStatement();
+ }
+ return stmt;
+ }
+
+}
diff --git a/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/CassandraGenericRecordSink.java b/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/CassandraGenericRecordSink.java
new file mode 100644
index 0000000000000..d986efbc6fb2b
--- /dev/null
+++ b/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 {
+
+ @Override
+ RecordWrapper wrapRecord(Record record) {
+ return new GenericRecordWrapper(record.getValue());
+ }
+}
diff --git a/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/CassandraSinkConfig.java b/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/CassandraSinkConfig.java
index 1dfc69b4d11a2..ac57f73b1a036 100644
--- a/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/CassandraSinkConfig.java
+++ b/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());
diff --git a/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/CassandraStringSink.java b/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/CassandraStringSink.java
index 789e7c2b73b73..ed90113d50f25 100644
--- a/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/CassandraStringSink.java
+++ b/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 {
-/**
- * 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 {
@Override
- public KeyValue extractKeyValue(Record record) {
- String key = record.getKey().orElseGet(() -> new String(record.getValue()));
- return new KeyValue<>(key, new String(record.getValue()));
+ RecordWrapper wrapRecord(Record record) {
+ return new StringRecordWrapper(record.getValue());
}
-}
\ No newline at end of file
+}
diff --git a/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/BoundStatementProvider.java b/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/BoundStatementProvider.java
new file mode 100644
index 0000000000000..983bedf7564ed
--- /dev/null
+++ b/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);
+ }
+
+}
diff --git a/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/CassandraConnector.java b/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/CassandraConnector.java
new file mode 100644
index 0000000000000..ad9eb329b2174
--- /dev/null
+++ b/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/CassandraConnector.java
@@ -0,0 +1,137 @@
+/*
+ * 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.Cluster;
+import com.datastax.driver.core.ColumnMetadata;
+import com.datastax.driver.core.Metadata;
+import com.datastax.driver.core.PreparedStatement;
+import com.datastax.driver.core.Session;
+import com.datastax.driver.core.TableMetadata;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.pulsar.io.cassandra.CassandraSinkConfig;
+
+public class CassandraConnector implements AutoCloseable {
+
+ private Cluster cluster;
+ private Session session;
+ private PreparedStatement statement;
+ private List tableFields;
+ private final CassandraSinkConfig config;
+
+ public CassandraConnector(CassandraSinkConfig config) {
+ this.config = config;
+ }
+
+ public void connect() {
+ session = getCluster().connect(config.getKeyspace());
+ }
+
+ public synchronized Session getSession() {
+ if (session == null || session.isClosed()) {
+ this.connect();
+ }
+ return session;
+ }
+
+ public Metadata getTableMetadata() {
+ return getCluster().getMetadata();
+ }
+
+ public PreparedStatement getPreparedStatement() {
+ if (statement == null) {
+ List fields = getTableFields();
+
+ StringBuilder sb = new StringBuilder("INSERT INTO ")
+ .append(config.getKeyspace() + "." + config.getColumnFamily() + " (");
+
+ for (int idx = 0; idx < fields.size(); idx++) {
+ sb.append(fields.get(idx));
+ if (idx < fields.size() - 1) {
+ sb.append(", ");
+ }
+ }
+
+ sb.append(") VALUES (");
+
+ for (int idx = 0; idx < fields.size(); idx++) {
+ sb.append("?");
+ if (idx < fields.size() - 1) {
+ sb.append(", ");
+ }
+ }
+
+ sb.append(")");
+ statement = getSession().prepare(sb.toString());
+ }
+
+ return statement;
+ }
+
+ List getTableFields() {
+
+ if (tableFields == null) {
+
+ TableMetadata meta = getCluster().getMetadata()
+ .getKeyspace(config.getKeyspace())
+ .getTable(config.getColumnFamily());
+
+ tableFields = new ArrayList(meta.getColumns().size());
+
+ for (ColumnMetadata col : meta.getColumns()) {
+ tableFields.add(col.getName());
+ }
+ }
+ return tableFields;
+ }
+
+ private synchronized Cluster getCluster() {
+ if (cluster == null) {
+ String[] hosts = config.getRoots().split(",");
+
+ Cluster.Builder builder = Cluster.builder().withoutMetrics();
+
+ for (int i = 0; i < hosts.length; ++i) {
+ String[] hostPort = hosts[i].split(":");
+ builder.addContactPoint(hostPort[0]);
+ if (hostPort.length > 1) {
+ builder.withPort(Integer.parseInt(hostPort[1]));
+ }
+ }
+
+ // Authenticate if credentials have been provided
+ if (config.getUserName() != null
+ && config.getPassword() != null) {
+ builder.withCredentials(
+ config.getUserName(),
+ config.getPassword()
+ );
+ }
+ cluster = builder.build();
+ }
+
+ return cluster;
+ }
+
+ public void close() {
+ getSession().close();
+ getCluster().close();
+ }
+}
diff --git a/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/GenericRecordWrapper.java b/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/GenericRecordWrapper.java
new file mode 100644
index 0000000000000..f7f725be366e6
--- /dev/null
+++ b/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/GenericRecordWrapper.java
@@ -0,0 +1,38 @@
+/*
+ * 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 org.apache.pulsar.client.api.schema.GenericRecord;
+
+public class GenericRecordWrapper extends RecordWrapper {
+
+ public GenericRecordWrapper(GenericRecord value) {
+ super(value);
+ }
+
+ @Override
+ public Object get(TableMetadataProvider.ColumnId column) {
+ return getValueAsExpectedType(recordValue.getField(column.getName()), column);
+ }
+
+ @Override
+ public boolean containsKey(String name) {
+ return this.recordValue.getField(name) != null;
+ }
+}
diff --git a/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/RecordWrapper.java b/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/RecordWrapper.java
new file mode 100644
index 0000000000000..628f0c430f1dc
--- /dev/null
+++ b/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/RecordWrapper.java
@@ -0,0 +1,48 @@
+/*
+ * 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 org.apache.commons.beanutils.converters.IntegerConverter;
+import org.apache.commons.beanutils.converters.NumberConverter;
+
+public abstract class RecordWrapper {
+
+ T recordValue;
+ NumberConverter converter = new IntegerConverter();
+
+ public RecordWrapper(T value) {
+ this.recordValue = value;
+ }
+
+ public abstract Object get(TableMetadataProvider.ColumnId column);
+
+ public abstract boolean containsKey(String name);
+
+ Object getValueAsExpectedType(Object value, TableMetadataProvider.ColumnId column) {
+ switch (column.getType().getName()) {
+ case FLOAT: return converter.convert(Float.class, value);
+ case INT: return converter.convert(Integer.class, value);
+ case DOUBLE: return converter.convert(Double.class, value);
+ case TEXT: return value.toString();
+ default: return value;
+ }
+
+ }
+
+}
diff --git a/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/StringRecordWrapper.java b/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/StringRecordWrapper.java
new file mode 100644
index 0000000000000..6ca16e6d54069
--- /dev/null
+++ b/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/StringRecordWrapper.java
@@ -0,0 +1,48 @@
+/*
+ * 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.fasterxml.jackson.databind.ObjectMapper;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.util.Map;
+
+public class StringRecordWrapper extends RecordWrapper {
+ private static final ObjectMapper MAPPER = new ObjectMapper();
+ private Map valuesMap;
+
+ public StringRecordWrapper(String jsonString) {
+ super(jsonString);
+ try {
+ valuesMap = MAPPER.readValue(jsonString, Map.class);
+ } catch (IOException e) {
+ throw new UncheckedIOException(e);
+ }
+ }
+
+ @Override
+ public Object get(TableMetadataProvider.ColumnId column) {
+ return getValueAsExpectedType(valuesMap.get(column.getName()), column);
+ }
+
+ @Override
+ public boolean containsKey(String name) {
+ return valuesMap.containsKey(name);
+ }
+}
diff --git a/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/TableMetadataProvider.java b/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/TableMetadataProvider.java
new file mode 100644
index 0000000000000..d978dcf31396a
--- /dev/null
+++ b/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/TableMetadataProvider.java
@@ -0,0 +1,107 @@
+/*
+ * 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.ColumnMetadata;
+import com.datastax.driver.core.DataType;
+import com.datastax.driver.core.Metadata;
+import com.datastax.driver.core.TableMetadata;
+import com.google.common.collect.Lists;
+import java.util.List;
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.Setter;
+import lombok.ToString;
+
+public class TableMetadataProvider {
+
+ @Data(staticConstructor = "of")
+ public static class TableId {
+ private final String keyspaceName;
+ private final String columnFamily;
+ }
+
+ @Data(staticConstructor = "of")
+ public static class ColumnId {
+ private final TableId tableId;
+ private final String name;
+ private final DataType type;
+ }
+
+ @Setter
+ @Getter
+ @EqualsAndHashCode
+ @ToString
+ public static class TableDefinition {
+ private final TableId tableId;
+ private final List columns;
+ private final List partitionKeyColumns;
+ private final List primaryKeyColumns;
+
+ private TableDefinition(TableId tableId, List columns) {
+ this(tableId, columns, null, null);
+ }
+ private TableDefinition(TableId tableId, List columns,
+ List partitionKeyColumns, List primaryKeyColumns) {
+ this.tableId = tableId;
+ this.columns = columns;
+ this.partitionKeyColumns = partitionKeyColumns;
+ this.primaryKeyColumns = primaryKeyColumns;
+ }
+
+ public static TableDefinition of(TableId tableId, List columns) {
+ return new TableDefinition(tableId, columns);
+ }
+
+ public static TableDefinition of(TableId tableId, List columns,
+ List nonKeyColumns, List keyColumns) {
+ return new TableDefinition(tableId, columns, nonKeyColumns, keyColumns);
+ }
+
+ }
+
+ public static TableDefinition getTableDefinition(Metadata clusterMetadata, String keyspace, String columnFamily) {
+
+ TableId tableId = TableId.of(keyspace, columnFamily);
+ TableDefinition table = TableDefinition.of(tableId,
+ Lists.newArrayList(), Lists.newArrayList(), Lists.newArrayList());
+
+ TableMetadata meta = clusterMetadata
+ .getKeyspace(keyspace)
+ .getTable(columnFamily);
+
+ for (ColumnMetadata col : meta.getPrimaryKey()) {
+ ColumnId columnId = ColumnId.of(tableId, col.getName(), col.getType());
+ table.getPrimaryKeyColumns().add(columnId);
+ }
+
+ for (ColumnMetadata col : meta.getPartitionKey()) {
+ ColumnId columnId = ColumnId.of(tableId, col.getName(), col.getType());
+ table.getPartitionKeyColumns().add(columnId);
+ }
+
+ for (ColumnMetadata col : meta.getColumns()) {
+ ColumnId columnId = ColumnId.of(tableId, col.getName(), col.getType());
+ table.getColumns().add(columnId);
+ }
+
+ return table;
+ }
+}
diff --git a/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/package-info.java b/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/package-info.java
new file mode 100644
index 0000000000000..904579e812163
--- /dev/null
+++ b/pulsar-io/cassandra/src/main/java/org/apache/pulsar/io/cassandra/util/package-info.java
@@ -0,0 +1,19 @@
+/*
+ * 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;
\ No newline at end of file
diff --git a/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/CassandraSinkExec.java b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/CassandraSinkExec.java
new file mode 100644
index 0000000000000..613c21ad0cbf5
--- /dev/null
+++ b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/CassandraSinkExec.java
@@ -0,0 +1,99 @@
+/*
+ * 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 java.io.File;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import org.apache.pulsar.common.io.SinkConfig;
+import org.apache.pulsar.functions.LocalRunner;
+import org.apache.pulsar.io.cassandra.producers.InputTopicProducerThread;
+import org.apache.pulsar.io.cassandra.producers.ReadingSchemaRecordProducer;
+import org.yaml.snakeyaml.Yaml;
+
+/**
+ * Useful for testing within IDE.
+ *
+ */
+@SuppressWarnings({"unchecked", "rawtypes"})
+public class CassandraSinkExec {
+
+ public static final String BROKER_URL = "pulsar://localhost:6650";
+ public static final String INPUT_TOPIC = "persistent://public/default/air-quality-reading-generic";
+
+ public static final String CONFIG_FILE = "cassandra-sink-config.yaml";
+
+ public static void main(String[] args) throws Exception {
+
+ SinkConfig config = getSinkConfig();
+
+ final LocalRunner localRunner =
+ LocalRunner.builder()
+ .brokerServiceUrl(BROKER_URL)
+ .sinkConfig(config)
+ .build();
+
+ localRunner.start(false);
+
+ sendData();
+ TimeUnit.MINUTES.sleep(10);
+
+ localRunner.stop();
+
+ System.exit(0);
+ }
+
+ private static SinkConfig getSinkConfig() throws IOException {
+ SinkConfig sinkConfig = SinkConfig.builder()
+ .autoAck(true)
+ .cleanupSubscription(Boolean.TRUE)
+ .configs(getConfigs())
+ .className(CassandraGenericRecordSink.class.getName())
+ .inputs(Collections.singletonList(INPUT_TOPIC))
+ .name("CassandraSink")
+ .build();
+
+ return sinkConfig;
+ }
+
+ private static Map getConfigs() throws IOException {
+ Map configs = new HashMap();
+
+ ClassLoader classLoader = Thread.currentThread().getContextClassLoader();
+ File file = new File(classLoader.getResource(CONFIG_FILE).getFile());
+
+ try (FileInputStream fis = new FileInputStream(file)) {
+ configs = new Yaml().load(fis);
+ } catch (IOException ex) {
+ throw ex;
+ }
+
+ return configs;
+ }
+
+ private static void sendData() throws InterruptedException {
+ TimeUnit.SECONDS.sleep(10);
+ InputTopicProducerThread writer = new ReadingSchemaRecordProducer(BROKER_URL, INPUT_TOPIC);
+ writer.run();
+ }
+}
diff --git a/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/producers/AbstractGenericRecordProducer.java b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/producers/AbstractGenericRecordProducer.java
new file mode 100644
index 0000000000000..96c341c46f125
--- /dev/null
+++ b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/producers/AbstractGenericRecordProducer.java
@@ -0,0 +1,42 @@
+/*
+ * 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.producers;
+
+import org.apache.pulsar.client.api.Schema;
+import org.apache.pulsar.client.api.schema.GenericRecord;
+import org.apache.pulsar.common.schema.SchemaInfo;
+
+@SuppressWarnings({"unchecked", "rawtypes"})
+public abstract class AbstractGenericRecordProducer extends InputTopicProducerThread {
+
+ public AbstractGenericRecordProducer(String brokerUrl, String inputTopic) {
+ super(brokerUrl, inputTopic);
+ }
+
+ @Override
+ Schema getSchema() {
+ return Schema.generic(getGenericSchemaInfo());
+ }
+
+ @Override
+ abstract GenericRecord getValue();
+
+ abstract SchemaInfo getGenericSchemaInfo();
+
+ }
diff --git a/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/producers/InputTopicProducerThread.java b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/producers/InputTopicProducerThread.java
new file mode 100644
index 0000000000000..260cbbc1db602
--- /dev/null
+++ b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/producers/InputTopicProducerThread.java
@@ -0,0 +1,83 @@
+/*
+ * 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.producers;
+
+import java.util.Random;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.pulsar.client.api.Producer;
+import org.apache.pulsar.client.api.PulsarClient;
+import org.apache.pulsar.client.api.PulsarClientException;
+import org.apache.pulsar.client.api.Schema;
+
+@SuppressWarnings({"unchecked", "rawtypes"})
+@Slf4j
+public abstract class InputTopicProducerThread implements Runnable {
+
+ private Random rnd = new Random();
+ final String inputTopic;
+ final String brokerUrl;
+ PulsarClient client;
+ Producer producer;
+
+ public InputTopicProducerThread(String brokerUrl, String inputTopic) {
+ this.brokerUrl = brokerUrl;
+ this.inputTopic = inputTopic;
+ }
+
+ @Override
+ public void run() {
+ for (int idx = 0; idx < 100; idx++) {
+ try {
+ getProducer().newMessage().key(getKey()).value(getValue()).send();
+ } catch (PulsarClientException e) {
+ log.error("Unable to connect to Pulsar", e);
+ }
+ }
+ }
+
+ String getKey() {
+ Integer i = Integer.valueOf(rnd.nextInt(999999));
+ return i.toString();
+ }
+
+ abstract T getValue();
+
+ abstract Schema getSchema();
+
+ private PulsarClient getPulsarClient() throws PulsarClientException {
+ if (client == null) {
+ client = PulsarClient.builder()
+ .serviceUrl(brokerUrl)
+ .build();
+ }
+
+ return client;
+ }
+
+ private Producer getProducer() throws PulsarClientException {
+ if (producer == null) {
+ producer = getPulsarClient().newProducer(getSchema())
+ .topic(inputTopic)
+ .create();
+ }
+
+ return producer;
+ }
+
+}
diff --git a/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/producers/InputTopicStringProducer.java b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/producers/InputTopicStringProducer.java
new file mode 100644
index 0000000000000..7ada099728ea0
--- /dev/null
+++ b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/producers/InputTopicStringProducer.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 org.apache.pulsar.io.cassandra.producers;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import java.io.UncheckedIOException;
+import java.util.Random;
+import java.util.SortedMap;
+import java.util.TreeMap;
+import org.apache.pulsar.client.api.Schema;
+
+@SuppressWarnings({"unchecked", "rawtypes"})
+public class InputTopicStringProducer extends InputTopicProducerThread {
+
+ private final Random rnd = new Random();
+ int lastReadingId = rnd.nextInt(900000);
+
+ public InputTopicStringProducer(String brokerUrl, String inputTopic) {
+ super(brokerUrl, inputTopic);
+ }
+
+ @Override
+ String getValue() {
+
+ SortedMap elements = new TreeMap();
+ elements.put("readingid", lastReadingId++ + "");
+ elements.put("avg_ozone", rnd.nextDouble());
+ elements.put("min_ozone", rnd.nextDouble());
+ elements.put("max_ozone", rnd.nextDouble());
+ elements.put("avg_pm10", rnd.nextDouble());
+ elements.put("min_pm10", rnd.nextDouble());
+ elements.put("max_pm10", rnd.nextDouble());
+ elements.put("avg_pm25", rnd.nextDouble());
+ elements.put("min_pm25", rnd.nextDouble());
+ elements.put("max_pm25", rnd.nextDouble());
+ elements.put("local_time_zone", "PST");
+ elements.put("state_code", "CA");
+ elements.put("reporting_area", lastReadingId + "");
+ elements.put("hour_observed", rnd.nextInt(24));
+ elements.put("date_observed", "2022-06-18");
+ elements.put("latitude", 40.021f);
+ elements.put("longitude", -122.33f);
+
+ ObjectMapper objectMapper = new ObjectMapper();
+ try {
+ return objectMapper.writeValueAsString(elements);
+ } catch (JsonProcessingException e) {
+ throw new UncheckedIOException(e);
+ }
+ }
+
+ @Override
+ Schema getSchema() {
+ return Schema.STRING;
+ }
+
+}
diff --git a/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/producers/ObservationSchemaRecordProducer.java b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/producers/ObservationSchemaRecordProducer.java
new file mode 100644
index 0000000000000..e77ebf72b70b3
--- /dev/null
+++ b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/producers/ObservationSchemaRecordProducer.java
@@ -0,0 +1,60 @@
+/*
+ * 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.producers;
+
+import org.apache.pulsar.client.api.Schema;
+import org.apache.pulsar.client.api.schema.GenericRecord;
+import org.apache.pulsar.client.api.schema.RecordSchemaBuilder;
+import org.apache.pulsar.client.api.schema.SchemaBuilder;
+import org.apache.pulsar.common.schema.SchemaInfo;
+import org.apache.pulsar.common.schema.SchemaType;
+
+@SuppressWarnings({"unchecked", "rawtypes"})
+public class ObservationSchemaRecordProducer extends AbstractGenericRecordProducer {
+
+ public ObservationSchemaRecordProducer(String brokerUrl, String inputTopic) {
+ super(brokerUrl, inputTopic);
+ }
+
+ @Override
+ GenericRecord getValue() {
+ String val = "Some random string";
+
+ GenericRecord record = Schema.generic(getGenericSchemaInfo()).newRecordBuilder()
+ .set("key", getKey())
+ .set("observed", val)
+ .build();
+
+ return record;
+ }
+
+ @Override
+ SchemaInfo getGenericSchemaInfo() {
+ RecordSchemaBuilder recordSchemaBuilder =
+ SchemaBuilder.record("airquality.observation");
+
+ recordSchemaBuilder.field("key")
+ .type(SchemaType.STRING).required();
+
+ recordSchemaBuilder.field("observed")
+ .type(SchemaType.STRING).optional();
+
+ return recordSchemaBuilder.build(SchemaType.AVRO);
+ }
+}
diff --git a/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/producers/ReadingSchemaRecordProducer.java b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/producers/ReadingSchemaRecordProducer.java
new file mode 100644
index 0000000000000..e533a32ef6b23
--- /dev/null
+++ b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/producers/ReadingSchemaRecordProducer.java
@@ -0,0 +1,92 @@
+/*
+ * 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.producers;
+
+import java.util.Random;
+import org.apache.pulsar.client.api.Schema;
+import org.apache.pulsar.client.api.schema.GenericRecord;
+import org.apache.pulsar.client.api.schema.RecordSchemaBuilder;
+import org.apache.pulsar.client.api.schema.SchemaBuilder;
+import org.apache.pulsar.common.schema.SchemaInfo;
+import org.apache.pulsar.common.schema.SchemaType;
+
+@SuppressWarnings({"unchecked", "rawtypes"})
+public class ReadingSchemaRecordProducer extends AbstractGenericRecordProducer {
+
+ private Random rnd = new Random();
+
+ int lastReadingId = rnd.nextInt(50000);
+
+ public ReadingSchemaRecordProducer(String brokerUrl, String inputTopic) {
+ super(brokerUrl, inputTopic);
+ }
+
+ @Override
+ GenericRecord getValue() {
+ GenericRecord record = Schema.generic(getGenericSchemaInfo())
+ .newRecordBuilder()
+ .set("readingid", lastReadingId++ + "")
+ .set("avg_ozone", rnd.nextDouble())
+ .set("min_ozone", rnd.nextDouble())
+ .set("max_ozone", rnd.nextDouble())
+ .set("avg_pm10", rnd.nextDouble())
+ .set("min_pm10", rnd.nextDouble())
+ .set("max_pm10", rnd.nextDouble())
+ .set("avg_pm25", rnd.nextDouble())
+ .set("min_pm25", rnd.nextDouble())
+ .set("max_pm25", rnd.nextDouble())
+ .set("local_time_zone", "PST")
+ .set("state_code", "CA")
+ .set("reporting_area", lastReadingId + "")
+ .set("hour_observed", rnd.nextInt(24))
+ .set("date_observed", "2022-06-18")
+ .set("latitude", Float.valueOf(40.021f))
+ .set("longitude", Float.valueOf(-122.33f))
+ .build();
+
+ return record;
+ }
+
+ @Override
+ SchemaInfo getGenericSchemaInfo() {
+ RecordSchemaBuilder schemaBuilder =
+ SchemaBuilder.record("airquality.reading");
+
+ schemaBuilder.field("readingid").type(SchemaType.STRING).required();
+ schemaBuilder.field("avg_ozone").type(SchemaType.DOUBLE);
+ schemaBuilder.field("min_ozone").type(SchemaType.DOUBLE);
+ schemaBuilder.field("max_ozone").type(SchemaType.DOUBLE);
+ schemaBuilder.field("avg_pm10").type(SchemaType.DOUBLE);
+ schemaBuilder.field("min_pm10").type(SchemaType.DOUBLE);
+ schemaBuilder.field("max_pm10").type(SchemaType.DOUBLE);
+ schemaBuilder.field("avg_pm25").type(SchemaType.DOUBLE);
+ schemaBuilder.field("min_pm25").type(SchemaType.DOUBLE);
+ schemaBuilder.field("max_pm25").type(SchemaType.DOUBLE);
+
+ schemaBuilder.field("local_time_zone").type(SchemaType.STRING);
+ schemaBuilder.field("state_code").type(SchemaType.STRING);
+ schemaBuilder.field("reporting_area").type(SchemaType.STRING).required();
+ schemaBuilder.field("hour_observed").type(SchemaType.INT32);
+ schemaBuilder.field("date_observed").type(SchemaType.STRING);
+ schemaBuilder.field("latitude").type(SchemaType.FLOAT);
+ schemaBuilder.field("longitude").type(SchemaType.FLOAT);
+
+ return schemaBuilder.build(SchemaType.AVRO);
+ }
+}
diff --git a/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/util/AbstractCassandraTest.java b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/util/AbstractCassandraTest.java
new file mode 100644
index 0000000000000..cf88a334e1d65
--- /dev/null
+++ b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/util/AbstractCassandraTest.java
@@ -0,0 +1,53 @@
+/*
+ * 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 java.net.InetSocketAddress;
+import org.apache.pulsar.io.cassandra.CassandraSinkConfig;
+import org.testcontainers.cassandra.CassandraContainer;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+
+public class AbstractCassandraTest {
+ protected CassandraSinkConfig config;
+ protected CassandraContainer cassandraContainer;
+
+ @BeforeClass
+ public void startCassandraContainer() {
+ cassandraContainer = new CassandraContainer("cassandra:3.11");
+ cassandraContainer.withInitScript("init.cql");
+ cassandraContainer.start();
+ }
+
+ @AfterClass(alwaysRun = true)
+ public void stopCassandraContainer() {
+ if (cassandraContainer != null) {
+ cassandraContainer.stop();
+ cassandraContainer = null;
+ }
+ }
+
+ protected void createSinkConfig() {
+ config = new CassandraSinkConfig();
+ InetSocketAddress contactPoint = cassandraContainer.getContactPoint();
+ config.setRoots(contactPoint.getHostString() + ":" + contactPoint.getPort());
+ config.setUserName(cassandraContainer.getUsername());
+ config.setPassword(cassandraContainer.getPassword());
+ }
+}
diff --git a/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/util/CassandraConnectorTest.java b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/util/CassandraConnectorTest.java
new file mode 100644
index 0000000000000..079f112180d67
--- /dev/null
+++ b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/util/CassandraConnectorTest.java
@@ -0,0 +1,64 @@
+/*
+ * 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 static org.testng.AssertJUnit.assertEquals;
+import static org.testng.AssertJUnit.assertNotNull;
+import lombok.Cleanup;
+import org.testng.annotations.Test;
+
+public class CassandraConnectorTest extends AbstractCassandraTest {
+
+ @Test
+ public final void securedTest() {
+ createSinkConfig();
+
+ @Cleanup
+ CassandraConnector connector = new CassandraConnector(config);
+ connector.connect();
+ assertNotNull(connector.getSession());
+ }
+
+ @Test
+ public final void getObservationPreparedStatementTest() {
+ createSinkConfig();
+ config.setColumnFamily("observation");
+ config.setKeyspace("airquality");
+
+ @Cleanup
+ CassandraConnector connector = new CassandraConnector(config);
+ assertEquals("INSERT INTO airquality.observation (key, observed) VALUES (?, ?)",
+ connector.getPreparedStatement().getQueryString());
+ }
+
+ @Test
+ public final void getReadingPreparedStatementTest() {
+ createSinkConfig();
+ config.setColumnFamily("reading");
+ config.setKeyspace("airquality");
+
+ @Cleanup
+ CassandraConnector connector = new CassandraConnector(config);
+ assertEquals("INSERT INTO airquality.reading "
+ + "(reporting_area, date_observed, hour_observed, avg_ozone, avg_pm10, avg_pm25, latitude, "
+ + "local_time_zone, longitude, max_ozone, max_pm10, max_pm25, min_ozone, min_pm10, min_pm25, "
+ + "readingid, state_code) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
+ connector.getPreparedStatement().getQueryString());
+ }
+}
diff --git a/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/util/TableMetadataProviderTest.java b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/util/TableMetadataProviderTest.java
new file mode 100644
index 0000000000000..57ef5cb96024c
--- /dev/null
+++ b/pulsar-io/cassandra/src/test/java/org/apache/pulsar/io/cassandra/util/TableMetadataProviderTest.java
@@ -0,0 +1,48 @@
+/*
+ * 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 static org.testng.Assert.assertNotNull;
+import static org.testng.AssertJUnit.assertEquals;
+import lombok.Cleanup;
+import org.testng.annotations.Test;
+
+public class TableMetadataProviderTest extends AbstractCassandraTest {
+
+ @Test
+ public final void getTableDefinitionTest() {
+
+ createSinkConfig();
+ config.setColumnFamily("observation");
+ config.setKeyspace("airquality");
+
+ @Cleanup
+ CassandraConnector connector = new CassandraConnector(config);
+
+ TableMetadataProvider.TableDefinition table =
+ TableMetadataProvider.getTableDefinition(
+ connector.getTableMetadata(),
+ "airquality", "reading");
+
+ assertNotNull(table);
+ assertEquals(17, table.getColumns().size());
+ assertEquals(1, table.getPartitionKeyColumns().size());
+ assertEquals(3, table.getPrimaryKeyColumns().size());
+ }
+}
diff --git a/pulsar-io/cassandra/src/test/resources/cassandra-sink-config.yaml b/pulsar-io/cassandra/src/test/resources/cassandra-sink-config.yaml
new file mode 100644
index 0000000000000..0ba3bdb878728
--- /dev/null
+++ b/pulsar-io/cassandra/src/test/resources/cassandra-sink-config.yaml
@@ -0,0 +1,25 @@
+#
+# 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.
+#
+
+roots: "localhost"
+keyspace : "airquality"
+columnFamily: "reading"
+userName : "cassandra"
+password : "cassandra"
+
diff --git a/pulsar-io/cassandra/src/test/resources/init.cql b/pulsar-io/cassandra/src/test/resources/init.cql
new file mode 100644
index 0000000000000..8ad23b3cd92a4
--- /dev/null
+++ b/pulsar-io/cassandra/src/test/resources/init.cql
@@ -0,0 +1,60 @@
+-- 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.
+
+-- Create keyspace for air quality data
+CREATE KEYSPACE IF NOT EXISTS airquality
+WITH replication = {
+ 'class': 'SimpleStrategy',
+ 'replication_factor': 1
+};
+
+-- Use the airquality keyspace
+USE airquality;
+
+-- Create observation table
+CREATE TABLE IF NOT EXISTS observation (
+ key text PRIMARY KEY,
+ observed text
+);
+
+-- Create reading table with air quality measurements
+CREATE TABLE IF NOT EXISTS reading (
+ reporting_area text,
+ date_observed text,
+ hour_observed int,
+ readingid text,
+ avg_ozone double,
+ min_ozone double,
+ max_ozone double,
+ avg_pm10 double,
+ min_pm10 double,
+ max_pm10 double,
+ avg_pm25 double,
+ min_pm25 double,
+ max_pm25 double,
+ local_time_zone text,
+ state_code text,
+ latitude float,
+ longitude float,
+ PRIMARY KEY ((reporting_area), date_observed, hour_observed)
+) WITH CLUSTERING ORDER BY (date_observed DESC, hour_observed DESC);
+
+-- Create index on readingid for queries by reading ID
+CREATE INDEX IF NOT EXISTS reading_id_idx ON reading (readingid);
+
+-- Create index on state_code for queries by state
+CREATE INDEX IF NOT EXISTS state_code_idx ON reading (state_code);