github-ds postgres version of GitHub::KV

initializer.rb

module BoxOutSports
  def self.kv
    @kv ||= KeyValueJsonb.new(use_local_time: Rails.env.test?) {
      ActiveRecord::Base.connection
    }
  end
end

key_value_jsonb.rb

require "github/result"
require "github/sql"

class KeyValueJsonb
  MAX_KEY_LENGTH = 255
  MAX_VALUE_LENGTH = 65535

  # Private: Valid types that can be stored for key.
  VALID_TYPES = [
    Array,
    Hash,
    String,
  ].freeze

  KeyLengthError = Class.new(StandardError)
  ValueLengthError = Class.new(StandardError)
  UnavailableError = Class.new(StandardError)
  InvalidValueError = Class.new(StandardError)

  class MissingConnectionError < StandardError; end

  # Initialize a new KeyValueJsonb instance.
  #
  # encapsulated_errors - An Array of Exception subclasses that, when raised,
  #                       will be replaced with UnavailableError.
  # use_local_time:     - Whether to use Ruby's `Time.now` instead of databases
  #                       `NOW()` function. This is mostly useful in testing
  #                       where time needs to be modified (eg. Timecop).
  #                       Default false.
  # table_name:         - The table name to use if you have more than one table.
  #                       Default "key_values".
  # &conn_block         - A block to call to open a new database connection.
  #
  # Returns nothing.
  def initialize(encapsulated_errors: [SystemCallError], use_local_time: false, table_name: "key_values", &conn_block)
    @encapsulated_errors = encapsulated_errors
    @use_local_time = use_local_time
    @table_name = table_name
    @conn_block = conn_block
  end

  def connection
    @conn_block.try(:call) || (raise MissingConnectionError, "KeyValueJsonb must be initialized with a block that returns a connection")
  end

  # get :: String -> Result<String | nil>
  #
  # Gets the value of the specified key.
  #
  # Example:
  #
  #   kv.get("foo")
  #     # => #<Result value: "bar">
  #
  #   kv.get("octocat")
  #     # => #<Result value: nil>
  #
  def get(key)
    validate_key(key)

    mget([key]).map { |values| values[0] }
  end

  # mget :: [String] -> Result<[String | nil]>
  #
  # Gets the values of all specified keys. Values will be returned in the
  # same order as keys are specified. nil will be returned in place of a
  # String for keys which do not exist.
  #
  # Example:
  #
  #   kv.mget(["foo", "octocat"])
  #     # => #<Result value: ["bar", nil]
  #
  def mget(keys)
    validate_key_array(keys)

    GitHub::Result.new {
      kvs = GitHub::SQL.results(<<-SQL, keys: keys, now: now, connection: connection).to_h
        SELECT key, value FROM #{@table_name} WHERE key IN :keys AND (expires_at IS NULL OR expires_at > :now)
      SQL

      keys.map do |key|
        if value = kvs[key]
          JSON.parse(kvs[key])
        end
      end
    }
  end

  # set :: String, String, expires: Time? -> nil
  #
  # Sets the specified key to the specified value. Returns nil. Raises on
  # error.
  #
  # Example:
  #
  #   kv.set("foo", "bar")
  #     # => nil
  #
  def set(key, value, expires: nil)
    validate_key(key)
    validate_value(value)

    mset({ key => value }, expires: expires)
  end

  # mset :: { String => String }, expires: Time? -> nil
  #
  # Sets the specified hash keys to their associated values, setting them to
  # expire at the specified time. Returns nil. Raises on error.
  #
  # Example:
  #
  #   kv.mset({ "foo" => "bar", "baz" => "quux" })
  #     # => nil
  #
  #   kv.mset({ "expires" => "soon" }, expires: 1.hour.from_now)
  #     # => nil
  #
  def mset(kvs, expires: nil)
    validate_key_value_hash(kvs)
    validate_expires(expires) if expires

    rows = kvs.map { |key, value|
      [key, JSON.generate(value), now, now, expires || GitHub::SQL::NULL]
    }

    encapsulate_error do
      GitHub::SQL.run(<<-SQL, rows: GitHub::SQL::ROWS(rows), connection: connection)
        INSERT INTO #{@table_name} (key, value, created_at, updated_at, expires_at)
        VALUES :rows
        ON CONFLICT (key) DO UPDATE SET
          value = EXCLUDED.value,
          updated_at = EXCLUDED.updated_at,
          expires_at = EXCLUDED.expires_at
      SQL
    end

    nil
  end

  # exists :: String -> Result<Boolean>
  #
  # Checks for existence of the specified key.
  #
  # Example:
  #
  #   kv.exists("foo")
  #     # => #<Result value: true>
  #
  #   kv.exists("octocat")
  #     # => #<Result value: false>
  #
  def exists(key)
    validate_key(key)

    mexists([key]).map { |values| values[0] }
  end

  # mexists :: [String] -> Result<[Boolean]>
  #
  # Checks for existence of all specified keys. Booleans will be returned in
  # the same order as keys are specified.
  #
  # Example:
  #
  #   kv.mexists(["foo", "octocat"])
  #     # => #<Result value: [true, false]>
  #
  def mexists(keys)
    validate_key_array(keys)

    GitHub::Result.new {
      existing_keys = GitHub::SQL.values(<<-SQL, keys: keys, now: now, connection: connection).to_set
        SELECT key FROM #{@table_name} WHERE key IN :keys AND (expires_at IS NULL OR expires_at > :now)
      SQL

      keys.map { |key| existing_keys.include?(key) }
    }
  end

  # del :: String -> nil
  #
  # Deletes the specified key. Returns nil. Raises on error.
  #
  # Example:
  #
  #   kv.del("foo")
  #     # => nil
  #
  def del(key)
    validate_key(key)

    mdel([key])
  end

  # mdel :: String -> nil
  #
  # Deletes the specified keys. Returns nil. Raises on error.
  #
  # Example:
  #
  #   kv.mdel(["foo", "octocat"])
  #     # => nil
  #
  def mdel(keys)
    validate_key_array(keys)

    encapsulate_error do
      GitHub::SQL.run(<<-SQL, keys: keys, connection: connection)
        DELETE FROM #{@table_name} WHERE key IN :keys
      SQL
    end

    nil
  end

  # ttl :: String -> Result<[Time | nil]>
  #
  # Returns the expires_at time for the specified key or nil.
  #
  # Example:
  #
  #  kv.ttl("foo")
  #    # => #<Result value: 2018-04-23 11:34:54 +0200>
  #
  #  kv.ttl("foo")
  #    # => #<Result value: nil>
  #
  def ttl(key)
    validate_key(key)

    GitHub::Result.new {
      GitHub::SQL.value(<<-SQL, key: key, now: now, connection: connection)
        SELECT expires_at FROM #{@table_name}
        WHERE key = :key AND (expires_at IS NULL OR expires_at > :now)
      SQL
    }
  end

  private

  def now
    @use_local_time ? Time.now : GitHub::SQL::NOW
  end

  def validate_key(key, error_message: nil)
    unless key.is_a?(String)
      raise TypeError, error_message || "key must be a String in #{self.class.name}, but was #{key.class}"
    end

    validate_key_length(key)
  end

  def validate_value(value, error_message: nil)
    unless VALID_TYPES.any? { |type| value.is_a?(type) }
      raise TypeError, error_message || "value must be a Array | Hash | String in #{self.class.name}, but was #{value.class}"
    end
  end

  def validate_key_array(keys)
    unless keys.is_a?(Array)
      raise TypeError, "keys must be a [String] in #{self.class.name}, but was #{keys.class}"
    end

    keys.each do |key|
      unless key.is_a?(String)
        raise TypeError, "keys must be a [String] in #{self.class.name}, but also saw at least one #{key.class}"
      end

      validate_key_length(key)
    end
  end

  def validate_key_value_hash(kvs)
    unless kvs.is_a?(Hash)
      raise TypeError, "kvs must be a {String => #{VALID_TYPES.join(" | ")}} in #{self.class.name}, but was #{kvs.class}"
    end

    kvs.each do |key, value|
      validate_key(key, error_message: "kvs must be a {String => [#{VALID_TYPES.join(" | ")}]} in #{self.class.name}, but also saw at least one key of type #{key.class}")
      validate_value(value, error_message: "kvs must be a {String => [#{VALID_TYPES.join(" | ")}]} in #{self.class.name}, but also saw at least one value of type #{value.class}")
    end
  end

  def validate_key_length(key)
    if key.length > MAX_KEY_LENGTH
      raise KeyLengthError, "key of length #{key.length} exceeds maximum key length of #{MAX_KEY_LENGTH}\n\nkey: #{key.inspect}"
    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}"
    end
  end

  def encapsulate_error
    yield
  rescue *@encapsulated_errors => error
    raise UnavailableError, "#{error.class}: #{error.message}"
  end
