Add Cluster#protocol_version

Watson1978 · Feb 6, 2024 · e7fa838 · e7fa838
1 parent 88a2f09
commit e7fa838
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 0 deletions.
diff --git a/ext/ilios/cluster.c b/ext/ilios/cluster.c
@@ -151,6 +151,23 @@ static VALUE cluster_connect_timeout(VALUE self, VALUE timeout_ms)
     return self;
 }
 
+/**
+ * Sets the protocol version. The driver will automatically downgrade to the lowest supported protocol version.
+ * Default is +PROTOCOL_VERSION_V4+.
+ *
+ * @param timeout_ms [Integer] A connect timeout in milliseconds.
+ * @return [Cassandra::Cluster] self.
+ */
+static VALUE cluster_protocol_version(VALUE self, VALUE version)
+{
+    CassandraCluster *cassandra_cluster;
+
+    GET_CLUSTER(self, cassandra_cluster);
+    cass_cluster_set_protocol_version(cassandra_cluster->cluster, NUM2INT(version));
+
+    return self;
+}
+
 /**
  * Sets the timeout for waiting for a response from a node.
  * Default is +12000+ milliseconds.
@@ -242,9 +259,17 @@ void Init_cluster(void)
     rb_define_method(cCluster, "hosts", cluster_hosts, 1);
     rb_define_method(cCluster, "port", cluster_port, 1);
     rb_define_method(cCluster, "keyspace", cluster_keyspace, 1);
+    rb_define_method(cCluster, "protocol_version", cluster_protocol_version, 1);
     rb_define_method(cCluster, "connect_timeout", cluster_connect_timeout, 1);
     rb_define_method(cCluster, "request_timeout", cluster_request_timeout, 1);
     rb_define_method(cCluster, "resolve_timeout", cluster_resolve_timeout, 1);
     rb_define_method(cCluster, "constant_speculative_execution_policy", cluster_constant_speculative_execution_policy, 2);
 
+    rb_define_const(cCluster, "PROTOCOL_VERSION_V1", INT2NUM(CASS_PROTOCOL_VERSION_V1));
+    rb_define_const(cCluster, "PROTOCOL_VERSION_V2", INT2NUM(CASS_PROTOCOL_VERSION_V2));
+    rb_define_const(cCluster, "PROTOCOL_VERSION_V3", INT2NUM(CASS_PROTOCOL_VERSION_V3));
+    rb_define_const(cCluster, "PROTOCOL_VERSION_V4", INT2NUM(CASS_PROTOCOL_VERSION_V4));
+    rb_define_const(cCluster, "PROTOCOL_VERSION_V5", INT2NUM(CASS_PROTOCOL_VERSION_V5));
+    rb_define_const(cCluster, "PROTOCOL_VERSION_DSEV1", INT2NUM(CASS_PROTOCOL_VERSION_DSEV1));
+    rb_define_const(cCluster, "PROTOCOL_VERSION_DSEV2", INT2NUM(CASS_PROTOCOL_VERSION_DSEV2));
 }
diff --git a/sig/ilios.rbs b/sig/ilios.rbs
@@ -7,6 +7,7 @@ module Ilios
       def hosts: (Array[String]) -> self
       def port: (Integer) -> self
       def keyspace: (String) -> self
+      def protocol_version: (Integer) -> self
       def connect_timeout: (Integer) -> self
       def request_timeout: (Integer) -> self
       def resolve_timeout: (Integer) -> self

diff --git a/test/test_cluster.rb b/test/test_cluster.rb
@@ -26,6 +26,13 @@ def test_keyspace
     assert_kind_of(Ilios::Cassandra::Cluster, cluster.keyspace('ilios'))
   end
 
+  def test_protocol_version
+    cluster = Ilios::Cassandra::Cluster.new
+
+    assert_raises(TypeError) { cluster.keyspace(Object.new) }
+    assert_kind_of(Ilios::Cassandra::Cluster, cluster.protocol_version(Ilios::Cassandra::Cluster::PROTOCOL_VERSION_V4))
+  end
+
   def test_connect_timeout
     cluster = Ilios::Cassandra::Cluster.new