Add config option for binary protocol version

src/driver/src/main/java/com/edgedb/driver/EdgeDBClientConfig.java

@@ -1,10 +1,13 @@
 package com.edgedb.driver;
 
+import com.edgedb.driver.binary.protocol.ProtocolProvider;
 import com.edgedb.driver.namingstrategies.NamingStrategy;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.time.Duration;
 import java.time.temporal.ChronoUnit;
+import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 
 public class EdgeDBClientConfig {
@@ -28,6 +31,15 @@ public class EdgeDBClientConfig {
     private ClientType clientType = ClientType.TCP;
     private int clientAvailability = 10;
     private Duration clientMaxAge = Duration.of(10, ChronoUnit.MINUTES);
+    private @Nullable ProtocolVersion forcedProtocolVersion;
+
+    /**
+     * Gets the forced protocol version to use with the client.
+     * @return an @{@linkplain Optional} wrapping a {@linkplain ProtocolVersion} if one is specified.
+     */
+    public Optional<ProtocolVersion> getForcedProtocolVersion() {
+        return Optional.ofNullable(this.forcedProtocolVersion);
+    }
 
     /**
      * Gets the number of attempts to try to connect.
@@ -190,6 +202,29 @@ public static final class Builder {
         private ClientType clientType = DEFAULT.clientType;
         private int clientAvailability = DEFAULT.clientAvailability;
         private Duration clientMaxAge = DEFAULT.clientMaxAge;
+        private @Nullable ProtocolVersion forcedProtocolVersion = DEFAULT.forcedProtocolVersion;
+
+        /**
+         *  Sets the forced protocol version of the current builder.
+         * @param version The version to set.
+         * @return The current builder.
+         * @throws IllegalArgumentException The provided version has no implementation.
+         */
+        public @NotNull Builder withForcedProtocolVersion(@Nullable ProtocolVersion version)
+                throws IllegalArgumentException {
+            if(version != null && !ProtocolProvider.supports(version)) {
+                throw new IllegalArgumentException(
+                        String.format(
+                                "The provided protocol version '%s' has no implementation",
+                                version
+                        )
+                );
+            }
+
+            this.forcedProtocolVersion = version;
+
+            return this;
+        }
 
         /**
          * Sets the pool size of the current builder.
@@ -365,6 +400,7 @@ public static final class Builder {
             edgeDBClientConfig.implicitTypeIds = this.implicitTypeIds;
             edgeDBClientConfig.retryMode = this.retryMode;
             edgeDBClientConfig.messageTimeout = this.messageTimeout;
+            edgeDBClientConfig.forcedProtocolVersion = this.forcedProtocolVersion;
             return edgeDBClientConfig;
         }
     }

src/driver/src/main/java/com/edgedb/driver/binary/protocol/ProtocolVersion.java → src/driver/src/main/java/com/edgedb/driver/ProtocolVersion.java

@@ -1,4 +1,4 @@
-package com.edgedb.driver.binary.protocol;
+package com.edgedb.driver;
 
 
 import org.joou.UShort;
@@ -7,15 +7,29 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
+/**
+ * Represents a <a href="https://www.edgedb.com/docs/reference/protocol/index#binary-protocol">EdgeDB Binary Protocol</a>
+ * version, used to negotiate and specify which protocol version to use.
+ */
 public final class ProtocolVersion {
+    /**
+     * Gets the default {@linkplain ProtocolVersion} used by this version of the binding.
+     */
     public static final ProtocolVersion BINARY_PROTOCOL_DEFAULT_VERSION = of(2, 0);
 
     private static final ConcurrentMap<Integer, ProtocolVersion> VERSIONS = new ConcurrentHashMap<>() {{
         put(keyOf(1, 0), new ProtocolVersion(1, 0));
         put(keyOf(2, 0), new ProtocolVersion(2, 0));
     }};
 
+    /**
+     * The major part of the version
+     */
     public final short major;
+
+    /**
+     * The minor part of the version.
+     */
     public final short minor;
 
     private ProtocolVersion(int major, int minor) {
@@ -28,12 +42,24 @@ private ProtocolVersion(int version) {
         this.minor = (short)(version & 0xFFFF);
     }
 
+    /**
+     * Gets a protocol version representing the specified components.
+     * @param major The major component of the version.
+     * @param minor The minor component of the version.
+     * @return a protocol version matching the provided components.
+     */
     public static ProtocolVersion of(int major, int minor) {
         return VERSIONS != null
                 ? VERSIONS.computeIfAbsent(keyOf(major, minor), ProtocolVersion::new)
                 : new ProtocolVersion(major, minor);
     }
 
+    /**
+     * Gets a protocol version representing the specified components.
+     * @param major The major component of the version.
+     * @param minor The minor component of the version.
+     * @return a protocol version matching the provided components.
+     */
     public static ProtocolVersion of(UShort major, UShort minor) {
         return of(major.intValue(), minor.intValue());
     }
@@ -42,19 +68,37 @@ private static int keyOf(int major, int minor) {
         return ((major & 0xFFFF) << 16) + (minor & 0xFFFF);
     }
 
+    /**
+     * Determines if the provided components equal this {@linkplain ProtocolVersion}
+     * @param major The major component
+     * @param minor The minor component
+     * @return {@code true} if the components equal this {@linkplain ProtocolVersion}; otherwise {@code false}
+     */
     public boolean equals(UShort major, UShort minor) {
         return equals(major.intValue(), minor.intValue());
     }
 
+    /**
+     * Determines if the provided components equal this {@linkplain ProtocolVersion}
+     * @param major The major component
+     * @param minor The minor component
+     * @return {@code true} if the components equal this {@linkplain ProtocolVersion}; otherwise {@code false}
+     */
     public boolean equals(int major, int minor) {
         return this.major == major && this.minor == minor;
     }
 
+    /**
+     * {@inheritDoc}
+     */
     @Override
     public String toString() {
         return String.format("%d.%d", major, minor);
     }
 
+    /**
+     * {@inheritDoc}
+     */
     @Override
     public boolean equals(Object obj) {
         if(!(obj instanceof ProtocolVersion)) {
@@ -66,6 +110,9 @@ public boolean equals(Object obj) {
         return this.major == other.major && this.minor == other.minor;
     }
 
+    /**
+     * {@inheritDoc}
+     */
     @Override
     public int hashCode() {
         return Objects.hash(major, minor);

src/driver/src/main/java/com/edgedb/driver/binary/builders/CodecBuilder.java

@@ -7,7 +7,7 @@
 import com.edgedb.driver.binary.codecs.scalars.complex.DateTimeCodec;
 import com.edgedb.driver.binary.codecs.scalars.complex.RelativeDurationCodec;
 import com.edgedb.driver.binary.protocol.ProtocolProvider;
-import com.edgedb.driver.binary.protocol.ProtocolVersion;
+import com.edgedb.driver.ProtocolVersion;
 import com.edgedb.driver.binary.protocol.TypeDescriptorInfo;
 import com.edgedb.driver.binary.protocol.common.Cardinality;
 import com.edgedb.driver.binary.protocol.common.IOFormat;

src/driver/src/main/java/com/edgedb/driver/binary/protocol/ProtocolProvider.java

@@ -1,6 +1,7 @@
 package com.edgedb.driver.binary.protocol;
 
 import com.edgedb.driver.EdgeDBConnection;
+import com.edgedb.driver.ProtocolVersion;
 import com.edgedb.driver.binary.PacketReader;
 import com.edgedb.driver.binary.codecs.Codec;
 import com.edgedb.driver.binary.protocol.v1.V1ProtocolProvider;
@@ -26,13 +27,24 @@ public interface ProtocolProvider {
        put(ProtocolVersion.of(2, 0), V2ProtocolProvider::new);
     }};
 
+    static boolean supports(ProtocolVersion version) {
+        return PROVIDERS.containsKey(version);
+    }
+
     static ProtocolProvider getProvider(EdgeDBBinaryClient client) {
         return PROVIDERS_FACTORY.computeIfAbsent(
                 client.getConnectionArguments(),
                 ignored -> PROVIDERS.get(ProtocolVersion.BINARY_PROTOCOL_DEFAULT_VERSION)
         ).apply(client);
     }
 
+    static ProtocolProvider getProvider(EdgeDBBinaryClient client, ProtocolVersion version) {
+        return PROVIDERS_FACTORY.compute(
+                client.getConnectionArguments(),
+                (connection, old) -> PROVIDERS.get(version)
+        ).apply(client);
+    }
+
     static void updateProviderFor(EdgeDBBinaryClient client, ProtocolProvider provider) {
         PROVIDERS_FACTORY.put(client.getConnectionArguments(), PROVIDERS.get(provider.getVersion()));
     }

src/driver/src/main/java/com/edgedb/driver/binary/protocol/v1/V1ProtocolProvider.java

@@ -2,6 +2,7 @@
 
 import com.edgedb.driver.Capabilities;
 import com.edgedb.driver.ErrorCode;
+import com.edgedb.driver.ProtocolVersion;
 import com.edgedb.driver.binary.PacketReader;
 import com.edgedb.driver.binary.builders.CodecBuilder;
 import com.edgedb.driver.binary.codecs.*;

src/driver/src/main/java/com/edgedb/driver/binary/protocol/v2/V2ProtocolProvider.java

@@ -4,7 +4,7 @@
 import com.edgedb.driver.binary.builders.CodecBuilder;
 import com.edgedb.driver.binary.codecs.*;
 import com.edgedb.driver.binary.protocol.ProtocolProvider;
-import com.edgedb.driver.binary.protocol.ProtocolVersion;
+import com.edgedb.driver.ProtocolVersion;
 import com.edgedb.driver.binary.protocol.TypeDescriptor;
 import com.edgedb.driver.binary.protocol.TypeDescriptorInfo;
 import com.edgedb.driver.binary.protocol.v1.V1ProtocolProvider;

src/driver/src/main/java/com/edgedb/driver/clients/EdgeDBBinaryClient.java

@@ -52,7 +52,12 @@ public EdgeDBBinaryClient(EdgeDBConnection connection, EdgeDBClientConfig config
         this.querySemaphore = new Semaphore(1);
         this.readyPromise = new CompletableFuture<>();
         this.stateDescriptorId = CodecBuilder.INVALID_CODEC_ID;
-        this.protocolProvider = ProtocolProvider.getProvider(this);
+
+        this.protocolProvider = config.getForcedProtocolVersion()
+                .map(v -> ProtocolProvider.getProvider(this, v))
+                .orElseGet(() -> ProtocolProvider.getProvider(this));
+
+        logger.debug("Picked protocol provider {}", this.protocolProvider);
     }
 
     public abstract Duplexer getDuplexer();