redcord 0.0.3 → 0.1.3

data/lib/redcord/migration/ttl.rb

@@ -2,10 +2,15 @@
 module Redcord::Migration::TTL
   extend T::Sig
 
-  # This won't change ttl until we call update on a record
+  sig { params(model: T.class_of(Redcord::Base)).returns(T.untyped) }
+  def _get_ttl(model)
+    model.class_variable_get(:@@ttl) || -1
+  end
+
   sig { params(model: T.class_of(Redcord::Base)).void }
-  def change_ttl_passive(model)
-    ttl = model.class_variable_get(:@@ttl)
-    model.redis.set("#{model.model_key}:ttl", ttl ? ttl : -1)
+  def change_ttl_active(model)
+    model.redis.scan_each_shard("#{model.model_key}:id:*") do |key|
+      model.redis.expire(key, _get_ttl(model))
+    end
   end
 end

data/lib/redcord/railtie.rb

@@ -29,7 +29,6 @@ class Redcord::Railtie < Rails::Railtie
       )
     end
 
-    Redcord::PreparedRedis.load_server_scripts!
     Redcord._after_initialize!
   end
 end

data/lib/redcord/redis.rb

@@ -0,0 +1,200 @@
+# typed: true
+require 'digest'
+require 'redis'
+require 'securerandom'
+
+class Redcord::Redis < Redis
+  extend T::Sig
+
+  sig do
+    params(
+      key: T.any(String, Symbol),
+      args: T::Hash[T.untyped, T.untyped],
+      ttl: T.nilable(Integer),
+      index_attrs: T::Array[Symbol],
+      range_index_attrs: T::Array[Symbol],
+      custom_index_attrs: T::Hash[Symbol, T::Array],
+      hash_tag: T.nilable(String),
+    ).returns(String)
+  end
+  def create_hash_returning_id(key, args, ttl:, index_attrs:, range_index_attrs:, custom_index_attrs:, hash_tag: nil)
+    id = "#{SecureRandom.uuid}#{hash_tag}"
+    custom_index_attrs_flat = custom_index_attrs.inject([]) do |result, (index_name, attrs)|
+      result << index_name
+      result << attrs.size
+      result + attrs
+    end
+    run_script(
+      :create_hash,
+      keys: [id, hash_tag],
+      argv: [key, ttl, index_attrs.size, range_index_attrs.size, custom_index_attrs_flat.size] +
+        index_attrs + range_index_attrs + custom_index_attrs_flat + args.to_a.flatten,
+    )
+    id
+  end
+
+  sig do
+    params(
+      model: String,
+      id: String,
+      args: T::Hash[T.untyped, T.untyped],
+      ttl: T.nilable(Integer),
+      index_attrs: T::Array[Symbol],
+      range_index_attrs: T::Array[Symbol],
+      custom_index_attrs: T::Hash[Symbol, T::Array],
+      hash_tag: T.nilable(String),
+    ).void
+  end
+  def update_hash(model, id, args, ttl:, index_attrs:, range_index_attrs:, custom_index_attrs:, hash_tag:)
+    custom_index_attrs_flat = custom_index_attrs.inject([]) do |result, (index_name, attrs)|
+      if !(args.keys.to_set & attrs.to_set).empty?
+        result << index_name
+        result << attrs.size
+        result + attrs
+      else
+        result
+      end
+    end
+    run_script(
+      :update_hash,
+      keys: [id, hash_tag],
+      argv: [model, ttl, index_attrs.size, range_index_attrs.size, custom_index_attrs_flat.size] +
+        index_attrs + range_index_attrs + custom_index_attrs_flat + args.to_a.flatten,
+    )
+  end
+
+  sig do
+    params(
+      model: String,
+      id: String,
+      index_attrs: T::Array[Symbol],
+      range_index_attrs: T::Array[Symbol],
+      custom_index_attrs: T::Hash[Symbol, T::Array],
+    ).returns(Integer)
+  end
+  def delete_hash(model, id, index_attrs:, range_index_attrs:, custom_index_attrs:)
+    custom_index_names = custom_index_attrs.keys
+    run_script(
+      :delete_hash,
+      keys: [id, id.match(/\{.*\}$/)&.send(:[], 0)],
+      argv: [model, index_attrs.size, range_index_attrs.size] + index_attrs + range_index_attrs + custom_index_names,
+    )
+  end
+
+  sig do
+    params(
+      model: String,
+      query_conditions: T::Hash[T.untyped, T.untyped],
+      index_attrs: T::Array[Symbol],
+      range_index_attrs: T::Array[Symbol],
+      select_attrs: T::Set[Symbol],
+      custom_index_attrs: T::Array[Symbol],
+      hash_tag: T.nilable(String),
+      custom_index_name: T.nilable(Symbol),
+    ).returns(T::Hash[Integer, T::Hash[T.untyped, T.untyped]])
+  end
+  def find_by_attr(
+        model,
+        query_conditions,
+        select_attrs: Set.new,
+        index_attrs:,
+        range_index_attrs:,
+        custom_index_attrs: Array.new,
+        hash_tag: nil,
+        custom_index_name: nil
+      )
+    conditions = flatten_with_partial_sort(query_conditions.clone, custom_index_attrs)
+    res = run_script(
+      :find_by_attr,
+      keys: [hash_tag],
+      argv: [model, custom_index_name, index_attrs.size, range_index_attrs.size, custom_index_attrs.size, conditions.size] + 
+        index_attrs + range_index_attrs + custom_index_attrs + conditions + select_attrs.to_a.flatten
+    )
+    # The Lua script will return this as a flattened array.
+    # Convert the result into a hash of {id -> model hash}
+    res_hash = res.each_slice(2)
+    res_hash.map { |key, val| [key, val.each_slice(2).to_h] }.to_h
+  end
+
+  sig do
+    params(
+      model: String,
+      query_conditions: T::Hash[T.untyped, T.untyped],
+      index_attrs: T::Array[Symbol],
+      range_index_attrs: T::Array[Symbol],
+      custom_index_attrs: T::Array[Symbol],
+      hash_tag: T.nilable(String),
+      custom_index_name: T.nilable(Symbol),
+    ).returns(Integer)
+  end
+  def find_by_attr_count(
+        model,
+        query_conditions,
+        index_attrs:,
+        range_index_attrs:,
+        custom_index_attrs: Array.new,
+        hash_tag: nil,
+        custom_index_name: nil
+      )
+    conditions = flatten_with_partial_sort(query_conditions.clone, custom_index_attrs)
+    run_script(
+      :find_by_attr_count,
+      keys: [hash_tag],
+      argv: [model, custom_index_name, index_attrs.size, range_index_attrs.size, custom_index_attrs.size] +
+        index_attrs + range_index_attrs + custom_index_attrs + conditions
+    )
+  end
+
+  def scan_each_shard(key, count: 1000, &blk)
+    clients = instance_variable_get(:@client)
+      &.instance_variable_get(:@node)
+      &.instance_variable_get(:@clients)
+      &.values
+
+    if clients.nil?
+      scan_each(match: key, count: count, &blk)
+    else
+      clients.each do |client|
+        cursor = 0
+        loop do
+          cursor, keys = client.call([:scan, cursor, 'match', key, 'count', count])
+          keys.each(&blk)
+          break if cursor == "0"
+        end
+      end
+    end
+  end
+
+  private
+
+  def run_script(script_name, *args)
+    # Use EVAL when a redis shard has not loaded the script before
+    hash_var_name = :"@script_sha_#{script_name}"
+    hash = instance_variable_get(hash_var_name)
+
+    begin
+      return evalsha(hash, *args) if hash
+    rescue Redis::CommandError => e
+      if e.message != 'NOSCRIPT No matching script. Please use EVAL.'
+        raise e
+      end
+    end
+
+    script_content = Redcord::LuaScriptReader.read_lua_script(script_name.to_s)
+    instance_variable_set(hash_var_name, Digest::SHA1.hexdigest(script_content))
+    self.eval(script_content, *args)
+  end
+
+  # When using custom index: On Lua side script expects query conditions sorted 
+  # in the order of appearance of attributes in specified index
+  sig { params(query_conditions: T::Hash[T.untyped, T.untyped], partial_order: T::Array[Symbol]).returns(T::Array[T.untyped]) }
+  def flatten_with_partial_sort(query_conditions, partial_order)
+    conditions = partial_order.inject([]) do |result, attr|
+      if !query_conditions[attr].nil?
+        result << attr << query_conditions.delete(attr)
+      end
+      result.flatten
+    end
+    conditions += query_conditions.to_a.flatten
+  end
+end

