Support CQL Vector type (upgrade Core Driver to 4.16.0 and DSBulk to 1.11.0-SNAPSHOT)

Author: eolivelli

common/src/main/java/com/datastax/oss/common/sink/config/TopicConfig.java

@@ -16,6 +16,7 @@
 package com.datastax.oss.common.sink.config;
 
 import com.datastax.oss.common.sink.ConfigException;
+import com.datastax.oss.driver.internal.core.type.codec.registry.DefaultCodecRegistry;
 import com.datastax.oss.driver.shaded.guava.common.base.Splitter;
 import com.datastax.oss.dsbulk.codecs.api.ConversionContext;
 import com.datastax.oss.dsbulk.codecs.api.ConvertingCodecFactory;
@@ -139,7 +140,7 @@ public String toString() {
   }
 
   @NonNull
-  public ConvertingCodecFactory createCodecFactory() {
+  public ConvertingCodecFactory createCodecFactory(DefaultCodecRegistry defaultCodecRegistry) {
     ConversionContext context =
         new TextConversionContext()
             .setLocale(
@@ -150,7 +151,7 @@ public ConvertingCodecFactory createCodecFactory() {
             .setTimeZone(ZoneId.of(getString(getTopicSettingPath(topicName, TIMEZONE_OPT))))
             .setTimeUnit(
                 TimeUnit.valueOf(getString(getTopicSettingPath(topicName, TIME_UNIT_OPT))));
-    return new ConvertingCodecFactory(context);
+    return new ConvertingCodecFactory(defaultCodecRegistry, context);
   }
 
   /**

common/src/main/java/com/datastax/oss/common/sink/state/LifeCycleManager.java

@@ -53,16 +53,23 @@
 import com.datastax.oss.driver.api.core.metadata.schema.ColumnMetadata;
 import com.datastax.oss.driver.api.core.metadata.schema.KeyspaceMetadata;
 import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata;
+import com.datastax.oss.driver.api.core.type.CqlVectorType;
+import com.datastax.oss.driver.api.core.type.DataType;
 import com.datastax.oss.driver.api.core.type.DataTypes;
+import com.datastax.oss.driver.api.core.type.codec.TypeCodec;
+import com.datastax.oss.driver.api.core.type.reflect.GenericType;
 import com.datastax.oss.driver.internal.core.auth.PlainTextAuthProvider;
 import com.datastax.oss.driver.internal.core.config.typesafe.DefaultDriverConfigLoader;
 import com.datastax.oss.driver.internal.core.config.typesafe.DefaultProgrammaticDriverConfigLoaderBuilder;
+import com.datastax.oss.driver.internal.core.type.codec.CqlVectorCodec;
+import com.datastax.oss.driver.internal.core.type.codec.registry.DefaultCodecRegistry;
 import com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting;
 import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
 import com.datastax.oss.dsbulk.codecs.api.ConvertingCodecFactory;
 import com.typesafe.config.Config;
 import com.typesafe.config.ConfigFactory;
 import edu.umd.cs.findbugs.annotations.NonNull;
+import edu.umd.cs.findbugs.annotations.Nullable;
 import java.net.InetSocketAddress;
 import java.nio.file.Path;
 import java.util.Collection;
@@ -89,6 +96,21 @@ public class LifeCycleManager {
   private static final ConcurrentMap<String, InstanceState> INSTANCE_STATES =
       new ConcurrentHashMap<>();
   private static MetricRegistry metricRegistry = new MetricRegistry();
+  private static final DefaultCodecRegistry CODEC_REGISTRY =
+      new DefaultCodecRegistry("default-registry") {
+
+        protected TypeCodec<?> createCodec(
+            @Nullable DataType cqlType,
+            @Nullable GenericType<?> javaType,
+            boolean isJavaCovariant) {
+          if (cqlType instanceof CqlVectorType) {
+            log.info("Automatically Registering codec for CqlVectorType {}", cqlType);
+            CqlVectorType vectorType = (CqlVectorType) cqlType;
+            return new CqlVectorCodec<>(vectorType, codecFor(vectorType.getSubtype()));
+          }
+          return super.createCodec(cqlType, javaType, isJavaCovariant);
+        }
+      };
 
   /** This is a utility class that no one should instantiate. */
   private LifeCycleManager() {}
@@ -420,7 +442,8 @@ private static InstanceState buildInstanceState(CqlSession session, CassandraSin
             .stream()
             .map(
                 topicConfig -> {
-                  ConvertingCodecFactory codecFactory = topicConfig.createCodecFactory();
+                  ConvertingCodecFactory codecFactory =
+                      topicConfig.createCodecFactory(CODEC_REGISTRY);
                   TopicState topicState = new TopicState(codecFactory);
                   topicStates.put(topicConfig.getTopicName(), topicState);
 
@@ -485,6 +508,7 @@ public static CqlSession buildCqlSession(
     SslConfig sslConfig = config.getSslConfig();
     CqlSessionBuilder builder =
         new SessionBuilder(sslConfig)
+            .withCodecRegistry(CODEC_REGISTRY)
             .withApplicationVersion(version)
             .withApplicationName(applicationName)
             .withClientId(generateClientId(config.getInstanceName()));

common/src/test/java/com/datastax/oss/common/sink/RecordMapperTest.java

@@ -59,11 +59,11 @@
 import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList;
 import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
 import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableSet;
+import com.datastax.oss.dsbulk.codecs.api.format.temporal.CqlTemporalFormat;
+import com.datastax.oss.dsbulk.codecs.api.format.temporal.TemporalFormat;
+import com.datastax.oss.dsbulk.codecs.api.format.temporal.ZonedTemporalFormat;
 import com.datastax.oss.dsbulk.codecs.api.util.CodecUtils;
-import com.datastax.oss.dsbulk.codecs.api.util.CqlTemporalFormat;
 import com.datastax.oss.dsbulk.codecs.api.util.OverflowStrategy;
-import com.datastax.oss.dsbulk.codecs.api.util.TemporalFormat;
-import com.datastax.oss.dsbulk.codecs.api.util.ZonedTemporalFormat;
 import com.datastax.oss.dsbulk.codecs.text.string.StringToIntegerCodec;
 import com.datastax.oss.dsbulk.codecs.text.string.StringToLongCodec;
 import com.datastax.oss.protocol.internal.response.result.ColumnSpec;

pom.xml

@@ -40,8 +40,8 @@
     <java.release.version>8</java.release.version>
     <kafka.connect.version>2.4.0</kafka.connect.version>
     <caffeine.version>2.6.2</caffeine.version>
-    <oss.driver.version>4.6.0</oss.driver.version>
-    <dsbulk.version>1.6.0</dsbulk.version>
+    <oss.driver.version>4.16.0</oss.driver.version>
+    <dsbulk.version>1.11.0-SNAPSHOT</dsbulk.version>
     <reactive-streams.version>1.0.3</reactive-streams.version>
     <guava.version>25.1-jre</guava.version>
     <slf4j.version>1.7.25</slf4j.version>