Merge pull request #45 from github/fletchto99/kv-increment

kv: Add increment functionality

Author: zerowidth

.travis.yml

@@ -13,3 +13,5 @@ env:
   - RAILS_VERSION=5.1.5
   - RAILS_VERSION=5.0.6
   - RAILS_VERSION=4.2.10
+services:
+  - mysql

github-ds.gemspec

@@ -39,4 +39,6 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency "activesupport"
   spec.add_development_dependency "mysql2"
   spec.add_development_dependency "mocha", "~> 1.2.1"
+  spec.add_development_dependency "minitest-focus", "~> 1.1.2"
+  spec.add_development_dependency "pry", "~> 0.12.2"
 end

lib/github/ds/version.rb

@@ -1,5 +1,5 @@
 module GitHub
   module DS
-    VERSION = "0.2.11"
+    VERSION = "0.3.0"
   end
 end

lib/github/kv.rb

@@ -48,6 +48,7 @@ class KV
     KeyLengthError = Class.new(StandardError)
     ValueLengthError = Class.new(StandardError)
     UnavailableError = Class.new(StandardError)
+    InvalidValueError = Class.new(StandardError)
 
     class MissingConnectionError < StandardError; end
 
@@ -252,6 +253,102 @@ def setnx(key, value, expires: nil)
       }
     end
 
+    # increment :: String, Integer, expires: Time? -> Integer
+    #
+    # Increment the key's value by an amount.
+    #
+    # key             - The key to increment.
+    # amount          - The amount to increment the key's value by.
+    #                   The user can increment by both positive and
+    #                   negative values
+    # expires         - When the key should expire.
+    # touch_on_insert - Only when expires is specified. When true
+    #                   the expires value is only touched upon
+    #                   inserts. Otherwise the record is always
+    #                   touched.
+    #
+    # Returns the key's value after incrementing.
+    def increment(key, amount: 1, expires: nil, touch_on_insert: false)
+      validate_key(key)
+      validate_amount(amount) if amount
+      validate_expires(expires) if expires
+      validate_touch(touch_on_insert, expires)
+
+      expires ||= GitHub::SQL::NULL
+
+      # This query uses a few MySQL "hacks" to ensure that the incrementing
+      # is done atomically and the value is returned. The first trick is done
+      # using the `LAST_INSERT_ID` function. This allows us to manually set
+      # the LAST_INSERT_ID returned by the query. Here we are able to set it
+      # to the new value when an increment takes place, essentially allowing us
+      # to do: `UPDATE...;SELECT value from key_value where key=:key` in a
+      # single step.
+      #
+      # However the `LAST_INSERT_ID` trick is only used when the value is
+      # updated. Upon a fresh insert we know the amount is going to be set
+      # to the amount specified.
+      #
+      # Lastly we only do these tricks when the value at the key is an integer.
+      # If the value is not an integer the update ensures the values remain the
+      # same and we raise an error.
+      encapsulate_error {
+        sql = GitHub::SQL.run(<<-SQL, key: key, amount: amount, now: now, expires: expires, touch: !touch_on_insert, connection: connection)
+          INSERT INTO key_values (`key`, `value`, `created_at`, `updated_at`, `expires_at`)
+          VALUES(:key, :amount, :now, :now, :expires)
+          ON DUPLICATE KEY UPDATE
+            `value`=IF(
+              concat('',`value`*1) = `value`,
+              LAST_INSERT_ID(IF(
+                `expires_at` IS NULL OR `expires_at`>=:now,
+                `value`+:amount,
+                :amount
+              )),
+              `value`
+            ),
+            `updated_at`=IF(
+              concat('',`value`*1) = `value`,
+              :now,
+              `updated_at`
+            ),
+            `expires_at`=IF(
+              concat('',`value`*1) = `value`,
+              IF(
+                :touch,
+                :expires,
+                `expires_at`
+              ),
+              `expires_at`
+            )
+        SQL
+
+        # The ordering of these statements is extremely important if we are to
+        # support incrementing a negative amount. The checks occur in this order:
+        # 1. Check if an update with new values occurred? If so return the result
+        #    This could potentially result in `sql.last_insert_id` with a value
+        #    of 0, thus it must be before the second check.
+        # 2. Check if an update took place but nothing changed (I.E. no new value
+        #    was set)
+        # 3. Check if an insert took place.
+        #
+        # See https://dev.mysql.com/doc/refman/8.0/en/insert-on-duplicate.html for
+        # more information (NOTE: CLIENT_FOUND_ROWS is set)
+        if sql.affected_rows == 2
+          # An update took place in which data changed. We use a hack to set
+          # the last insert ID to be the new value.
+          sql.last_insert_id
+        elsif sql.affected_rows == 0 || (sql.affected_rows == 1 && sql.last_insert_id == 0)
+          # No insert took place nor did any update occur. This means that
+          # the value was not an integer thus not incremented.
+          raise InvalidValueError
+        elsif sql.affected_rows == 1
+          # If the number of affected_rows is 1 then a new value was inserted
+          # thus we can just return the amount given to us since that is the
+          # value at the key
+          amount
+        end
+      }
+    end
+
     # del :: String -> nil
     #
     # Deletes the specified key. Returns nil. Raises on error.