data/lib/redcord/redis_connection.rb

@@ -5,12 +5,15 @@
 require 'rails'
 
 require 'redcord/lua_script_reader'
-require 'redcord/prepared_redis'
+require 'redcord/redis'
+require 'redcord/connection_pool'
 
 module Redcord::RedisConnection
   extend T::Sig
   extend T::Helpers
 
+  RedcordClientType = T.type_alias { T.any(Redcord::Redis, Redcord::ConnectionPool) }
+
   @connections = T.let(nil, T.nilable(T::Hash[String, T.untyped]))
   @procs_to_prepare = T.let([], T::Array[Proc])
 
@@ -29,17 +32,17 @@ module Redcord::RedisConnection
       (env_config[name.underscore] || env_config['default']).symbolize_keys
     end
 
-    sig { returns(Redcord::PreparedRedis) }
+    sig { returns(RedcordClientType) }
     def redis
       Redcord::RedisConnection.connections[name.underscore] ||= prepare_redis!
     end
 
-    sig { returns(Redcord::PreparedRedis) }
+    sig { returns(RedcordClientType) }
     def establish_connection
       Redcord::RedisConnection.connections[name.underscore] = prepare_redis!
     end
 
-    sig { params(redis: Redis).returns(Redcord::PreparedRedis) }
+    sig { params(redis: Redis).returns(RedcordClientType) }
     def redis=(redis)
       Redcord::RedisConnection.connections[name.underscore] =
         prepare_redis!(redis)