end

key_value_jsonb_test.rb

require "rails_helper"

describe KeyValueJsonb do
  before do
    @kv = KeyValueJsonb.new { ActiveRecord::Base.connection }
  end

  it "is preconfigured for the application" do
    expect(BoxOutSports.kv.get("foo").value!).to be(nil)
    expect(BoxOutSports.kv.instance_variable_get("@use_local_time")).to be(true)
  end

  it "initialize without connection" do
    kv = KeyValueJsonb.new
    expect { kv.get("foo").value! }.to raise_error(KeyValueJsonb::MissingConnectionError)
  end

  it "get and set" do
    expect(@kv.get("foo").value!).to be(nil)

    @kv.set("foo", "bar")

    expect(@kv.get("foo").value!).to eq("bar")
  end

  it "get and set with array value" do
    expect(@kv.get("foo").value!).to be(nil)

    @kv.set("foo", [1, 2, 3])

    expect(@kv.get("foo").value!).to eq([1, 2, 3])
  end

  it "get and set with hash value" do
    expect(@kv.get("foo").value!).to be(nil)

    @kv.set("foo", "token" => "asdf", "grant_type" => "booyeah")

    expect(@kv.get("foo").value!).to eq("token" => "asdf", "grant_type" => "booyeah")
  end

  it "mget and mset" do
    expect(@kv.mget(["a", "b"]).value!).to eq([nil, nil])

    @kv.mset("a" => "1", "b" => "2")

    expect(@kv.mget(["a", "b"]).value!).to eq(["1", "2"])
    expect(@kv.mget(["b", "a"]).value!).to eq(["2", "1"])
  end

  it "mget and mset with array values" do
    expect(@kv.mget(["a", "b"]).value!).to eq([nil, nil])

    @kv.mset("a" => [1, 1], "b" => [2, 2, 2])

    expect(@kv.mget(["a", "b"]).value!).to eq([[1, 1], [2, 2, 2]])
    expect(@kv.mget(["b", "a"]).value!).to eq([[2, 2, 2], [1, 1]])
  end

  it "get with failure" do
    allow(ActiveRecord::Base.connection).to receive(:select_all).and_raise(Errno::ECONNRESET)

    result = @kv.get("foo")

    expect(result.ok?).to be(false)
  end

  it "set with failure" do
    allow(ActiveRecord::Base.connection).to receive(:insert).and_raise(Errno::ECONNRESET)

    expect { @kv.set("foo", "bar") }.to raise_error(KeyValueJsonb::UnavailableError)
  end

  it "exists" do
    expect(@kv.exists("foo").value!). to be(false)

    @kv.set("foo", "bar")

    expect(@kv.exists("foo").value!). to be(true)
  end

  it "mexists" do
    @kv.set("foo", "bar")

    expect(@kv.mexists(["foo", "notfoo"]).value!).to eq([true, false])
    expect(@kv.mexists(["notfoo", "foo"]).value!).to eq([false, true])
  end

  it "del" do
    @kv.set("foo", "bar")
    @kv.del("foo")

    expect(@kv.get("foo").value!).to be(nil)
  end

  it "del with failure" do
    allow(ActiveRecord::Base.connection).to receive(:delete).and_raise(Errno::ECONNRESET)

    expect {
      @kv.del("foo")
    }.to raise_error(KeyValueJsonb::UnavailableError)
  end

  it "mdel" do
    @kv.set("foo", "bar")
    @kv.mdel(["foo", "notfoo"])

    expect(@kv.get("foo").value!).to be(nil)
    expect(@kv.get("notfoo").value!).to be(nil)
  end

  it "set with expiry" do
    expires = Time.at(1.hour.from_now.to_i).utc

    @kv.set("foo", "bar", expires: expires)

    actual = GitHub::SQL.value(<<-SQL)
      SELECT expires_at FROM key_values WHERE key = 'foo'
    SQL
    expect(actual).to eq(expires)
  end

  it "get respects expiry" do
    @kv.set("foo", "bar", expires: 1.hour.from_now)

    expect(@kv.get("foo").value!).to eq("bar")

    @kv.set("foo", "bar", expires: 1.hour.ago)

    expect(@kv.get("foo").value!).to be(nil)
  end

  it "exists respects expiry" do
    @kv.set("foo", "bar", expires: 1.hour.from_now)

    expect(@kv.exists("foo").value!).to be(true)

    @kv.set("foo", "bar", expires: 1.hour.ago)

    expect(@kv.exists("foo").value!).to be(false)
  end

  it "set resets expiry" do
    @kv.set("foo", "bar", expires: 1.hour.from_now)
    @kv.set("foo", "bar")

    expect(GitHub::SQL.value(<<-SQL)).to be(nil)
      SELECT expires_at FROM key_values WHERE key = 'foo'
    SQL
  end

  it "ttl" do
    expect(@kv.ttl("foo-ttl").value!).to be(nil)

    # the Time.at dance is necessary because db column does not support sub-second
    expires = Time.at(1.hour.from_now.to_i).utc
    @kv.set("foo-ttl", "bar", expires: expires)

    expect(@kv.ttl("foo-ttl").value!).to eq(expires)
  end

  it "ttl for key that exists but is expired" do
    @kv.set("foo-ttl", "bar", expires: 1.hour.ago)

    row_count = GitHub::SQL.value <<-SQL, key: "foo-ttl"
      SELECT count(*) FROM key_values WHERE key = :key
    SQL
    expect(row_count).to be(1)
    expect(@kv.ttl("foo-ttl").value!).to be(nil)
  end

  it "get type checks key" do
    expect { @kv.get(0) }.to raise_error(TypeError)
    expect { @kv.mget([0, 1]) }.to raise_error(TypeError)
  end

  it "get length checks key" do
    expect { @kv.get("A" * 256) }.to raise_error(KeyValueJsonb::KeyLengthError)
    expect { @kv.mget(["A" * 256]) }.to raise_error(KeyValueJsonb::KeyLengthError)
  end

  it "set type checks value" do
    expect { @kv.set("foo", 1) }.to raise_error(TypeError)
  end

  it "works with timecop if using local time option" do
    @kv = KeyValueJsonb.new(use_local_time: true) {
      ActiveRecord::Base.connection
    }

    Timecop.freeze(1.month.ago) do
      # set/get
      @kv.set("foo", "bar", expires: 1.day.from_now.utc)

      expect(@kv.get("foo").value!).to eq("bar")

      # exists
      expect(@kv.exists("foo").value!).to be(true)

      # ttl
      expect(@kv.ttl("foo").value!.to_i).to be(1.day.from_now.to_i)

      # mset/mget
      @kv.mset({"foo" => "baz"}, expires: 1.day.from_now.utc)
      expect(@kv.mget(["foo"]).value!).to eq(["baz"])
    end
  end
end

migration.rb

class CreateKeyValues < ActiveRecord::Migration[5.2]
  def change
    create_table :key_values do |t|
      t.string :key
      t.jsonb :value, default: {}, null: false
      t.datetime :expires_at, null: true
      t.timestamps
    end

    add_index :key_values, :key, unique: true
  end
end