@@ -311,6 +408,28 @@ def ttl(key)
       }
     end
 
+    # mttl :: [String] -> Result<[Time | nil]>
+    #
+    # Returns the expires_at time for the specified key or nil.
+    #
+    # Example:
+    #
+    #  kv.mttl(["foo", "octocat"])
+    #    # => #<Result value: [2018-04-23 11:34:54 +0200, nil]>
+    #
+    def mttl(keys)
+      validate_key_array(keys)
+
+      Result.new {
+        kvs = GitHub::SQL.results(<<-SQL, :keys => keys, :now => now, :connection => connection).to_h
+          SELECT `key`, expires_at FROM key_values
+          WHERE `key` in :keys AND (expires_at IS NULL OR expires_at > :now)
+        SQL
+
+        keys.map { |key| kvs[key] }
+      }
+    end
+
   private
     def now
       use_local_time ? Time.now : GitHub::SQL::NOW
@@ -369,6 +488,19 @@ def validate_value_length(value)
       end
     end
 
+    def validate_amount(amount)
+      raise ArgumentError.new("The amount specified must be an integer") unless amount.is_a? Integer
+      raise ArgumentError.new("The amount specified cannot be 0") if amount == 0
+    end
+
+    def validate_touch(touch, expires)
+      raise ArgumentError.new("touch_on_insert must be a boolean value") unless [true, false].include?(touch)
+
+      if touch && expires.nil?
+        raise ArgumentError.new("Please specify an expires value if you wish to touch on insert")
+      end
+    end
+
     def validate_expires(expires)
       unless expires.respond_to?(:to_time)
         raise TypeError, "expires must be a time of some sort (Time, DateTime, ActiveSupport::TimeWithZone, etc.), but was #{expires.class}"

test/github/kv_test.rb

@@ -54,6 +54,159 @@ def test_set_failure
     end
   end
 