@@ -50,27 +53,23 @@ module Redcord::RedisConnection
     # definitions in each Redis query.
     #
     # TODO: Replace this with Redcord migrations
-    sig { params(client: T.nilable(Redis)).returns(Redcord::PreparedRedis) }
+    sig { params(client: T.nilable(Redis)).returns(RedcordClientType) }
     def prepare_redis!(client = nil)
-      return client if client.is_a?(Redcord::PreparedRedis)
-
-      client = Redcord::PreparedRedis.new(
-        **(
-          if client.nil?
-            connection_config
-          else
-            client.instance_variable_get(:@options)
-          end
-        ),
-        logger: Redcord::Logger.proxy,
-      )
-
-      client.pipelined do
-        Redcord::RedisConnection.procs_to_prepare.each do |proc_to_prepare|
-          proc_to_prepare.call(client)
+      return client if client.is_a?(Redcord::Redis) || client.is_a?(Redcord::ConnectionPool)
+
+      options = client.nil? ? connection_config : client.instance_variable_get(:@options)
+      client =
+        if options[:pool]
+          Redcord::ConnectionPool.new(
+            pool_size: options[:pool],
+            timeout: options[:connection_timeout] || 1.0,
+            **options
+          )
+        else
+          Redcord::Redis.new(**options, logger: Redcord::Logger.proxy)
         end
-      end
 
+      client.ping
       client
     end
   end
@@ -78,7 +77,7 @@ module Redcord::RedisConnection
   module InstanceMethods
     extend T::Sig
 
-    sig { returns(Redcord::PreparedRedis) }
+    sig { returns(RedcordClientType) }
     def redis
       self.class.redis
     end
@@ -108,3 +107,10 @@ module Redcord::RedisConnection
 
   mixes_in_class_methods(ClassMethods)
 end
+
+module Redcord
+  sig { void }
+  def self.establish_connections
+    Redcord::Base.descendants.select(&:name).each(&:establish_connection)
+  end
+end

data/lib/redcord/relation.rb

@@ -5,42 +5,62 @@
 require 'active_support/core_ext/array'
 require 'active_support/core_ext/module'
 
+module Redcord
+  class InvalidQuery < StandardError; end
+end
+
 class Redcord::Relation
   extend T::Sig
 
   sig { returns(T.class_of(Redcord::Base)) }
   attr_reader :model
 
-  sig { returns(T::Hash[Symbol, T.untyped]) }
-  attr_reader :query_conditions
-
   sig { returns(T::Set[Symbol]) }
   attr_reader :select_attrs
 
+  sig { returns(T.nilable(Symbol)) }
+  attr_reader :custom_index_name
+
+  sig { returns(T::Hash[Symbol, T.untyped]) }
+  attr_reader :regular_index_query_conditions
+
+  sig { returns(T::Hash[Symbol, T.untyped]) }
+  attr_reader :custom_index_query_conditions
+
   sig do
     params(
       model: T.class_of(Redcord::Base),
-      query_conditions: T::Hash[Symbol, T.untyped],
+      regular_index_query_conditions: T::Hash[Symbol, T.untyped],
+      custom_index_query_conditions: T::Hash[Symbol, T.untyped],
       select_attrs: T::Set[Symbol],
+      custom_index_name: T.nilable(Symbol)
     ).void
   end
   def initialize(
     model,
-    query_conditions = {},
-    select_attrs = Set.new
+    regular_index_query_conditions = {},
+    custom_index_query_conditions = {},
+    select_attrs = Set.new,
+    custom_index_name: nil
   )
     @model = model
-    @query_conditions = query_conditions
+    @regular_index_query_conditions = regular_index_query_conditions
+    @custom_index_query_conditions = custom_index_query_conditions
     @select_attrs = select_attrs
+    @custom_index_name = custom_index_name
   end
 
   sig { params(args: T::Hash[Symbol, T.untyped]).returns(Redcord::Relation) }
   def where(args)
     encoded_args = args.map do |attr_key, attr_val|
-      encoded_val = model.validate_and_encode_query(attr_key, attr_val)
+      encoded_val = model.validate_types_and_encode_query(attr_key, attr_val)
       [attr_key, encoded_val]
     end
