RUBY-332: Add support for duration type (#262)

Author: absurdfarce

Rakefile

@@ -43,6 +43,8 @@ end
 
 Rake::TestTask.new(integration: :compile) do |t|
   t.libs.push 'lib'
+  t.libs.push 'support'
+  t.libs.push 'integration'
   t.test_files = FileList['integration/*_test.rb',
                           'integration/security/*_test.rb',
                           'integration/load_balancing/*_test.rb',

integration/integration_test_case.rb

@@ -16,11 +16,11 @@
 # limitations under the License.
 #++
 
-require File.dirname(__FILE__) + '/../support/ccm.rb'
-require File.dirname(__FILE__) + '/../support/retry.rb'
-require File.dirname(__FILE__) + '/schema_change_listener.rb'
-require 'minitest/unit'
+require 'ccm.rb'
+require 'retry.rb'
+require 'schema_change_listener.rb'
 require 'minitest/autorun'
+require 'minitest/unit'
 require 'cassandra'
 require 'delorean'
 require 'ansi/code'

integration/types/duration_test.rb

@@ -0,0 +1,54 @@
+# encoding: utf-8
+
+#--
+# Copyright DataStax, Inc.
+#
+# Licensed 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.
+#++
+
+require 'integration_test_case'
+
+class DurationTest < IntegrationTestCase
+
+  def setup
+    @@ccm_cluster.setup_schema("CREATE KEYSPACE foo WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}")
+  end
+
+  def test_can_insert_duration
+    skip("Duration type was added in DSE 5.1/C* 3.10") if (CCM.dse_version < '5.1' || CCM.cassandra_version < '3.10')
+
+    cluster = Cassandra.cluster
+    session = cluster.connect "foo"
+
+    durations = {
+      'week' => Cassandra::Types.duration.new(0,7,0),
+      'year' => Cassandra::Types.duration.new(12,0,0),
+      'day' => Cassandra::Types.duration.new(0,1,0)
+    }
+    session.execute 'CREATE TABLE bar ("name" varchar, "dur" duration, primary key ("name"))'
+
+    insert = Retry.with_attempts(5) { session.prepare "INSERT INTO foo.bar (name,dur) VALUES (?,?)" }    
+    durations.each_pair do |key,val|
+      Retry.with_attempts(5) { session.execute insert, arguments: [key,val] }
+    end
+
+    result = session.execute("SELECT * FROM foo.bar")
+    result.each do |row|
+      duration = row["dur"]
+      duration_name = row["name"]
+      assert_equal durations[duration_name], duration
+    end
+  ensure
+    cluster && cluster.close
+  end
+end

lib/cassandra.rb

@@ -777,6 +777,8 @@ def self.validate_and_massage_options(options)
 require 'cassandra/time'
 
 require 'cassandra/types'
+require 'cassandra/custom_data'
+require 'cassandra/duration'
 
 require 'cassandra/errors'
 require 'cassandra/compression'
@@ -788,7 +790,6 @@ def self.validate_and_massage_options(options)
 require 'cassandra/executors'
 require 'cassandra/future'
 require 'cassandra/cluster'
-require 'cassandra/custom_data'
 require 'cassandra/driver'
 require 'cassandra/host'
 require 'cassandra/session'

lib/cassandra/cluster/schema/cql_type_parser.rb

@@ -60,6 +60,7 @@ def lookup_type(node, types)
           when 'smallint'          then Cassandra::Types.smallint
           when 'time'              then Cassandra::Types.time
           when 'tinyint'           then Cassandra::Types.tinyint
+          when 'duration'          then Cassandra::Types.duration
           when 'map'               then
             Cassandra::Types.map(*node.children.map { |t| lookup_type(t, types)})
           when 'set'               then

lib/cassandra/duration.rb

@@ -0,0 +1,100 @@
+# encoding: utf-8
+
+#--
+# Copyright DataStax, Inc.
+#
+# Licensed 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.
+#++
+module Cassandra
+  module Types
+
+    class Duration < Type
+      include CustomData
+
+      @@four_byte_max = 2 ** 32
+      @@eight_byte_max = 2 ** 64
+
+      # @private
+      attr_reader :months, :days, :nanos
+
+      # @private
+      def initialize(months, days, nanos)
+        super(:duration)
+        @months = months
+        @days = days
+        @nanos = nanos
+      end
+
+      def new(*values)
+        Util.assert_size(3, values, "Duration type expects three values, #{values.size} were provided")
+        values.each { |v| Util.assert_type(Int, v) }
+        Util.assert (Util.encode_zigzag32(values[0]) < @@four_byte_max), "Months value must be a valid 32-bit integer"
+        Util.assert (Util.encode_zigzag32(values[1]) < @@four_byte_max), "Days value must be a valid 32-bit integer"
+        Util.assert (Util.encode_zigzag64(values[2]) < @@eight_byte_max), "Nanos value must be a valid 64-bit integer"
+        all_positive = values.all? {|i| i >= 0 }
+        all_negative = values.all? {|i| i <= 0 }
+        Util.assert (all_positive or all_negative), "Values in a duration must be uniformly positive or negative"
+        Duration.new *values
+      end
+
+      def assert(value, message = nil, &block)
+        Util.assert_instance_of(Duration, value, message, &block)
+      end
+
+      def to_s
+        "Duration: months => #{@months}, days => #{@days}, nanos => #{@nanos}"
+      end
+
+      def hash
+        @hash ||= begin
+          h = 17
+          h = 31 * h + @months.hash
+          h = 31 * h + @days.hash
+          h = 31 * h + @nanos.hash
+          h
+        end
+      end
+
+      def eql?(other)
+        other.is_a?(Duration) &&
+          @months == other.months &&
+          @days == other.days &&
+          @nanos == other.nanos
+      end
+
+      alias == eql?
+
+      def self.cql_type
+        Type.new(@kind)        
+      end
+
+      # Requirements for CustomData module
+      def self.deserialize(bytestr)
+        buffer = Cassandra::Protocol::CqlByteBuffer.new.append(bytestr)
+        Cassandra::Types::Duration.new(buffer.read_signed_vint,buffer.read_signed_vint,buffer.read_signed_vint)
+      end
+
+      def self.type
+        Cassandra::Types::Custom.new('org.apache.cassandra.db.marshal.DurationType')
+      end
+
+      def serialize
+        rv = Cassandra::Protocol::CqlByteBuffer.new
+        rv.append_signed_vint32(@months)
+        rv.append_signed_vint32(@days)
+        rv.append_signed_vint64(@nanos)
+        rv
+      end
+    end
+  end
+end

lib/cassandra/protocol/cql_byte_buffer.rb

@@ -287,6 +287,29 @@ def read_tinyint
               "Not enough bytes available to decode a tinyint: #{e.message}", e.backtrace
       end
 
+      def read_vint
+        n = read_byte
+
+        # Bits are indexed in Integer in little-endian order
+        bytes_to_read = 7.downto(0).take_while {|i| n[i] == 1}.size
+        return n unless bytes_to_read > 0
+
+        rv = n & (0xff >> bytes_to_read)
+        1.upto(bytes_to_read) do |idx|
+          new_byte = read_byte
+          rv <<= 8
+          rv |= (new_byte & 0xff)
+        end
+
+        rv
+      rescue RangeError => e
+        raise Errors::DecodingError, e.message, e.backtrace
+      end
+
+      def read_signed_vint
+        Util.decode_zigzag(read_vint)
+      end
+
       def append_tinyint(n)
         append([n].pack(Formats::CHAR_FORMAT))
       end
@@ -409,6 +432,25 @@ def append_float(n)
         append([n].pack(Formats::FLOAT_FORMAT))
       end
 
+      def append_vint(n)
+        send_bytes = Util.to_min_byte_array(n)
+        send_cnt = send_bytes.length
+
+        raise Errors::EncodingError, "Too many bytes (#{bytes_to_send.length}) to send!" if send_cnt > 8
+
+        send_cnt_byte = (0xff << (8 - send_cnt)) & 0xff
+        append([send_cnt_byte].pack(Formats::BYTES_FORMAT))
+        append(send_bytes.pack(Formats::BYTES_FORMAT))
+      end
+
+      def append_signed_vint32(n)
+        append_vint(Util.encode_zigzag32(n))
+      end
+
+      def append_signed_vint64(n)
+        append_vint(Util.encode_zigzag64(n))
+      end
+
       def eql?(other)
         other.eql?(to_str)
       end

lib/cassandra/protocol/v4.rb

@@ -61,6 +61,9 @@ def write_parameters(buffer, params, types, names = EMPTY_LIST)
 
       class Decoder
         def initialize(handler, compressor = nil, custom_type_handlers = {})
+          # In v4 the duration type is represented as a custom type so we always want to have
+          # a handler included here.  This handler can be overridden via the connection options.
+          custom_type_handlers[Cassandra::Types::Duration.type] ||= Cassandra::Types::Duration
           @handler    = handler
           @compressor = compressor
           @state      = :initial

lib/cassandra/types.rb

@@ -1700,5 +1700,9 @@ def udt(keyspace, name, *fields)
     def custom(name)
       Custom.new(name)
     end
+
+    def duration
+      Duration.new 0,0,0
+    end
   end
 end

lib/cassandra/util.rb

@@ -302,6 +302,26 @@ def assert_equal(expected, actual, message = nil, &block)
       end
     end
 
+    def to_byte_array(n)
+      [n].pack("Q>").unpack("C*")
+    end
+
+    def to_min_byte_array(n)
+      to_byte_array(n).drop_while {|i| i == 0}
+    end
+
+    def decode_zigzag(n)
+      (n >> 1) ^ -(n & 1)
+    end
+
+    def encode_zigzag32(n)
+      (n >> 31) ^ (n << 1)
+    end
+
+    def encode_zigzag64(n)
+      (n >> 63) ^ (n << 1)
+    end
+
     # @private
     LOWERCASE_REGEXP = /[[:lower:]\_]*/
     # @private