+  def test_increment_failure
+    ActiveRecord::Base.connection.stubs(:insert).raises(Errno::ECONNRESET)
+
+    assert_raises GitHub::KV::UnavailableError do
+      @kv.increment("foo")
+    end
+  end
+
+  def test_increment_default_value
+    result = @kv.increment("foo")
+
+    assert_equal 1, result
+    assert_nil @kv.ttl("foo").value!
+  end
+
+  def test_increment_large_value
+    result = @kv.increment("foo", amount: 10000)
+
+    assert_equal 10000, result
+  end
+
+  def test_increment_negative
+    result = @kv.increment("foo", amount: -1)
+
+    assert_equal -1, result
+  end
+
+  def test_increment_negative_to_0
+    @kv.set("foo", "1")
+    result = @kv.increment("foo", amount: -1)
+
+    assert_equal 0, result
+  end
+
+  def test_increment_multiple
+    @kv.increment("foo")
+    result = @kv.increment("foo")
+
+    assert_equal 2, result
+  end
+
+  def test_increment_multiple_different_values
+    @kv.increment("foo")
+    result = @kv.increment("foo", amount: 2)
+
+    assert_equal 3, result
+  end
+
+  def test_increment_existing
+    @kv.set("foo", "1")
+
+    result = @kv.increment("foo")
+
+    assert_equal 2, result
+  end
+
+  def test_increment_overwrites_expired_value
+    @kv.set("foo", "100", expires: 1.hour.ago)
+    expires = 1.hour.from_now.utc
+    result = @kv.increment("foo", expires: expires)
+
+    assert_equal expires.to_i, @kv.ttl("foo").value!.to_i
+    assert_equal 1, result
+  end
+
+  def test_increment_overwrites_expired_value_from_incrementing
+    @kv.increment("foo", expires: 1.hour.ago)
+    expires = 1.hour.from_now.utc
+    result = @kv.increment("foo", expires: expires)
+
+    assert_equal 1, result
+    assert_equal expires.to_i, @kv.ttl("foo").value!.to_i
+  end
+
+  def test_increment_overwrites_expired_with_nil
+    @kv.increment("foo", expires: 1.hour.ago)
+    result = @kv.increment("foo")
+
+    assert_equal 1, result
+    assert_nil @kv.ttl("foo").value!
+  end
+
+  def test_increment_sets_expires
+    expires = 1.hour.from_now.utc
+    @kv.increment("foo", expires: expires)
+
+    assert_equal expires.to_i, @kv.ttl("foo").value!.to_i
+  end
+
+  def test_increment_sets_expires_only_on_insert
+    expires = 1.hour.from_now.utc
+    result = @kv.increment("foo", expires: expires, touch_on_insert: true)
+    assert_equal 1, result
+
+    result = @kv.increment("foo", expires: 3.hours.from_now.utc, touch_on_insert: true)
+    assert_equal 2, result
+
+    assert_equal expires.to_i, @kv.ttl("foo").value!.to_i
+  end
+
+  def test_increment_sets_expires_only_on_insert_for_existing
+    expires = 1.hour.from_now.utc
+    @kv.set("foo", "100", expires: expires)
+
+    result = @kv.increment("foo", expires: 3.hours.from_now.utc, touch_on_insert: true)
+    assert_equal 101, result
+
+    assert_equal expires.to_i, @kv.ttl("foo").value!.to_i
+  end
+
+  def test_increment_updates_expires
+    expires = 2.hours.from_now.utc
+
+    @kv.set("foo", "100", expires: 1.hour.from_now)
+    result = @kv.increment("foo", expires: expires)
+
+    assert_equal 101, result
+    assert_equal expires.to_i, @kv.ttl("foo").value!.to_i
+  end
+
+  def test_increment_non_integer_key_value
+    @kv.set("foo", "bar")
+
+    assert_raises GitHub::KV::InvalidValueError do
+      @kv.increment("foo")
+    end
+    assert_equal "bar", @kv.get("foo").value!
+  end
+
+  def test_increment_only_accepts_integer_amounts
+    assert_raises ArgumentError do
+      @kv.increment("foo", amount: "bar")
+    end
+  end
+
+  def test_increment_only_accepts_integer_amounts
+    assert_raises ArgumentError do
+      @kv.increment("foo", amount: 0)
+    end
+  end
+
+  def test_increment_with_touch_expects_expires
+    assert_raises ArgumentError do
+      @kv.increment("foo", touch_on_insert: true)
+    end
+  end
+
+  def test_increment_with_touch_must_be_boolean
+    assert_raises ArgumentError do
+      @kv.increment("foo", touch_on_insert: "blahblah")
+    end
+  end
+
   def test_exists
     assert_equal false, @kv.exists("foo").value!
 
@@ -187,6 +340,18 @@ def test_ttl_for_key_that_exists_but_is_expired
     assert_nil @kv.ttl("foo-ttl").value!
   end
 
+  def test_mttl
+    assert_equal [nil, nil], @kv.mttl(["foo-ttl", "bar-ttl"]).value!
+
+    # the Time.at dance is necessary because MySQL does not support sub-second
+    # precision in DATETIME values
+    expires = Time.at(1.hour.from_now.to_i).utc
+    @kv.set("foo-ttl", "bar", expires: expires)
+
+    assert_equal [expires, nil], @kv.mttl(["foo-ttl", "bar-ttl"]).value!
+    assert_equal [nil, expires], @kv.mttl(["bar-ttl", "foo-ttl"]).value!
+  end
+
   def test_type_checks_key
     assert_raises TypeError do
       @kv.get(0)
@@ -240,6 +405,10 @@ def test_timecop
         # mset/mget
         @kv.mset({"foo" => "baz"}, expires: 1.day.from_now.utc)
         assert_equal ["baz"], @kv.mget(["foo"]).value!
+
+        # increment
+        @kv.increment("foo-increment", expires: 1.day.from_now.utc)
+        assert_equal 1, @kv.get("foo-increment").value!.to_i
       end
     end
   end

test/test_helper.rb

@@ -8,6 +8,8 @@
 require "timecop"
 require "minitest/autorun"
 require "mocha/mini_test"
+require "minitest/focus"
+require "pry"
 
 ActiveRecord::Base.configurations = {
   "without_database" => {