-    query_conditions.merge!(encoded_args.to_h)
+
+    regular_index_query_conditions.merge!(encoded_args.to_h)
+    if custom_index_name
+      with_index(custom_index_name)
+    end
     self
   end
 
@@ -68,7 +88,29 @@ class Redcord::Relation
 
   sig { returns(Integer) }
   def count
-    redis.find_by_attr_count(model.model_key, query_conditions)
+    Redcord::Base.trace(
+     'redcord_relation_count',
+     model_name: model.name,
+    ) do
+      model.validate_index_attributes(query_conditions.keys, custom_index_name: custom_index_name)
+      redis.find_by_attr_count(
+        model.model_key,
+        extract_query_conditions!,
+        index_attrs: model._script_arg_index_attrs,
+        range_index_attrs: model._script_arg_range_index_attrs,
+        custom_index_attrs: model._script_arg_custom_index_attrs[custom_index_name],
+        hash_tag: extract_hash_tag!,
+        custom_index_name: custom_index_name
+      )
+    end
+  end
+
+  sig { params(index_name: T.nilable(Symbol)).returns(Redcord::Relation) }
+  def with_index(index_name)
+    @custom_index_name = index_name
+    adjusted_query_conditions = model.validate_and_adjust_custom_index_query_conditions(regular_index_query_conditions)
+    custom_index_query_conditions.merge!(adjusted_query_conditions)
+    self
   end
 
   delegate(
@@ -137,17 +179,52 @@ class Redcord::Relation
 
   private
 
+  sig { returns(T.nilable(String)) }
+  def extract_hash_tag!
+    attr = model.shard_by_attribute
+    return nil if attr.nil?
+
+    if !query_conditions.keys.include?(attr)
+      raise(
+        Redcord::InvalidQuery,
+        "Queries must contain attribute '#{attr}' since model #{model.name} is sharded by this attribute"
+      )
+    end
+
+    # Query conditions on custom index are always in form of range, even when query is by value condition is [value_x, value_x]
+    # When in fact query is by value, range is trasformed to a single value to pass the validation.
+    condition = query_conditions[attr]
+    if custom_index_name and condition.first == condition.last
+      condition = condition.first
+    end
+    case condition
+    when Integer, String
+      "{#{condition}}"
+    else
+      raise(
+        Redcord::InvalidQuery,
+        "Does not support query condition #{condition} on a Redis Cluster",
+      )
+    end
+  end
+
   sig { returns(T::Array[T.untyped]) }
   def execute_query
     Redcord::Base.trace(
      'redcord_relation_execute_query',
      model_name: model.name,
     ) do
+      model.validate_index_attributes(query_conditions.keys, custom_index_name: custom_index_name)
       if !select_attrs.empty?
         res_hash = redis.find_by_attr(
           model.model_key,
-          query_conditions,
-          select_attrs,
+          extract_query_conditions!,
+          select_attrs: select_attrs,
+          index_attrs: model._script_arg_index_attrs,
+          range_index_attrs: model._script_arg_range_index_attrs,
+          custom_index_attrs: model._script_arg_custom_index_attrs[custom_index_name],
+          hash_tag: extract_hash_tag!,
+          custom_index_name: custom_index_name
         )
 
         res_hash.map do |id, args|
@@ -158,7 +235,12 @@ class Redcord::Relation
       else
         res_hash = redis.find_by_attr(
           model.model_key,
-          query_conditions,
+          extract_query_conditions!,
+          index_attrs: model._script_arg_index_attrs,
+          range_index_attrs: model._script_arg_range_index_attrs,
+          custom_index_attrs: model._script_arg_custom_index_attrs[custom_index_name],
+          hash_tag: extract_hash_tag!,
+          custom_index_name: custom_index_name
         )
 
         res_hash.map { |id, args| model.coerce_and_set_id(args, id) }
@@ -166,8 +248,24 @@ class Redcord::Relation
     end
   end
 
-  sig { returns(Redcord::PreparedRedis) }
+  sig { returns(Redcord::RedisConnection::RedcordClientType) }
   def redis
     model.redis
   end
+
+  sig { returns(T::Hash[Symbol, T.untyped]) }
+  def query_conditions
+    custom_index_name ? custom_index_query_conditions : regular_index_query_conditions
+  end
+
+  sig { returns(T::Hash[Symbol, T.untyped]) }
+  def extract_query_conditions!
+    attr = model.shard_by_attribute
+    return query_conditions if attr.nil?
+
+    cond = query_conditions.reject { |key| key == attr }
+    raise Redcord::InvalidQuery, "Cannot query only by shard_by_attribute: #{attr}" if cond.empty?
+
+    cond
+  end
 end